mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Merge c244e5fca4 into c8635ecc67
This commit is contained in:
commit
43fafbc42a
5 changed files with 1078 additions and 43 deletions
|
|
@ -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
|
||||
|
|
@ -19,7 +20,7 @@ from collections.abc import AsyncGenerator, Mapping, Sequence
|
|||
from datetime import datetime, timezone
|
||||
from itertools import accumulate, groupby
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Optional, cast
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, NoReturn, Optional, cast
|
||||
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
|
|
@ -37,6 +38,7 @@ from litellm.litellm_core_utils.litellm_logging import (
|
|||
_get_masked_values, # pyright: ignore[reportPrivateUsage] # the shared header-masking helper has no public name
|
||||
)
|
||||
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.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
effective_scan_only_tool_results_for_guardrail,
|
||||
|
|
@ -62,16 +64,24 @@ from litellm.types.guardrails import (
|
|||
GuardrailEventHooks,
|
||||
LitellmParams,
|
||||
)
|
||||
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,
|
||||
)
|
||||
|
|
@ -142,6 +152,20 @@ _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"})
|
||||
|
||||
# AWS image-filter limits, none of which litellm checked before. Anything other than
|
||||
# png/jpeg has no representation in the payload and cannot be scanned at all.
|
||||
# https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails-mmfilter.html
|
||||
_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)."""
|
||||
|
|
@ -361,28 +385,256 @@ 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
|
||||
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) 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: Final = 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:
|
||||
"""Pull the url out of an image_url part. The caller owns the type dispatch."""
|
||||
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
|
||||
|
||||
@classmethod
|
||||
def _image_urls_in(cls, messages: "Sequence[AllMessageValues] | None") -> frozenset[str]:
|
||||
"""Normalized image urls already carried by these messages."""
|
||||
found: Final[set[str]] = set() # mutable-ok: accumulator, frozen on return
|
||||
for message in messages or ():
|
||||
content = message.get("content")
|
||||
if not isinstance(content, list):
|
||||
continue
|
||||
for part in content:
|
||||
if not isinstance(part, dict) or part.get("type") != "image_url":
|
||||
continue
|
||||
url = cls._get_image_url(item=part)
|
||||
if url is not None:
|
||||
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) -> NoReturn:
|
||||
"""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
|
||||
"""
|
||||
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.
|
||||
_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 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:
|
||||
"""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
|
||||
source outright, so a check reading it never fires while the provider still
|
||||
forwards the file. The cost is that the scope flags no longer narrow this
|
||||
check, and over-refusing is the safer error here.
|
||||
"""
|
||||
if input_type != "request":
|
||||
return
|
||||
found: Final = self._file_backed_image_count(request_data.get("messages"))
|
||||
if not found:
|
||||
return
|
||||
self._handle_unscannable_image(
|
||||
reason=f"{found} image(s) reference a provider file id, whose bytes are not available here"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _file_backed_image_count(cls, messages: object) -> int:
|
||||
"""Count Anthropic `{"type": "image", "source": {"type": "file"}}` parts.
|
||||
|
||||
Matches that one shape rather than comparing counts against
|
||||
`inputs["images"]`, whose length the skip and scope flags already narrow, so
|
||||
a mismatch there is not by itself evidence of a dropped image.
|
||||
"""
|
||||
if not isinstance(messages, list):
|
||||
return 0
|
||||
return sum(cls._file_backed_parts(message.get("content")) for message in messages if isinstance(message, dict))
|
||||
|
||||
@classmethod
|
||||
def _file_backed_parts(cls, content: object) -> int:
|
||||
"""Count file-backed images in one content list, descending into tool_result.
|
||||
|
||||
The extractor pulls scannable images out of a tool_result's nested blocks, so
|
||||
a file source sitting there has to be refused for the same reason a top-level
|
||||
one is: nothing else in the request will surface it.
|
||||
"""
|
||||
if not isinstance(content, list):
|
||||
return 0
|
||||
found = 0 # rebind-ok: running count over the content list
|
||||
for part in content:
|
||||
if not isinstance(part, dict):
|
||||
continue
|
||||
if part.get("type") == "tool_result":
|
||||
found += cls._file_backed_parts(part.get("content"))
|
||||
continue
|
||||
if part.get("type") != "image":
|
||||
continue
|
||||
source = part.get("source")
|
||||
if isinstance(source, dict) and source.get("type") == "file":
|
||||
found += 1
|
||||
return found
|
||||
|
||||
async def _build_image_content_item(self, image_url: str) -> BedrockContentItem:
|
||||
"""Decode an inline image into an ApplyGuardrail image block.
|
||||
|
||||
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.
|
||||
"""
|
||||
if image_url.startswith(("http://", "https://")):
|
||||
self._handle_unscannable_image(reason="remote image URLs are not supported")
|
||||
|
||||
try:
|
||||
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}")
|
||||
|
||||
image_block: Final = block.get("image")
|
||||
image_format: Final = (
|
||||
_APPLY_GUARDRAIL_IMAGE_FORMATS.get(str(image_block.get("format"))) if image_block else None
|
||||
)
|
||||
image_source: Final = image_block.get("source") if image_block else None
|
||||
image_bytes: Final = image_source.get("bytes") if image_source else None
|
||||
if image_format is None or not image_bytes:
|
||||
self._handle_unscannable_image(reason="attachment is not a png/jpeg image")
|
||||
|
||||
# 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 BedrockContentItem(
|
||||
image=BedrockImageContent(
|
||||
format=image_format,
|
||||
source=BedrockGuardrailImageSource(bytes=image_bytes),
|
||||
)
|
||||
)
|
||||
|
||||
def _create_bedrock_output_content_request(
|
||||
self,
|
||||
response: object,
|
||||
|
|
@ -429,7 +681,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,
|
||||
|
|
@ -444,12 +696,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
Returns:
|
||||
BedrockRequest: The bedrock request object.
|
||||
"""
|
||||
bedrock_request: BedrockRequest = BedrockRequest(source=source)
|
||||
if source == "INPUT":
|
||||
bedrock_request = 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
|
||||
return await self._create_bedrock_input_content_request(messages=messages)
|
||||
if source == "OUTPUT":
|
||||
return self._create_bedrock_output_content_request(response=response, messages=messages)
|
||||
return BedrockRequest(source=source)
|
||||
|
||||
def get_content_items_for_message(self, message: AllMessageValues) -> list[QualifiedTextBlock] | None:
|
||||
"""
|
||||
|
|
@ -882,7 +1133,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:
|
||||
|
|
@ -3124,7 +3375,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
|
||||
|
|
@ -3132,10 +3384,28 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
# NOTE: Use `or []` to handle case where inputs["texts"] is explicitly None.
|
||||
# dict.get("texts", []) would return None if the key exists with a None value.
|
||||
texts: Final = inputs.get("texts") or []
|
||||
# Images only exist on the request side; ApplyGuardrail's OUTPUT source takes
|
||||
# model-generated text. The endpoint translations already extract them:
|
||||
# OpenAIChatCompletionsHandler from `image_url` parts, AnthropicMessagesHandler
|
||||
# from `image`/`source` blocks. Five other guardrails already consume this
|
||||
# field; Bedrock was the one that dropped it on the floor.
|
||||
image_urls: Final = tuple(inputs.get("images") or ()) if input_type == "request" else ()
|
||||
# Before the shortcuts below, so a file-backed image cannot be skipped either.
|
||||
# Both conditions live in the callee: apply_guardrail sits one branch under
|
||||
# ruff-strict's complexity ceiling, and two more here would cross it.
|
||||
self._refuse_file_backed_images(request_data=request_data, input_type=input_type)
|
||||
try:
|
||||
verbose_proxy_logger.debug("Bedrock Guardrail: Applying guardrail to %s text(s)", len(texts))
|
||||
verbose_proxy_logger.debug(
|
||||
"Bedrock Guardrail: Applying guardrail to %s text(s) and %s image(s)", len(texts), len(image_urls)
|
||||
)
|
||||
|
||||
if input_type == "request":
|
||||
# Both of the optimizations below decide what to scan by looking at
|
||||
# `texts` alone, so an image rides along unscanned and the proxy still
|
||||
# reports the guardrail as run. Neither tracks images in its session
|
||||
# cache either, so "already seen" cannot be established for them.
|
||||
# Presence of an image therefore forces the full image-aware path:
|
||||
# correctness over the optimization, and it fails closed.
|
||||
if input_type == "request" and not image_urls:
|
||||
incremental_result: Final = await self._apply_incremental_request_scan(
|
||||
texts=texts,
|
||||
inputs=inputs,
|
||||
|
|
@ -3152,14 +3422,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)
|
||||
|
|
@ -3182,7 +3455,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
),
|
||||
finish_reason="stop",
|
||||
)
|
||||
for _idx, _msg in enumerate(filtered_messages)
|
||||
# `or ()`: skip_scan is now bypassed when an image is present,
|
||||
# and images exist on the request side only, so a response scan
|
||||
# still always has messages here. Spelled out rather than left
|
||||
# leaning on that indirection.
|
||||
for _idx, _msg in enumerate(filtered_messages or ())
|
||||
]
|
||||
)
|
||||
bedrock_response = await self.make_bedrock_api_request(
|
||||
|
|
@ -3192,9 +3469,33 @@ 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 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 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)
|
||||
for normalized in (self._normalize_image_input(url) for url in image_urls)
|
||||
if normalized not in already_scanned
|
||||
]
|
||||
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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
from typing import Literal
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
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: TypeAlias = 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
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@
|
|||
Unit tests for Bedrock Guardrails
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
|
@ -2408,12 +2410,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 +2965,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)
|
||||
|
||||
|
|
@ -5276,6 +5278,723 @@ async def test_terminal_failure_logs_usage_and_cost_of_prior_passed_chunks(monke
|
|||
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="
|
||||
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_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_twenty_inline_images_are_accepted(self):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"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"]) == 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",
|
||||
[
|
||||
pytest.param(None, id="no content"),
|
||||
pytest.param(123, id="content is not a list"),
|
||||
pytest.param([123], id="part is not a mapping"),
|
||||
pytest.param([{"type": "image_url"}], id="image part with no url"),
|
||||
pytest.param([{"type": "image_url", "image_url": 123}], id="url is not a string or mapping"),
|
||||
pytest.param([{"type": "image_url", "image_url": {"url": 123}}], id="url value is not a string"),
|
||||
pytest.param([{"type": "input_audio"}], id="part carries neither image nor text"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_malformed_part_never_becomes_scannable_content(self, content):
|
||||
"""Default-deny is the point of this PR, so pin it rather than trust it.
|
||||
|
||||
Each shape below is one a caller can put on the wire. None of them may turn
|
||||
into a content item: an unrecognised part that fell through to the text
|
||||
branch would be reported to the operator as scanned when it was not.
|
||||
"""
|
||||
request = await self._guardrail().convert_to_bedrock_format(
|
||||
source="INPUT", messages=[{"role": "user", "content": content}]
|
||||
)
|
||||
|
||||
assert request["content"] == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_bare_string_part_is_scanned_as_text(self):
|
||||
"""A content list may hold plain strings, not only typed parts."""
|
||||
request = await self._guardrail().convert_to_bedrock_format(
|
||||
source="INPUT", messages=[{"role": "user", "content": ["just text"]}]
|
||||
)
|
||||
|
||||
assert request["content"] == [{"text": {"text": "just text"}}]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_url_given_as_a_plain_string_is_accepted(self):
|
||||
"""OpenAI accepts `image_url` as a bare string as well as `{"url": ...}`.
|
||||
|
||||
Both reach the model as an image, so both have to reach the scan.
|
||||
"""
|
||||
request = await self._guardrail().convert_to_bedrock_format(
|
||||
source="INPUT",
|
||||
messages=[{"role": "user", "content": [{"type": "image_url", "image_url": self._PNG_DATA_URI}]}],
|
||||
)
|
||||
|
||||
kinds = [k for item in request["content"] for k in item]
|
||||
assert kinds == ["image"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_unrecognized_payload_is_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, 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.
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await self._guardrail().apply_guardrail(
|
||||
inputs={"texts": [], "images": ["R0lGODlhAQABAAAAACw="]},
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert "could not be read" in str(exc_info.value.detail) or "not a png/jpeg" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_remote_image_url_is_rejected_without_being_fetched(self):
|
||||
"""A remote url is named as its own rejection rather than left to the decoder.
|
||||
|
||||
Fetching one safely (size cap, SSRF/redirect validation) is separate work; this
|
||||
PR only scans inline images, so a url has to fail closed rather than be ignored
|
||||
or silently forwarded to the model unscanned.
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "image_url", "image_url": {"url": "https://example.com/pic.png"}}],
|
||||
}
|
||||
]
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
|
||||
|
||||
assert "remote image URLs are not supported" in str(exc_info.value.detail)
|
||||
|
||||
def test_file_backed_part_counting_skips_non_mapping_entries(self):
|
||||
"""A content list may mix plain strings in with typed parts.
|
||||
|
||||
`_file_backed_image_count` walks the raw request rather than a normalized
|
||||
shape, so a non-dict entry must be skipped rather than raise or be miscounted.
|
||||
"""
|
||||
content = [
|
||||
"just a string",
|
||||
{"type": "text", "text": "hi"},
|
||||
{"type": "image", "source": {"type": "file", "file_id": "file_abc"}},
|
||||
]
|
||||
|
||||
assert BedrockGuardrail._file_backed_parts(content) == 1
|
||||
|
||||
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
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_file_backed_image_is_refused_rather_than_ignored(self):
|
||||
"""`{"type": "file"}` carries no bytes, so nothing reaches inputs["images"].
|
||||
|
||||
The provider still forwards the file to the model, so ignoring it is exactly
|
||||
the silent pass this path exists to remove. Documented is not the same as
|
||||
safe; under the default policy the request is refused.
|
||||
"""
|
||||
inputs = {"texts": ["what does this say?"], "images": []}
|
||||
request_data = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "what does this say?"},
|
||||
{"type": "image", "source": {"type": "file", "file_id": "file_abc"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
g = self._guardrail()
|
||||
sent: list = []
|
||||
|
||||
async def spy(**kwargs):
|
||||
sent.append(kwargs["messages"])
|
||||
return {"action": "NONE", "outputs": []}
|
||||
|
||||
# Stubbed so that without the refusal this request would simply succeed:
|
||||
# the failure mode being pinned is a silent pass, not an AWS error.
|
||||
with patch.object(g, "make_bedrock_api_request", new=spy):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await g.apply_guardrail(inputs=inputs, request_data=request_data, input_type="request")
|
||||
|
||||
assert "file id" in str(exc_info.value.detail)
|
||||
assert sent == [], "refused before any scan was attempted"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_file_backed_image_is_refused_through_the_real_translation(self):
|
||||
"""Drive the /v1/messages handler instead of hand-building its output.
|
||||
|
||||
The handler translates to OpenAI spec before filling structured_messages, and
|
||||
that translation drops a file source, so a check reading structured_messages
|
||||
passes every hand-written fixture and never fires in production.
|
||||
"""
|
||||
from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
|
||||
|
||||
data = {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "what does this say?"},
|
||||
{"type": "image", "source": {"type": "file", "file_id": "file_abc"}},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
assert not self._file_parts_in(AnthropicMessagesHandler().get_structured_messages(data)), (
|
||||
"the translation is expected to drop the file source; that is why this test exists"
|
||||
)
|
||||
|
||||
g = self._guardrail()
|
||||
with patch.object(g, "make_bedrock_api_request", new=AsyncMock(return_value={"action": "NONE"})):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=g)
|
||||
|
||||
assert "file id" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_file_backed_image_inside_a_tool_result_is_refused(self):
|
||||
"""The extractor pulls scannable images out of a tool_result's nested blocks.
|
||||
|
||||
A file source sitting there is invisible to both: it yields no bytes to scan
|
||||
and, until the count descended into tool_result, no refusal either.
|
||||
"""
|
||||
g = self._guardrail()
|
||||
sent: list = []
|
||||
|
||||
async def spy(**kwargs):
|
||||
sent.append(kwargs["messages"])
|
||||
return {"action": "NONE", "outputs": []}
|
||||
|
||||
with patch.object(g, "make_bedrock_api_request", new=spy):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await g.apply_guardrail(
|
||||
inputs={"texts": ["what does this say?"], "images": []},
|
||||
request_data={
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "what does this say?"},
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "tu_1",
|
||||
"content": [
|
||||
{"type": "image", "source": {"type": "file", "file_id": "file_abc"}}
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert "file id" in str(exc_info.value.detail)
|
||||
assert sent == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_scannable_image_inside_a_tool_result_is_not_refused(self):
|
||||
"""Descending into tool_result must not start refusing what can be scanned."""
|
||||
g = self._guardrail()
|
||||
sent: list = []
|
||||
|
||||
async def spy(**kwargs):
|
||||
sent.append(kwargs["messages"])
|
||||
return {"action": "NONE", "outputs": []}
|
||||
|
||||
with patch.object(g, "make_bedrock_api_request", new=spy):
|
||||
await g.apply_guardrail(
|
||||
inputs={"texts": ["hello"], "images": [self._PNG_DATA_URI]},
|
||||
request_data={
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "hello"},
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "tu_1",
|
||||
"content": [
|
||||
{"type": "image", "source": {"type": "base64", "data": "AAAA"}}
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert sent, "a scannable nested image still has to be scanned, not refused"
|
||||
|
||||
@staticmethod
|
||||
def _file_parts_in(messages) -> int:
|
||||
return sum(
|
||||
1
|
||||
for message in messages or ()
|
||||
if isinstance(message, dict)
|
||||
for part in (message.get("content") if isinstance(message.get("content"), list) else ())
|
||||
if isinstance(part, dict) and part.get("type") == "image"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_latest_message_only_does_not_scan_the_same_image_twice(self):
|
||||
"""The selected message carries its own image parts into the scan payload.
|
||||
|
||||
`inputs["images"]` holds that same url, so appending it again would fetch and
|
||||
bill the image twice.
|
||||
"""
|
||||
url = self._PNG_DATA_URI
|
||||
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": ["hello"],
|
||||
"images": [url],
|
||||
"structured_messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "hello"},
|
||||
{"type": "image_url", "image_url": {"url": url}},
|
||||
],
|
||||
}
|
||||
],
|
||||
},
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
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_the_scannable_source_shapes_are_not_refused(self):
|
||||
"""The refusal has to be specific to the shape that cannot be read.
|
||||
|
||||
structured_messages is already narrowed by the skip and scope flags, so a
|
||||
count mismatch against inputs["images"] is not evidence of a dropped image.
|
||||
Blocking a legitimate request would be worse than the gap being closed.
|
||||
"""
|
||||
g = self._guardrail()
|
||||
sent: list = []
|
||||
|
||||
async def spy(**kwargs):
|
||||
sent.append(await g.convert_to_bedrock_format(source="INPUT", messages=kwargs["messages"]))
|
||||
return {"action": "NONE", "outputs": []}
|
||||
|
||||
with patch.object(g, "make_bedrock_api_request", new=spy):
|
||||
await g.apply_guardrail(
|
||||
inputs={"texts": ["hello"], "images": [self._PNG_DATA_URI]},
|
||||
request_data={
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "hello"},
|
||||
{"type": "image", "source": {"type": "base64", "data": "AAAA"}},
|
||||
{"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
kinds = [k for item in sent[0]["content"] for k in item]
|
||||
assert "image" in kinds
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_malformed_structured_message_does_not_derail_the_file_check(self):
|
||||
"""structured_messages comes from the caller, so its shape is not guaranteed.
|
||||
|
||||
The scan must keep walking past an entry it cannot read rather than throwing
|
||||
or giving up, or a single junk element would hide a file image sitting after
|
||||
it -- turning a defensive guard into the bypass it was meant to prevent.
|
||||
"""
|
||||
g = self._guardrail()
|
||||
sent: list = []
|
||||
|
||||
async def spy(**kwargs):
|
||||
sent.append(kwargs["messages"])
|
||||
return {"action": "NONE", "outputs": []}
|
||||
|
||||
with patch.object(g, "make_bedrock_api_request", new=spy):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await g.apply_guardrail(
|
||||
inputs={"texts": ["hello"], "images": []},
|
||||
request_data={
|
||||
"messages": [
|
||||
"not a message",
|
||||
123,
|
||||
{"role": "user", "content": "a plain string, not a list"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "image", "source": {"type": "file", "file_id": "file_abc"}}],
|
||||
},
|
||||
]
|
||||
},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert "file id" in str(exc_info.value.detail)
|
||||
assert sent == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_file_backed_image_on_the_response_side_is_not_refused(self):
|
||||
"""Images are a request-side concern; an OUTPUT scan takes generated text."""
|
||||
g = self._guardrail()
|
||||
|
||||
async def spy(**kwargs):
|
||||
return {"action": "NONE", "outputs": []}
|
||||
|
||||
with patch.object(g, "make_bedrock_api_request", new=spy):
|
||||
result = await g.apply_guardrail(
|
||||
inputs={
|
||||
"texts": ["the model said this"],
|
||||
"structured_messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "image", "source": {"type": "file", "file_id": "file_abc"}}],
|
||||
}
|
||||
],
|
||||
},
|
||||
request_data={},
|
||||
input_type="response",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_scans_images_from_inputs(self):
|
||||
"""The proxy reaches BedrockGuardrail through `apply_guardrail`, not the native hook.
|
||||
|
||||
ProxyLogging._execute_guardrail_hook routes any guardrail that defines
|
||||
`apply_guardrail` through unified_guardrail, so a fix that only touches
|
||||
async_pre_call_hook never runs on a real request. The endpoint translations
|
||||
already put images in inputs["images"]; this asserts they reach the payload.
|
||||
"""
|
||||
g = self._guardrail()
|
||||
sent: list = []
|
||||
|
||||
async def spy(**kwargs):
|
||||
sent.append(await g.convert_to_bedrock_format(source="INPUT", messages=kwargs["messages"]))
|
||||
return {"action": "NONE", "outputs": []}
|
||||
|
||||
with patch.object(g, "make_bedrock_api_request", new=spy):
|
||||
await g.apply_guardrail(
|
||||
inputs={"texts": ["what does this say?"], "images": [self._PNG_DATA_URI]},
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
kinds = [k for item in sent[0]["content"] for k in item]
|
||||
assert "image" in kinds, f"image never reached the payload: {kinds}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_scans_bare_base64_images(self):
|
||||
"""Anthropic's translation drops media_type and passes bare base64.
|
||||
|
||||
`_image_sources` returns source["data"] only, so the entry is not a data URI.
|
||||
Without sniffing the format back it would be rejected as unreadable and 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_an_unknown_source_yields_an_empty_request(self):
|
||||
"""The dispatch went from seed-then-assign to early returns when INPUT became
|
||||
async. The trailing fallback preserves what the seeded request used to return
|
||||
for a source that is neither INPUT nor OUTPUT, so nothing starts scanning an
|
||||
unrecognized source as if it were input.
|
||||
"""
|
||||
from typing import Literal, cast
|
||||
|
||||
# cast-ok: pins the runtime fallback the Literal forbids at type-check time
|
||||
source = cast(Literal["INPUT", "OUTPUT"], "SOMETHING_ELSE")
|
||||
request = await self._guardrail().convert_to_bedrock_format(
|
||||
source=source,
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
)
|
||||
|
||||
assert request == {"source": "SOMETHING_ELSE"}
|
||||
|
||||
@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_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
|
||||
|
|
|
|||
|
|
@ -916,7 +916,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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue