From 735e560c1160523a3bfb93d9f066adfb761e0dee Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Wed, 26 Aug 2026 19:03:44 +0800 Subject: [PATCH] fix(guardrails): scan images in apply_guardrail, the path the proxy actually uses The payload fix alone never runs on a real request. ProxyLogging._execute_guardrail_hook (proxy/utils.py:1216) routes any guardrail that defines `apply_guardrail` through unified_guardrail unless it sets `use_native_lifecycle_hooks`: has_apply_guardrail = "apply_guardrail" in type(callback).__dict__ and not getattr( callback, "use_native_lifecycle_hooks", False) target = unified_guardrail if has_apply_guardrail else callback BedrockGuardrail defines apply_guardrail and does not set that flag, so a /v1/chat/completions request reaches apply_guardrail, which read inputs["texts"] only. Its own docstring said "images unchanged". Measured against a live guardrail with the IMAGE modality enabled, same messages, same account: guardrail.async_pre_call_hook png imageUnits=1 gif blocked 400 ProxyLogging.pre_call_hook png imageUnits=0 gif passed through The second row is what a proxy user gets, and it matches the imageUnits: 0 the reporter measured in #35332. Nothing new is needed on the extraction side. The endpoint translations already populate inputs["images"]: OpenAIChatCompletionsHandler from `image_url` parts (openai/chat/guardrail_translation/handler.py:253), AnthropicMessagesHandler from `image`/`source` blocks. Five guardrails already consume that field (vigil_guard, custom_code, deepkeep, straiker, generic_guardrail_api), so the contract exists and Bedrock was the one dropping it. apply_guardrail now appends the images as one extra user message and lets the normal message path build the payload, so the unified and native routes share the decoding, the png/jpeg check and on_unscannable_image instead of drifting apart. Requests that carry an image but no text are no longer skipped. _normalize_image_input handles the two shapes that field carries. The OpenAI translation appends the caller's image_url verbatim, already a data: URI or an https URL. The Anthropic one returns source["data"] only, dropping media_type, so the entry is bare base64 that the decoder would reject as unreadable and, under on_unscannable_image=block, turn a legitimate /v1/messages call into a 400. The format is sniffed back from the base64 prefix. Images are only read for input_type == "request"; the OUTPUT source scans model-generated text. Three tests, all failing before this change with "image never reached the payload: ['text']". Co-Authored-By: Claude Opus 5 (1M context) --- .../guardrail_hooks/bedrock_guardrails.py | 72 +++++++++++++++++-- .../test_bedrock_guardrails.py | 71 ++++++++++++++++++ 2 files changed, 137 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 3f02d512e08..fa9b3567638 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -56,7 +56,12 @@ from litellm.proxy.guardrails.anthropic_sse import ( ) from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import BedrockChecksConfigModel, GuardrailEventHooks -from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionImageObject, + ChatCompletionImageUrlObject, + ChatCompletionUserMessage, +) from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockChecksMessage, BedrockChecksViolation, @@ -424,6 +429,37 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): reason, ) + #: base64 magic-byte prefixes for the formats ApplyGuardrail accepts. + _BASE64_IMAGE_PREFIXES: ClassVar[tuple[tuple[str, str], ...]] = ( + ("iVBORw0KGgo", "image/png"), + ("/9j/", "image/jpeg"), + ) + + @staticmethod + def _image_content_part(url: str) -> ChatCompletionImageObject: + """One OpenAI-format image content part, so the payload builder handles it.""" + return ChatCompletionImageObject(type="image_url", image_url=ChatCompletionImageUrlObject(url=url)) + + @classmethod + def _normalize_image_input(cls, value: str) -> str: + """Return a data URI or URL that `_build_image_content_item` can consume. + + `GenericGuardrailAPIInputs["images"]` is not a single shape. The OpenAI chat + translation appends the caller's `image_url` verbatim, so entries are already a + `data:` URI or an `https://` URL. The Anthropic translation's `_image_sources` + returns `source["data"]` only, which is bare base64 with the `media_type` + dropped. Sniff the format back from the base64 prefix so both shapes reach the + same decoder instead of the bare-base64 one failing as unreadable. + """ + if value.startswith(("data:", "http://", "https://")): + return value + for prefix, media_type in cls._BASE64_IMAGE_PREFIXES: + if value.startswith(prefix): + return f"data:{media_type};base64,{value}" + # Unrecognized: hand it over as-is and let the decoder reject it, so the + # on_unscannable_image policy decides rather than this helper. + return value + async def _build_image_content_item(self, image_url: str) -> BedrockContentItem | None: """Decode or fetch an image part into an ApplyGuardrail image block. @@ -3170,7 +3206,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): logging_obj: Optional logging object Returns: - GenericGuardrailAPIInputs - processed_texts may be masked, images unchanged + GenericGuardrailAPIInputs - processed_texts may be masked, images are + scanned but returned unchanged (ApplyGuardrail does not rewrite images) Raises: Exception: If content is blocked by Bedrock guardrail @@ -3178,8 +3215,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # NOTE: Use `or []` to handle case where inputs["texts"] is explicitly None. # dict.get("texts", []) would return None if the key exists with a None value. texts: Final = inputs.get("texts") or [] + # Images only exist on the request side; ApplyGuardrail's OUTPUT source takes + # model-generated text. The endpoint translations already extract them: + # OpenAIChatCompletionsHandler from `image_url` parts, AnthropicMessagesHandler + # from `image`/`source` blocks. Five other guardrails already consume this + # field; Bedrock was the one that dropped it on the floor. + image_urls: Final = tuple(inputs.get("images") or ()) if input_type == "request" else () try: - verbose_proxy_logger.debug("Bedrock Guardrail: Applying guardrail to %s text(s)", len(texts)) + verbose_proxy_logger.debug( + "Bedrock Guardrail: Applying guardrail to %s text(s) and %s image(s)", len(texts), len(image_urls) + ) if input_type == "request": incremental_result: Final = await self._apply_incremental_request_scan( @@ -3204,8 +3249,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): scanned_slice: Final = selection.scanned_slice scanned_role_subset: Final = selection.scanned_role_subset - # Bedrock will throw an error if there is no text to process - if filtered_messages: + # Bedrock rejects an empty content list, so only skip when there is + # neither text nor an image to scan. + if filtered_messages or image_urls: _log_hook = GuardrailEventHooks.pre_call if input_type == "request" else GuardrailEventHooks.post_call # Map the abstract input_type to the Bedrock source parameter. # "request" -> INPUT (scan user-supplied content) @@ -3238,9 +3284,23 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): logging_event_type=_log_hook, ) else: + # Append the images as one extra user message. Reusing the normal + # message path means `_create_bedrock_input_content_request` does the + # decoding, format check and on_unscannable_image handling, so the + # unified and native lifecycle paths cannot drift apart. + image_parts: Final = [ # mutable-ok: OpenAI message content is a list in the wire format + self._image_content_part(self._normalize_image_input(url)) for url in image_urls + ] + image_message: Final = ( + (ChatCompletionUserMessage(role="user", content=image_parts),) if image_parts else () + ) + scan_messages: Final = [ # mutable-ok: make_bedrock_api_request takes a list of messages + *filtered_messages, + *image_message, + ] bedrock_response = await self.make_bedrock_api_request( source="INPUT", - messages=filtered_messages, + messages=scan_messages, request_data=request_data, logging_event_type=_log_hook, ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 542cdab64ff..931a4f30d24 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5467,6 +5467,77 @@ class TestBedrockGuardrailImageInput: get.assert_not_awaited() assert request["content"] == [] + @pytest.mark.asyncio + async def test_apply_guardrail_scans_images_from_inputs(self): + """The proxy reaches BedrockGuardrail through `apply_guardrail`, not the native hook. + + ProxyLogging._execute_guardrail_hook routes any guardrail that defines + `apply_guardrail` through unified_guardrail, so a fix that only touches + async_pre_call_hook never runs on a real request. The endpoint translations + already put images in inputs["images"]; this asserts they reach the payload. + """ + g = self._guardrail() + sent: list = [] + + async def spy(**kwargs): + sent.append(await g.convert_to_bedrock_format(source="INPUT", messages=kwargs["messages"])) + return {"action": "NONE", "outputs": []} + + with patch.object(g, "make_bedrock_api_request", new=spy): + await g.apply_guardrail( + inputs={"texts": ["what does this say?"], "images": [self._PNG_DATA_URI]}, + request_data={}, + input_type="request", + ) + + kinds = [k for item in sent[0]["content"] for k in item] + assert "image" in kinds, f"image never reached the payload: {kinds}" + + @pytest.mark.asyncio + async def test_apply_guardrail_scans_bare_base64_images(self): + """Anthropic's translation drops media_type and passes bare base64. + + `_image_sources` returns source["data"] only, so the entry is not a data URI. + Without sniffing the format back it would be rejected as unreadable and, under + on_unscannable_image=block, turn a legitimate /v1/messages call into a 400. + """ + g = self._guardrail() + sent: list = [] + + async def spy(**kwargs): + sent.append(await g.convert_to_bedrock_format(source="INPUT", messages=kwargs["messages"])) + return {"action": "NONE", "outputs": []} + + bare_base64 = self._PNG_DATA_URI.split(",")[1] + with patch.object(g, "make_bedrock_api_request", new=spy): + await g.apply_guardrail( + inputs={"texts": [], "images": [bare_base64]}, + request_data={}, + input_type="request", + ) + + kinds = [k for item in sent[0]["content"] for k in item] + assert "image" in kinds, f"bare base64 image never reached the payload: {kinds}" + + @pytest.mark.asyncio + async def test_apply_guardrail_ignores_images_on_the_response_side(self): + """ApplyGuardrail's OUTPUT source scans model-generated text, not input images.""" + g = self._guardrail() + sent: list = [] + + async def spy(**kwargs): + sent.append(kwargs) + return {"action": "NONE", "outputs": []} + + with patch.object(g, "make_bedrock_api_request", new=spy): + await g.apply_guardrail( + inputs={"texts": ["some model output"], "images": [self._PNG_DATA_URI]}, + request_data={}, + input_type="response", + ) + + assert sent and sent[0]["source"] == "OUTPUT" + def test_masking_keeps_image_parts_in_the_request(self): """Masking rewrites text in place; the image must survive to reach the model.""" image_part = {"type": "image_url", "image_url": {"url": self._PNG_DATA_URI}}