From bb3b77cbebc6b4f5beeb02d3e95a0c99eaee4f9a Mon Sep 17 00:00:00 2001 From: YutaSaito <36355491+uc4w6c@users.noreply.github.com> Date: Fri, 17 Oct 2025 13:50:30 +0900 Subject: [PATCH] feat: add guardrail for image generation (#15619) * feat: add guardrail for image generation * fix: use get_str_from_messages --- .../guardrail_hooks/azure/prompt_shield.py | 53 ++++----- .../guardrail_hooks/azure/text_moderation.py | 37 +++--- .../guardrails_ai/guardrails_ai.py | 49 ++++---- .../guardrails/guardrail_hooks/presidio.py | 71 +++++------- litellm/proxy/image_endpoints/endpoints.py | 17 +++ .../proxy/image_endpoints/test_endpoints.py | 106 ++++++++++++++++++ 6 files changed, 217 insertions(+), 116 deletions(-) create mode 100644 tests/test_litellm/proxy/image_endpoints/test_endpoints.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py index 689777ad5c..7486bd85f6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py +++ b/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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py index 4c6f8d335a..d02a9751bc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py +++ b/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( diff --git a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py index db22f630ad..84cbef0526 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py +++ b/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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index b77e802c71..23c73098c4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/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 diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 1ab9d14484..16aa8f1657 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/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, diff --git a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py new file mode 100644 index 0000000000..a3b6a9c602 --- /dev/null +++ b/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