diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 86cfbf70255..ab60e827943 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -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,26 @@ 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) + # Forward max_bytes only when the caller set one. Passing it + # unconditionally would reach every override and test stub written + # against the previous signature, so an additive parameter would + # break them; omitting it keeps the call byte-for-byte as it was. + capped: Final = {} if max_bytes is None else {"max_bytes": max_bytes} # mutable-ok: kwargs for one call + fetched: Final = await BedrockImageProcessor.get_image_details_async(image_url, **capped) + img_bytes, mime_type = fetched # rebind-ok: mime_type is overridden below by `format` image_format = mime_type.split("/")[1] else: raise ValueError("Unsupported image type. Expected either image url or base64 encoded string") diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 0a59eaa75d3..50cdc679f9b 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -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, Any], # 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: Final[list[bytes]] = [] # mutable-ok: accumulator for the capped body + total = 0 # rebind-ok: running byte count for the cap + 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, Any], # 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, Any] = { # 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 diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 721a6653597..a09048be684 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -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": ""} + {"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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index c70a2ee8a74..fc083010147 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -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,94 @@ _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 + +# A per-image cap does not bound a request. Content items are built by gathering +# over every message and then every part, so N urls are fetched concurrently, and +# the image-count limit above is not reached until _bin_pack_bedrock_content runs +# on items that are already resident. 500 urls is 500 fetches before anything says +# stop. These two bound the request itself: how much may be fetched in total, and +# how much of it may be in flight at once. +# +# The total is what a single ApplyGuardrail call would accept anyway (20 images at +# 4 MB), so no request the API would take in one call is refused. A conversation +# chunked across several calls can exceed it, and images past the budget are then +# unscannable and left to on_unscannable_image -- blocked by default. +_MAX_TOTAL_IMAGE_FETCH_BYTES: Final = _MAX_IMAGE_BYTES * _MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL +_MAX_CONCURRENT_IMAGE_FETCHES: Final = 4 + + +class _ImageFetchBudget: + """Bytes still fetchable for one guardrail request, and a concurrency gate. + + Held for the lifetime of a single content-request build and passed down rather + than kept on the guardrail, which is a callback instance shared by every + request. A fetch reserves one whole image's worth up front and hands back what + the decoded image did not take, so the worst-case resident size is the budget + plus whatever the in-flight fetches have pulled, not the sum of every url a + caller listed. See `claim` for why the reservation is all or nothing. + """ + + def __init__(self, total: int = _MAX_TOTAL_IMAGE_FETCH_BYTES) -> None: + self._remaining = total + self._gate = asyncio.Semaphore(_MAX_CONCURRENT_IMAGE_FETCHES) + + def claim(self) -> int: + """Reserve one image's worth of budget. 0 means exhausted. + + All or nothing rather than handing out whatever is left. A partial grant + would cap the fetch below the per-image limit, and the rejection then + surfaces as "over ApplyGuardrail's 4 MB limit" while naming a few hundred + bytes -- blaming AWS for this request having spent its own budget. The two + failures stay separately legible at the cost of up to one image's worth of + headroom going unused at the tail. + """ + if self._remaining < _MAX_IMAGE_BYTES: + return 0 + self._remaining -= _MAX_IMAGE_BYTES + return _MAX_IMAGE_BYTES + + def give_back(self, unused: int) -> None: + self._remaining += unused + + def gate(self) -> "asyncio.Semaphore": + return self._gate + + +def _retained_image_bytes(item: "BedrockContentItem | None") -> int: + """Approximate what a built image item holds, for budget accounting. + + Measured from the base64 payload rather than decoding it a second time; the + ratio is exact enough for a quota and costs nothing. + """ + if item is None: + return 0 + image: Final = item.get("image") + if not image: + return 0 + encoded: Final = image.get("source", {}).get("bytes") # mutable-ok: {} is a .get default, never mutated + return len(encoded) * 3 // 4 if isinstance(encoded, str) else 0 + + +_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 +320,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 +329,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 +420,276 @@ 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 + budget: Final = _ImageFetchBudget() + per_message: Final = await asyncio.gather( + *(self._build_input_content_items(message=message, budget=budget) 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, budget: "_ImageFetchBudget | None" = None + ) -> 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) + ) + # A direct caller gets a fresh budget rather than an unbounded fetch. + request_budget: Final = budget if budget is not None else _ImageFetchBudget() + items: Final = await asyncio.gather( + *(self._build_input_content_item(item=item, budget=request_budget) 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, budget: "_ImageFetchBudget | None" = None + ) -> 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, budget=budget) + 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 + + def _refuse_file_backed_images(self, inputs: "GenericGuardrailAPIInputs", input_type: str) -> None: + """Hand a file-backed image to on_unscannable_image rather than ignoring it. + + 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")) + if not found: + return + self._handle_unscannable_image( + reason=f"{found} image(s) reference a provider file id, whose bytes are not available here" + ) + + @staticmethod + def _file_backed_image_count(structured_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. + + 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. + + 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. + """ + if not isinstance(structured_messages, list): + return 0 + found = 0 # rebind-ok: running count over the message list + for message in structured_messages: + if not isinstance(message, dict): + continue + content = message.get("content") + if not isinstance(content, list): + continue + for part in content: + if not isinstance(part, dict) or part.get("type") != "image": + continue + source = part.get("source") + if isinstance(source, dict) and source.get("type") == "file": + found += 1 + return found + + async def _build_image_content_item( + self, image_url: str, budget: "_ImageFetchBudget | None" = None + ) -> 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 + """ + is_remote: Final = not image_url.startswith("data:") + if is_remote 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 + + if not is_remote: + # Already in the request body the proxy accepted; nothing is fetched, + # so it draws on neither the byte budget nor the concurrency gate. + return await self._decode_image_content_item(image_url=image_url, max_bytes=None) + + request_budget: Final = budget if budget is not None else _ImageFetchBudget() + granted: Final = request_budget.claim() + if granted <= 0: + self._handle_unscannable_image( + reason="remote image skipped: this request already used its image download budget" + ) + return None + async with request_budget.gate(): + item: Final = await self._decode_image_content_item(image_url=image_url, max_bytes=granted) + + # Refund only what a usable image did not take. Returning the whole + # reservation would make the budget bound in-flight bytes alone, while + # decoded images stay resident in the request being assembled. + # + # A response that produced nothing is charged in full rather than refunded: + # the transfer still happened, and refunding it would let a url serving + # megabytes of unusable bytes be repeated down the whole list for free -- + # the exact shape of the exhaustion this guards against. + # + # No try/finally: the only escape from the line above is the HTTPException + # _handle_unscannable_image raises under the block policy, which ends the + # request and takes this request-scoped budget with it. + request_budget.give_back(granted - _retained_image_bytes(item) if item is not None else 0) + return item + + async def _decode_image_content_item(self, image_url: str, max_bytes: int | None) -> BedrockContentItem | None: + """Turn a data URI or a fetched url into an ApplyGuardrail image block.""" + try: + block: Final = await BedrockImageProcessor.process_image_async( + image_url=image_url, format=None, max_bytes=max_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 +738,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 +755,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 +1191,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 +1307,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 +1740,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 +3285,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 +3454,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 +3463,28 @@ 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 () + # 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) 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 +3501,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) @@ -3105,7 +3534,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ), finish_reason="stop", ) - for _idx, _msg in enumerate(filtered_messages) + # `or ()`: skip_scan is now bypassed when an image is present, + # and images exist on the request side only, so a response scan + # still always has messages here. Spelled out rather than left + # leaning on that indirection. + for _idx, _msg in enumerate(filtered_messages or ()) ] ) bedrock_response = await self.make_bedrock_api_request( @@ -3115,9 +3548,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, ) diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 0d23e19f88d..4ba23474d6c 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -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, diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index c7cdfaad780..0223d651d4a 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -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): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 8d66b624341..e194f4f1132 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -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): diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index 43d088268eb..bb958bcc304 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -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 ) diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 6265779b90d..036656f3251 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -2,12 +2,16 @@ import base64 import json import logging import os +import socket from typing import Final from unittest.mock import MagicMock, patch import pytest +import httpx + import litellm +from litellm.litellm_core_utils.url_utils import PayloadTooLargeError from litellm.litellm_core_utils.prompt_templates.factory import ( BAD_MESSAGE_ERROR_STR, BEDROCK_DOCUMENT_PLACEHOLDER_TEXT, @@ -3578,3 +3582,71 @@ async def test_bedrock_converse_pdf_only_user_message_gets_text_block_async(): assert len(result) == 1 assert any("document" in block for block in result[0]["content"]) assert _text_blocks(result[0]) == [BEDROCK_DOCUMENT_PLACEHOLDER_TEXT] + + +def _resolve_to_public(host, port, *args, **kwargs): + """Keep validate_url's DNS lookup off the network without faking the fetch.""" + return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", port or 443))] + + +class TestBedrockImageProcessorMaxBytes: + """`max_bytes` is threaded from the caller down to the fetch. + + The Bedrock guardrail is the only caller that sets it. Everything else, the + model-call image paths included, must keep the previous unbounded fetch, and the + keyword has to be absent from the call rather than merely defaulted -- an + override or stub written against the old signature would otherwise break. + """ + + _REMOTE_URL = "https://93.184.216.34/a.png" + + @staticmethod + def _fake_stream(chunks): + import contextlib + + @contextlib.asynccontextmanager + async def _stream(self, method, url, **kwargs): + async def aiter_bytes(): + for chunk in chunks: + yield chunk + + response = MagicMock() + response.status_code = 200 + response.headers = httpx.Headers({"content-type": "image/png"}) + response.request = httpx.Request("GET", str(url)) + response.aiter_bytes = aiter_bytes + yield response + + return _stream + + @pytest.mark.asyncio + async def test_a_remote_fetch_is_capped_when_max_bytes_is_given(self, monkeypatch): + monkeypatch.setattr(socket, "getaddrinfo", _resolve_to_public, raising=False) + + with patch.object(httpx.AsyncClient, "stream", new=self._fake_stream([b"\0" * 8192])): + with pytest.raises(PayloadTooLargeError): + await BedrockImageProcessor.get_image_details_async(self._REMOTE_URL, max_bytes=1024) + + @pytest.mark.asyncio + async def test_omitting_max_bytes_leaves_the_call_as_it_was(self, monkeypatch): + """A stub written against the previous one-parameter signature still works. + + This is what test_url_with_format_param asserts through the model path; here + it is pinned on the helper itself so the plumbing cannot start passing the + keyword unconditionally again. + """ + monkeypatch.setattr(socket, "getaddrinfo", _resolve_to_public, raising=False) + seen: list = [] + + async def one_parameter_stub(image_url): + seen.append(image_url) + return "ZmFrZQ==", "image/png" + + monkeypatch.setattr( + BedrockImageProcessor, "get_image_details_async", staticmethod(one_parameter_stub) + ) + + block = await BedrockImageProcessor.process_image_async(image_url=self._REMOTE_URL, format=None) + + assert seen == [self._REMOTE_URL] + assert block["image"]["source"]["bytes"] == "ZmFrZQ==" diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index aaaa43a0dc4..e285f749ebc 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -1,11 +1,17 @@ +import contextlib import socket +from types import SimpleNamespace +from unittest.mock import MagicMock, patch +import httpx import pytest import litellm from litellm.litellm_core_utils import url_utils from litellm.litellm_core_utils.url_utils import ( + PayloadTooLargeError, SSRFError, + _underlying_httpx_client, _is_blocked_ip, assert_same_origin, encode_url_path_segment, @@ -535,3 +541,91 @@ def test_assert_same_origin_error_message_does_not_leak_hostnames(): detail = str(exc.value) assert "attacker.example.com" not in detail assert "api.internal-corp.example" not in detail + + +class TestCappedFetch: + """`async_safe_get(max_bytes=...)` streams and aborts past the cap. + + `client.get` buffers the whole body first, so a caller-supplied url serving an + arbitrarily large or indefinitely chunked response is an unbounded allocation. + """ + + def test_a_client_with_no_httpx_client_is_rejected(self): + """Streaming needs the wrapped httpx client. + + AsyncHTTPHandler forwards get/post but not stream, so the wrapped `.client` + is what gets used. Something with neither is a programming error and says so, + rather than failing later inside httpx with nothing pointing back here. + """ + with pytest.raises(TypeError) as exc: + _underlying_httpx_client(object()) + + assert "no httpx client" in str(exc.value) + + def test_a_raw_httpx_client_is_used_as_is(self): + client = httpx.AsyncClient() + + assert _underlying_httpx_client(client) is client + + def test_a_wrapped_client_resolves_to_the_one_it_wraps(self): + inner = httpx.AsyncClient() + wrapper = SimpleNamespace(client=inner) + + assert _underlying_httpx_client(wrapper) is inner + + @pytest.mark.asyncio + async def test_the_body_is_cut_off_once_it_passes_the_cap(self, mock_dns_public): + """Asserting on how much was pulled is what separates a streamed abort from + buffering everything and rejecting afterwards.""" + served: list[int] = [] + chunks = [b"\0" * 1024 for _ in range(100)] + + @contextlib.asynccontextmanager + async def fake_stream(self, method, url, **kwargs): + async def aiter_bytes(): + for chunk in chunks: + served.append(len(chunk)) + yield chunk + + response = MagicMock() + response.status_code = 200 + response.headers = httpx.Headers({"content-type": "image/png"}) + response.request = httpx.Request("GET", str(url)) + response.aiter_bytes = aiter_bytes + yield response + + client = httpx.AsyncClient() + with patch.object(httpx.AsyncClient, "stream", new=fake_stream): + with pytest.raises(PayloadTooLargeError): + await url_utils.async_safe_get(client, "https://93.184.216.34/a.png", max_bytes=4096) + + assert sum(served) <= 5 * 1024, f"pulled {sum(served)} bytes past a 4096 byte cap" + assert len(served) < len(chunks), "the whole body was read before rejecting it" + + @pytest.mark.asyncio + async def test_a_body_inside_the_cap_comes_back_whole(self, mock_dns_public): + """The rebuilt response drops content-encoding and content-length: aiter_bytes + yields decoded bytes, so carrying those over would describe the body wrongly.""" + + @contextlib.asynccontextmanager + async def fake_stream(self, method, url, **kwargs): + async def aiter_bytes(): + yield b"tiny-image" + + response = MagicMock() + response.status_code = 200 + response.headers = httpx.Headers( + {"content-type": "image/png", "content-length": "999", "content-encoding": "gzip"} + ) + response.request = httpx.Request("GET", str(url)) + response.aiter_bytes = aiter_bytes + yield response + + client = httpx.AsyncClient() + with patch.object(httpx.AsyncClient, "stream", new=fake_stream): + result = await url_utils.async_safe_get(client, "https://93.184.216.34/a.png", max_bytes=4096) + + assert result.content == b"tiny-image" + assert result.headers.get("content-type") == "image/png" + assert "content-encoding" not in result.headers + assert result.headers.get("content-length") == str(len(b"tiny-image")) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 2b392456763..ed25ada891f 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -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. 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 dd339d4e51f..f51190ef683 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 @@ -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 @@ -20,6 +24,7 @@ from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockContentChunkResult, BedrockGuardrail, _redact_pii_matches, + _retained_image_bytes, ) from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks @@ -2408,12 +2413,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 +2968,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 +5279,834 @@ 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_a_request_full_of_urls_is_bounded_in_total_not_just_per_image(self): + """A per-image cap does not bound a request. + + Content items are gathered over every message and every part, so the urls + are fetched concurrently, and the 20-image limit is not applied until + _bin_pack_bedrock_content runs on items that are already resident. Without + a request-wide budget, 200 urls at 4 MB is 800 MB the caller chose. + """ + # 200 parts, each serving a 1 MB image, against a 20 x 4 MB budget. + served: list[int] = [] + one_mb: list[bytes] = [b"\0" * (1024 * 1024)] + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": f"{self._REMOTE_IMAGE_URL}?i={i}"}} for i in range(200) + ], + } + ] + + with patch.object(httpx.AsyncClient, "stream", new=self._fake_stream(one_mb, served)): + request = await self._guardrail(on_unscannable_image="allow").convert_to_bedrock_format( + source="INPUT", messages=messages + ) + + fetched: int = sum(served) + assert fetched <= 20 * 4 * 1024 * 1024, f"fetched {fetched} bytes for one request" + assert len(request["content"]) < 200, "every url was kept despite the budget" + assert request["content"], "the budget swallowed the whole request" + + @pytest.mark.asyncio + async def test_an_oversized_remote_image_is_dropped_under_the_allow_policy(self): + """The transfer is cut off, and then the request has to carry on. + + `block` raises out of the size rejection, so this is the only path that + reaches its fall-through. An operator who set `allow` asked for the image to + go unscanned, not for the whole request to die on it. + """ + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "look"}, + {"type": "image_url", "image_url": {"url": self._REMOTE_IMAGE_URL}}, + ], + } + ] + served: list[int] = [] + chunks: list[bytes] = [b"\0" * (1024 * 1024) for _ in range(8)] + + with patch.object(httpx.AsyncClient, "stream", new=self._fake_stream(chunks, served)): + request = await self._guardrail(on_unscannable_image="allow").convert_to_bedrock_format( + source="INPUT", messages=messages + ) + + assert request["content"] == [{"text": {"text": "look"}}] + assert served, "the capped stream path did not run" + assert len(served) < len(chunks), "the whole body was pulled before dropping it" + + def test_the_budget_grants_a_whole_image_or_nothing(self): + """A partial grant would cap a fetch below the per-image limit. + + The rejection then reads "over ApplyGuardrail's 4 MB limit" while naming a + few hundred bytes, blaming AWS for this request having spent its own budget. + Keeping the two failures separately legible is worth leaving one image's + worth of headroom unused at the tail. + """ + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + _MAX_IMAGE_BYTES, + _ImageFetchBudget, + ) + + budget = _ImageFetchBudget(total=_MAX_IMAGE_BYTES + 100) + + assert budget.claim() == _MAX_IMAGE_BYTES + assert budget.claim() == 0, "100 bytes left must read as exhausted, not as a 100 byte cap" + + @pytest.mark.asyncio + async def test_inline_images_do_not_draw_on_the_download_budget(self): + """Base64 arrives in the request body the proxy already accepted. + + Nothing is fetched for it, so charging it against a download quota would + refuse inline images for no reason. + """ + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": self._PNG_DATA_URI}} for _ in range(40)], + } + ] + + request = await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + + assert len(request["content"]) == 40 + + @pytest.mark.parametrize( + "content", + [ + pytest.param(None, id="no content"), + pytest.param(123, id="content is not a list"), + pytest.param([123], id="part is not a mapping"), + pytest.param([{"type": "image_url"}], id="image part with no url"), + pytest.param([{"type": "image_url", "image_url": 123}], id="url is not a string or mapping"), + pytest.param([{"type": "image_url", "image_url": {"url": 123}}], id="url value is not a string"), + pytest.param([{"type": "input_audio"}], id="part carries neither image nor text"), + ], + ) + @pytest.mark.asyncio + async def test_a_malformed_part_never_becomes_scannable_content(self, content): + """Default-deny is the point of this PR, so pin it rather than trust it. + + Each shape below is one a caller can put on the wire. None of them may turn + into a content item: an unrecognised part that fell through to the text + branch would be reported to the operator as scanned when it was not. + """ + request = await self._guardrail(on_unscannable_image="allow").convert_to_bedrock_format( + source="INPUT", messages=[{"role": "user", "content": content}] + ) + + assert request["content"] == [] + + @pytest.mark.asyncio + async def test_a_bare_string_part_is_scanned_as_text(self): + """A content list may hold plain strings, not only typed parts.""" + request = await self._guardrail().convert_to_bedrock_format( + source="INPUT", messages=[{"role": "user", "content": ["just text"]}] + ) + + assert request["content"] == [{"text": {"text": "just text"}}] + + @pytest.mark.asyncio + async def test_image_url_given_as_a_plain_string_is_accepted(self): + """OpenAI accepts `image_url` as a bare string as well as `{"url": ...}`. + + Both reach the model as an image, so both have to reach the scan. + """ + request = await self._guardrail().convert_to_bedrock_format( + source="INPUT", + messages=[{"role": "user", "content": [{"type": "image_url", "image_url": self._PNG_DATA_URI}]}], + ) + + kinds = [k for item in request["content"] for k in item] + assert kinds == ["image"] + + @pytest.mark.asyncio + async def test_an_unrecognized_payload_is_left_to_the_unscannable_policy(self): + """_normalize_image_input sniffs png and jpeg out of bare base64. + + Anything else is handed to the decoder as-is rather than guessed at, so the + rejection comes from on_unscannable_image and not from a helper deciding + quietly on its own. + """ + # Reached through apply_guardrail: bare base64 arrives in inputs["images"], + # which is the only caller that normalizes before decoding. + with pytest.raises(HTTPException) as exc_info: + await self._guardrail().apply_guardrail( + inputs={"texts": [], "images": ["R0lGODlhAQABAAAAACw="]}, + request_data={}, + input_type="request", + ) + + assert "could not be read" in str(exc_info.value.detail) or "not a png/jpeg" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_an_oversized_inline_image_is_dropped_under_the_allow_policy(self): + """The allow policy has to survive the size rejection, not just the format one. + + Under `block` the oversized branch raises and never returns, so this is the + only path that reaches its fall-through. + """ + oversized_png = base64.b64encode(b"\x89PNG\r\n\x1a\n" + b"\x00" * (5 * 1024 * 1024)).decode() + + request = await self._guardrail(on_unscannable_image="allow").convert_to_bedrock_format( + source="INPUT", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "look"}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{oversized_png}"}}, + ], + } + ], + ) + + assert request["content"] == [{"text": {"text": "look"}}] + + def test_the_url_and_budget_helpers_guard_their_own_inputs(self): + """Both are reached only through callers that already checked the shape. + + Exercised directly so the guards are not silently dropped in a refactor that + gives either one a second caller. + """ + assert BedrockGuardrail._get_image_url(item={"type": "text", "text": "hi"}) is None + + assert _retained_image_bytes(None) == 0 + assert _retained_image_bytes({"text": {"text": "not an image"}}) == 0 + assert _retained_image_bytes({"image": {"format": "png", "source": {"bytes": 123}}}) == 0 + assert _retained_image_bytes({"image": {"format": "png", "source": {"bytes": "AAAA"}}}) == 3 + + @pytest.mark.asyncio + async def test_a_file_backed_image_is_refused_rather_than_ignored(self): + """`{"type": "file"}` carries no bytes, so nothing reaches inputs["images"]. + + The provider still forwards the file to the model, so ignoring it is exactly + 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": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "what does this say?"}, + {"type": "image", "source": {"type": "file", "file_id": "file_abc"}}, + ], + } + ], + } + + g = self._guardrail() + sent: list = [] + + async def spy(**kwargs): + sent.append(kwargs["messages"]) + return {"action": "NONE", "outputs": []} + + # Stubbed so that without the refusal this request would simply succeed: + # 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") + + 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_let_through_under_the_allow_policy(self): + """An operator who would rather serve it unscanned can still say so.""" + g = self._guardrail(on_unscannable_image="allow") + sent: list = [] + + async def spy(**kwargs): + sent.append(kwargs["messages"]) + return {"action": "NONE", "outputs": []} + + with patch.object(g, "make_bedrock_api_request", new=spy): + await g.apply_guardrail( + inputs={ + "texts": ["hello"], + "images": [], + "structured_messages": [ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "file", "file_id": "file_abc"}}], + } + ], + }, + request_data={}, + input_type="request", + ) + + assert sent, "the text alongside the file image still has to be scanned" + + @pytest.mark.asyncio + async def test_the_scannable_source_shapes_are_not_refused(self): + """The refusal has to be specific to the shape that cannot be read. + + structured_messages is already narrowed by the skip and scope flags, so a + count mismatch against inputs["images"] is not evidence of a dropped image. + Blocking a legitimate request would be worse than the gap being closed. + """ + 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": ["hello"], + "images": [self._PNG_DATA_URI], + "structured_messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "image", "source": {"type": "base64", "data": "AAAA"}}, + {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}, + ], + } + ], + }, + request_data={}, + input_type="request", + ) + + kinds = [k for item in sent[0]["content"] for k in item] + assert "image" in kinds + + @pytest.mark.asyncio + async def test_a_malformed_structured_message_does_not_derail_the_file_check(self): + """structured_messages comes from the caller, so its shape is not guaranteed. + + The scan must keep walking past an entry it cannot read rather than throwing + or giving up, or a single junk element would hide a file image sitting after + it -- turning a defensive guard into the bypass it was meant to prevent. + """ + g = self._guardrail() + sent: list = [] + + async def spy(**kwargs): + sent.append(kwargs["messages"]) + return {"action": "NONE", "outputs": []} + + 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": [ + "not a message", + 123, + {"role": "user", "content": "a plain string, not a list"}, + { + "role": "user", + "content": [{"type": "image", "source": {"type": "file", "file_id": "file_abc"}}], + }, + ], + }, + request_data={}, + input_type="request", + ) + + assert "file id" in str(exc_info.value.detail) + assert sent == [] + + @pytest.mark.asyncio + async def test_a_file_backed_image_on_the_response_side_is_not_refused(self): + """Images are a request-side concern; an OUTPUT scan takes generated text.""" + g = self._guardrail() + + async def spy(**kwargs): + return {"action": "NONE", "outputs": []} + + with patch.object(g, "make_bedrock_api_request", new=spy): + result = await g.apply_guardrail( + inputs={ + "texts": ["the model said this"], + "structured_messages": [ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "file", "file_id": "file_abc"}}], + } + ], + }, + request_data={}, + input_type="response", + ) + + assert result is not None + + @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, + ] diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 45f5afef1bc..d6f7d5305be 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -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, diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index 82363302d2e..70004847d59 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -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"