This commit is contained in:
EarthFeng 2026-08-27 04:05:16 +00:00 committed by GitHub
commit 97e4c2899e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 1057 additions and 58 deletions

View file

@ -3365,14 +3365,14 @@ class BedrockImageProcessor:
return base64_bytes, content_type
@staticmethod
async def get_image_details_async(image_url) -> tuple[str, str]:
async def get_image_details_async(image_url, max_bytes: int | None = None) -> tuple[str, str]:
try:
client: Final = get_async_httpx_client(
llm_provider=httpxSpecialProvider.PromptFactory,
params={"concurrent_limit": 1},
)
# Send a GET request to the image URL
response: Final[httpx.Response] = await async_safe_get(client, image_url)
response: Final[httpx.Response] = await async_safe_get(client, image_url, max_bytes=max_bytes)
response.raise_for_status() # Raise an exception for HTTP errors
return BedrockImageProcessor._post_call_image_processing(response, image_url)
@ -3553,13 +3553,20 @@ class BedrockImageProcessor:
return cls._create_bedrock_block(img_bytes, mime_type, image_format)
@classmethod
async def process_image_async(cls, image_url: str, format: str | None) -> BedrockContentBlock:
"""Asynchronous image processing."""
async def process_image_async(
cls, image_url: str, format: str | None, max_bytes: int | None = None
) -> BedrockContentBlock:
"""Asynchronous image processing.
``max_bytes`` caps a remote fetch and is ignored for a base64 data URI,
whose size the caller already knows before calling. Omitting it keeps the
previous unbounded fetch.
"""
if "base64" in image_url:
img_bytes, mime_type, image_format = cls._parse_base64_image(image_url)
elif "http://" in image_url or "https://" in image_url:
img_bytes, mime_type = await BedrockImageProcessor.get_image_details_async(image_url)
img_bytes, mime_type = await BedrockImageProcessor.get_image_details_async(image_url, max_bytes=max_bytes)
image_format = mime_type.split("/")[1]
else:
raise ValueError("Unsupported image type. Expected either image url or base64 encoded string")

View file

@ -44,6 +44,14 @@ class SSRFError(ValueError):
"""Raised when a URL targets a blocked network."""
class PayloadTooLargeError(ValueError):
"""Raised when a fetched body exceeds the caller's byte cap.
A ``ValueError`` subclass, like :class:`SSRFError`, so callers that already
treat a malformed remote response as a rejected fetch need no new except arm.
"""
def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") -> str:
"""Percent-encode one user-controlled URL path segment.
@ -412,21 +420,88 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> Any:
raise SSRFError("Too many redirects")
async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any:
"""Async version of safe_get."""
# Headers that describe the wire encoding of the body. `aiter_bytes` yields
# decoded bytes, so carrying these onto the rebuilt response would describe it
# wrongly (a gzip label on already-inflated bytes, a length from before decoding).
_TRANSFER_ENCODING_HEADERS: Final = frozenset({"content-encoding", "content-length"})
def _underlying_httpx_client(client: object) -> httpx.AsyncClient:
"""Return the object exposing ``stream``.
``AsyncHTTPHandler``/``HTTPHandler`` wrap an httpx client and forward only
``get``/``post``/..., so streaming has to go through the wrapped ``.client``.
A raw httpx client is returned unchanged.
"""
inner: Final = getattr(client, "client", client)
if not isinstance(inner, httpx.AsyncClient):
raise TypeError(f"cannot stream from {type(client).__name__}: no httpx client to stream with")
return inner
async def _async_get_capped(
client: object,
url: str,
max_bytes: int,
request_kwargs: dict[str, object], # mutable-ok: forwarded straight to httpx as **kwargs
) -> httpx.Response:
"""GET ``url``, aborting the transfer once the body exceeds ``max_bytes``.
``client.get`` buffers the whole body before returning, so a caller-supplied
URL serving an arbitrarily large or indefinitely chunked response is an
unbounded allocation. Streaming makes the cap effective during the transfer
rather than after it.
Returns a fully-read response so callers keep using ``.content`` as before.
"""
async with _underlying_httpx_client(client).stream("GET", url, **request_kwargs) as response:
chunks: list[bytes] = [] # mutable-ok: accumulator for the capped body
total = 0
async for chunk in response.aiter_bytes():
total += len(chunk)
if total > max_bytes:
raise PayloadTooLargeError(f"remote body exceeds {max_bytes} bytes")
chunks.append(chunk)
kept_headers: Final = [ # mutable-ok: httpx.Response takes the header pairs as a list
(k, v) for k, v in response.headers.multi_items() if k.lower() not in _TRANSFER_ENCODING_HEADERS
]
return httpx.Response(
status_code=response.status_code,
headers=kept_headers,
content=b"".join(chunks),
request=response.request,
)
async def async_safe_get(client: Any, url: str, max_bytes: int | None = None, **kwargs: Any) -> Any:
"""Async version of safe_get.
``max_bytes`` caps the response body, rejecting an oversized transfer with
:class:`PayloadTooLargeError` while it is still in flight. Omitting it keeps
the previous unbounded buffering, so existing callers are unaffected.
"""
async def _issue(
target_url: str,
request_kwargs: dict[str, object], # mutable-ok: forwarded straight to httpx as **kwargs
) -> httpx.Response:
if max_bytes is None:
return await client.get(target_url, **request_kwargs)
return await _async_get_capped(client, target_url, max_bytes, request_kwargs)
if not getattr(litellm, "user_url_validation", True):
kwargs.setdefault("follow_redirects", True)
return await client.get(url, **kwargs)
return await _issue(url, kwargs)
kwargs.pop("follow_redirects", None)
caller_headers: Final = kwargs.pop("headers", {})
for _ in range(_MAX_REDIRECTS):
validated_url, original_host = validate_url(url)
response = await client.get(
validated_url,
headers={**caller_headers, "Host": original_host},
follow_redirects=False,
hop_kwargs: dict[str, object] = { # mutable-ok: a fresh per-hop kwargs dict, consumed by this call
**kwargs,
)
"headers": {**caller_headers, "Host": original_host}, # mutable-ok: httpx takes headers as a dict
"follow_redirects": False,
}
response = await _issue(validated_url, hop_kwargs)
if not response.is_redirect:
return response
# Resolve the next hop against the ORIGINAL (pre-rewrite) URL so

View file

@ -788,12 +788,40 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _image_sources(block: Mapping[str, object]) -> tuple[str, ...]:
"""Normalize an Anthropic image block into strings a guardrail can read.
`source` is one of three shapes (types/llms/anthropic.py:259):
{"type": "base64", "media_type": "image/png", "data": "<b64>"}
{"type": "url", "url": "https://..."}
{"type": "file", "file_id": "..."}
base64 is returned as a data URI rather than the bare payload: consumers of
``GenericGuardrailAPIInputs["images"]`` otherwise have no way to know the
format, and an API like Bedrock's ApplyGuardrail requires it. url is passed
through so the consumer can fetch it under its own SSRF policy.
file is not resolvable here (the bytes live behind the Files API), so it
yields nothing. That is a silent gap for any consumer that treats a missing
entry as "no image to scan"; scanning a file_id needs a fetch this extractor
has no client for.
"""
source: Final = block.get("source")
if not isinstance(source, Mapping):
return ()
# Could be base64 or url
source_type: Final = source.get("type")
if source_type == "url":
url: Final = source.get("url")
return (url,) if isinstance(url, str) and url else ()
data: Final = source.get("data")
return (data,) if data else ()
if not isinstance(data, str) or not data:
return ()
media_type: Final = source.get("media_type")
if isinstance(media_type, str) and media_type:
return (f"data:{media_type};base64,{data}",)
return (data,)
async def _apply_guardrail_responses_to_input(
self,

View file

@ -10,6 +10,7 @@ import sys
sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path
import asyncio
import binascii
import copy
import json
import re
@ -32,6 +33,8 @@ from litellm.exceptions import ModifyResponseException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import bedrock_guardrail_cost
from litellm.litellm_core_utils.prompt_templates.factory import BedrockImageProcessor
from litellm.litellm_core_utils.url_utils import PayloadTooLargeError, SSRFError
from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_scan_only_tool_results_for_guardrail,
@ -53,16 +56,24 @@ 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,
BedrockContentItem,
BedrockGuardrailChecksResponse,
BedrockGuardrailImageFormat,
BedrockGuardrailImageSource,
BedrockGuardrailOutput,
BedrockGuardrailQualifier,
BedrockGuardrailResponse,
BedrockGuardrailUsage,
BedrockImageContent,
BedrockRequest,
BedrockTextContent,
)
@ -133,6 +144,27 @@ _CONTENT_TYPE_TO_QUALIFIER: Final[dict[str, BedrockGuardrailQualifier]] = {
# must not be graded against as if it were the application's own source material.
_GROUNDING_SOURCE_TRUSTED_ROLES: Final = frozenset({"system", "developer"})
# ApplyGuardrail only accepts png/jpeg image blocks. Anything else (gif, webp, ...)
# has no representation in the payload, so it cannot be scanned at all.
# AWS's hard limits for image content filters:
# https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails-mmfilter.html
# 4 MB per image, 20 images per request, 8000x8000, PNG/JPEG only, 25 images/second,
# and only the first 100 words of text inside an image are evaluated.
# Nothing in litellm checked these. BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS is the
# text-side quota and says nothing about images, and the strings
# _BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS matches are all text-shaped, so an image-count
# rejection would not even reach the chunking fallback.
_MAX_IMAGE_BYTES: Final = 4 * 1024 * 1024
_MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL: Final = 20
_APPLY_GUARDRAIL_IMAGE_FORMATS: Final[
dict[str, BedrockGuardrailImageFormat]
] = { # mutable-ok: module-level lookup table, never mutated
"png": "png",
"jpeg": "jpeg",
"jpg": "jpeg",
}
class QualifiedTextBlock(NamedTuple):
"""A piece of message text paired with its Bedrock grounding qualifier (if any)."""
@ -221,6 +253,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
prompt_attack_threshold: float | None = 0.5,
pii_confidence_threshold: float | None = 0.5,
chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS,
on_unscannable_image: Literal["block", "allow"] = "block",
**kwargs,
):
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
@ -229,6 +262,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self.guardrail_provider = "bedrock"
self.chunk_budget_chars = chunk_budget_chars
self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only"))
# What to do with an image part ApplyGuardrail cannot scan (non png/jpeg, or a
# remote url we refuse to fetch). Defaults to blocking: the image reaches the
# model regardless, so allowing it would be a silent guardrail bypass.
self.on_unscannable_image: Literal["block", "allow"] = on_unscannable_image
# Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks`
# routes the guardrail to InvokeGuardrailChecks; absent => ApplyGuardrail.
@ -316,28 +353,180 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
return cleaned or None
def _create_bedrock_input_content_request(self, messages: list[AllMessageValues] | None) -> BedrockRequest:
"""
Create a bedrock request for the input content - the LLM request.
"""
async def _create_bedrock_input_content_request(self, messages: list[AllMessageValues] | None) -> BedrockRequest:
"""Create a bedrock request for the input content - the LLM request."""
bedrock_request: Final[BedrockRequest] = BedrockRequest(source="INPUT")
bedrock_request_content: Final[list[BedrockContentItem]] = []
if messages is None:
return bedrock_request
for message in messages:
blocks = self.get_content_items_for_message(message=message)
if blocks is None:
continue
for block in blocks:
# INPUT scans send plain text only. Grounding qualifiers are attached
# exclusively when assembling the OUTPUT request, so a caller cannot use
# a grounding_source/query tag to change how input-safety policies treat
# their content (which would be an input-guardrail bypass).
bedrock_request_content.append(BedrockContentItem(text=BedrockTextContent(text=block.text)))
bedrock_request["content"] = bedrock_request_content
per_message = await asyncio.gather(*(self._build_input_content_items(message=message) for message in messages))
# mutable-ok: BedrockRequest["content"] is a list in the AWS wire format
bedrock_request["content"] = [item for items in per_message for item in items]
return bedrock_request
async def _build_input_content_items(self, message: AllMessageValues) -> tuple[BedrockContentItem, ...]:
"""Flatten one request message into ApplyGuardrail INPUT content items.
Grounding qualifiers are attached only when assembling the OUTPUT request, so a
grounding_source/query tag cannot change how input-safety policies treat content
"""
content = message.get("content")
if content is None:
return ()
if isinstance(content, str):
return (BedrockContentItem(text=BedrockTextContent(text=content)),)
if not isinstance(content, list):
return ()
parts: Final = cast( # cast-ok: AllMessageValues content is a union of part TypedDicts
tuple[object, ...], tuple(content)
)
items: Final = await asyncio.gather(*(self._build_input_content_item(item=item) for item in parts))
return tuple(item for item in items if item is not None)
async def _build_input_content_item(self, item: object) -> BedrockContentItem | None:
if isinstance(item, str):
return BedrockContentItem(text=BedrockTextContent(text=item))
if not isinstance(item, dict):
return None
part: Final = cast(Mapping[str, object], item) # cast-ok: narrowed to dict on the line above
# Provider transformations branch on `type`, so an image_url part reaches the
# model as an image even when it also carries `text`. Reading `text` first would
# scan that decoy and forward the image unscanned
if part.get("type") == "image_url":
image_url: Final = self._get_image_url(item=part)
if image_url is None:
return None
return await self._build_image_content_item(image_url=image_url)
text: Final = part.get("text")
if isinstance(text, str):
return BedrockContentItem(text=BedrockTextContent(text=text))
return None
@staticmethod
def _get_image_url(item: Mapping[str, object]) -> str | None:
if item.get("type") != "image_url":
return None
image_url: Final = item.get("image_url")
if isinstance(image_url, str):
return image_url
if isinstance(image_url, dict):
url: Final = cast(Mapping[str, object], image_url).get("url") # cast-ok: narrowed to dict on the line above
return url if isinstance(url, str) else None
return None
def _handle_unscannable_image(self, reason: str) -> None:
"""Block or warn for an image part ApplyGuardrail cannot scan.
The image reaches the model either way, so skipping it silently would let a
caller defeat an IMAGE-modality guardrail by picking a format the API rejects
"""
if self.on_unscannable_image == "block":
raise HTTPException(
status_code=400,
detail={ # mutable-ok: HTTPException detail payload, serialized immediately
"error": "Violated guardrail policy",
"bedrock_guardrail_response": (
f"Request contains an image the guardrail cannot scan ({reason}). "
"ApplyGuardrail accepts png/jpeg images only. Set "
"'on_unscannable_image: allow' on this guardrail to send such "
"requests to the model unscanned."
),
"guardrail_name": self.guardrail_name,
},
)
verbose_proxy_logger.warning(
"Bedrock Guardrail %s: image part will not be scanned (%s); on_unscannable_image=allow, forwarding it to the model anyway",
self.guardrail_name,
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.
With `user_url_validation` off, `async_safe_get` degrades to an unrestricted,
redirect-following GET on a caller-supplied URL, so it is not fetched at all
"""
if not image_url.startswith("data:") and not getattr(litellm, "user_url_validation", True):
self._handle_unscannable_image(
reason=f"remote image url not fetched because litellm.user_url_validation is disabled: {image_url}"
)
return None
try:
# Cap the fetch itself. The decoded-size check below runs after the
# bytes are already resident, so on its own it does not stop a caller
# from pointing the proxy at an arbitrarily large or chunked response
# -- and `_build_input_content_items` gathers these concurrently, so
# one request with several URLs multiplies the allocation.
block: Final = await BedrockImageProcessor.process_image_async(
image_url=image_url, format=None, max_bytes=_MAX_IMAGE_BYTES
)
except PayloadTooLargeError as e:
# Named before the ValueError arm it subclasses, so the operator sees
# "too large" rather than "could not be read" for a size rejection.
self._handle_unscannable_image(reason=f"remote image over ApplyGuardrail's 4 MB limit: {e}")
return None
except (httpx.HTTPError, SSRFError, ValueError, TypeError, KeyError, binascii.Error) as e:
self._handle_unscannable_image(reason=f"image content could not be read: {e}")
return None
image_block: Final = block.get("image")
image_format: Final = (
_APPLY_GUARDRAIL_IMAGE_FORMATS.get(str(image_block.get("format"))) if image_block else None
)
# mutable-ok: {} is an empty default for .get, never mutated
image_bytes: Final = image_block.get("source", {}).get("bytes") if image_block else None
if image_format is None or not image_bytes:
self._handle_unscannable_image(reason="attachment is not a png/jpeg image")
return None
# base64 decodes to roughly 3/4 of its length; estimate rather than decode the
# whole image a second time just to measure it.
decoded_size: Final = len(image_bytes) * 3 // 4
if decoded_size > _MAX_IMAGE_BYTES:
self._handle_unscannable_image(
reason=f"image is {decoded_size / 1024 / 1024:.1f} MB, over ApplyGuardrail's 4 MB limit"
)
return None
return BedrockContentItem(
image=BedrockImageContent(
format=image_format,
source=BedrockGuardrailImageSource(bytes=image_bytes),
)
)
def _create_bedrock_output_content_request(
self,
response: Any | ModelResponse,
@ -386,7 +575,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
items.append(self._build_content_item(block))
return items
def convert_to_bedrock_format(
async def convert_to_bedrock_format(
self,
source: Literal["INPUT", "OUTPUT"],
messages: list[AllMessageValues] | None = None,
@ -403,7 +592,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
"""
bedrock_request: BedrockRequest = BedrockRequest(source=source)
if source == "INPUT":
bedrock_request = self._create_bedrock_input_content_request(messages=messages)
bedrock_request = await self._create_bedrock_input_content_request(messages=messages)
elif source == "OUTPUT":
bedrock_request = self._create_bedrock_output_content_request(response=response, messages=messages)
return bedrock_request
@ -839,7 +1028,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
"""
start_time: Final = datetime.now(timezone.utc)
bedrock_request_data: Final[dict] = dict(
self.convert_to_bedrock_format(source=source, messages=messages, response=response)
await self.convert_to_bedrock_format(source=source, messages=messages, response=response)
)
api_key: str | None = None
if request_data:
@ -955,7 +1144,48 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
guardrail block on any (sub-)chunk raises immediately
-- callers must not lose that signal by continuing to post the remaining
chunks.
Images are the one case split up front rather than reactively. Chunking here
is a recovery path: the content goes out as a single call and is only
re-batched once AWS rejects it AND `_is_input_too_large_error` matches. Those
substrings ("text unit", "too long", ...) are all text-shaped, so a rejection
for exceeding 20 images per request may never reach this fallback at all, and
bisection cannot rescue it either -- `_split_bedrock_content` reads
`item["text"]` to halve a lone item, which is empty for an image, so it gives
up and re-raises the original error. Splitting before the first call keeps a
many-image request inside the documented limit instead.
"""
image_count: Final = sum(1 for item in content if "image" in item)
if allow_chunking and image_count > _MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL:
preemptive_batches: Final = self._bin_pack_bedrock_content(content, budget=self.chunk_budget_chars)
if len(preemptive_batches) > 1:
verbose_proxy_logger.warning(
"Bedrock Guardrail: %d image(s) exceeds ApplyGuardrail's limit of %d per request; "
"splitting into %d calls before sending",
image_count,
_MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL,
len(preemptive_batches),
)
preemptive_results: Final = [ # mutable-ok: await needs a list comprehension; frozen to a tuple below
await self._apply_guardrail_content_with_chunking(
content=batch,
base_request_data=base_request_data,
credentials=credentials,
aws_region_name=aws_region_name,
api_key=api_key,
request_data=request_data,
event_type=event_type,
start_time=start_time,
# Safe to keep enabled: every batch _bin_pack_bedrock_content
# returns holds at most 20 images, so the recursive call falls
# straight through this branch and text chunking still applies.
allow_chunking=allow_chunking,
completed_chunk_usages=completed_chunk_usages,
)
for batch in preemptive_batches
]
return tuple(result for results in preemptive_results for result in results)
try:
response: Final = await self._post_apply_guardrail_content_with_retry(
content=content,
@ -1347,15 +1577,27 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
if not content:
return (tuple(content),)
lengths: Final = tuple(len((item.get("text") or BedrockTextContent()).get("text") or "") for item in content)
# An image item has no `text`, so it measures 0 against `budget` and
# `used + 0 <= budget` always holds: without a second dimension every image
# lands in whichever batch is open, however many there are. That measurement
# was complete when a content item could only be text. Packing also has to
# respect ApplyGuardrail's limit of 20 images per request, which is a count
# rather than a character budget, so the two are carried separately: image
# bytes are deliberately not charged against `budget`, which is the text-unit
# quota.
measured: Final = tuple(
(len((item.get("text") or BedrockTextContent()).get("text") or ""), 1 if "image" in item else 0)
for item in content
)
def assign(carried: tuple[int, int], length: int) -> tuple[int, int]:
batch_index, used = carried
if used + length <= budget:
return batch_index, used + length
return batch_index + 1, length
def assign(carried: tuple[int, int, int], item: tuple[int, int]) -> tuple[int, int, int]:
batch_index, used, images = carried
length, is_image = item
if used + length <= budget and images + is_image <= _MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL:
return batch_index, used + length, images + is_image
return batch_index + 1, length, is_image
batch_numbers: Final = (index for index, _ in tuple(accumulate(lengths, assign, initial=(0, 0)))[1:])
batch_numbers: Final = (index for index, _, _ in tuple(accumulate(measured, assign, initial=(0, 0, 0)))[1:])
return tuple(
tuple(item for _, item in group)
for _, group in groupby(zip(batch_numbers, content), key=lambda pair: pair[0])
@ -2880,6 +3122,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
masking_index += 1
if item is not None:
new_content.append(item)
else:
new_content.append(item)
return new_content, masking_index
@ -3047,7 +3291,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
@ -3055,10 +3300,24 @@ 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":
# Both of the optimizations below decide what to scan by looking at
# `texts` alone, so an image rides along unscanned and the proxy still
# reports the guardrail as run. Neither tracks images in its session
# cache either, so "already seen" cannot be established for them.
# Presence of an image therefore forces the full image-aware path:
# correctness over the optimization, and it fails closed.
if input_type == "request" and not image_urls:
incremental_result: Final = await self._apply_incremental_request_scan(
texts=texts,
inputs=inputs,
@ -3075,14 +3334,17 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
request_data=request_data,
input_type=input_type,
)
if selection.skip_scan:
# `experimental_use_latest_role_message_only` marks a latest message with
# no text as skip_scan; an image-only message is exactly that shape.
if selection.skip_scan and not image_urls:
return inputs
filtered_messages: Final = selection.filtered_messages
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)
@ -3115,9 +3377,26 @@ 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 ()
)
# `filtered_messages` is Optional; the guard above only proves one of
# it and `image_urls` is truthy, so it can still be None here when the
# request carries an image and no text.
scan_messages: Final = [ # mutable-ok: make_bedrock_api_request takes a list of messages
*(filtered_messages or ()),
*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,
)

View file

@ -21,6 +21,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
prompt_attack_threshold=litellm_params.prompt_attack_threshold,
pii_confidence_threshold=litellm_params.pii_confidence_threshold,
chunk_budget_chars=litellm_params.chunk_budget_chars,
on_unscannable_image=litellm_params.on_unscannable_image,
default_on=litellm_params.default_on,
disable_exception_on_block=litellm_params.disable_exception_on_block,
mask_request_content=litellm_params.mask_request_content,

View file

@ -535,6 +535,13 @@ class BedrockGuardrailConfigModel(BaseModel):
"still rejects is bisected automatically, so this value only trades round trips against "
"batch size and cannot fail a request on its own.",
)
on_unscannable_image: Literal["block", "allow"] = Field(
default="block",
description="What to do with an image the guardrail cannot scan - ApplyGuardrail "
"accepts png/jpeg only, and remote image URLs are not fetched while "
"litellm.user_url_validation is disabled. 'block' (default) rejects the request; "
"'allow' logs a warning and sends the image to the model unscanned.",
)
class LakeraV2GuardrailConfigModel(BaseModel):

View file

@ -1,6 +1,6 @@
from typing import Literal
from typing_extensions import TypedDict
from typing_extensions import ReadOnly, TypedDict
# Bedrock contextual grounding tags each content block so the guardrail knows
# which text is the reference source, the user question, and the content to grade.
@ -12,8 +12,21 @@ class BedrockTextContent(TypedDict, total=False):
qualifiers: list[BedrockGuardrailQualifier]
BedrockGuardrailImageFormat = Literal["png", "jpeg"]
class BedrockGuardrailImageSource(TypedDict, total=False):
bytes: ReadOnly[str]
class BedrockImageContent(TypedDict, total=False):
format: ReadOnly[BedrockGuardrailImageFormat]
source: ReadOnly[BedrockGuardrailImageSource]
class BedrockContentItem(TypedDict, total=False):
text: BedrockTextContent
image: ReadOnly[BedrockImageContent]
class BedrockRequest(TypedDict, total=False):

View file

@ -945,7 +945,7 @@ async def test_convert_to_bedrock_format_input_source():
]
# Call the method
result = guardrail.convert_to_bedrock_format(source="INPUT", messages=mock_messages)
result = await guardrail.convert_to_bedrock_format(source="INPUT", messages=mock_messages)
# Verify the result structure
assert isinstance(result, dict)
@ -1007,7 +1007,7 @@ async def test_convert_to_bedrock_format_output_source():
)
# Call the method
result = guardrail.convert_to_bedrock_format(
result = await guardrail.convert_to_bedrock_format(
source="OUTPUT", response=mock_response
)

View file

@ -1472,6 +1472,92 @@ class MockCanaryMaskingGuardrail(CustomGuardrail):
return inputs
class TestAnthropicMessagesImageSources:
"""An Anthropic image block has three source shapes (types/llms/anthropic.py:259).
Only the base64 one carries "data", so reading that key alone drops url images
entirely -- for every guardrail consuming GenericGuardrailAPIInputs["images"],
not just Bedrock.
"""
def _data(self, messages):
return {"model": "claude-sonnet-4-5", "messages": messages}
async def _images_seen(self, content) -> list[str]:
handler = AnthropicMessagesHandler()
class ImageRecordingGuardrail(MockCanaryMaskingGuardrail):
def __init__(self):
super().__init__()
self.seen_images: list[str] = [] # mutable-ok: accumulator for the assertion
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
self.seen_images.extend(inputs.get("images") or [])
return await super().apply_guardrail(inputs, request_data, input_type, logging_obj)
guardrail = ImageRecordingGuardrail()
# The text block is what gets the guardrail invoked at all: a message with
# no text gives the handler nothing to scan, so it never reaches the
# guardrail and every source shape would look equally "dropped".
await handler.process_input_messages(
data=self._data([{"role": "user", "content": [{"type": "text", "text": "describe it"}, *content]}]),
guardrail_to_apply=guardrail,
)
return guardrail.seen_images
@pytest.mark.asyncio
async def test_url_source_reaches_the_guardrail(self):
"""A url source has no "data" key, so it used to yield nothing at all."""
seen = await self._images_seen(
[{"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}]
)
assert seen == ["https://example.com/a.png"]
@pytest.mark.asyncio
async def test_base64_source_carries_its_media_type(self):
"""Bare base64 leaves the consumer no way to recover the format.
An API like Bedrock's ApplyGuardrail needs it to build the request, so the
media_type travels with the payload as a data URI.
"""
seen = await self._images_seen(
[{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AAAA"}}]
)
assert seen == ["data:image/png;base64,AAAA"]
@pytest.mark.asyncio
async def test_base64_source_without_a_media_type_is_passed_through(self):
"""There is no format to attach, so the payload goes through unchanged."""
seen = await self._images_seen([{"type": "image", "source": {"type": "base64", "data": "AAAA"}}])
assert seen == ["AAAA"]
@pytest.mark.asyncio
async def test_file_source_yields_nothing(self):
"""The bytes live behind the Files API and this extractor has no client.
Documented as a known gap rather than silently handed on as a file_id string,
which a consumer would try to decode as an image.
"""
seen = await self._images_seen([{"type": "image", "source": {"type": "file", "file_id": "file_abc"}}])
assert seen == []
@pytest.mark.asyncio
async def test_a_malformed_source_is_dropped_rather_than_passed_on(self):
seen = await self._images_seen(
[
{"type": "image", "source": {"type": "base64"}},
{"type": "image", "source": {"type": "url"}},
{"type": "image", "source": {"type": "base64", "data": ""}},
]
)
assert seen == []
class TestAnthropicMessagesToolResultScanning:
"""LIT-5251: tool_result blocks carry whatever a client's local tool fetched, so
they are the request-path payload an indirect prompt injection actually arrives in.

View file

@ -2,8 +2,12 @@
Unit tests for Bedrock Guardrails
"""
import asyncio
import base64
import contextlib
import json
import sys
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@ -2408,12 +2412,14 @@ _GUARD_BLOCK = {"text": {"text": _GROUNDING_RESPONSE_TEXT, "qualifiers": ["guard
def _input_request(messages: list) -> dict:
"""Arrange a guardrail and act: build the Bedrock INPUT payload."""
return _grounding_guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
return asyncio.run(_grounding_guardrail().convert_to_bedrock_format(source="INPUT", messages=messages))
def _output_request(messages: list, response=None) -> dict:
"""Arrange a guardrail and act: build the Bedrock OUTPUT payload."""
return _grounding_guardrail().convert_to_bedrock_format(source="OUTPUT", response=response, messages=messages)
return asyncio.run(
_grounding_guardrail().convert_to_bedrock_format(source="OUTPUT", response=response, messages=messages)
)
def test_grounding_input_strips_grounding_and_query_qualifiers():
@ -2961,9 +2967,7 @@ async def test_streaming_hook_reraises_guardrail_service_failures():
guardrail = _sse_guardrail()
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
mock_api.side_effect = HTTPException(
status_code=500, detail="Bedrock guardrail throttle retries exhausted"
)
mock_api.side_effect = HTTPException(status_code=500, detail="Bedrock guardrail throttle retries exhausted")
with pytest.raises(HTTPException) as exc:
await _drain_streaming_hook(guardrail)
@ -5274,3 +5278,466 @@ async def test_terminal_failure_logs_usage_and_cost_of_prior_passed_chunks(monke
assert logged["guardrail_cost"] == pytest.approx(0.0003)
assert logged["guardrail_response"]["usage"] == {"contentPolicyUnits": 2, "wordPolicyUnits": 1}
assert "error" in logged["guardrail_response"]
class TestBedrockGuardrailImageInput:
"""Image parts must reach ApplyGuardrail, not just the text sitting next to them."""
_PNG_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUg=="
_GIF_DATA_URI = "data:image/gif;base64,R0lGODlhAQABAAAAACw="
_JPEG_BYTES = b"\xff\xd8\xff\xdb"
# A literal, globally-routable IP keeps validate_url's getaddrinfo off DNS; the
# transport is faked in every test below, so no request leaves the process.
_REMOTE_IMAGE_URL = "https://93.184.216.34/a.jpg"
def _jpeg_response(self, url: str) -> httpx.Response:
return httpx.Response(
200,
content=self._JPEG_BYTES,
headers={"content-type": "image/jpeg"},
request=httpx.Request("GET", url),
)
@staticmethod
def _fake_stream(chunks: list[bytes], served: list[int] | None = None):
"""Stand in for httpx.AsyncClient.stream, serving `chunks` one at a time.
The guardrail caps the transfer, so the fetch goes through `stream` rather
than `get`. `served` counts the chunks actually pulled, which is how a test
tells "stopped mid-transfer" apart from "read everything, then rejected".
"""
@contextlib.asynccontextmanager
async def _stream(self, method: str, url, **kwargs):
async def _aiter_bytes():
for chunk in chunks:
if served is not None:
served.append(len(chunk))
yield chunk
response = MagicMock()
response.status_code = 200
response.headers = httpx.Headers({"content-type": "image/jpeg"})
response.request = httpx.Request("GET", str(url))
response.aiter_bytes = _aiter_bytes
yield response
return _stream
def _guardrail(self, **kwargs) -> BedrockGuardrail:
return BedrockGuardrail(
guardrail_name="bedrock-image",
guardrailIdentifier="gr-image",
guardrailVersion="DRAFT",
**kwargs,
)
@pytest.mark.asyncio
async def test_inline_image_is_sent_for_scanning(self):
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "what does this say?"},
{"type": "image_url", "image_url": {"url": self._PNG_DATA_URI}},
],
}
]
request = await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
assert request["content"] == [
{"text": {"text": "what does this say?"}},
{"image": {"format": "png", "source": {"bytes": self._PNG_DATA_URI.split(",")[1]}}},
]
@pytest.mark.asyncio
async def test_image_part_carrying_a_text_field_is_still_scanned_as_an_image(self):
"""A part tagged image_url reaches the model as an image, text field or not.
Provider transformations branch on `type`, so reading `text` first would scan
the decoy and forward the image unscanned - the exact bypass this change closes.
"""
messages = [
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {"url": self._PNG_DATA_URI},
"text": "just a friendly note",
}
],
}
]
request = await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
assert request["content"] == [
{"image": {"format": "png", "source": {"bytes": self._PNG_DATA_URI.split(",")[1]}}}
]
@pytest.mark.asyncio
async def test_unscannable_image_part_carrying_a_text_field_still_blocks(self):
"""The decoy text must not turn an unscannable image into a scanned request."""
messages = [
{
"role": "user",
"content": [{"type": "image_url", "image_url": {"url": self._GIF_DATA_URI}, "text": "hello"}],
}
]
with pytest.raises(HTTPException) as exc_info:
await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_unscannable_image_blocks_the_request_by_default(self):
"""ApplyGuardrail takes png/jpeg only, and the image reaches the model either way.
Skipping it silently would let a caller defeat an IMAGE-modality guardrail by
sending a gif, so the default is to reject the request.
"""
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "hello"},
{"type": "image_url", "image_url": {"url": self._GIF_DATA_URI}},
],
}
]
with pytest.raises(HTTPException) as exc_info:
await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
assert exc_info.value.status_code == 400
assert exc_info.value.detail["error"] == "Violated guardrail policy"
@pytest.mark.asyncio
async def test_undecodable_image_blocks_the_request_by_default(self):
messages = [
{
"role": "user",
"content": [{"type": "image_url", "image_url": {"url": "not-an-image"}}],
}
]
with pytest.raises(HTTPException) as exc_info:
await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_unscannable_image_is_skipped_when_explicitly_allowed(self):
"""on_unscannable_image: allow restores the permissive behavior, opt-in only."""
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "hello"},
{"type": "image_url", "image_url": {"url": self._GIF_DATA_URI}},
{"type": "image_url", "image_url": {"url": "not-an-image"}},
],
}
]
request = await self._guardrail(on_unscannable_image="allow").convert_to_bedrock_format(
source="INPUT", messages=messages
)
assert request["content"] == [{"text": {"text": "hello"}}]
@pytest.mark.asyncio
async def test_remote_image_is_fetched_and_scanned(self):
messages = [
{
"role": "user",
"content": [{"type": "image_url", "image_url": {"url": self._REMOTE_IMAGE_URL}}],
}
]
with patch.object(httpx.AsyncClient, "stream", new=self._fake_stream([self._JPEG_BYTES])):
request = await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
assert request["content"] == [
{
"image": {
"format": "jpeg",
"source": {"bytes": base64.b64encode(self._JPEG_BYTES).decode()},
}
}
]
@pytest.mark.asyncio
async def test_remote_image_is_not_fetched_when_url_validation_is_disabled(self, monkeypatch):
"""With validation off, async_safe_get is an unrestricted redirect-following GET.
The url comes straight from the caller, so fetching it here would make the
guardrail an SSRF primitive. Treat the image as unscannable instead, and make
no request at all.
"""
monkeypatch.setattr(litellm, "user_url_validation", False, raising=False)
url = "http://169.254.169.254/latest/meta-data/"
messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": url}}]}]
get = AsyncMock(return_value=self._jpeg_response(url))
served: list[int] = []
with (
patch.object(httpx.AsyncClient, "get", new=get),
patch.object(httpx.AsyncClient, "stream", new=self._fake_stream([self._JPEG_BYTES], served)),
):
with pytest.raises(HTTPException):
await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
request = await self._guardrail(on_unscannable_image="allow").convert_to_bedrock_format(
source="INPUT", messages=messages
)
# Both transports are stubbed: neither the plain nor the capped fetch ran.
get.assert_not_awaited()
assert served == []
assert request["content"] == []
@pytest.mark.asyncio
async def test_incremental_scan_does_not_skip_a_request_carrying_an_image(self):
"""`only_scan_new_messages` decides what to scan from `texts` alone.
With a session id and every text segment already seen, the incremental path
returns before the image-aware code runs, so benign text plus a violating
image is never scanned -- and the proxy still reports the guardrail as run.
Nothing caches images either, so "already seen" cannot cover them.
"""
g = self._guardrail(only_scan_new_messages=True)
sent: list = []
async def spy(**kwargs):
sent.append(await g.convert_to_bedrock_format(source="INPUT", messages=kwargs["messages"]))
return {"action": "NONE", "outputs": []}
# Every text is already in the session cache: the exact state that made the
# incremental path return early.
with (
patch.object(g, "filter_new_texts_for_session", new=AsyncMock(return_value=[])),
patch.object(g, "make_bedrock_api_request", new=spy),
):
await g.apply_guardrail(
inputs={"texts": ["look at this"], "images": [self._PNG_DATA_URI]},
request_data={"litellm_session_id": "sess-1"},
input_type="request",
)
assert sent, "image-carrying request was skipped entirely by the incremental path"
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_latest_message_only_does_not_skip_an_image_only_message(self):
"""`experimental_use_latest_role_message_only` skips a latest message with no text.
An image-only user message is exactly that shape, so the whole request was
dropped from the scan.
"""
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": [],
"images": [self._PNG_DATA_URI],
"structured_messages": [
{"role": "user", "content": [{"type": "image_url", "image_url": {"url": self._PNG_DATA_URI}}]}
],
},
request_data={},
input_type="request",
)
assert sent, "image-only latest message was skipped entirely"
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_oversized_remote_image_is_cut_off_during_the_transfer(self):
"""A caller-supplied url can serve an unbounded or indefinitely chunked body.
The decoded-size check runs once the bytes are already resident, so the cap
has to apply while the transfer is in flight. Asserting on how much was
pulled is what separates that from buffering it all and rejecting after.
"""
messages = [
{
"role": "user",
"content": [{"type": "image_url", "image_url": {"url": self._REMOTE_IMAGE_URL}}],
}
]
# 8 MB offered one MB at a time against a 4 MB cap.
chunks: list[bytes] = [b"\0" * (1024 * 1024) for _ in range(8)]
served: list[int] = []
with patch.object(httpx.AsyncClient, "stream", new=self._fake_stream(chunks, served)):
with pytest.raises(HTTPException):
await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
# Without the cap the fetch buffers through `get` instead, so `served` stays
# empty and the HTTPException above would come from an unstubbed transport
# rather than from the size rejection. Assert the capped path actually ran.
assert served, "the fetch did not go through the capped stream path"
assert sum(served) <= 5 * 1024 * 1024, f"read {sum(served)} bytes past a 4 MB cap"
assert len(served) < len(chunks), "the whole body was pulled before rejecting it"
@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"
@pytest.mark.asyncio
async def test_oversized_image_is_rejected_before_sending(self):
"""ApplyGuardrail caps images at 4 MB; AWS's rejection is not worth a round trip.
Built as a real data URI rather than a patched decoder so the size check runs
against what the decoder actually produces.
"""
oversized_png = base64.b64encode(b"\x89PNG\r\n\x1a\n" + b"\x00" * (5 * 1024 * 1024)).decode()
with pytest.raises(HTTPException) as exc_info:
await self._guardrail().convert_to_bedrock_format(
source="INPUT",
messages=[
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{oversized_png}"}}
],
}
],
)
assert "4 MB limit" in str(exc_info.value.detail)
def test_bin_packing_respects_the_twenty_image_limit(self):
"""An image measures 0 against the character budget, so count them separately."""
image = {"image": {"format": "png", "source": {"bytes": "AAAA"}}}
batches = BedrockGuardrail._bin_pack_bedrock_content([image] * 45, budget=25_000)
sizes = [len(batch) for batch in batches]
assert all(size <= 20 for size in sizes), sizes
assert sum(sizes) == 45, sizes
def test_bin_packing_still_splits_on_the_text_budget(self):
text = {"text": {"text": "x" * 20_000}}
batches = BedrockGuardrail._bin_pack_bedrock_content([text] * 3, budget=25_000)
assert [len(batch) for batch in batches] == [1, 1, 1]
@pytest.mark.asyncio
async def test_many_images_are_split_before_the_first_call(self):
"""Chunking is reactive; an image-count rejection may never match the too-large check."""
g = self._guardrail()
sent: list = []
async def fake_post(content, **kwargs):
sent.append(sum(1 for item in content if "image" in item))
return {"action": "NONE", "outputs": []}
image = {"image": {"format": "png", "source": {"bytes": "AAAA"}}}
with patch.object(g, "_post_apply_guardrail_content_with_retry", new=fake_post):
await g._apply_guardrail_content_with_chunking(
content=[image] * 45,
base_request_data={},
credentials=None,
aws_region_name="us-west-2",
api_key=None,
request_data=None,
event_type=GuardrailEventHooks.pre_call,
start_time=datetime.now(timezone.utc),
allow_chunking=True,
completed_chunk_usages=[],
)
assert sent == [20, 20, 5], sent
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}}
messages = [
{
"role": "user",
"content": [{"type": "text", "text": "my ssn is 123-45-6789"}, image_part],
}
]
updated = self._guardrail()._apply_masking_to_messages(messages=messages, masked_texts=["my ssn is {SSN}"])
assert updated[0]["content"] == [
{"type": "text", "text": "my ssn is {SSN}"},
image_part,
]

View file

@ -896,7 +896,9 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key():
guardrail_hook.async_handler, "post", AsyncMock(return_value=mock_response)
),
patch.object(guardrail_hook, "_load_credentials") as mock_load_creds,
patch.object(guardrail_hook, "convert_to_bedrock_format") as mock_convert,
patch.object(
guardrail_hook, "convert_to_bedrock_format", new_callable=AsyncMock
) as mock_convert,
patch.object(
guardrail_hook, "get_guardrail_dynamic_request_body_params"
) as mock_get_params,

View file

@ -118,3 +118,37 @@ def test_initialize_guardrail_sets_run_in_parallel(config_value, expected):
custom_guardrail = guardrail_handler.guardrail_id_to_custom_guardrail[result["guardrail_id"]]
assert custom_guardrail.run_in_parallel is expected
def test_initialize_bedrock_forwards_on_unscannable_image():
"""Regression: `on_unscannable_image` set in config.yaml must reach the guardrail.
Same shape as chunk_budget_chars above: the field lives on
BedrockGuardrailConfigModel so LitellmParams parses it, but initialize_bedrock
enumerates its kwargs explicitly. Dropped here, an operator who opted into
`allow` would keep getting 400s on unscannable images with no indication why.
"""
import litellm
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail
test_guardrail = {
"guardrail_name": "test_bedrock_unscannable_image",
"litellm_params": {
"guardrail": SupportedGuardrailIntegrations.BEDROCK.value,
"mode": "pre_call",
"guardrailIdentifier": "test-guardrail",
"guardrailVersion": "DRAFT",
"on_unscannable_image": "allow",
},
}
guardrail_handler = InMemoryGuardrailHandler()
guardrail_handler.initialize_guardrail(guardrail=test_guardrail)
initialized = [
callback
for callback in litellm.callbacks
if isinstance(callback, BedrockGuardrail) and callback.guardrail_name == "test_bedrock_unscannable_image"
]
assert initialized, "bedrock guardrail was not registered as a callback"
assert initialized[-1].on_unscannable_image == "allow"