scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(guardrails): make the file-backed image refusal actually fire

Browse source 
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.
This commit is contained in:
feng.tsai 2026-09-01 12:14:33 +08:00
parent c2db05f0a5
commit e4eca8c64c
2 changed files with 144 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

66
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

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

121
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

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