scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(guardrails): scan images in apply_guardrail, the path the proxy actually uses

Browse source 
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) <noreply@anthropic.com>
This commit is contained in:
samtsai15 2026-08-26 19:03:44 +08:00
parent a19dd2d92e
commit 735e560c11
2 changed files with 137 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

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