diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 4b228cdd187..b8f7eda7a91 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -14,6 +14,7 @@ from pydantic import InstanceOf, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, @@ -61,6 +62,7 @@ if TYPE_CHECKING: _CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) _INVOKE_STREAM_ACTION: Final = "invoke-with-response-stream" _INVOKE_ACTIONS: Final = frozenset({"invoke", _INVOKE_STREAM_ACTION}) +_BLOCK_TYPES_THE_ANTHROPIC_HANDLER_DOES_NOT_SCAN: Final = frozenset({"document", "search_result"}) _EVENT_STREAM_CONTENT_TYPE: Final = "vnd.amazon.eventstream" _EVENT_STREAM_MEDIA_TYPE: Final = "application/vnd.amazon.eventstream" @@ -87,9 +89,31 @@ def _anthropic_invoke_body( invoke_body: Final = _BODY_BY_REFERENCE.validate_python(body) if not isinstance(invoke_body.get("anthropic_version"), str) or not isinstance(invoke_body.get("messages"), list): return None + blocks: Final = _anthropic_message_blocks(_ARRAY_BY_REFERENCE.validate_python(invoke_body["messages"])) + if any(block.get("type") in _BLOCK_TYPES_THE_ANTHROPIC_HANDLER_DOES_NOT_SCAN for block in blocks): + return None return invoke_body +def _anthropic_message_blocks(messages: Sequence[object]) -> tuple[Mapping[str, object], ...]: + fields: Final = (as_str_mapping(message) for message in messages) + return tuple(chain.from_iterable(_anthropic_content_blocks(field.get("content")) for field in fields if field)) + + +def _anthropic_content_blocks(content: object) -> tuple[Mapping[str, object], ...]: + if not isinstance(content, list): + return () + blocks: Final = tuple( + block + for block in (as_str_mapping(item) for item in _ARRAY_BY_REFERENCE.validate_python(content)) + if block is not None + ) + nested: Final = ( + _anthropic_content_blocks(block.get("content")) for block in blocks if block.get("type") == "tool_result" + ) + return (*blocks, *chain.from_iterable(nested)) + + def _anthropic_invoke_reply( response: object, request_data: object, diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index 17d7971e1f5..9d1168327f6 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -3981,6 +3981,45 @@ async def test_invoke_body_without_anthropic_version_is_scanned_whole(): assert "structured_messages" not in guardrail.received[-1] +_PLAIN_TEXT_DOCUMENT: Final = { + "type": "document", + "source": {"type": "text", "media_type": "text/plain", "data": "Card 4111 1111 1111 1111"}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "content", + [ + pytest.param([{"type": "text", "text": "Summarize this"}, _PLAIN_TEXT_DOCUMENT], id="document"), + pytest.param( + [ + { + "type": "search_result", + "source": "https://example.com/cards", + "title": "Cards", + "content": [{"type": "text", "text": "Card 4111 1111 1111 1111"}], + } + ], + id="search_result", + ), + pytest.param( + [{"type": "tool_result", "tool_use_id": "toolu_1", "content": [_PLAIN_TEXT_DOCUMENT]}], + id="document_in_tool_result", + ), + ], +) +async def test_anthropic_invoke_body_with_blocks_the_messages_handler_does_not_scan_is_scanned_whole(content: list): + invoke = _anthropic_invoke_data() + data = {**invoke, "data": {**invoke["data"], "messages": [{"role": "user", "content": content}]}} + guardrail = _replying({}) + + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail) + + assert json.loads(guardrail.received[-1]["texts"][0])["data"] == data["data"] + assert "structured_messages" not in guardrail.received[-1] + + @pytest.mark.asyncio async def test_anthropic_invoke_response_masks_are_written_back(): response = {