fix(llm_guard): apply sanitized prompt returned by moderation API to request (#33331)

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

enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py

@@ -7,7 +7,8 @@
 #  Thank you users! We ❤️ you! - Krrish & Ishaan
 ## This provides an LLM Guard Integration for content moderation on the proxy
 
-from typing import Literal, Optional
+import asyncio
+from typing import Optional
 
 import aiohttp
 from fastapi import HTTPException
@@ -18,7 +19,6 @@ from litellm.integrations.custom_logger import CustomLogger
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.utils import CallTypesLiteral
-from litellm.utils import get_formatted_prompt
 
 
 class _ENTERPRISE_LLMGuard(CustomLogger):
@@ -46,45 +46,44 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
         except Exception:
             pass
 
-    async def moderation_check(self, text: str):
+    async def moderation_check(self, text: str) -> str:
         """
+        Runs the LLM Guard moderation check on ``text``.
+
+        Raises an HTTPException when the content violates the safety policy;
+        otherwise returns the sanitized prompt from LLM Guard, falling back to
+        the original text when the API does not provide one.
+
         [TODO] make this more performant for high-throughput scenario
         """
         try:
-            async with aiohttp.ClientSession() as session:
-                if self.mock_redacted_text is not None:
-                    redacted_text = self.mock_redacted_text
-                else:
-                    # Make the first request to /analyze
-                    analyze_url = f"{self.llm_guard_api_base}analyze/prompt"
-                    verbose_proxy_logger.debug("Making request to: %s", analyze_url)
-                    analyze_payload = {"prompt": text}
-                    redacted_text = None
+            if self.mock_redacted_text is not None:
+                redacted_text = self.mock_redacted_text
+            else:
+                analyze_url = f"{self.llm_guard_api_base}analyze/prompt"
+                verbose_proxy_logger.debug("Making request to: %s", analyze_url)
+                async with aiohttp.ClientSession() as session:
                     async with session.post(
-                        analyze_url, json=analyze_payload
+                        analyze_url, json={"prompt": text}
                     ) as response:
                         redacted_text = await response.json()
-                verbose_proxy_logger.debug(
-                    f"LLM Guard: Received response - {redacted_text}"
+            verbose_proxy_logger.debug(
+                f"LLM Guard: Received response - {redacted_text}"
+            )
+            if redacted_text is None:
+                raise HTTPException(
+                    status_code=500,
+                    detail={
+                        "error": f"Invalid content moderation response: {redacted_text}"
+                    },
                 )
-                if redacted_text is not None:
-                    if (
-                        redacted_text.get("is_valid", None) is not None
-                        and redacted_text["is_valid"] is False
-                    ):
-                        raise HTTPException(
-                            status_code=400,
-                            detail={"error": "Violated content safety policy"},
-                        )
-                    else:
-                        pass
-                else:
-                    raise HTTPException(
-                        status_code=500,
-                        detail={
-                            "error": f"Invalid content moderation response: {redacted_text}"
-                        },
-                    )
+            if redacted_text.get("is_valid", None) is False:
+                raise HTTPException(
+                    status_code=400,
+                    detail={"error": "Violated content safety policy"},
+                )
+            sanitized_prompt = redacted_text.get("sanitized_prompt")
+            return sanitized_prompt if isinstance(sanitized_prompt, str) else text
         except Exception as e:
             verbose_proxy_logger.exception(
                 "litellm.enterprise.enterprise_hooks.llm_guard::moderation_check - Exception occurred - {}".format(
@@ -138,23 +137,75 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
             return
 
         self.print_verbose("Makes LLM Guard Check")
-        try:
-            assert call_type in [
-                "completion",
-                "embeddings",
-                "image_generation",
-                "moderation",
-                "audio_transcription",
-            ]
-        except Exception:
+        if call_type not in [
+            "completion",
+            "embeddings",
+            "image_generation",
+            "moderation",
+            "audio_transcription",
+        ]:
             self.print_verbose(
                 f"Call Type - {call_type}, not in accepted list - ['completion','embeddings','image_generation','moderation','audio_transcription']"
             )
             return data
 
-        formatted_prompt = get_formatted_prompt(data=data, call_type=call_type)  # type: ignore
-        self.print_verbose(f"LLM Guard, formatted_prompt: {formatted_prompt}")
-        return await self.moderation_check(text=formatted_prompt)
+        return await self._moderate_request(data=data)
+
+    async def _moderate_request(self, data: dict) -> dict:
+        """
+        Sanitizes the request in place using the prompt returned by LLM Guard so
+        the provider-bound request carries the redacted content, then returns it.
+        """
+        messages = data.get("messages")
+        if messages is not None:
+            data["messages"] = list(
+                await asyncio.gather(
+                    *(self._moderate_message(message) for message in messages)
+                )
+            )
+            return data
+
+        input_ = data.get("input")
+        if input_ is not None:
+            data["input"] = await self._moderate_input(input_)
+            return data
+
+        prompt = data.get("prompt")
+        if isinstance(prompt, str):
+            data["prompt"] = await self.moderation_check(text=prompt)
+        return data
+
+    async def _moderate_message(self, message: dict) -> dict:
+        content = message.get("content")
+        if isinstance(content, str):
+            return {**message, "content": await self.moderation_check(text=content)}
+        if isinstance(content, list):
+            return {
+                **message,
+                "content": list(
+                    await asyncio.gather(
+                        *(self._moderate_content_part(part) for part in content)
+                    )
+                ),
+            }
+        return message
+
+    async def _moderate_content_part(self, part: dict) -> dict:
+        if part.get("type") == "text" and isinstance(part.get("text"), str):
+            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):
+            return [
+                await self.moderation_check(text=item)
+                if isinstance(item, str)
+                else item
+                for item in input_
+            ]
+        return input_
 
     async def async_post_call_streaming_hook(
         self, user_api_key_dict: UserAPIKeyAuth, response: str

tests/local_testing/test_llm_guard.py

@@ -28,8 +28,11 @@ from litellm.caching.caching import DualCache
 @pytest.mark.asyncio
 async def test_llm_guard_valid_response():
     """
-    Tests to see llm guard raises an error for a flagged response
+    A valid (is_valid=True) LLM Guard response must apply the returned
+    sanitized_prompt back onto the request data so the provider receives the
+    redacted content.
     """
+    litellm.llm_guard_mode = "all"
     input_a_anonymizer_results = {
         "sanitized_prompt": "hello world",
         "is_valid": True,
@@ -44,21 +47,65 @@ async def test_llm_guard_valid_response():
     user_api_key_dict = UserAPIKeyAuth(api_key=_api_key)
     local_cache = DualCache()
 
-    try:
-        await llm_guard.async_moderation_hook(
-            data={
-                "messages": [
-                    {
-                        "role": "user",
-                        "content": "hello world, my name is Jane Doe. My number is: 23r323r23r2wwkl",
-                    }
-                ]
-            },
-            user_api_key_dict=user_api_key_dict,
-            call_type="completion",
-        )
-    except Exception as e:
-        pytest.fail(f"An exception occurred - {str(e)}")
+    data = {
+        "messages": [
+            {
+                "role": "user",
+                "content": "hello world, my name is Jane Doe. My number is: 23r323r23r2wwkl",
+            }
+        ]
+    }
+
+    result = await llm_guard.async_moderation_hook(
+        data=data,
+        user_api_key_dict=user_api_key_dict,
+        call_type="completion",
+    )
+
+    assert result is data
+    assert data["messages"][0]["content"] == "hello world"
+
+
+@pytest.mark.asyncio
+async def test_llm_guard_sanitizes_multimodal_and_input():
+    """
+    Sanitization must reach text parts of multimodal message content and the
+    ``input`` field (embeddings/moderation) while leaving non-text parts intact.
+    """
+    litellm.llm_guard_mode = "all"
+    llm_guard = _ENTERPRISE_LLMGuard(
+        mock_testing=True,
+        mock_redacted_text={
+            "sanitized_prompt": "email: [REDACTED]",
+            "is_valid": True,
+            "scanners": {"Regex": 0.0},
+        },
+    )
+    user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-12345"))
+
+    image_part = {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}
+    data = {
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "email: person@example.com"},
+                    image_part,
+                ],
+            }
+        ]
+    }
+    result = await llm_guard.async_moderation_hook(
+        data=data, user_api_key_dict=user_api_key_dict, call_type="completion"
+    )
+    assert result["messages"][0]["content"][0]["text"] == "email: [REDACTED]"
+    assert result["messages"][0]["content"][1] == image_part
+
+    input_data = {"input": ["email: person@example.com", "another prompt"]}
+    input_result = await llm_guard.async_moderation_hook(
+        data=input_data, user_api_key_dict=user_api_key_dict, call_type="embeddings"
+    )
+    assert input_result["input"] == ["email: [REDACTED]", "email: [REDACTED]"]
 
 
 @pytest.mark.asyncio