fix(llmguard): scan list valued completion prompts

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

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",
     (