diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py index 9c8537e6820..7338352106a 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py @@ -177,12 +177,12 @@ class _ENTERPRISE_LLMGuard(CustomLogger): input_ = data.get("input") if input_ is not None: - data["input"] = await self._moderate_input(input_) + data["input"] = await self._moderate_text_or_list(input_) return data prompt = data.get("prompt") - if isinstance(prompt, str): - data["prompt"] = await self.moderation_check(text=prompt) + if prompt is not None: + data["prompt"] = await self._moderate_text_or_list(prompt) return data async def _moderate_message(self, message: dict) -> dict: @@ -205,17 +205,17 @@ class _ENTERPRISE_LLMGuard(CustomLogger): return {**part, "text": await self.moderation_check(text=part["text"])} return part - async def _moderate_input(self, input_: object) -> object: - if isinstance(input_, str): - return await self.moderation_check(text=input_) - if isinstance(input_, list): + async def _moderate_text_or_list(self, value: object) -> object: + if isinstance(value, str): + return await self.moderation_check(text=value) + if isinstance(value, list): return [ await self.moderation_check(text=item) if isinstance(item, str) else item - for item in input_ + for item in value ] - return input_ + return value async def async_post_call_streaming_hook( self, user_api_key_dict: UserAPIKeyAuth, response: str diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py index dcb14e176e9..ef2aa96c36f 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py @@ -68,6 +68,32 @@ async def test_llm_guard_call_type_aliases( ) +@pytest.mark.parametrize("call_type", ("text_completion", "atext_completion")) +@pytest.mark.parametrize("is_valid", (True, False)) +@pytest.mark.asyncio +async def test_llm_guard_scans_list_prompt( + call_type: CallTypesLiteral, is_valid: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "llm_guard_mode", "all") + llm_guard: Final = _ENTERPRISE_LLMGuard( + mock_testing=True, + mock_redacted_text={"sanitized_prompt": "[REDACTED]", "is_valid": is_valid}, + ) + data: Final = {"prompt": ["email: person@example.com", "say ok", [1, 2, 3]]} + + if not is_valid: + with pytest.raises(HTTPException) as exc_info: + await llm_guard.async_moderation_hook(data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type) + assert exc_info.value.status_code == 400 + return + + result: Final = await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type + ) + assert result is data + assert data["prompt"] == ["[REDACTED]", "[REDACTED]", [1, 2, 3]] + + @pytest.mark.parametrize( "call_type", (