This commit is contained in:
NotZhu 2026-11-05 05:44:30 +00:00 • committed by GitHub
commit 8b56987eba
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 746 additions and 87 deletions

View file

@ -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

View file

@ -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

View file

@ -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(

View file

@ -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,
}

View file

@ -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.

View file

@ -0,0 +1,171 @@
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] = False
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://")):
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"},
)
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

View file

@ -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(

View file

@ -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"

View file

@ -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):
"""One base64 image for the Azure Content Safety image analysis endpoint
Remote URLs are rejected by the guardrail because the bytes Azure scans and
the bytes the model later fetches can differ.
"""
content: 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"

View file

@ -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):

View file

@ -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"""

View file

@ -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):

View file

@ -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"""

View file

@ -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"""

View file

@ -0,0 +1,231 @@
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_rejects_remote_image_urls():
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": [REMOTE_IMAGE_URL]},
request_data={},
input_type="request",
)
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
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 "image must be a base64 data 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_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(
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"]