From 333841b8f3960ea5f9f456fb64c9aaee78026bca Mon Sep 17 00:00:00 2001 From: NotZhu <3495159203@qq.com> Date: Mon, 2 Nov 2026 17:18:12 +0800 Subject: [PATCH 1/7] refactor(guardrails): share Azure severity threshold check --- .../guardrails/guardrail_hooks/azure/base.py | 66 ++++++++++++++++++- .../guardrail_hooks/azure/text_moderation.py | 58 +++------------- .../azure/azure_text_moderation.py | 10 +-- .../guardrails/guardrail_hooks/azure/base.py | 8 +++ 4 files changed, 83 insertions(+), 59 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py index 5f57969c310..a3c16b4e7d4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py @@ -1,8 +1,9 @@ import re -from collections.abc import Mapping -from typing import Any, Final, cast +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Literal, cast import httpx +from fastapi import HTTPException from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -16,9 +17,15 @@ from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import AllMessageValues, ResponseInputParam from litellm.types.utils import CallTypes, CallTypesLiteral +if TYPE_CHECKING: + from litellm.types.proxy.guardrails.guardrail_hooks.azure.base import AzureCategoriesAnalysis + # Azure Content Safety APIs have a 10,000 character limit per request. AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH: Final = 10000 +# Azure Content Safety accepts images up to 4 MB after base64 decoding. +AZURE_CONTENT_SAFETY_MAX_IMAGE_BYTES: Final = 4 * 1024 * 1024 + # Azure Content Safety bills text in 1,000-character "text records"; a submitted # chunk of N characters consumes ceil(N / 1000) text records. AZURE_CONTENT_SAFETY_TEXT_RECORD_LENGTH: Final = 1000 @@ -37,6 +44,61 @@ def resolve_content_safety_api_version(configured: str | None) -> str: return configured +def check_categories_against_thresholds( + categories_analysis: "Sequence[AzureCategoriesAnalysis]", + *, + severity_threshold: int | None, + severity_threshold_by_category: "Mapping[str, int] | None", + default_severity_threshold: int, + guardrail_label: str, +) -> Literal[True]: + """Raise ``HTTPException(400)`` when any category reaches its configured severity. + + Shared by the Azure text and image moderation guardrails so both accept the + same ``severity_threshold`` / ``severity_threshold_by_category`` config and + return the same error shape. + """ + if severity_threshold_by_category: + for category in categories_analysis: + if (threshold := severity_threshold_by_category.get(category["category"])) is not None and ( + category["severity"] >= threshold + ): + raise HTTPException( + status_code=400, + detail={ + "error": "{}: {} crossed severity {}, Got severity: {}".format( + guardrail_label, category["category"], threshold, category["severity"] + ) + }, + ) + if severity_threshold: + for category in categories_analysis: + if category["severity"] >= severity_threshold: + raise HTTPException( + status_code=400, + detail={ + "error": "{}: {} crossed severity {}, Got severity: {}".format( + guardrail_label, category["category"], severity_threshold, category["severity"] + ) + }, + ) + if severity_threshold is None and severity_threshold_by_category is None: + for category in categories_analysis: + if category["severity"] >= default_severity_threshold: + raise HTTPException( + status_code=400, + detail={ + "error": "{}: {} crossed severity {}, Got severity: {}".format( + guardrail_label, + category["category"], + default_severity_threshold, + category["severity"], + ) + }, + ) + return True + + class AzureGuardrailBase: """ Base class for Azure guardrails. diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py index 27f36f2724b..d78d6c250de 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py @@ -20,6 +20,7 @@ from .base import ( # noqa: F401 # legacy module exports _RESPONSES_API_CALL_TYPES, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export RESPONSES_API_CALL_TYPES, AzureGuardrailBase, + check_categories_against_thresholds, ) if TYPE_CHECKING: @@ -167,56 +168,13 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr return inputs def check_severity_threshold(self, response: "AzureTextModerationGuardrailResponse") -> Literal[True]: - """ - - Check if threshold set by category - - Check if general severity threshold set - - If both none, use default_severity_threshold - """ - - if self.severity_threshold_by_category: - for category in response["categoriesAnalysis"]: - severity_category_threshold_item = self.severity_threshold_by_category.get(category["category"]) - if ( - severity_category_threshold_item is not None - and category["severity"] >= severity_category_threshold_item - ): - raise HTTPException( - status_code=400, - detail={ - "error": "Azure Content Safety Guardrail: {} crossed severity {}, Got severity: {}".format( - category["category"], - self.severity_threshold_by_category.get(category["category"]), - category["severity"], - ) - }, - ) - if self.severity_threshold: - for category in response["categoriesAnalysis"]: - if category["severity"] >= self.severity_threshold: - raise HTTPException( - status_code=400, - detail={ - "error": "Azure Content Safety Guardrail: {} crossed severity {}, Got severity: {}".format( - category["category"], - self.severity_threshold, - category["severity"], - ) - }, - ) - if self.severity_threshold is None 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, - detail={ - "error": "Azure Content Safety Guardrail: {} crossed severity {}, Got severity: {}".format( - category["category"], - self.default_severity_threshold, - category["severity"], - ) - }, - ) - return True + return check_categories_against_thresholds( + response["categoriesAnalysis"], + severity_threshold=self.severity_threshold, + severity_threshold_by_category=self.severity_threshold_by_category, + default_severity_threshold=self.default_severity_threshold, + guardrail_label="Azure Content Safety Guardrail", + ) @log_guardrail_information async def async_pre_call_hook( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py index 6ee5741baf5..135490690e5 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py @@ -1,4 +1,4 @@ -from typing import Final, Literal +from typing import Final, Literal, TypeAlias from pydantic import Field from typing_extensions import Required, TypedDict @@ -6,7 +6,7 @@ from typing_extensions import Required, TypedDict from litellm.types.llms.base import LiteLLMBaseModel from ..base import GuardrailConfigModel -from .base import AzureContentSafetyConfigModel +from .base import AzureCategoriesAnalysis, AzureContentSafetyConfigModel # Shared Azure Content Safety categories AZURE_CONTENT_SAFETY_CATEGORIES: Final = ["Hate", "SelfHarm", "Sexual", "Violence"] @@ -27,11 +27,7 @@ class AzureTextModerationGuardrailRequestBody(AzureTextModerationRequestBodyOpti text: Required[str] -class AzureTextModerationGuardrailResponseCategoriesAnalysis(TypedDict): - """Response from the Azure Text Moderation guardrail""" - - category: str - severity: int +AzureTextModerationGuardrailResponseCategoriesAnalysis: TypeAlias = AzureCategoriesAnalysis class AzureTextModerationGuardrailResponse(TypedDict): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/azure/base.py b/litellm/types/proxy/guardrails/guardrail_hooks/azure/base.py index cc09e2b679d..0afbfb4dc8a 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/azure/base.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/azure/base.py @@ -1,8 +1,16 @@ from pydantic import Field +from typing_extensions import TypedDict from litellm.types.llms.base import LiteLLMBaseModel +class AzureCategoriesAnalysis(TypedDict): + """Per-category severity returned by Azure Content Safety analyse endpoints""" + + category: str + severity: int + + class AzureContentSafetyConfigModel(LiteLLMBaseModel): """Configuration parameters for the Azure Content Safety Prompt Shield guardrail""" From c2c9b9786d1be9ba929cfa142a4b80d625950a54 Mon Sep 17 00:00:00 2001 From: NotZhu <3495159203@qq.com> Date: Mon, 2 Nov 2026 17:18:17 +0800 Subject: [PATCH 2/7] feat(guardrails): add Azure Content Safety image moderation --- .../guardrail_hooks/azure/__init__.py | 18 +- .../guardrail_hooks/azure/image_moderation.py | 164 +++++++++++++ litellm/types/guardrails.py | 1 + .../azure/azure_image_moderation.py | 65 +++++ .../azure/test_azure_image_moderation.py | 225 ++++++++++++++++++ 5 files changed, 472 insertions(+), 1 deletion(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/azure/image_moderation.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_image_moderation.py create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_image_moderation.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py index de77cd55671..486d3334954 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Final from litellm.types.guardrails import SupportedGuardrailIntegrations +from .image_moderation import AzureContentSafetyImageModerationGuardrail from .prompt_shield import AzureContentSafetyPromptShieldGuardrail from .text_moderation import AzureContentSafetyTextModerationGuardrail @@ -25,7 +26,9 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" if azure_guardrail == "prompt_shield": azure_content_safety_guardrail: ( - AzureContentSafetyPromptShieldGuardrail | AzureContentSafetyTextModerationGuardrail + AzureContentSafetyPromptShieldGuardrail + | AzureContentSafetyTextModerationGuardrail + | AzureContentSafetyImageModerationGuardrail ) = AzureContentSafetyPromptShieldGuardrail( guardrail_name=guardrail_name, **{ @@ -47,6 +50,17 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" "event_hook": litellm_params.mode, }, ) + elif azure_guardrail == "image_moderations": + azure_content_safety_guardrail = AzureContentSafetyImageModerationGuardrail( + guardrail_name=guardrail_name, + **{ + **litellm_params.model_dump(exclude_none=True), + "api_key": litellm_params.api_key, + "api_base": litellm_params.api_base, + "default_on": litellm_params.default_on, + "event_hook": litellm_params.mode, + }, + ) else: raise ValueError(f"Azure Content Safety: {azure_guardrail} is not a valid guardrail") @@ -57,10 +71,12 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" guardrail_initializer_registry: Final = { SupportedGuardrailIntegrations.AZURE_PROMPT_SHIELD.value: initialize_guardrail, SupportedGuardrailIntegrations.AZURE_TEXT_MODERATIONS.value: initialize_guardrail, + SupportedGuardrailIntegrations.AZURE_IMAGE_MODERATIONS.value: initialize_guardrail, } guardrail_class_registry: Final = { SupportedGuardrailIntegrations.AZURE_PROMPT_SHIELD.value: AzureContentSafetyPromptShieldGuardrail, SupportedGuardrailIntegrations.AZURE_TEXT_MODERATIONS.value: AzureContentSafetyTextModerationGuardrail, + SupportedGuardrailIntegrations.AZURE_IMAGE_MODERATIONS.value: AzureContentSafetyImageModerationGuardrail, } diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/image_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/image_moderation.py new file mode 100644 index 00000000000..69e2c17bb32 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/image_moderation.py @@ -0,0 +1,164 @@ +import base64 +import binascii +from typing import TYPE_CHECKING, ClassVar, Final, Literal, cast + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # decorator is untyped +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + +from .base import ( + AZURE_CONTENT_SAFETY_MAX_IMAGE_BYTES, + AzureGuardrailBase, + check_categories_against_thresholds, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_image_moderation import ( + AzureContentSafetyImageModerationConfigModel, + AzureImageModerationGuardrailResponse, + AzureImageModerationImage, + ) + +GUARDRAIL_LABEL: Final = "Azure Content Safety Guardrail" + + +class AzureContentSafetyImageModerationGuardrail(AzureGuardrailBase, CustomGuardrail): + """LiteLLM guardrail for Azure Content Safety image analysis. + + Sends request and response images to the Azure ``image:analyze`` endpoint and + applies the same per-category severity thresholds as the text moderation + guardrail. + """ + + use_native_lifecycle_hooks: ClassVar[bool] = True + + default_severity_threshold: int = 2 + + def __init__( + self, + guardrail_name: str, + api_key: object, + api_base: object, + severity_threshold: object = None, + severity_threshold_by_category: object = None, + categories: object = None, + **kwargs: object, + ): + from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_image_moderation import ( + AZURE_CONTENT_SAFETY_IMAGE_CATEGORIES, + ) + + kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) + super().__init__( + api_key=cast("str", api_key), + api_base=cast("str", api_base), + guardrail_name=guardrail_name, + **kwargs, + ) + + raw_severity_threshold: Final = cast("int | float | str | None", severity_threshold) + self.severity_threshold = int(raw_severity_threshold) if raw_severity_threshold else None + self.severity_threshold_by_category = cast("dict[str, int] | None", severity_threshold_by_category) + configured_categories: Final = cast("list[object]", categories) if isinstance(categories, list) else None + self.categories = ( + [str(category) for category in configured_categories] + if configured_categories + else list(AZURE_CONTENT_SAFETY_IMAGE_CATEGORIES) + ) + + verbose_proxy_logger.info("Initialized Azure Image Moderation Guardrail: %s", guardrail_name) + + @classmethod + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: + return [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ] + + @staticmethod + def get_config_model() -> type["AzureContentSafetyImageModerationConfigModel"] | None: + from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_image_moderation import ( + AzureContentSafetyImageModerationConfigModel, + ) + + return AzureContentSafetyImageModerationConfigModel + + def _to_azure_image(self, image: str) -> "AzureImageModerationImage": + if image.startswith("data:"): + header, separator, payload = image.partition(",") + if not separator or ";base64" not in header: + raise HTTPException( + status_code=400, + detail={ + "error": f"{GUARDRAIL_LABEL}: image must be a base64 data URL, got {header[:32]}", + }, + ) + stripped_payload: Final = payload.strip() + try: + decoded: Final = base64.b64decode(stripped_payload, validate=True) + except (binascii.Error, ValueError) as e: + raise HTTPException( + status_code=400, + detail={"error": f"{GUARDRAIL_LABEL}: image data URL is not valid base64"}, + ) from e + if len(decoded) > AZURE_CONTENT_SAFETY_MAX_IMAGE_BYTES: + raise HTTPException( + status_code=413, + detail={ + "error": ( + f"{GUARDRAIL_LABEL}: image is {len(decoded)} bytes, " + f"the limit is {AZURE_CONTENT_SAFETY_MAX_IMAGE_BYTES} bytes" + ), + }, + ) + return {"content": stripped_payload} + if image.startswith(("http://", "https://")): + return {"blobUrl": image} + raise HTTPException( + status_code=400, + detail={"error": f"{GUARDRAIL_LABEL}: image must be a base64 data URL or an http(s) URL"}, + ) + + async def async_make_request(self, image: str) -> "AzureImageModerationGuardrailResponse": + from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_image_moderation import ( + AzureImageModerationGuardrailRequestBody, + ) + + request_body: Final[AzureImageModerationGuardrailRequestBody] = { + "image": self._to_azure_image(image), + "categories": list(self.categories), + "outputType": "FourSeverityLevels", + } + response_json: Final = await self._post_to_content_safety( + "image:analyze", cast(dict[str, object], request_body) + ) + response: Final = cast("AzureImageModerationGuardrailResponse", response_json) + + check_categories_against_thresholds( + response["categoriesAnalysis"], + severity_threshold=self.severity_threshold, + severity_threshold_by_category=self.severity_threshold_by_category, + default_severity_threshold=self.default_severity_threshold, + guardrail_label=GUARDRAIL_LABEL, + ) + return response + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # pyright: ignore[reportMissingTypeArgument, reportUnknownParameterType] # matches CustomGuardrail + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + for image in inputs.get("images") or (): + if image: + await self.async_make_request(image=image) + return inputs diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 718a651bdee..3744010dfdc 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -119,6 +119,7 @@ class SupportedGuardrailIntegrations(Enum): CISCO_AI_DEFENSE = "cisco_ai_defense" AZURE_PROMPT_SHIELD = "azure/prompt_shield" AZURE_TEXT_MODERATIONS = "azure/text_moderations" + AZURE_IMAGE_MODERATIONS = "azure/image_moderations" MODEL_ARMOR = "model_armor" OPENAI_MODERATION = "openai_moderation" NOMA = "noma" diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_image_moderation.py b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_image_moderation.py new file mode 100644 index 00000000000..33ee2b28da3 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_image_moderation.py @@ -0,0 +1,65 @@ +from typing import Final, Literal + +from pydantic import Field +from typing_extensions import Required, TypedDict + +from litellm.types.llms.base import LiteLLMBaseModel + +from ..base import GuardrailConfigModel +from .azure_text_moderation import AzureHarmCategories +from .base import AzureCategoriesAnalysis, AzureContentSafetyConfigModel + +AZURE_CONTENT_SAFETY_IMAGE_CATEGORIES: Final = ["Hate", "SelfHarm", "Sexual", "Violence"] + + +class AzureImageModerationImage(TypedDict, total=False): + """One image for the Azure Content Safety image analysis endpoint. + + Exactly one of ``content`` (base64) or ``blobUrl`` (public URL) must be set. + """ + + content: str + blobUrl: str + + +class AzureImageModerationGuardrailRequestBody(TypedDict, total=False): + """Request body for the Azure Content Safety image analysis endpoint""" + + image: Required[AzureImageModerationImage] + categories: list[str] + outputType: Literal["FourSeverityLevels"] + + +class AzureImageModerationGuardrailResponse(TypedDict): + """Response from the Azure Content Safety image analysis endpoint""" + + categoriesAnalysis: list[AzureCategoriesAnalysis] + + +class AzureImageModerationOptionalParams(LiteLLMBaseModel): + severity_threshold: int | None = Field( + default=None, + description="Severity threshold for the Azure Content Safety Image Moderation guardrail across all categories", + ) + severity_threshold_by_category: dict[AzureHarmCategories, int] | None = Field( + default=None, + description="Severity threshold by category for the Azure Content Safety Image Moderation guardrail. See list of categories - https://learn.microsoft.com/en-us/azure/ai-services/content-safety/concepts/harm-categories?tabs=warning", + ) + categories: list[AzureHarmCategories] | None = Field( + default=None, + description="Categories to scan for the Azure Content Safety Image Moderation guardrail. See list of categories - https://learn.microsoft.com/en-us/azure/ai-services/content-safety/concepts/harm-categories?tabs=warning", + ) + + +class AzureContentSafetyImageModerationConfigModel( + AzureContentSafetyConfigModel, + GuardrailConfigModel[AzureImageModerationOptionalParams], +): + optional_params: AzureImageModerationOptionalParams | None = Field( + default=None, + description="Optional parameters for the Azure Content Safety Image Moderation guardrail", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Azure Content Safety Image Moderation" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_image_moderation.py b/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_image_moderation.py new file mode 100644 index 00000000000..30a49a9484f --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_image_moderation.py @@ -0,0 +1,225 @@ +import base64 +from typing import Final +from unittest.mock import Mock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy.guardrails.guardrail_hooks.azure import ( + AzureContentSafetyImageModerationGuardrail, + guardrail_class_registry, + guardrail_initializer_registry, +) +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler +from litellm.types.guardrails import SupportedGuardrailIntegrations + +IMAGE_BASE64: Final = base64.b64encode(b"small-image-bytes").decode() +IMAGE_DATA_URL: Final = f"data:image/png;base64,{IMAGE_BASE64}" +REMOTE_IMAGE_URL: Final = "https://example.com/cat.png" +AZURE_API_BASE: Final = "https://example.cognitiveservices.azure.com" + + +def _image_response(severity: int, category: str = "Violence") -> Mock: + response = Mock() + response.json.return_value = { + "categoriesAnalysis": [ + {"category": "Hate", "severity": severity if category == "Hate" else 0}, + {"category": "Sexual", "severity": severity if category == "Sexual" else 0}, + {"category": "SelfHarm", "severity": severity if category == "SelfHarm" else 0}, + {"category": "Violence", "severity": severity if category == "Violence" else 0}, + ] + } + return response + + +def _image_guardrail(**kwargs) -> AzureContentSafetyImageModerationGuardrail: + return AzureContentSafetyImageModerationGuardrail( + guardrail_name="azure_image_moderation", + api_key="azure_image_moderation_api_key", + api_base=AZURE_API_BASE, + **kwargs, + ) + + +@pytest.mark.asyncio +async def test_image_moderation_posts_base64_content(): + guardrail = _image_guardrail() + + with patch.object(guardrail.async_handler, "post", return_value=_image_response(0)) as mock_post: + result = await guardrail.apply_guardrail( + inputs={"images": [IMAGE_DATA_URL]}, + request_data={}, + input_type="request", + ) + + assert result == {"images": [IMAGE_DATA_URL]} + assert mock_post.call_count == 1 + assert mock_post.call_args.kwargs["url"] == (f"{AZURE_API_BASE}/contentsafety/image:analyze?api-version=2024-09-01") + body = mock_post.call_args.kwargs["json"] + assert body["image"] == {"content": IMAGE_BASE64} + assert body["categories"] == ["Hate", "SelfHarm", "Sexual", "Violence"] + assert body["outputType"] == "FourSeverityLevels" + + +@pytest.mark.asyncio +async def test_image_moderation_uses_blob_url_for_remote_images(): + guardrail = _image_guardrail() + + with patch.object(guardrail.async_handler, "post", return_value=_image_response(0)) as mock_post: + await guardrail.apply_guardrail( + inputs={"images": [REMOTE_IMAGE_URL]}, + request_data={}, + input_type="request", + ) + + body = mock_post.call_args.kwargs["json"] + assert body["image"] == {"blobUrl": REMOTE_IMAGE_URL} + assert "content" not in body["image"] + + +@pytest.mark.asyncio +async def test_image_moderation_raises_on_violation(): + guardrail = _image_guardrail() + + with patch.object(guardrail.async_handler, "post", return_value=_image_response(4)): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"images": [IMAGE_DATA_URL]}, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 400 + assert "Violence crossed severity 2" in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_image_moderation_honours_per_category_threshold(): + guardrail = _image_guardrail(severity_threshold_by_category={"Hate": 4}) + + with patch.object(guardrail.async_handler, "post", return_value=_image_response(4, category="Violence")): + await guardrail.apply_guardrail( + inputs={"images": [IMAGE_DATA_URL]}, + request_data={}, + input_type="request", + ) + + with patch.object(guardrail.async_handler, "post", return_value=_image_response(4, category="Hate")): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"images": [IMAGE_DATA_URL]}, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 400 + assert "Hate crossed severity 4" in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_image_moderation_skips_empty_images(): + guardrail = _image_guardrail() + + with patch.object(guardrail.async_handler, "post") as mock_post: + result = await guardrail.apply_guardrail(inputs={"images": ["", ""]}, request_data={}, input_type="request") + + mock_post.assert_not_called() + assert result == {"images": ["", ""]} + + +@pytest.mark.asyncio +async def test_image_moderation_rejects_plain_paths(): + guardrail = _image_guardrail() + + with patch.object(guardrail.async_handler, "post") as mock_post: + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"images": ["file.png"]}, + request_data={}, + input_type="request", + ) + + mock_post.assert_not_called() + assert exc_info.value.status_code == 400 + assert "base64 data URL or an http(s) URL" in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_image_moderation_rejects_data_url_without_base64(): + guardrail = _image_guardrail() + + with patch.object(guardrail.async_handler, "post") as mock_post: + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"images": ["data:image/png,not-base64"]}, + request_data={}, + input_type="request", + ) + + mock_post.assert_not_called() + assert exc_info.value.status_code == 400 + assert "must be a base64 data URL" in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_image_moderation_rejects_invalid_base64_payload(): + guardrail = _image_guardrail() + + with patch.object(guardrail.async_handler, "post") as mock_post: + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"images": ["data:image/png;base64,not-base64!"]}, + request_data={}, + input_type="request", + ) + + mock_post.assert_not_called() + assert exc_info.value.status_code == 400 + assert "not valid base64" in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_image_moderation_rejects_oversized_images(): + guardrail = _image_guardrail() + oversized = "data:image/png;base64," + base64.b64encode(b"x" * (4 * 1024 * 1024 + 1)).decode() + + with patch.object(guardrail.async_handler, "post") as mock_post: + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail(inputs={"images": [oversized]}, request_data={}, input_type="request") + + mock_post.assert_not_called() + assert exc_info.value.status_code == 413 + assert "the limit is 4194304 bytes" in exc_info.value.detail["error"] + + +def test_image_moderation_registers_initializer_and_class(): + integration = SupportedGuardrailIntegrations.AZURE_IMAGE_MODERATIONS.value + assert integration in guardrail_initializer_registry + assert guardrail_class_registry[integration] is AzureContentSafetyImageModerationGuardrail + + +def test_image_moderation_event_hooks(): + assert AzureContentSafetyImageModerationGuardrail.get_supported_event_hooks() == [ + "pre_call", + "post_call", + ] + + +def test_image_moderation_loads_from_config(): + handler = InMemoryGuardrailHandler() + registered = handler.initialize_guardrail( + guardrail={ + "guardrail_name": "azure-image-moderation", + "litellm_params": { + "guardrail": "azure/image_moderations", + "mode": "pre_call", + "api_key": "azure_image_moderation_api_key", + "api_base": AZURE_API_BASE, + }, + } + ) + + assert registered is not None + guardrail = handler.guardrail_id_to_custom_guardrail[registered["guardrail_id"]] + assert isinstance(guardrail, AzureContentSafetyImageModerationGuardrail) + assert guardrail.categories == ["Hate", "SelfHarm", "Sexual", "Violence"] From 3a24eeb515b6058c3b2f256e0179f606b9845203 Mon Sep 17 00:00:00 2001 From: NotZhu <3495159203@qq.com> Date: Mon, 2 Nov 2026 17:18:23 +0800 Subject: [PATCH 3/7] fix(guardrails): scan image-only chat requests in guardrails --- .../llms/openai/chat/guardrail_translation/handler.py | 4 ++-- .../test_openai_guardrail_handler.py | 9 +++++---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 5c5204ef2f4..113dc73d3b8 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -139,8 +139,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): scan_only_tool_results=scan_only_tool_results, ) - # Step 2: Apply guardrail to all texts and tool calls in batch - if texts_to_check or tool_calls_to_check: + # Step 2: Apply guardrail to all texts, images and tool calls in batch + if texts_to_check or tool_calls_to_check or images_to_check: inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: inputs["images"] = images_to_check diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 5d71cbe9afc..c729ff78d2f 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -2048,9 +2048,8 @@ class TestNoScannableContentRecordsNotRun: assert all(e.get("guardrail_status") != "not_run" for e in self._recorded_entries(data)) @pytest.mark.asyncio - async def test_image_only_content_is_not_reported_as_not_run(self): - """Images are only scanned alongside text, so an image-only request is a - pre-existing scan gap, not a message-scoping skip, and must not be labelled one""" + async def test_image_only_content_is_scanned(self): + """An image in a scanned role reaches apply_guardrail, not a pre-existing scan gap""" handler = OpenAIChatCompletionsHandler() guardrail = MockGuardrail(guardrail_name="image-guardrail") guardrail.skip_system_message_in_guardrail = True @@ -2066,7 +2065,9 @@ class TestNoScannableContentRecordsNotRun: await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) - assert self._recorded_entries(data) == [] + assert guardrail.last_inputs is not None + assert guardrail.last_inputs.get("images") == ["https://example.com/cat.png"] + assert all(entry.get("guardrail_status") != "not_run" for entry in self._recorded_entries(data)) @pytest.mark.asyncio async def test_scoped_out_image_only_message_is_not_reported_as_not_run(self): From cdc2bbb1a8930844c458c280d9cb750867ceba8d Mon Sep 17 00:00:00 2001 From: NotZhu <3495159203@qq.com> Date: Mon, 2 Nov 2026 17:26:04 +0800 Subject: [PATCH 4/7] fix(guardrails): dispatch image moderation through apply_guardrail --- .../guardrails/guardrail_hooks/azure/image_moderation.py | 2 +- .../guardrail_hooks/azure/test_azure_image_moderation.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/image_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/image_moderation.py index 69e2c17bb32..51885506ba4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/image_moderation.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/image_moderation.py @@ -37,7 +37,7 @@ class AzureContentSafetyImageModerationGuardrail(AzureGuardrailBase, CustomGuard guardrail. """ - use_native_lifecycle_hooks: ClassVar[bool] = True + use_native_lifecycle_hooks: ClassVar[bool] = False default_severity_threshold: int = 2 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_image_moderation.py b/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_image_moderation.py index 30a49a9484f..ff64a1b4b5a 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_image_moderation.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_image_moderation.py @@ -205,6 +205,11 @@ def test_image_moderation_event_hooks(): ] +def test_image_moderation_uses_unified_lifecycle_dispatch(): + """Native lifecycle hooks are inherited no-ops, so pre and post call must go through apply_guardrail.""" + assert AzureContentSafetyImageModerationGuardrail.use_native_lifecycle_hooks is False + + def test_image_moderation_loads_from_config(): handler = InMemoryGuardrailHandler() registered = handler.initialize_guardrail( From 1c6e0f85774aec506971167d608d51f9e7cfb13b Mon Sep 17 00:00:00 2001 From: NotZhu <3495159203@qq.com> Date: Mon, 2 Nov 2026 18:16:33 +0800 Subject: [PATCH 5/7] fix(guardrails): reject mutable remote image URLs --- .../guardrail_hooks/azure/image_moderation.py | 11 +++++++-- .../azure/azure_image_moderation.py | 8 +++---- .../azure/test_azure_image_moderation.py | 23 ++++++++++--------- 3 files changed, 25 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/image_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/image_moderation.py index 51885506ba4..228461a0a42 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/image_moderation.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/image_moderation.py @@ -120,10 +120,17 @@ class AzureContentSafetyImageModerationGuardrail(AzureGuardrailBase, CustomGuard ) return {"content": stripped_payload} if image.startswith(("http://", "https://")): - return {"blobUrl": image} + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"{GUARDRAIL_LABEL}: remote image URLs are not supported, send the image as a base64 data URL" + ), + }, + ) raise HTTPException( status_code=400, - detail={"error": f"{GUARDRAIL_LABEL}: image must be a base64 data URL or an http(s) URL"}, + detail={"error": f"{GUARDRAIL_LABEL}: image must be a base64 data URL"}, ) async def async_make_request(self, image: str) -> "AzureImageModerationGuardrailResponse": diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_image_moderation.py b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_image_moderation.py index 33ee2b28da3..c59a7d3913f 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_image_moderation.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_image_moderation.py @@ -12,14 +12,14 @@ from .base import AzureCategoriesAnalysis, AzureContentSafetyConfigModel AZURE_CONTENT_SAFETY_IMAGE_CATEGORIES: Final = ["Hate", "SelfHarm", "Sexual", "Violence"] -class AzureImageModerationImage(TypedDict, total=False): - """One image for the Azure Content Safety image analysis endpoint. +class AzureImageModerationImage(TypedDict): + """One base64 image for the Azure Content Safety image analysis endpoint - Exactly one of ``content`` (base64) or ``blobUrl`` (public URL) must be set. + Remote URLs are rejected by the guardrail because the bytes Azure scans and + the bytes the model later fetches can differ. """ content: str - blobUrl: str class AzureImageModerationGuardrailRequestBody(TypedDict, total=False): diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_image_moderation.py b/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_image_moderation.py index ff64a1b4b5a..fb284ab10de 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_image_moderation.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_image_moderation.py @@ -62,19 +62,20 @@ async def test_image_moderation_posts_base64_content(): @pytest.mark.asyncio -async def test_image_moderation_uses_blob_url_for_remote_images(): +async def test_image_moderation_rejects_remote_image_urls(): guardrail = _image_guardrail() - with patch.object(guardrail.async_handler, "post", return_value=_image_response(0)) as mock_post: - await guardrail.apply_guardrail( - inputs={"images": [REMOTE_IMAGE_URL]}, - request_data={}, - input_type="request", - ) + with patch.object(guardrail.async_handler, "post") as mock_post: + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"images": [REMOTE_IMAGE_URL]}, + request_data={}, + input_type="request", + ) - body = mock_post.call_args.kwargs["json"] - assert body["image"] == {"blobUrl": REMOTE_IMAGE_URL} - assert "content" not in body["image"] + mock_post.assert_not_called() + assert exc_info.value.status_code == 400 + assert "remote image URLs are not supported" in exc_info.value.detail["error"] @pytest.mark.asyncio @@ -141,7 +142,7 @@ async def test_image_moderation_rejects_plain_paths(): mock_post.assert_not_called() assert exc_info.value.status_code == 400 - assert "base64 data URL or an http(s) URL" in exc_info.value.detail["error"] + assert "image must be a base64 data URL" in exc_info.value.detail["error"] @pytest.mark.asyncio From dda979bbbb4c29558434ca0e26693adb4a4ba731 Mon Sep 17 00:00:00 2001 From: NotZhu <3495159203@qq.com> Date: Mon, 2 Nov 2026 18:16:33 +0800 Subject: [PATCH 6/7] fix(guardrails): scan Responses API native input images --- .../guardrail_translation/handler.py | 26 ++++++-- ...test_openai_responses_guardrail_handler.py | 61 +++++++++++++++++++ 2 files changed, 81 insertions(+), 6 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index a905bff3521..b8c8cd1a3eb 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -810,7 +810,7 @@ class OpenAIResponsesHandler(BaseTranslation): skip_tool=skip_tool, scan_only=scan_only, ) - if not extracted.inputs.get("texts"): + if not extracted.inputs.get("texts") and not extracted.inputs.get("images"): return data if scoped_structured_messages: extracted.inputs["structured_messages"] = scoped_structured_messages @@ -1003,6 +1003,11 @@ class OpenAIResponsesHandler(BaseTranslation): texts_to_check.append(text) task_mappings.append((msg_idx, content_idx)) return + if message.get("type") == "input_image": + image_url = message.get("image_url") + if isinstance(image_url, str) and image_url: + images_to_check.append(image_url) + return content: Final = message.get("content", None) if content is None: return @@ -1023,11 +1028,20 @@ class OpenAIResponsesHandler(BaseTranslation): task_mappings.append((msg_idx, int(content_idx))) # Extract images - if content_item.get("type") == "image_url": - image_url = content_item.get("image_url", {}) - if isinstance(image_url, dict): - url = image_url.get("url") - if url: + image_item = cast("Mapping[str, object]", content_item) + if image_item.get("type") == "image_url": + image_url = image_item.get("image_url", {}) + if isinstance(image_url, Mapping): + url = cast("Mapping[str, object]", image_url).get("url") + if isinstance(url, str) and url: + images_to_check.append(url) + elif image_item.get("type") == "input_image": + image_url = image_item.get("image_url") + if isinstance(image_url, str) and image_url: + images_to_check.append(image_url) + elif isinstance(image_url, Mapping): + url = cast("Mapping[str, object]", image_url).get("url") + if isinstance(url, str) and url: images_to_check.append(url) async def _apply_guardrail_responses_to_input( diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index d80a5e835c9..5d7a7d94cd8 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -95,6 +95,24 @@ class MockGuardrail(CustomGuardrail): return inputs +class ImageRecordingGuardrail(CustomGuardrail): + """Mock guardrail that records the images it was asked to scan""" + + def __init__(self, guardrail_name: str = "test") -> None: + super().__init__(guardrail_name=guardrail_name) + self.seen_images: list[list[str]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.seen_images.append(list(inputs.get("images", []))) + return inputs + + class RecordingMaskingGuardrail(MockGuardrail): """MockGuardrail that also records the texts and structured message contents it was shown""" @@ -376,6 +394,49 @@ class TestOpenAIResponsesHandlerInputProcessing: # Image URL should remain unchanged assert result["input"][0]["content"][1]["image_url"]["url"] == "https://example.com/image.jpg" + @pytest.mark.asyncio + async def test_process_input_extracts_native_input_image(self): + handler = OpenAIResponsesHandler() + guardrail = ImageRecordingGuardrail() + + data = { + "input": [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Describe this image"}, + {"type": "input_image", "image_url": "https://example.com/cat.png"}, + ], + "type": "message", + } + ], + "model": "gpt-4", + } + + await handler.process_input_messages(data, guardrail) + + assert guardrail.seen_images == [["https://example.com/cat.png"]] + + @pytest.mark.asyncio + async def test_process_input_image_only_dispatches_guardrail(self): + handler = OpenAIResponsesHandler() + guardrail = ImageRecordingGuardrail() + + data = { + "input": [ + { + "role": "user", + "content": [{"type": "input_image", "image_url": "https://example.com/cat.png"}], + "type": "message", + } + ], + "model": "gpt-4", + } + + await handler.process_input_messages(data, guardrail) + + assert guardrail.seen_images == [["https://example.com/cat.png"]] + @pytest.mark.asyncio async def test_process_input_with_empty_content(self): """Test processing input with empty or None content""" From 0cbec4b14c330af37640398ec58223a7720f8d6c Mon Sep 17 00:00:00 2001 From: NotZhu <3495159203@qq.com> Date: Mon, 2 Nov 2026 18:43:18 +0800 Subject: [PATCH 7/7] fix(guardrails): scan generated images in image responses --- .../guardrail_translation/handler.py | 54 +++++++++++++------ ...test_image_generation_guardrail_handler.py | 51 ++++++++++++++++++ 2 files changed, 90 insertions(+), 15 deletions(-) diff --git a/litellm/llms/openai/image_generation/guardrail_translation/handler.py b/litellm/llms/openai/image_generation/guardrail_translation/handler.py index cce9272017d..a3a82655b6a 100644 --- a/litellm/llms/openai/image_generation/guardrail_translation/handler.py +++ b/litellm/llms/openai/image_generation/guardrail_translation/handler.py @@ -6,7 +6,7 @@ The handler processes the 'prompt' parameter for guardrails. """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, cast from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation @@ -81,6 +81,16 @@ class OpenAIImageGenerationHandler(BaseTranslation): return data + @staticmethod + def _image_field(item: object, name: str) -> object: + fields = cast("Mapping[str, object]", item) if isinstance(item, Mapping) else None + if fields is not None: + return fields.get(name) + return cast( + "object", + getattr(item, name, None), # pyright: ignore[reportUnknownArgumentType] # typed image item + ) + async def process_output_response( self, response: "ImageResponse", @@ -89,21 +99,35 @@ class OpenAIImageGenerationHandler(BaseTranslation): user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, ) -> "ImageResponse": + """Scan generated images with the guardrail before the response is returned. + + ``b64_json`` outputs are passed as base64 data URLs so the bytes the + guardrail sees are exactly the bytes the caller receives. URL outputs + are passed through as URLs and image guardrails that cannot guarantee + the same bytes should reject them. """ - Process output response - typically not needed for image generation. + data_items: Final = cast("list[object]", getattr(response, "data", None) or []) + images: Final[list[str]] = [] + for item in data_items: + b64_json = self._image_field(item, "b64_json") + url = self._image_field(item, "url") + if isinstance(b64_json, str) and b64_json: + images.append(b64_json if b64_json.startswith("data:") else f"data:image/png;base64,{b64_json}") + elif isinstance(url, str) and url: + images.append(url) - Image responses don't contain text to apply guardrails to, so this - method returns the response unchanged. This is provided for completeness - and can be overridden if needed for custom image metadata processing. + if not images: + verbose_proxy_logger.debug("OpenAI Image Generation: No images found in response for guardrail") + return response - Args: - response: Image generation response object - guardrail_to_apply: The guardrail instance to apply - litellm_logging_obj: Optional logging object (unused) - user_api_key_dict: User API key metadata (unused) - - Returns: - Unmodified response (images don't need text guardrails) - """ - verbose_proxy_logger.debug("OpenAI Image Generation: Output processing not needed for image responses") + inputs: Final = GenericGuardrailAPIInputs(images=images) + model: Final = getattr(response, "model", None) + if isinstance(model, str) and model: + inputs["model"] = model + await guardrail_to_apply.apply_guardrail( # pyright: ignore[reportUnknownMemberType] # guardrail hook is untyped + inputs=inputs, + request_data=request_data or {}, + input_type="response", + logging_obj=litellm_logging_obj, + ) return response diff --git a/tests/unit/llms/openai/image_generation/test_image_generation_guardrail_handler.py b/tests/unit/llms/openai/image_generation/test_image_generation_guardrail_handler.py index 0d699b1ec95..a7c4942f882 100644 --- a/tests/unit/llms/openai/image_generation/test_image_generation_guardrail_handler.py +++ b/tests/unit/llms/openai/image_generation/test_image_generation_guardrail_handler.py @@ -25,6 +25,20 @@ class MockGuardrail(CustomGuardrail): return {"texts": [f"{text} [GUARDRAILED]" for text in texts]} +class ImageRecordingGuardrail(CustomGuardrail): + """Mock guardrail that records the images it was asked to scan""" + + def __init__(self, guardrail_name: str = "test") -> None: + super().__init__(guardrail_name=guardrail_name) + self.seen_images: list[list[str]] = [] + + async def apply_guardrail( + self, inputs: dict, request_data: dict, input_type: str, **kwargs + ) -> dict: + self.seen_images.append(list(inputs.get("images", []))) + return inputs + + class TestHandlerDiscovery: """Test that the handler is properly discovered""" @@ -130,6 +144,43 @@ class TestOutputProcessing: assert result.created == 1589478378 assert len(result.data) == 1 + @pytest.mark.asyncio + async def test_process_output_scans_base64_images(self): + handler = OpenAIImageGenerationHandler() + guardrail = ImageRecordingGuardrail() + response = ImageResponse( + created=1589478378, + data=[ImageObject(b64_json="aGVsbG8=")], + ) + + result = await handler.process_output_response(response, guardrail) + + assert result == response + assert guardrail.seen_images == [["data:image/png;base64,aGVsbG8="]] + + @pytest.mark.asyncio + async def test_process_output_passes_url_images_to_guardrail(self): + handler = OpenAIImageGenerationHandler() + guardrail = ImageRecordingGuardrail() + response = ImageResponse( + created=1589478378, + data=[ImageObject(url="https://example.com/image.png")], + ) + + await handler.process_output_response(response, guardrail) + + assert guardrail.seen_images == [["https://example.com/image.png"]] + + @pytest.mark.asyncio + async def test_process_output_without_images_skips_guardrail(self): + handler = OpenAIImageGenerationHandler() + guardrail = ImageRecordingGuardrail() + response = ImageResponse(created=1589478378, data=[]) + + await handler.process_output_response(response, guardrail) + + assert guardrail.seen_images == [] + class TestPIIMaskingScenario: """Test real-world scenario: PII masking in prompts"""