Merge 7f0d38a586 into 9007e0fa70

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

@@ -596,9 +596,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:
@@ -610,7 +613,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(
@@ -682,7 +687,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(
@@ -690,7 +694,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

@@ -59,6 +59,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})
 _EVENT_STREAM_CONTENT_TYPE: Final = "vnd.amazon.eventstream"
 _EVENT_STREAM_MEDIA_TYPE: Final = "application/vnd.amazon.eventstream"
 
@@ -72,6 +73,33 @@ 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
+    return invoke_body
+
+
+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
@@ -1284,6 +1312,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,
@@ -1409,6 +1447,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

@@ -1393,6 +1393,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,
@@ -1402,6 +1403,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)
 
 
@@ -1428,16 +1430,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(
@@ -1456,7 +1472,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)),
     )
@@ -3840,3 +3856,144 @@ class TestInvokeStreamGuardrailRelay:
 
         assert received_key is not None and received_key.texts == ("In Paris",)
         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]
+
+
+@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"}]}