From e4eca8c64ca4ba71ba1239ad0219c832f1ad6416 Mon Sep 17 00:00:00 2001 From: "feng.tsai" Date: Tue, 1 Sep 2026 12:14:33 +0800 Subject: [PATCH] fix(guardrails): make the file-backed image refusal actually fire The refusal read inputs["structured_messages"], which the /v1/messages handler fills by translating to OpenAI spec. That translation drops a file source outright, so the count was always zero and the default block never applied while the provider still forwarded the file to the model. Read the raw request instead. Also stop scanning an image twice: under experimental_use_latest_role_message_only the selected message carries its own image parts into the scan payload, and inputs["images"] holds the same url, so it was fetched and billed once per copy. The tests that covered the refusal hand-built structured_messages in a shape the handler never produces, so they passed throughout. They now feed the raw request, and a new case drives the real handler end to end. --- .../guardrail_hooks/bedrock_guardrails.py | 66 +++++++--- .../test_bedrock_guardrails.py | 121 ++++++++++++++---- 2 files changed, 144 insertions(+), 43 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 50f66b79849..9a92a884437 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -520,6 +520,22 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return url if isinstance(url, str) else None return None + @classmethod + def _image_urls_in(cls, messages: "Sequence[AllMessageValues] | None") -> frozenset[str]: + """Normalized image urls already carried by these messages.""" + found: set[str] = set() # mutable-ok: accumulator, frozen on return + for message in messages or (): + content = message.get("content") + if not isinstance(content, list): + continue + for part in content: + if not isinstance(part, dict) or part.get("type") != "image_url": + continue + url = cls._get_image_url(item=part) + if url is not None: + found.add(cls._normalize_image_input(url)) + return frozenset(found) + def _handle_unscannable_image(self, reason: str) -> None: """Block or warn for an image part ApplyGuardrail cannot scan. @@ -577,15 +593,20 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # on_unscannable_image policy decides rather than this helper. return value - def _refuse_file_backed_images(self, inputs: "GenericGuardrailAPIInputs", input_type: str) -> None: + def _refuse_file_backed_images(self, request_data: Mapping[str, object], input_type: str) -> None: """Hand a file-backed image to on_unscannable_image rather than ignoring it. + Read from the raw request rather than inputs["structured_messages"]: the + /v1/messages handler translates to OpenAI spec before populating that field, + and the translation drops a file source entirely, so the block never survives + to be counted. The provider still forwards it to the model. + Images are a request-side concern; an OUTPUT scan takes generated text, so a file reference sitting in the conversation history is not this scan's problem. """ if input_type != "request": return - found: Final = self._file_backed_image_count(inputs.get("structured_messages")) + found: Final = self._file_backed_image_count(request_data.get("messages")) if not found: return self._handle_unscannable_image( @@ -593,32 +614,32 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) @staticmethod - def _file_backed_image_count(structured_messages: object) -> int: + def _file_backed_image_count(messages: object) -> int: """Count image parts whose bytes live behind a provider Files API. An Anthropic `{"type": "image", "source": {"type": "file", "file_id": ...}}` block carries no data, so the guardrail translation yields nothing for it while the provider still forwards the file to the model. Left alone that is an image the policy never sees, which is the failure this whole path exists - to remove -- so it is counted here and refused rather than documented. + to remove, so it is counted here and refused rather than documented. - Detects that one shape rather than comparing counts against - `inputs["images"]`: structured_messages is already narrowed by the - skip/scope flags, so a mismatch is not by itself evidence of a dropped - image, and blocking a legitimate request is worse than the gap. + Counts this one shape rather than comparing totals against + `inputs["images"]`, whose length is narrowed by the skip and scope flags; a + mismatch there is not by itself evidence of a dropped image, and refusing a + legitimate request is worse than the gap. - That narrowing cuts both ways and this check inherits it. `images` is - extracted from every message while structured_messages holds only the - scoped subset (guardrail_translation/handler.py builds them from different - lists), so a file source in a message the scope excluded is not seen here. - Reading the unscoped list instead would refuse requests for content the - operator's skip flags deliberately took out of scanning, which is a - different wrong answer. + Deliberately reads the whole request rather than the scanned subset. Under + `experimental_use_latest_role_message_only` that means a file source in an + older turn is refused even though the operator scoped scanning to the latest + one. The alternative is the scoped list, which for /v1/messages has already + been translated to OpenAI spec with the file source dropped, so the check + would never fire at all. Over-refusing is the safer of the two errors here, + and `on_unscannable_image: allow` turns it off. """ - if not isinstance(structured_messages, list): + if not isinstance(messages, list): return 0 found = 0 # rebind-ok: running count over the message list - for message in structured_messages: + for message in messages: if not isinstance(message, dict): continue content = message.get("content") @@ -3517,7 +3538,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Before the shortcuts below, so a file-backed image cannot be skipped either. # Both conditions live in the callee: apply_guardrail sits one branch under # ruff-strict's complexity ceiling, and two more here would cross it. - self._refuse_file_backed_images(inputs=inputs, input_type=input_type) + self._refuse_file_backed_images(request_data=request_data, input_type=input_type) try: verbose_proxy_logger.debug( "Bedrock Guardrail: Applying guardrail to %s text(s) and %s image(s)", len(texts), len(image_urls) @@ -3597,8 +3618,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # 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. + # `experimental_use_latest_role_message_only` puts the selected + # message itself into filtered_messages, image parts included, and + # those go through the same builder below. Appending them again + # would fetch and bill each one twice. + already_scanned: Final = self._image_urls_in(filtered_messages) 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 + self._image_content_part(normalized) + for normalized in (self._normalize_image_input(url) for url in image_urls) + if normalized not in already_scanned ] image_message: Final = ( (ChatCompletionUserMessage(role="user", content=image_parts),) if image_parts else () 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 ad6c1b24d46..0131624f69d 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 @@ -5799,10 +5799,9 @@ class TestBedrockGuardrailImageInput: the silent pass this path exists to remove. Documented is not the same as safe; under the default policy the request is refused. """ - inputs = { - "texts": ["what does this say?"], - "images": [], - "structured_messages": [ + inputs = {"texts": ["what does this say?"], "images": []} + request_data = { + "messages": [ { "role": "user", "content": [ @@ -5810,7 +5809,7 @@ class TestBedrockGuardrailImageInput: {"type": "image", "source": {"type": "file", "file_id": "file_abc"}}, ], } - ], + ] } g = self._guardrail() @@ -5824,11 +5823,91 @@ class TestBedrockGuardrailImageInput: # the failure mode being pinned is a silent pass, not an AWS error. with patch.object(g, "make_bedrock_api_request", new=spy): with pytest.raises(HTTPException) as exc_info: - await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + await g.apply_guardrail(inputs=inputs, request_data=request_data, input_type="request") assert "file id" in str(exc_info.value.detail) assert sent == [], "refused before any scan was attempted" + @pytest.mark.asyncio + async def test_a_file_backed_image_is_refused_through_the_real_translation(self): + """Drive the /v1/messages handler instead of hand-building its output. + + The handler translates to OpenAI spec before filling structured_messages, and + that translation drops a file source, so a check reading structured_messages + passes every hand-written fixture and never fires in production. + """ + from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler + + data = { + "model": "claude-sonnet-4-5", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "what does this say?"}, + {"type": "image", "source": {"type": "file", "file_id": "file_abc"}}, + ], + } + ], + } + assert not self._file_parts_in(AnthropicMessagesHandler().get_structured_messages(data)), ( + "the translation is expected to drop the file source; that is why this test exists" + ) + + g = self._guardrail() + with patch.object(g, "make_bedrock_api_request", new=AsyncMock(return_value={"action": "NONE"})): + with pytest.raises(HTTPException) as exc_info: + await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=g) + + assert "file id" in str(exc_info.value.detail) + + @staticmethod + def _file_parts_in(messages) -> int: + return sum( + 1 + for message in messages or () + if isinstance(message, dict) + for part in (message.get("content") if isinstance(message.get("content"), list) else ()) + if isinstance(part, dict) and part.get("type") == "image" + ) + + @pytest.mark.asyncio + async def test_latest_message_only_does_not_scan_the_same_image_twice(self): + """The selected message carries its own image parts into the scan payload. + + `inputs["images"]` holds that same url, so appending it again would fetch and + bill the image twice. + """ + url = self._PNG_DATA_URI + g = self._guardrail(experimental_use_latest_role_message_only=True) + 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": ["hello"], + "images": [url], + "structured_messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "image_url", "image_url": {"url": url}}, + ], + } + ], + }, + request_data={}, + input_type="request", + ) + + images = [item for item in sent[0]["content"] if "image" in item] + assert len(images) == 1, f"the image was sent {len(images)} times" + @pytest.mark.asyncio async def test_a_file_backed_image_is_let_through_under_the_allow_policy(self): """An operator who would rather serve it unscanned can still say so.""" @@ -5841,17 +5920,15 @@ class TestBedrockGuardrailImageInput: with patch.object(g, "make_bedrock_api_request", new=spy): await g.apply_guardrail( - inputs={ - "texts": ["hello"], - "images": [], - "structured_messages": [ + inputs={"texts": ["hello"], "images": []}, + request_data={ + "messages": [ { "role": "user", "content": [{"type": "image", "source": {"type": "file", "file_id": "file_abc"}}], } - ], + ] }, - request_data={}, input_type="request", ) @@ -5874,10 +5951,9 @@ class TestBedrockGuardrailImageInput: with patch.object(g, "make_bedrock_api_request", new=spy): await g.apply_guardrail( - inputs={ - "texts": ["hello"], - "images": [self._PNG_DATA_URI], - "structured_messages": [ + inputs={"texts": ["hello"], "images": [self._PNG_DATA_URI]}, + request_data={ + "messages": [ { "role": "user", "content": [ @@ -5886,9 +5962,8 @@ class TestBedrockGuardrailImageInput: {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}, ], } - ], + ] }, - request_data={}, input_type="request", ) @@ -5913,10 +5988,9 @@ class TestBedrockGuardrailImageInput: with patch.object(g, "make_bedrock_api_request", new=spy): with pytest.raises(HTTPException) as exc_info: await g.apply_guardrail( - inputs={ - "texts": ["hello"], - "images": [], - "structured_messages": [ + inputs={"texts": ["hello"], "images": []}, + request_data={ + "messages": [ "not a message", 123, {"role": "user", "content": "a plain string, not a list"}, @@ -5924,9 +5998,8 @@ class TestBedrockGuardrailImageInput: "role": "user", "content": [{"type": "image", "source": {"type": "file", "file_id": "file_abc"}}], }, - ], + ] }, - request_data={}, input_type="request", )