mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
Merge 0e1ec44e35 into aedaf4d0b0
This commit is contained in:
commit
97e4c2899e
12 changed files with 1057 additions and 58 deletions
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue