feat: add guardrail for image generation (#15619)

* feat: add guardrail for image generation

* fix: use get_str_from_messages

litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py

@@ -135,37 +135,32 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
             "Azure Prompt Shield: Running pre-call prompt scan, on call_type: %s",
             call_type,
         )
-        if call_type == "acompletion":
-            new_messages: Optional[List[AllMessageValues]] = data.get("messages")
-            if new_messages is None:
-                verbose_proxy_logger.warning(
-                    "Lakera AI: not running guardrail. No messages in data"
-                )
-                return data
-            user_prompt = self.get_user_prompt(new_messages)
+        new_messages: Optional[List[AllMessageValues]] = data.get("messages")
+        if new_messages is None:
+            verbose_proxy_logger.warning(
+                "Lakera AI: not running guardrail. No messages in data"
+            )
+            return data
+        user_prompt = self.get_user_prompt(new_messages)
 
-            if user_prompt:
-                verbose_proxy_logger.debug(
-                    f"Azure Prompt Shield: User prompt: {user_prompt}"
-                )
-                azure_prompt_shield_response = await self.async_make_request(
-                    user_prompt=user_prompt,
-                )
-                if azure_prompt_shield_response["userPromptAnalysis"].get(
-                    "attackDetected"
-                ):
-                    verbose_proxy_logger.warning("Azure Prompt Shield: Attack detected")
-                    raise HTTPException(
-                        status_code=400,
-                        detail={
-                            "error": "Violated Azure Prompt Shield guardrail policy",
-                            "detection_message": f"Attack detected: {azure_prompt_shield_response['userPromptAnalysis']}",
-                        },
-                    )
-            else:
-                verbose_proxy_logger.warning(
-                    "Azure Prompt Shield: No user prompt found"
+        if user_prompt:
+            verbose_proxy_logger.debug(
+                f"Azure Prompt Shield: User prompt: {user_prompt}"
+            )
+            azure_prompt_shield_response = await self.async_make_request(
+                user_prompt=user_prompt,
+            )
+            if azure_prompt_shield_response["userPromptAnalysis"].get("attackDetected"):
+                verbose_proxy_logger.warning("Azure Prompt Shield: Attack detected")
+                raise HTTPException(
+                    status_code=400,
+                    detail={
+                        "error": "Violated Azure Prompt Shield guardrail policy",
+                        "detection_message": f"Attack detected: {azure_prompt_shield_response['userPromptAnalysis']}",
+                    },
                 )
+        else:
+            verbose_proxy_logger.warning("Azure Prompt Shield: No user prompt found")
         return None
 
     @log_guardrail_information

litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py

@@ -21,7 +21,6 @@ from litellm.proxy._types import UserAPIKeyAuth
 from .base import AzureGuardrailBase
 
 if TYPE_CHECKING:
-
     from litellm.proxy._types import UserAPIKeyAuth
     from litellm.types.llms.openai import AllMessageValues
     from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_text_moderation import (
@@ -189,7 +188,6 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr
             and self.severity_threshold_by_category is None
         ):
             for category in response["categoriesAnalysis"]:
-
                 if category["severity"] >= self.default_severity_threshold:
                     raise HTTPException(
                         status_code=400,
@@ -230,25 +228,24 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr
             "Azure Prompt Shield: Running pre-call prompt scan, on call_type: %s",
             call_type,
         )
-        if call_type == "acompletion":
-            new_messages: Optional[List[AllMessageValues]] = data.get("messages")
-            if new_messages is None:
-                verbose_proxy_logger.warning(
-                    "Lakera AI: not running guardrail. No messages in data"
-                )
-                return data
-            user_prompt = self.get_user_prompt(new_messages)
+        new_messages: Optional[List[AllMessageValues]] = data.get("messages")
+        if new_messages is None:
+            verbose_proxy_logger.warning(
+                "Lakera AI: not running guardrail. No messages in data"
+            )
+            return data
+        user_prompt = self.get_user_prompt(new_messages)
 
-            if user_prompt:
-                verbose_proxy_logger.info(
-                    f"Azure Text Moderation: User prompt: {user_prompt}"
-                )
-                azure_text_moderation_response = await self.async_make_request(
-                    text=user_prompt,
-                )
-                self.check_severity_threshold(response=azure_text_moderation_response)
-            else:
-                verbose_proxy_logger.warning("Azure Text Moderation: No text found")
+        if user_prompt:
+            verbose_proxy_logger.info(
+                f"Azure Text Moderation: User prompt: {user_prompt}"
+            )
+            azure_text_moderation_response = await self.async_make_request(
+                text=user_prompt,
+            )
+            self.check_severity_threshold(response=azure_text_moderation_response)
+        else:
+            verbose_proxy_logger.warning("Azure Text Moderation: No text found")
         return None
 
     async def async_post_call_success_hook(

litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py

@@ -121,7 +121,7 @@ class GuardrailsAI(CustomGuardrail):
     ) -> str:
         from httpx import URL
 
-        # This branch of code does not work with current version of GuardrailsAI API (as of July 2025), and it is unclear if it ever worked. 
+        # This branch of code does not work with current version of GuardrailsAI API (as of July 2025), and it is unclear if it ever worked.
         # Use guardrails_ai_api_input_format: "llmOutput" config line for all guardrails (which is the default anyway)
         # We can still use the "pre_call" mode to validate the inputs even if the API input format is technicallt "llmOutput"
 
@@ -162,30 +162,29 @@ class GuardrailsAI(CustomGuardrail):
         return response
 
     async def process_input(self, data: dict, call_type: str) -> dict:
-        if call_type == "acompletion" or call_type == "completion":
-            from litellm.litellm_core_utils.prompt_templates.common_utils import (
-                get_last_user_message,
-                set_last_user_message,
+        from litellm.litellm_core_utils.prompt_templates.common_utils import (
+            get_last_user_message,
+            set_last_user_message,
+        )
+
+        if "messages" not in data:  # invalid request
+            return data
+
+        text = get_last_user_message(data["messages"])
+        if text is None:
+            return data
+        if self.guardrails_ai_api_input_format == "inputs":
+            updated_text = await self.make_guardrails_ai_api_request_pre_call_request(
+                text_input=text, request_data=data
             )
-
-            if "messages" not in data:  # invalid request
-                return data
-
-            text = get_last_user_message(data["messages"])
-            if text is None:
-                return data
-            if self.guardrails_ai_api_input_format == "inputs":
-                updated_text = (
-                    await self.make_guardrails_ai_api_request_pre_call_request(
-                        text_input=text, request_data=data
-                    )
-                )
-            else:
-                _result = await self.make_guardrails_ai_api_request(
-                    llm_output=text, request_data=data
-                )
-                updated_text = _result.get("validatedOutput") or _result.get("rawLlmOutput") or text
-            data["messages"] = set_last_user_message(data["messages"], updated_text)
+        else:
+            _result = await self.make_guardrails_ai_api_request(
+                llm_output=text, request_data=data
+            )
+            updated_text = (
+                _result.get("validatedOutput") or _result.get("rawLlmOutput") or text
+            )
+        data["messages"] = set_last_user_message(data["messages"], updated_text)
 
         return data
 
@@ -209,13 +208,11 @@ class GuardrailsAI(CustomGuardrail):
     ) -> Optional[
         Union[Exception, str, dict]
     ]:  # raise exception if invalid, return a str for the user to receive - if rejected, or return a modified dictionary for passing into litellm
-
         return await self.process_input(data=data, call_type=call_type)
 
     async def async_logging_hook(
         self, kwargs: dict, result: Any, call_type: str
     ) -> Tuple[dict, Any]:
-
         if call_type == "acompletion" or call_type == "completion":
             kwargs = await self.process_input(data=kwargs, call_type=call_type)
 

litellm/proxy/guardrails/guardrail_hooks/presidio.py

@@ -43,7 +43,6 @@ from litellm.types.proxy.guardrails.guardrail_hooks.presidio import (
     PresidioAnalyzeRequest,
     PresidioAnalyzeResponseItem,
 )
-from litellm.types.utils import CallTypes as LitellmCallTypes
 from litellm.types.utils import GuardrailStatus
 from litellm.utils import (
     EmbeddingResponse,
@@ -68,7 +67,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
         output_parse_pii: Optional[bool] = False,
         presidio_ad_hoc_recognizers: Optional[str] = None,
         logging_only: Optional[bool] = None,
-        pii_entities_config: Optional[Dict[Union[PiiEntityType, str], PiiAction]] = None,
+        pii_entities_config: Optional[
+            Dict[Union[PiiEntityType, str], PiiAction]
+        ] = None,
         presidio_language: Optional[str] = None,
         **kwargs,
     ):
@@ -402,46 +403,34 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
             content_safety = data.get("content_safety", None)
             verbose_proxy_logger.debug("content_safety: %s", content_safety)
             presidio_config = self.get_presidio_settings_from_request_data(data)
-            if (
-                call_type
-                in [
-                    LitellmCallTypes.completion.value,
-                    LitellmCallTypes.acompletion.value,
-                ]
-                or call_type == "mcp_call"
-            ):
-                messages = data["messages"]
-                tasks = []
-                for m in messages:
-                    content = m.get("content", None)
-                    if content is None:
-                        continue
-                    if isinstance(content, str):
-                        tasks.append(
-                            self.check_pii(
-                                text=content,
-                                output_parse_pii=self.output_parse_pii,
-                                presidio_config=presidio_config,
-                                request_data=data,
-                            )
+            messages = data["messages"]
+            tasks = []
+            for m in messages:
+                content = m.get("content", None)
+                if content is None:
+                    continue
+                if isinstance(content, str):
+                    tasks.append(
+                        self.check_pii(
+                            text=content,
+                            output_parse_pii=self.output_parse_pii,
+                            presidio_config=presidio_config,
+                            request_data=data,
                         )
-                responses = await asyncio.gather(*tasks)
-                for index, r in enumerate(responses):
-                    content = messages[index].get("content", None)
-                    if content is None:
-                        continue
-                    if isinstance(content, str):
-                        messages[index][
-                            "content"
-                        ] = r  # replace content with redacted string
-                verbose_proxy_logger.debug(
-                    f"Presidio PII Masking: Redacted pii message: {data['messages']}"
-                )
-                data["messages"] = messages
-            else:
-                verbose_proxy_logger.debug(
-                    f"Not running async_pre_call_hook for call_type={call_type}"
-                )
+                    )
+            responses = await asyncio.gather(*tasks)
+            for index, r in enumerate(responses):
+                content = messages[index].get("content", None)
+                if content is None:
+                    continue
+                if isinstance(content, str):
+                    messages[index][
+                        "content"
+                    ] = r  # replace content with redacted string
+            verbose_proxy_logger.debug(
+                f"Presidio PII Masking: Redacted pii message: {data['messages']}"
+            )
+            data["messages"] = messages
             return data
         except Exception as e:
             raise e

litellm/proxy/image_endpoints/endpoints.py

@@ -8,10 +8,14 @@ from fastapi.responses import ORJSONResponse
 
 import litellm
 from litellm._logging import verbose_proxy_logger
+from litellm.litellm_core_utils.prompt_templates.common_utils import (
+    get_str_from_messages,
+)
 from litellm.proxy._types import *
 from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
 from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
 from litellm.proxy.route_llm_request import route_request
+from litellm.types.llms.openai import ChatCompletionUserMessage
 
 router = APIRouter()
 
@@ -108,10 +112,23 @@ async def image_generation(
             data["model"] = litellm.model_alias_map[data["model"]]
 
         ### CALL HOOKS ### - modify incoming data / reject request before calling the model
+        prompt_value = data.get("prompt")
+        if prompt_value is not None:
+            # Reformat the image prompt as a chat message so guardrails can process it.
+            user_message: ChatCompletionUserMessage = {
+                "role": "user",
+                "content": prompt_value,
+            }
+            data["messages"] = [user_message]
         data = await proxy_logging_obj.pre_call_hook(
             user_api_key_dict=user_api_key_dict, data=data, call_type="image_generation"
         )
 
+        messages = data.get("messages")
+        if isinstance(messages, list) and messages:
+            data["prompt"] = get_str_from_messages(messages)
+        data.pop("messages", None)
+
         ## ROUTE TO CORRECT ENDPOINT ##
         llm_call = await route_request(
             data=data,

tests/test_litellm/proxy/image_endpoints/test_endpoints.py

@@ -0,0 +1,106 @@
+import asyncio
+import copy
+from types import SimpleNamespace
+from typing import Any, Dict
+
+import orjson
+import pytest
+from starlette.requests import Request
+from starlette.responses import Response
+
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.image_endpoints import endpoints
+
+
+@pytest.mark.asyncio
+async def test_image_generation_prompt_rerouting(monkeypatch):
+    """Ensure image prompts are exposed to guardrails and restored afterwards."""
+
+    async def fake_add_litellm_data_to_request(**kwargs):
+        return kwargs["data"]
+
+    async def fake_update_request_status(**_: Any) -> None:
+        await asyncio.sleep(0)
+
+    proxy_logger_calls: Dict[str, Any] = {}
+
+    async def fake_pre_call_hook(*, user_api_key_dict, data, call_type):  # type: ignore[override]
+        proxy_logger_calls["pre_call_input"] = copy.deepcopy(data)
+        modified = {
+            **data,
+            "messages": [
+                {
+                    "role": "user",
+                    "content": "sanitized prompt",
+                }
+            ],
+        }
+        return modified
+
+    async def fake_post_call_failure_hook(**_: Any) -> None:
+        return None
+
+    fake_proxy_logger = SimpleNamespace(
+        pre_call_hook=fake_pre_call_hook,
+        update_request_status=fake_update_request_status,
+        post_call_failure_hook=fake_post_call_failure_hook,
+    )
+
+    captured_route_request_data: Dict[str, Any] = {}
+
+    async def fake_route_request(*, data, **kwargs):  # type: ignore[override]
+        captured_route_request_data.update(data)
+
+        async def _inner():
+            class FakeResponse(dict):
+                _hidden_params = {}
+
+            return FakeResponse(result="ok")
+
+        return _inner()
+
+    scope = {
+        "type": "http",
+        "method": "POST",
+        "path": "/v1/images/generations",
+        "headers": [],
+    }
+    body = orjson.dumps({"prompt": "original prompt"})
+
+    async def receive():
+        return {"type": "http.request", "body": body, "more_body": False}
+
+    request = Request(scope, receive)
+    response = Response()
+    user_api_key = UserAPIKeyAuth()
+
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.add_litellm_data_to_request",
+        fake_add_litellm_data_to_request,
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
+    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {})
+    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", fake_proxy_logger)
+    monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
+    monkeypatch.setattr("litellm.proxy.proxy_server.version", "test-version")
+    monkeypatch.setattr(
+        "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing.get_custom_headers",
+        classmethod(lambda *args, **kwargs: {}),
+    )
+    monkeypatch.setattr(
+        "litellm.proxy.image_endpoints.endpoints.route_request", fake_route_request
+    )
+
+    result = await endpoints.image_generation(
+        request=request,
+        fastapi_response=response,
+        user_api_key_dict=user_api_key,
+    )
+    await asyncio.sleep(0)
+
+    assert result == {"result": "ok"}
+    pre_call_input = proxy_logger_calls["pre_call_input"]
+    assert pre_call_input["messages"][0]["content"] == "original prompt"
+    assert captured_route_request_data["prompt"] == "sanitized prompt"
+    assert "messages" not in captured_route_request_data