Merge 4d16ebd414 into 640223719f

litellm/llms/anthropic/chat/guardrail_translation/handler.py

@@ -813,9 +813,12 @@ class AnthropicMessagesHandler(BaseTranslation):
         data: dict,
         guardrail_to_apply: "CustomGuardrail",
         litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
+        *,
+        request_data: dict[str, object] | None = None,  # mutable-ok: the request dict guardrails receive
     ) -> Mapping[str, object]:
         """
-        Process input messages by applying guardrails to text content.
+        Process input messages by applying guardrails to text content. ``request_data`` is what the guardrail sees as
+        the request when ``data`` is only its body, e.g. a provider pass-through route that nests the body
         """
         messages: Final = data.get("messages")
         if messages is None:
@@ -827,7 +830,9 @@ class AnthropicMessagesHandler(BaseTranslation):
 
         # The top-level prompt is translated on its own below so it can be hoisted in front of
         # any mid-turn system entries and scanned first, aligned with that structured position.
-        translation_source: Final = {key: value for key, value in data.items() if key != "system"}
+        request_model: Final = (request_data or {}).get("model")
+        model: Final = data.get("model") or (request_model if isinstance(request_model, str) else "")
+        translation_source: Final = {**{key: value for key, value in data.items() if key != "system"}, "model": model}
         chat_completion_compatible_request: Final = self._translate_to_openai(translation_source)
 
         full_structured_messages: Final = cast(
@@ -899,7 +904,6 @@ class AnthropicMessagesHandler(BaseTranslation):
             if structured_messages:
                 inputs["structured_messages"] = structured_messages
             # Include model information if available
-            model: Final = data.get("model")
             if model:
                 inputs["model"] = model
             original_tool_names: Final = _original_tool_names(
@@ -907,7 +911,7 @@ class AnthropicMessagesHandler(BaseTranslation):
             )
             guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
                 inputs=inputs,
-                request_data=data,
+                request_data=data if request_data is None else request_data,
                 input_type="request",
                 logging_obj=litellm_logging_obj,
             )

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,
@@ -60,6 +61,8 @@ 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"
 
@@ -73,6 +76,55 @@ def _is_invoke_stream_endpoint(endpoint: str) -> bool:
     return endpoint.rstrip("/").split("/")[-1] == _INVOKE_STREAM_ACTION
 
 
+def _anthropic_invoke_body(
+    data: Mapping[str, object],
+) -> dict[str, object] | None:  # mutable-ok: the body guardrails write back into
+    """The body of an InvokeModel call in Anthropic Messages format, which Bedrock requires to carry anthropic_version"""
+    endpoint: Final = data.get("endpoint")
+    body: Final = data.get("data")
+    if not isinstance(endpoint, str) or endpoint.rstrip("/").split("/")[-1] not in _INVOKE_ACTIONS:
+        return None
+    if not isinstance(body, dict):
+        return None
+    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,
+) -> dict[str, object] | None:  # mutable-ok: the reply guardrails write back into
+    if not isinstance(response, dict) or not isinstance(request_data, dict):
+        return None
+    if _anthropic_invoke_body(_OBJECT_BY_REFERENCE.validate_python(request_data)) is None:
+        return None
+    return _BODY_BY_REFERENCE.validate_python(response)
+
+
 def _generic_passthrough_handler() -> "PassThroughEndpointHandler":
     """
     Fallback for non-Converse Bedrock routes (e.g. invoke). The generic
@@ -1285,6 +1337,16 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
         body: Final = data.get("data")
 
         if not _is_converse_endpoint(endpoint):
+            passthrough_request: Final = _BODY_BY_REFERENCE.validate_python(data)
+            anthropic_body: Final = _anthropic_invoke_body(passthrough_request)
+            if anthropic_body is not None:
+                await AnthropicMessagesHandler().process_input_messages(  # pyright: ignore[reportUnknownMemberType]  # takes the untyped request dict
+                    data=anthropic_body,
+                    guardrail_to_apply=guardrail_to_apply,
+                    litellm_logging_obj=litellm_logging_obj,
+                    request_data=passthrough_request,
+                )
+                return passthrough_request
             return await _generic_passthrough_handler().process_input_messages(
                 data=data,
                 guardrail_to_apply=guardrail_to_apply,
@@ -1410,6 +1472,15 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
             and not _is_converse_endpoint(endpoint)
             and not _is_invoke_stream_endpoint(endpoint)
         ):
+            anthropic_reply: Final = _anthropic_invoke_reply(response, request_data)  # pyright: ignore[reportUnknownArgumentType]  # the base signature types request_data as a bare dict
+            if anthropic_reply is not None:
+                return await AnthropicMessagesHandler().process_output_response(  # pyright: ignore[reportUnknownMemberType]  # takes the untyped request dict
+                    response=anthropic_reply,  # pyright: ignore[reportArgumentType]  # the Anthropic reply dict, written back in place
+                    guardrail_to_apply=guardrail_to_apply,
+                    litellm_logging_obj=litellm_logging_obj,
+                    user_api_key_dict=user_api_key_dict,
+                    request_data=request_data,
+                )
             return await _generic_passthrough_handler().process_output_response(
                 response=response,
                 guardrail_to_apply=guardrail_to_apply,

tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py

@@ -1394,6 +1394,7 @@ class _ScriptedGuardrail(CustomGuardrail):
         self.skip_system_message_in_guardrail = skip_system
         self.scan_only_tool_results = scan_only_tool_results
         self.received: tuple[GenericGuardrailAPIInputs, ...] = ()
+        self.request_data: tuple[dict, ...] = ()
 
     async def apply_guardrail(
         self,
@@ -1403,6 +1404,7 @@ class _ScriptedGuardrail(CustomGuardrail):
         logging_obj: LiteLLMLoggingObj | None = None,
     ) -> Mapping[str, object]:
         self.received = (*self.received, copy.deepcopy(inputs))
+        self.request_data = (*self.request_data, request_data)
         return self.reply(inputs)
 
 
@@ -1429,16 +1431,30 @@ def _converse_request(body: dict) -> dict:
 
 
 async def _sent_to_bedrock(data: dict) -> dict:
-    """The Converse body Bedrock receives when the passthrough route sends ``data``."""
+    """The body Bedrock receives when the passthrough route sends ``data``."""
     received: list[httpx.Request] = []
 
+    action = data["endpoint"].rsplit("/", 1)[-1]
+
     def respond(request: httpx.Request) -> httpx.Response:
         received.append(request)
-        reply = {
-            "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}},
-            "stopReason": "end_turn",
-            "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2},
-        }
+        reply = (
+            {
+                "id": "msg_bdrk_1",
+                "type": "message",
+                "role": "assistant",
+                "model": "claude-haiku-4-5",
+                "content": [{"type": "text", "text": "ok"}],
+                "stop_reason": "end_turn",
+                "usage": {"input_tokens": 1, "output_tokens": 1},
+            }
+            if action == "invoke"
+            else {
+                "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}},
+                "stopReason": "end_turn",
+                "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2},
+            }
+        )
         return httpx.Response(200, json=reply, request=request)
 
     router = litellm.Router(
@@ -1457,7 +1473,7 @@ async def _sent_to_bedrock(data: dict) -> dict:
     await router.allm_passthrough_route(
         model="bedrock-model",
         method="POST",
-        endpoint="/model/bedrock-model/converse",
+        endpoint=f"/model/bedrock-model/{action}",
         data=data["data"],
         client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)),
     )
@@ -3843,6 +3859,186 @@ class TestInvokeStreamGuardrailRelay:
         assert rewritten_key == received_key
 
 
+def _anthropic_invoke_data(action: str = "invoke") -> dict:
+    return {
+        "endpoint": f"model/anthropic.claude-sonnet/{action}",
+        "custom_llm_provider": "bedrock",
+        "model": "anthropic.claude-sonnet",
+        "litellm_call_id": "call-1",
+        "data": {
+            "anthropic_version": "bedrock-2023-05-31",
+            "max_tokens": 512,
+            "system": "You are a travel assistant. Card on file 4111 1111 1111 1111",
+            "tools": [
+                {
+                    "name": "find_trips",
+                    "description": "Search trips",
+                    "input_schema": {"type": "object", "properties": {"port": {"type": "string"}}},
+                }
+            ],
+            "messages": [
+                {"role": "user", "content": "Pay with 4111 1111 1111 1111"},
+                {
+                    "role": "assistant",
+                    "content": [
+                        {"type": "tool_use", "id": "toolu_1", "name": "find_trips", "input": {"port": "Miami"}}
+                    ],
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "tool_result", "tool_use_id": "toolu_1", "content": "Card 4111 1111 1111 1111 charged"}
+                    ],
+                },
+            ],
+        },
+    }
+
+
+def _masked_anthropic_invoke_body() -> dict:
+    body = _anthropic_invoke_data()["data"]
+    return {
+        **body,
+        "system": "You are a travel assistant. Card on file [CARD] 1111 1111 1111",
+        "messages": [
+            {"role": "user", "content": "Pay with [CARD] 1111 1111 1111"},
+            body["messages"][1],
+            {
+                "role": "user",
+                "content": [
+                    {"type": "tool_result", "tool_use_id": "toolu_1", "content": "Card [CARD] 1111 1111 1111 charged"}
+                ],
+            },
+        ],
+    }
+
+
+def _rows_masked(inputs: GenericGuardrailAPIInputs) -> Mapping[str, object]:
+    rows = [
+        {**row, "content": row["content"].replace("4111", "[CARD]")} if isinstance(row.get("content"), str) else row
+        for row in inputs["structured_messages"]
+    ]
+    return {"structured_messages": rows}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"])
+async def test_anthropic_invoke_masks_are_written_back_into_the_body(action: str):
+    guardrail = _ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)})
+
+    result = await BedrockPassthroughGuardrailHandler().process_input_messages(
+        _anthropic_invoke_data(action), guardrail
+    )
+
+    assert result["data"] == _masked_anthropic_invoke_body()
+
+
+@pytest.mark.asyncio
+async def test_anthropic_invoke_masks_reach_bedrock():
+    data = _anthropic_invoke_data()
+
+    await BedrockPassthroughGuardrailHandler().process_input_messages(
+        data, _ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)})
+    )
+
+    assert await _sent_to_bedrock(data) == _masked_anthropic_invoke_body()
+
+
+@pytest.mark.asyncio
+async def test_anthropic_invoke_guardrail_gets_roles_tools_and_the_passthrough_request():
+    data = _anthropic_invoke_data()
+    guardrail = _replying({})
+
+    await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail)
+
+    inputs = guardrail.received[-1]
+    assert [row["role"] for row in inputs["structured_messages"]] == ["system", "user", "assistant", "tool"]
+    assert [tool["function"]["name"] for tool in inputs["tools"]] == ["find_trips"]
+    assert inputs["model"] == data["model"]
+    assert guardrail.request_data[-1] is data
+
+
+@pytest.mark.asyncio
+async def test_anthropic_invoke_structured_message_rewrites_are_written_back():
+    result = await BedrockPassthroughGuardrailHandler().process_input_messages(
+        _anthropic_invoke_data(), _ScriptedGuardrail(_rows_masked)
+    )
+
+    assert result["data"]["messages"][0] == {
+        "role": "user",
+        "content": [{"type": "text", "text": "Pay with [CARD] 1111 1111 1111"}],
+    }
+
+
+@pytest.mark.asyncio
+async def test_invoke_body_without_anthropic_version_is_scanned_whole():
+    data = {**_anthropic_invoke_data(), "data": {"messages": [{"role": "user", "content": "hello"}]}}
+    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]
+
+
+_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 = {
+        "id": "msg_bdrk_1",
+        "type": "message",
+        "role": "assistant",
+        "content": [{"type": "text", "text": "Charged 4111 1111 1111 1111"}],
+        "stop_reason": "end_turn",
+    }
+
+    result = await BedrockPassthroughGuardrailHandler().process_output_response(
+        response=response,
+        guardrail_to_apply=_ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)}),
+        request_data=_anthropic_invoke_data(),
+    )
+
+    assert result == {**response, "content": [{"type": "text", "text": "Charged [CARD] 1111 1111 1111"}]}
+
+
 def test_sse_events_keeps_a_data_line_whose_json_holds_unicode_line_separators() -> None:
     event: Final = {
         "type": "content_block_delta",