diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrail_attachments.py b/litellm/llms/bedrock/guardrail_attachments.py similarity index 86% rename from litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrail_attachments.py rename to litellm/llms/bedrock/guardrail_attachments.py index a5b90393c07..fbfd5842b50 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrail_attachments.py +++ b/litellm/llms/bedrock/guardrail_attachments.py @@ -11,7 +11,7 @@ import base64 import binascii from collections.abc import Callable, Mapping, Sequence from types import MappingProxyType -from typing import Final, Literal, NamedTuple +from typing import Final, Literal, NamedTuple, TypeGuard from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockContentItem, @@ -53,12 +53,13 @@ _CHAT_CALL_TYPES: Final = frozenset({CallTypes.completion.value, CallTypes.acomp _RESPONSES_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value}) _CHAT_UNSCANNABLE_TYPES: Final = frozenset({"file", "input_audio", "video_url", "audio_url", "document"}) _ANTHROPIC_UNSCANNABLE_TYPES: Final = frozenset({"document", "container_upload"}) +_RESPONSES_IMAGE_TYPES: Final = frozenset({"input_image", "computer_screenshot"}) _RESPONSES_UNSCANNABLE_TYPES: Final = frozenset({"input_file", "input_audio"}) _CONVERSE_UNSCANNABLE_KEYS: Final = ("document", "video") _MAX_IMAGE_BYTES: Final = 4 * 1024 * 1024 _CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) _TOOL_ROLES: Final = frozenset({"tool", "function"}) -_TOOL_OUTPUT_ITEM_TYPES: Final = frozenset({"function_call_output", "custom_tool_call_output"}) +_TOOL_OUTPUT_ITEM_TYPES: Final = frozenset({"function_call_output", "custom_tool_call_output", "computer_call_output"}) _MAX_LABEL_MIME_CHARS: Final = 40 @@ -96,7 +97,7 @@ def _messages_and_classifier( return _mappings(data.get("input")), _classify_responses_block, _no_nested_blocks if call_type == CallTypes.allm_passthrough_route.value and _is_bedrock_converse(data): body: Final = data.get("data") - messages: Final = _mappings(body.get("messages") if isinstance(body, Mapping) else None) + messages: Final = _mappings(body.get("messages") if _is_mapping(body) else None) return messages, _classify_converse_block, _converse_tool_result_blocks return (), _classify_nothing, _no_nested_blocks @@ -110,10 +111,18 @@ def _is_bedrock_converse(data: Mapping[str, object]) -> bool: ) +def _is_mapping(value: object) -> TypeGuard[Mapping[str, object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip + return isinstance(value, Mapping) + + +def _is_list(value: object) -> TypeGuard[list[object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip + return isinstance(value, list) + + def _mappings(value: object) -> tuple[Mapping[str, object], ...]: - if not isinstance(value, list): + if not _is_list(value): return () - return tuple(item for item in value if isinstance(item, Mapping)) + return tuple(item for item in value if _is_mapping(item)) def _latest_user_message(messages: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, object], ...]: @@ -122,7 +131,9 @@ def _latest_user_message(messages: Sequence[Mapping[str, object]]) -> tuple[Mapp def _message_blocks(message: Mapping[str, object], nested_tool_blocks: _NestedToolBlocks) -> tuple[_Block, ...]: if message.get("type") in _TOOL_OUTPUT_ITEM_TYPES: - return tuple(_Block(block, from_tool=True) for block in _mappings(message.get("output"))) + output: Final = message.get("output") + blocks: Final = (output,) if _is_mapping(output) else _mappings(output) + return tuple(_Block(block, from_tool=True) for block in blocks) from_tool_message: Final = message.get("role") in _TOOL_ROLES return tuple( entry @@ -148,7 +159,7 @@ def _anthropic_tool_result_blocks(block: Mapping[str, object]) -> tuple[Mapping[ def _converse_tool_result_blocks(block: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: tool_result: Final = block.get("toolResult") - return _mappings(tool_result.get("content")) if isinstance(tool_result, Mapping) else () + return _mappings(tool_result.get("content")) if _is_mapping(tool_result) else () def _classify_nothing(block: Mapping[str, object]) -> _Classified: @@ -159,7 +170,7 @@ def _classify_chat_block(block: Mapping[str, object]) -> _Classified: block_type: Final = block.get("type") if block_type == "image_url": image_url: Final = block.get("image_url") - url: Final = image_url.get("url") if isinstance(image_url, Mapping) else image_url + url: Final = image_url.get("url") if _is_mapping(image_url) else image_url return _classify_data_uri(url, "image_url") if isinstance(block_type, str) and block_type in _CHAT_UNSCANNABLE_TYPES: return _Unscannable(block_type) @@ -170,7 +181,7 @@ def _classify_anthropic_block(block: Mapping[str, object]) -> _Classified: block_type: Final = block.get("type") if block_type == "image": source: Final = block.get("source") - if isinstance(source, Mapping) and source.get("type") == "base64": + if _is_mapping(source) and source.get("type") == "base64": return _classify_base64(source.get("media_type"), source.get("data"), "image") return _Unscannable("image (url or file source)") if isinstance(block_type, str) and block_type in _ANTHROPIC_UNSCANNABLE_TYPES: @@ -180,8 +191,8 @@ def _classify_anthropic_block(block: Mapping[str, object]) -> _Classified: def _classify_responses_block(block: Mapping[str, object]) -> _Classified: block_type: Final = block.get("type") - if block_type == "input_image": - return _classify_data_uri(block.get("image_url"), "input_image") + if isinstance(block_type, str) and block_type in _RESPONSES_IMAGE_TYPES: + return _classify_data_uri(block.get("image_url"), block_type) if isinstance(block_type, str) and block_type in _RESPONSES_UNSCANNABLE_TYPES: return _Unscannable(block_type) return None @@ -189,10 +200,10 @@ def _classify_responses_block(block: Mapping[str, object]) -> _Classified: def _classify_converse_block(block: Mapping[str, object]) -> _Classified: image: Final = block.get("image") - if isinstance(image, Mapping): + if _is_mapping(image): image_format: Final = image.get("format") source: Final = image.get("source") - encoded: Final = source.get("bytes") if isinstance(source, Mapping) else None + encoded: Final = source.get("bytes") if _is_mapping(source) else None if encoded is None: return _Unscannable("image (no inline bytes)") mime: Final = f"image/{image_format}" if isinstance(image_format, str) else None diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 3d44315341b..a43006a992c 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -10364,7 +10364,7 @@ } ], "default": false, - "description": "Implemented by guardrail='model_armor'. When True, attachment references that carry no inline bytes (file_id, gs://, or http(s) URLs) pass through unscanned instead of blocking, while fail_on_error still governs real Model Armor API errors. Default False blocks them.", + "description": "Implemented by guardrail='model_armor' and guardrail='bedrock'. When True, attachments the guardrail cannot scan pass through unscanned instead of blocking. For Model Armor these are references with no inline bytes (file_id, gs://, or http(s) URLs), and fail_on_error still governs real Model Armor API errors. For Bedrock these are documents, files, audio, video, and images that are not inline PNG or JPEG up to 4 MB. Default False blocks them.", "title": "Skip Unscannable Attachments" }, "sticky_session_routing": { @@ -13172,7 +13172,7 @@ } ], "default": false, - "description": "Implemented by guardrail='model_armor'. When True, attachment references that carry no inline bytes (file_id, gs://, or http(s) URLs) pass through unscanned instead of blocking, while fail_on_error still governs real Model Armor API errors. Default False blocks them.", + "description": "Implemented by guardrail='model_armor' and guardrail='bedrock'. When True, attachments the guardrail cannot scan pass through unscanned instead of blocking. For Model Armor these are references with no inline bytes (file_id, gs://, or http(s) URLs), and fail_on_error still governs real Model Armor API errors. For Bedrock these are documents, files, audio, video, and images that are not inline PNG or JPEG up to 4 MB. Default False blocks them.", "title": "Skip Unscannable Attachments" }, "sticky_session_routing": { diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 5a58a8abbad..977e95ae01c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -46,6 +46,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_skip_tool_message_for_guardrail, ) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token, run_aws_signing +from litellm.llms.bedrock.guardrail_attachments import find_request_attachments from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -60,7 +61,6 @@ from litellm.proxy.guardrails.anthropic_sse import ( is_raw_sse_stream, model_response_text, ) -from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrail_attachments import find_request_attachments from litellm.types.guardrails import ( BedrockChecksConfigModel, BedrockGuardrailStreamingParams, @@ -77,6 +77,8 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrailQualifier, BedrockGuardrailResponse, BedrockGuardrailUsage, + BedrockImageContent, + BedrockImageSource, BedrockRequest, BedrockTextContent, ) @@ -225,6 +227,20 @@ def _redact_assessment_match_fields(assessments: list[dict]) -> list[dict]: _RESPONSES_API_CALL_TYPES: Final = frozenset({CallTypes.responses, CallTypes.aresponses}) +def _without_image_bytes(content: Sequence[BedrockContentItem]) -> list[BedrockContentItem]: + return [ + BedrockContentItem( + image=BedrockImageContent( + format=item["image"]["format"], + source=BedrockImageSource(bytes=f"<{len(item['image']['source']['bytes'])} base64 chars>"), + ) + ) + if "image" in item + else item + for item in content + ] + + def _is_responses_api_route(request_route: str | None) -> bool: if request_route is None: return False @@ -1263,7 +1279,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): headers_dict: Final = dict(prepared_request.headers) # mutable-ok: the masking helper requires a dict verbose_proxy_logger.debug( "Bedrock AI request body: %s, url %s, headers: %s", - bedrock_request_data, + {**bedrock_request_data, "content": _without_image_bytes(content)}, prepared_request.url, _get_masked_values(headers_dict), ) @@ -2642,6 +2658,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): event_type: Final[GuardrailEventHooks] = GuardrailEventHooks.pre_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: return data + await self.async_scan_request_attachments(data=data, call_type=call_type, event_type=event_type) new_messages: Final = self.get_guardrails_messages_for_call_type( call_type=cast(CallTypes, call_type), diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 579a3f6322f..7e708570a50 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -1037,9 +1037,11 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up skip_unscannable_attachments: bool | None = Field( default=False, description=( - "Implemented by guardrail='model_armor'. When True, attachment references that carry no " - "inline bytes (file_id, gs://, or http(s) URLs) pass through unscanned instead of blocking, " - "while fail_on_error still governs real Model Armor API errors. Default False blocks them." + "Implemented by guardrail='model_armor' and guardrail='bedrock'. When True, attachments the " + "guardrail cannot scan pass through unscanned instead of blocking. For Model Armor these are " + "references with no inline bytes (file_id, gs://, or http(s) URLs), and fail_on_error still " + "governs real Model Armor API errors. For Bedrock these are documents, files, audio, video, " + "and images that are not inline PNG or JPEG up to 4 MB. Default False blocks them." ), ) sanitize_error_detail: bool | None = Field( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 4d46d3594f0..1e566f491cf 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -1,6 +1,6 @@ from typing import Literal -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict # Bedrock contextual grounding tags each content block so the guardrail knows # which text is the reference source, the user question, and the content to grade. @@ -13,12 +13,12 @@ class BedrockTextContent(TypedDict, total=False): class BedrockImageSource(TypedDict): - bytes: str + bytes: ReadOnly[str] class BedrockImageContent(TypedDict): - format: Literal["png", "jpeg"] - source: BedrockImageSource + format: ReadOnly[Literal["png", "jpeg"]] + source: ReadOnly[BedrockImageSource] class BedrockContentItem(TypedDict, total=False): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 42bcdf678c1..4c78703ea38 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -6334,3 +6334,43 @@ async def test_attachment_scan_honors_scan_only_tool_results(): assert result is None assert mock_post.await_count == 0 assert data == _pdf_chat_request() + + +@pytest.mark.asyncio +async def test_native_pre_call_hook_blocks_unscannable_attachment(): + guardrail = _attachment_guardrail(event_hook=GuardrailEventHooks.pre_call, default_on=True) + + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post: + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(), + data=_pdf_chat_request(), + call_type="acompletion", + ) + + assert exc_info.value.status_code == 400 + assert "cannot scan 1 attachment(s) (file)" in exc_info.value.detail["reason"] + assert mock_post.await_count == 0 + + +@pytest.mark.asyncio +async def test_attachment_scan_debug_log_omits_image_bytes(): + guardrail = _attachment_guardrail() + post_patch, credentials_patch, prepare_patch = _patched_bedrock_post( + guardrail, _passing_bedrock_httpx_response("ok") + ) + + with ( + post_patch, + credentials_patch, + prepare_patch, + patch("litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug") as mock_debug, + ): + await guardrail.async_scan_request_attachments( + data=_image_only_chat_request(), call_type=CallTypes.acompletion.value + ) + + logged = " ".join(str(call.args) for call in mock_debug.call_args_list) + assert _ATTACHMENT_PNG_B64 not in logged + assert f"<{len(_ATTACHMENT_PNG_B64)} base64 chars>" in logged diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrail_attachments.py b/tests/unit/llms/bedrock/test_guardrail_attachments.py similarity index 93% rename from tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrail_attachments.py rename to tests/unit/llms/bedrock/test_guardrail_attachments.py index 26e8e66f430..5edf7f2bf83 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrail_attachments.py +++ b/tests/unit/llms/bedrock/test_guardrail_attachments.py @@ -2,7 +2,7 @@ import base64 import pytest -from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrail_attachments import find_request_attachments +from litellm.llms.bedrock.guardrail_attachments import find_request_attachments from litellm.types.utils import CallTypes PNG_B64 = base64.b64encode(b"\x89PNG\r\n\x1a\nfake-png").decode() @@ -237,6 +237,26 @@ TEXT = {"type": "text", "text": "hello"} ["input_file"], id="responses-function-call-output-image-and-file", ), + pytest.param( + { + "input": [ + { + "type": "computer_call_output", + "call_id": "call_1", + "output": {"type": "computer_screenshot", "image_url": f"data:image/png;base64,{PNG_B64}"}, + }, + { + "type": "computer_call_output", + "call_id": "call_2", + "output": {"type": "computer_screenshot", "file_id": "file-1"}, + }, + ] + }, + CallTypes.responses.value, + [_png_item()], + ["computer_screenshot (remote URL or file id)"], + id="responses-computer-screenshot", + ), pytest.param( _converse( { diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0500aeb95c8..22f7d186f22 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25548,7 +25548,7 @@ export interface components { skip_tool_message_in_guardrail?: boolean | null; /** * Skip Unscannable Attachments - * @description Implemented by guardrail='model_armor'. When True, attachment references that carry no inline bytes (file_id, gs://, or http(s) URLs) pass through unscanned instead of blocking, while fail_on_error still governs real Model Armor API errors. Default False blocks them. + * @description Implemented by guardrail='model_armor' and guardrail='bedrock'. When True, attachments the guardrail cannot scan pass through unscanned instead of blocking. For Model Armor these are references with no inline bytes (file_id, gs://, or http(s) URLs), and fail_on_error still governs real Model Armor API errors. For Bedrock these are documents, files, audio, video, and images that are not inline PNG or JPEG up to 4 MB. Default False blocks them. * @default false */ skip_unscannable_attachments: boolean | null; @@ -34775,7 +34775,7 @@ export interface components { skip_tool_message_in_guardrail?: boolean | null; /** * Skip Unscannable Attachments - * @description Implemented by guardrail='model_armor'. When True, attachment references that carry no inline bytes (file_id, gs://, or http(s) URLs) pass through unscanned instead of blocking, while fail_on_error still governs real Model Armor API errors. Default False blocks them. + * @description Implemented by guardrail='model_armor' and guardrail='bedrock'. When True, attachments the guardrail cannot scan pass through unscanned instead of blocking. For Model Armor these are references with no inline bytes (file_id, gs://, or http(s) URLs), and fail_on_error still governs real Model Armor API errors. For Bedrock these are documents, files, audio, video, and images that are not inline PNG or JPEG up to 4 MB. Default False blocks them. * @default false */ skip_unscannable_attachments: boolean | null;