mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
refactor(guardrails): narrow this PR to inline image scanning
Review raised two resource concerns that both trace back to remote url support rather than to image scanning itself: an image fetched for the guardrail is not provably the image the model sees, because the model request still carries the url (TOCTOU); and a byte budget does not bound the work a request can cause, because inline data URIs are free of it. Fetching a caller-supplied url safely is a separate piece of work -- a streamed size cap, SSRF and redirect validation, a per-request download budget, a concurrency gate, and handing the same bytes to both the guardrail and the model. Bundling it here makes the guardrail fix hostage to it, so it moves to its own change and this one covers inline images. Removed: - remote url fetching, `async_safe_get(max_bytes)`, the streamed body cap and `PayloadTooLargeError` - the per-request download budget and the concurrent-fetch gate - image-aware batching: the request-wide limit below caps a request at 20 images, so every batch text chunking produces already holds at most 20 and the image dimension can never split anything - `on_unscannable_image`, which configures behaviour this PR no longer has a second option for; it fails closed and the setting follows the remote support it was written for Changed: - the 20-image limit is now a request-wide rejection raised before any image is decoded, rather than a batch size applied after. A caller can list hundreds of tiny data URIs, and counting them first bounds the decode work and the ApplyGuardrail calls it would fan out into. - an unrecognized payload is no longer reported as a remote url. Only `http://`/`https://` is named as unsupported-because-remote; anything else reaches the decoder, which rejects it as unreadable.
This commit is contained in:
feng.tsai
parent
bb633fa9ef
commit
f04847f243
11 changed files with 109 additions and 855 deletions
|
|
@ -3378,14 +3378,14 @@ class BedrockImageProcessor:
|
|||
return base64_bytes, content_type
|
||||
|
||||
@staticmethod
|
||||
async def get_image_details_async(image_url, max_bytes: int | None = None) -> tuple[str, str]:
|
||||
async def get_image_details_async(image_url) -> 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, max_bytes=max_bytes)
|
||||
response: Final[httpx.Response] = await async_safe_get(client, image_url)
|
||||
response.raise_for_status() # Raise an exception for HTTP errors
|
||||
|
||||
return BedrockImageProcessor._post_call_image_processing(response, image_url)
|
||||
|
|
@ -3566,26 +3566,13 @@ 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, 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.
|
||||
"""
|
||||
async def process_image_async(cls, image_url: str, format: str | None) -> BedrockContentBlock:
|
||||
"""Asynchronous image processing."""
|
||||
|
||||
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:
|
||||
# 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`
|
||||
img_bytes, mime_type = await BedrockImageProcessor.get_image_details_async(image_url)
|
||||
image_format = mime_type.split("/")[1]
|
||||
else:
|
||||
raise ValueError("Unsupported image type. Expected either image url or base64 encoded string")
|
||||
|
|
|
|||
|
|
@ -92,14 +92,6 @@ 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: object, *, field_name: str = "path parameter") -> str:
|
||||
"""Percent-encode one user-controlled URL path segment.
|
||||
|
||||
|
|
@ -468,79 +460,11 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response:
|
|||
raise SSRFError("Too many redirects")
|
||||
|
||||
|
||||
# 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) -> httpx.Response:
|
||||
"""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: _AsyncUrlFetcher,
|
||||
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 target.get(target_url, **request_kwargs)
|
||||
return await _async_get_capped(target, target_url, max_bytes, request_kwargs)
|
||||
|
||||
async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response:
|
||||
"""Async version of safe_get."""
|
||||
if not getattr(litellm, "user_url_validation", True):
|
||||
kwargs.setdefault("follow_redirects", True)
|
||||
unvalidated: Final[_ResponseView] = {"response": await _issue(client, url, kwargs)}
|
||||
unvalidated: Final[_ResponseView] = {"response": await client.get(url, **kwargs)}
|
||||
return unvalidated["response"]
|
||||
fetcher_view: Final[_AsyncFetcherView] = {"fetcher": client}
|
||||
fetcher: Final = fetcher_view["fetcher"]
|
||||
|
|
@ -548,12 +472,12 @@ async def async_safe_get(client: Any, url: str, max_bytes: int | None = None, **
|
|||
headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})}
|
||||
for _ in range(_MAX_REDIRECTS):
|
||||
validated_url, original_host = validate_url(url)
|
||||
hop_kwargs: dict[str, Any] = { # mutable-ok: a fresh per-hop kwargs dict, consumed by this call
|
||||
response = await fetcher.get(
|
||||
validated_url,
|
||||
headers={**headers_view["headers"], "Host": original_host},
|
||||
follow_redirects=False,
|
||||
**kwargs,
|
||||
"headers": {**headers_view["headers"], "Host": original_host}, # mutable-ok: httpx takes headers as a dict
|
||||
"follow_redirects": False,
|
||||
}
|
||||
response = await _issue(fetcher, validated_url, hop_kwargs)
|
||||
)
|
||||
if not response.is_redirect:
|
||||
return response
|
||||
# Resolve the next hop against the ORIGINAL (pre-rewrite) URL so
|
||||
|
|
|
|||
|
|
@ -11810,16 +11810,6 @@
|
|||
"description": "Action to take when sensitive data is detected. 'block' raises an exception (default behavior). 'route' reroutes the request to the model specified in sensitive_data_route_to_model.",
|
||||
"title": "On Sensitive Data"
|
||||
},
|
||||
"on_unscannable_image": {
|
||||
"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.",
|
||||
"enum": [
|
||||
"block",
|
||||
"allow"
|
||||
],
|
||||
"title": "On Unscannable Image",
|
||||
"type": "string"
|
||||
},
|
||||
"on_violation": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -39,7 +39,6 @@ from litellm.litellm_core_utils.litellm_logging import (
|
|||
)
|
||||
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,
|
||||
|
|
@ -160,64 +159,6 @@ _GROUNDING_SOURCE_TRUSTED_ROLES: Final = frozenset({"system", "developer"})
|
|||
_MAX_IMAGE_BYTES: Final = 4 * 1024 * 1024
|
||||
_MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL: Final = 20
|
||||
|
||||
# A per-image cap does not bound a request: parts are gathered over every message,
|
||||
# so N urls are fetched concurrently long before the image-count limit is consulted.
|
||||
# The total is what one ApplyGuardrail call accepts anyway, 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 fall to on_unscannable_image.
|
||||
_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 _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
|
||||
|
|
@ -324,7 +265,6 @@ 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",
|
||||
streaming_buffer_until_moderated: bool | None = None,
|
||||
streaming_sampling_rate: int | None = None,
|
||||
streaming_end_of_stream_only: bool | None = None,
|
||||
|
|
@ -347,11 +287,6 @@ 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.
|
||||
self.checks: dict[str, object] | None = self._normalize_checks(checks)
|
||||
|
|
@ -456,17 +391,28 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
if messages is None:
|
||||
return bedrock_request
|
||||
|
||||
budget: Final = _ImageFetchBudget()
|
||||
image_count: Final = self._image_count_in(messages)
|
||||
if image_count > _MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={ # mutable-ok: HTTPException detail payload, serialized immediately
|
||||
"error": "Violated guardrail policy",
|
||||
"bedrock_guardrail_response": (
|
||||
f"Request contains {image_count} images; Bedrock ApplyGuardrail accepts at most "
|
||||
f"{_MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL} images per request"
|
||||
),
|
||||
"guardrail_name": self.guardrail_name,
|
||||
},
|
||||
)
|
||||
|
||||
per_message: Final = await asyncio.gather(
|
||||
*(self._build_input_content_items(message=message, budget=budget) for message in messages)
|
||||
*(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, budget: "_ImageFetchBudget | None" = None
|
||||
) -> tuple[BedrockContentItem, ...]:
|
||||
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
|
||||
|
|
@ -482,16 +428,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
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)
|
||||
)
|
||||
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, budget: "_ImageFetchBudget | None" = None
|
||||
) -> BedrockContentItem | 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):
|
||||
|
|
@ -504,7 +444,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
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)
|
||||
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))
|
||||
|
|
@ -537,30 +477,37 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
found.add(cls._normalize_image_input(url))
|
||||
return frozenset(found)
|
||||
|
||||
@classmethod
|
||||
def _image_count_in(cls, messages: "Sequence[AllMessageValues] | None") -> int:
|
||||
"""Count image occurrences in the exact messages sent to ApplyGuardrail."""
|
||||
count = 0 # rebind-ok: running count over request content
|
||||
for message in messages or ():
|
||||
content = message.get("content")
|
||||
if not isinstance(content, list):
|
||||
continue
|
||||
count += sum(
|
||||
1
|
||||
for part in content
|
||||
if isinstance(part, dict) and part.get("type") == "image_url" and cls._get_image_url(part) is not None
|
||||
)
|
||||
return count
|
||||
|
||||
def _handle_unscannable_image(self, reason: str) -> None:
|
||||
"""Block or warn for an image part ApplyGuardrail cannot scan.
|
||||
"""Block 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,
|
||||
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}). "
|
||||
"Bedrock ApplyGuardrail accepts inline png/jpeg images only"
|
||||
),
|
||||
"guardrail_name": self.guardrail_name,
|
||||
},
|
||||
)
|
||||
|
||||
#: base64 magic-byte prefixes for the formats ApplyGuardrail accepts.
|
||||
|
|
@ -590,12 +537,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
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.
|
||||
# Unrecognized inputs reach the decoder so the guardrail fails closed.
|
||||
return value
|
||||
|
||||
def _refuse_file_backed_images(self, request_data: Mapping[str, object], input_type: str) -> None:
|
||||
"""Hand a file-backed image to on_unscannable_image rather than ignoring it.
|
||||
"""Reject a file-backed image rather than ignoring it.
|
||||
|
||||
Reads the raw request, not inputs["structured_messages"]: the /v1/messages
|
||||
handler fills that field by translating to OpenAI spec, which drops a file
|
||||
|
|
@ -648,55 +594,22 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
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.
|
||||
async def _build_image_content_item(self, image_url: str) -> BedrockContentItem | None:
|
||||
"""Decode an inline image 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
|
||||
A remote url is named as its own rejection rather than left to the decoder:
|
||||
fetching one is a separate piece of work (size cap, SSRF, and handing the
|
||||
same bytes to the model), so the operator gets "not supported" instead of
|
||||
the decoder's "could not be read". Anything else that is not a data URI is
|
||||
an unrecognized payload and falls through to the decoder, which rejects it.
|
||||
"""
|
||||
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}"
|
||||
)
|
||||
if image_url.startswith(("http://", "https://")):
|
||||
self._handle_unscannable_image(reason="remote image URLs are not supported")
|
||||
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 what a usable image did not take, not the whole reservation: decoded
|
||||
# images stay resident in the request being assembled. A response that produced
|
||||
# nothing is charged in full, or one url serving unusable megabytes could be
|
||||
# repeated down the whole list for free.
|
||||
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:
|
||||
block: Final = await BedrockImageProcessor.process_image_async(image_url=image_url, format=None)
|
||||
except (ValueError, TypeError, KeyError, binascii.Error) as e:
|
||||
self._handle_unscannable_image(reason=f"image content could not be read: {e}")
|
||||
return None
|
||||
|
||||
|
|
@ -1343,47 +1256,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
-- 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,
|
||||
|
|
@ -1776,27 +1649,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
if not content:
|
||||
return (tuple(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
|
||||
lengths: Final = tuple(
|
||||
len((item.get("text") or BedrockTextContent()).get("text") or "") for item in content
|
||||
)
|
||||
|
||||
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
|
||||
def assign(carried: tuple[int, int], length: int) -> tuple[int, int]:
|
||||
batch_index, used = carried
|
||||
return (batch_index, used + length) if used + length <= budget else (batch_index + 1, length)
|
||||
|
||||
batch_numbers: Final = (index for index, _, _ in tuple(accumulate(measured, assign, initial=(0, 0, 0)))[1:])
|
||||
batch_numbers: Final = (index for index, _ in tuple(accumulate(lengths, assign, initial=(0, 0)))[1:])
|
||||
return tuple(
|
||||
tuple(item for _, item in group)
|
||||
for _, group in groupby(zip(batch_numbers, content), key=lambda pair: pair[0])
|
||||
|
|
@ -3616,12 +3477,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
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
|
||||
# decoding and fail-closed format checks, so the
|
||||
# unified and native lifecycle paths cannot drift apart.
|
||||
# `experimental_use_latest_role_message_only` puts the selected
|
||||
# message itself into filtered_messages, image parts included, and
|
||||
# those go through the same builder below. Appending them again
|
||||
# would fetch and bill each one twice.
|
||||
# would scan and bill each one twice.
|
||||
already_scanned: Final = self._image_urls_in(filtered_messages)
|
||||
image_parts: Final = [ # mutable-ok: OpenAI message content is a list in the wire format
|
||||
self._image_content_part(normalized)
|
||||
|
|
|
|||
|
|
@ -22,7 +22,6 @@ 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,
|
||||
|
|
|
|||
|
|
@ -551,13 +551,6 @@ 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 BedrockGuardrailStreamingParams(BaseModel):
|
||||
|
|
|
|||
|
|
@ -2,16 +2,12 @@ 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,
|
||||
|
|
@ -3606,73 +3602,6 @@ async def test_bedrock_converse_pdf_only_user_message_gets_text_block_async():
|
|||
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=="
|
||||
|
||||
def test_convert_to_anthropic_tool_result_keeps_tool_reference_blocks():
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import convert_to_anthropic_tool_result
|
||||
|
||||
|
|
|
|||
|
|
@ -1,17 +1,11 @@
|
|||
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,
|
||||
|
|
@ -541,91 +535,3 @@ 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"))
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ Unit tests for Bedrock Guardrails
|
|||
|
||||
import asyncio
|
||||
import base64
|
||||
import contextlib
|
||||
import json
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
|
|
@ -24,7 +23,6 @@ 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
|
||||
|
|
@ -5286,45 +5284,6 @@ class TestBedrockGuardrailImageInput:
|
|||
|
||||
_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",
|
||||
|
|
@ -5430,76 +5389,6 @@ class TestBedrockGuardrailImageInput:
|
|||
|
||||
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.
|
||||
|
|
@ -5564,129 +5453,37 @@ class TestBedrockGuardrailImageInput:
|
|||
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.
|
||||
"""
|
||||
async def test_twenty_inline_images_are_accepted(self):
|
||||
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)],
|
||||
"content": [{"type": "image_url", "image_url": {"url": self._PNG_DATA_URI}} for _ in range(20)],
|
||||
}
|
||||
]
|
||||
|
||||
request = await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
|
||||
|
||||
assert len(request["content"]) == 40
|
||||
assert len(request["content"]) == 20
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_twenty_one_duplicate_images_are_rejected_before_decode(self):
|
||||
"""Repeated tiny data URIs still consume image slots and must not amplify AWS calls."""
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "image_url", "image_url": {"url": self._PNG_DATA_URI}} for _ in range(21)],
|
||||
}
|
||||
]
|
||||
|
||||
with patch( # test-quality-ok: the decoder is the thing under assertion -- the cap must reject before it is ever awaited
|
||||
"litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.BedrockImageProcessor.process_image_async",
|
||||
new_callable=AsyncMock,
|
||||
) as decode:
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
|
||||
|
||||
decode.assert_not_awaited()
|
||||
assert "at most 20 images" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"content",
|
||||
|
|
@ -5708,7 +5505,7 @@ class TestBedrockGuardrailImageInput:
|
|||
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(
|
||||
request = await self._guardrail().convert_to_bedrock_format(
|
||||
source="INPUT", messages=[{"role": "user", "content": content}]
|
||||
)
|
||||
|
||||
|
|
@ -5738,12 +5535,11 @@ class TestBedrockGuardrailImageInput:
|
|||
assert kinds == ["image"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_unrecognized_payload_is_left_to_the_unscannable_policy(self):
|
||||
async def test_an_unrecognized_payload_is_rejected(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.
|
||||
Anything else is handed to the decoder as-is rather than guessed at, and the
|
||||
guardrail rejects it instead of forwarding the image unscanned.
|
||||
"""
|
||||
# Reached through apply_guardrail: bare base64 arrives in inputs["images"],
|
||||
# which is the only caller that normalizes before decoding.
|
||||
|
|
@ -5756,41 +5552,12 @@ class TestBedrockGuardrailImageInput:
|
|||
|
||||
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):
|
||||
def test_the_url_helper_guards_its_own_inputs(self):
|
||||
"""Exercised directly so the guards are not dropped in a later refactor."""
|
||||
assert BedrockGuardrail._get_image_url(item={"type": "image_url"}) is None
|
||||
assert BedrockGuardrail._get_image_url(item={"type": "image_url", "image_url": {"url": 7}}) is None
|
||||
assert BedrockGuardrail._get_image_url(item={"type": "image_url", "image_url": 7}) 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"].
|
||||
|
|
@ -5984,32 +5751,6 @@ class TestBedrockGuardrailImageInput:
|
|||
images = [item for item in sent[0]["content"] if "image" in item]
|
||||
assert len(images) == 1, f"the image was sent {len(images)} times"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_file_backed_image_is_let_through_under_the_allow_policy(self):
|
||||
"""An operator who would rather serve it unscanned can still say so."""
|
||||
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": []},
|
||||
request_data={
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "image", "source": {"type": "file", "file_id": "file_abc"}}],
|
||||
}
|
||||
]
|
||||
},
|
||||
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.
|
||||
|
|
@ -6138,8 +5879,8 @@ class TestBedrockGuardrailImageInput:
|
|||
"""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.
|
||||
Without sniffing the format back it would be rejected as unreadable and turn
|
||||
a legitimate /v1/messages call into a 400.
|
||||
"""
|
||||
g = self._guardrail()
|
||||
sent: list = []
|
||||
|
|
@ -6202,46 +5943,11 @@ class TestBedrockGuardrailImageInput:
|
|||
|
||||
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_load_credentials_assumes_role_with_external_id():
|
||||
"""A trust policy requiring sts:ExternalId must be satisfied by the guardrail's aws_external_id."""
|
||||
import datetime
|
||||
|
|
|
|||
|
|
@ -121,40 +121,6 @@ def test_initialize_guardrail_sets_run_in_parallel(config_value, expected):
|
|||
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"
|
||||
|
||||
|
||||
def test_initialize_presidio_forwards_analyze_chunk_size_bytes():
|
||||
"""Regression (LIT-4785): `presidio_analyze_chunk_size_bytes` set in
|
||||
config.yaml must reach the guardrail instance. The field lives on
|
||||
|
|
|
|||
7
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
7
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -30451,13 +30451,6 @@ export interface components {
|
|||
* @description Action to take when sensitive data is detected. 'block' raises an exception (default behavior). 'route' reroutes the request to the model specified in sensitive_data_route_to_model.
|
||||
*/
|
||||
on_sensitive_data?: ("block" | "route") | null;
|
||||
/**
|
||||
* On Unscannable Image
|
||||
* @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.
|
||||
* @default block
|
||||
* @enum {string}
|
||||
*/
|
||||
on_unscannable_image: "block" | "allow";
|
||||
/**
|
||||
* On Violation
|
||||
* @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue