Merge f06c6c94ae into f285229b51

litellm/proxy/guardrails/_content_utils.py

@@ -8,9 +8,13 @@ skip the other shapes — these helpers normalise that so every hook sees
 every text fragment.
 """
 
+import json
 from collections.abc import Callable, Iterator, Mapping, Sequence
 from typing import Any, Final
 
+from pydantic import JsonValue
+from pydantic_core import to_jsonable_python
+
 # Call types whose body carries free-form chat / prompt text that
 # text-content guardrails (banned keywords, content moderation, secret
 # detection, …) should inspect. The proxy ingress passes ``route_type``
@@ -307,3 +311,12 @@ def build_inspection_messages(data: dict[str, Any]) -> list[dict[str, str]]:
         role = message.get("role", "user") or "user"
         flattened.append({"role": role, "content": text})
     return flattened
+
+
+def as_json_value(value: object) -> JsonValue:
+    """Round-trips through the stdlib codec because pydantic's serializer turns anything nested
+    past 254 levels into "...", while this keeps about the depth the proxy's request parser accepts"""
+    parsed: Final[JsonValue] = json.loads(  # pyright: ignore[reportAny]  # untyped stdlib parse of json.dumps output
+        json.dumps(value, default=to_jsonable_python)
+    )
+    return parsed

litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py

@@ -11,6 +11,8 @@ from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional
 
 import httpx
+from pydantic import JsonValue
+from typing_extensions import TypeIs
 
 from litellm._logging import verbose_proxy_logger
 from litellm._version import version as litellm_version
@@ -20,9 +22,11 @@ from litellm.integrations.custom_guardrail import (
     log_guardrail_information,
 )
 from litellm.llms.custom_httpx.http_handler import (
+    AsyncHTTPHandler,
     get_async_httpx_client,
     httpxSpecialProvider,
 )
+from litellm.proxy.guardrails._content_utils import as_json_value
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
 from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
@@ -30,6 +34,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import
     GenericGuardrailAPIRequest,
     GenericGuardrailAPIResponse,
     GuardrailToolParam,
+    structured_messages_from_json,
 )
 from litellm.types.utils import GenericGuardrailAPIInputs
 
@@ -150,6 +155,28 @@ def _extract_inbound_headers(
     return None
 
 
+def _is_part_list(value: object) -> TypeIs[list[object]]:  # guard-ok: trivial isinstance narrowing
+    return isinstance(value, list)
+
+
+def _row_as_sent(dumped: JsonValue, caller: Mapping[str, object]) -> JsonValue:
+    """The request model validates list content lazily and dumps a list holding
+    any part it rejects as [], so such a row is sent with the caller's content."""
+    caller_content: Final = caller.get("content")
+    if not isinstance(dumped, dict) or not _is_part_list(caller_content):
+        return dumped
+    dumped_content: Final = dumped.get("content")
+    if isinstance(dumped_content, list) and len(dumped_content) == len(caller_content):
+        return dumped
+    return {**dumped, "content": as_json_value(caller_content)}
+
+
+def _rows_as_sent(dumped_rows: JsonValue, caller_rows: Sequence[Mapping[str, object]] | None) -> JsonValue:
+    if caller_rows is None or not isinstance(dumped_rows, list):
+        return dumped_rows
+    return [_row_as_sent(dumped, caller) for dumped, caller in zip(dumped_rows, caller_rows, strict=True)]
+
+
 def _structured_rows_to_write_back(
     original_rows: Sequence[AllMessageValues] | None,
     shown_rows: Sequence[AllMessageValues] | None,
@@ -204,9 +231,12 @@ class GenericGuardrailAPI(CustomGuardrail):
         streaming_end_of_stream_only: bool | None = None,
         streaming_sampling_rate: int | None = None,
         streaming_transform_mode: Literal["block_only", "incremental_diff"] | None = None,
+        async_handler: AsyncHTTPHandler | None = None,
         **kwargs,
     ):
-        self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
+        self.async_handler = async_handler or get_async_httpx_client(
+            llm_provider=httpxSpecialProvider.GuardrailCallback
+        )
         self.headers = headers or {}
         self.extra_headers = extra_headers or []
 
@@ -470,12 +500,14 @@ class GenericGuardrailAPI(CustomGuardrail):
             )
 
             headers: Final = self._build_request_headers()
+            # The model's list content is a lazy iterator that this dump consumes, so it cannot be read again
+            dumped: Final[Mapping[str, JsonValue]] = guardrail_request.model_dump(mode="json")
+            sent_messages: Final = _rows_as_sent(dumped.get("structured_messages"), structured_messages)
+            request_json: Final = {**dumped, "structured_messages": sent_messages}  # mutable-ok: JSON POST body
 
-            # Make the API request
-            # Use mode="json" to ensure all iterables are converted to lists
             response: Final = await self.async_handler.post(
                 url=self.api_base,
-                json=guardrail_request.model_dump(mode="json"),
+                json=request_json,
                 headers=headers,
             )
 
@@ -503,7 +535,7 @@ class GenericGuardrailAPI(CustomGuardrail):
                 images=images,
                 tools=tools,
                 structured_messages=structured_messages,
-                shown_messages=guardrail_request.structured_messages,
+                shown_messages=structured_messages_from_json(sent_messages),
                 guardrail_response=guardrail_response,
             )
 

litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py

@@ -159,7 +159,7 @@ def coerce_stream_holdback_value(value: Any) -> int:
         return 0
 
 
-def structured_messages_from_response(value: object) -> Sequence[AllMessageValues] | None:
+def structured_messages_from_json(value: object) -> Sequence[AllMessageValues] | None:
     if not isinstance(value, list):
         return None
     if not all(isinstance(message, Mapping) and isinstance(message.get("role"), str) for message in value):
@@ -212,5 +212,5 @@ class GenericGuardrailAPIResponse:
             images=data.get("images"),
             tools=data.get("tools"),
             stream_holdback_chars=stream_holdback_chars,
-            structured_messages=structured_messages_from_response(data.get("structured_messages")),
+            structured_messages=structured_messages_from_json(data.get("structured_messages")),
         )

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py

@@ -5,16 +5,23 @@ This test file tests the Generic Guardrail API implementation,
 specifically focusing on metadata extraction and passing.
 """
 
+import json
 import os
+from collections.abc import Callable, Mapping
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
 import pytest
+from pydantic import JsonValue
 
 import litellm
 from litellm import ModelResponse
 from litellm._version import version as litellm_version
 from litellm.exceptions import GuardrailRaisedException, Timeout
+from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
     GenericGuardrailAPI,
@@ -22,6 +29,8 @@ from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
 from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import (
     _HEADER_PRESENT_PLACEHOLDER,
 )
+from litellm.types.llms.anthropic import AllAnthropicMessageValues
+from litellm.types.llms.openai import AllMessageValues, ChatCompletionImageObject
 from litellm.types.utils import Choices, Message
 
 
@@ -721,6 +730,287 @@ class TestStructuredMessagesInResponse:
         assert guardrailed_inputs["texts"] == ["[REDACTED]"]
 
 
+_SSN: Final = "123-45-6789"
+
+
+def _image_part() -> ChatCompletionImageObject:
+    return {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}
+
+
+GuardrailAnswer = Callable[[Mapping[str, JsonValue]], Mapping[str, JsonValue]]
+
+
+def _guardrail_answering(answer: GuardrailAnswer) -> GenericGuardrailAPI:
+    def serve(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(200, json=answer(json.loads(request.content)))
+
+    return GenericGuardrailAPI(
+        api_base="https://guardrail.test/beta/litellm_basic_guardrail_api",
+        guardrail_name="pii-masker",
+        event_hook="pre_call",
+        default_on=True,
+        async_handler=AsyncHTTPHandler(transport=httpx.MockTransport(serve)),
+    )
+
+
+def _masked(text: str) -> str:
+    return text.replace(_SSN, "[SSN]")
+
+
+def _echo_every_row_and_mask_texts(request_json: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
+    return {
+        "action": "GUARDRAIL_INTERVENED",
+        "structured_messages": request_json["structured_messages"],
+        "texts": [_masked(text) for text in request_json["texts"]],
+    }
+
+
+def _echo_first_row_and_mask_the_rest(request_json: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
+    first_row, *other_rows = request_json["structured_messages"]
+    return {
+        "action": "GUARDRAIL_INTERVENED",
+        "structured_messages": [first_row, *({**row, "content": _masked(row["content"])} for row in other_rows)],
+    }
+
+
+def _masked_part(part: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
+    return {**part, "text": _masked(part["text"])} if part["type"] == "text" else part
+
+
+def _masked_row(row: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
+    content: Final = row["content"]
+    if isinstance(content, str):
+        return {**row, "content": _masked(content)}
+    return {**row, "content": [_masked_part(part) for part in content]}
+
+
+def _mask_every_row_and_text(request_json: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
+    return {
+        "action": "GUARDRAIL_INTERVENED",
+        "texts": [_masked(text) for text in request_json["texts"]],
+        "structured_messages": [_masked_row(row) for row in request_json["structured_messages"]],
+    }
+
+
+def _guarded_text_part() -> Mapping[str, JsonValue]:
+    return {"type": "guarded_text", "text": "keep this guarded"}
+
+
+def _pdf_document_part() -> Mapping[str, JsonValue]:
+    return {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0="}}
+
+
+def _nested(depth: int) -> JsonValue:
+    return {"leaf": "x"} if depth == 0 else {"nested": _nested(depth - 1)}
+
+
+async def _llm_bound_messages(guardrail: GenericGuardrailAPI, messages: list[AllMessageValues]) -> object:
+    data: Final = await OpenAIChatCompletionsHandler().process_input_messages(
+        data={"model": "gpt-5.6", "messages": messages}, guardrail_to_apply=guardrail
+    )
+    return data["messages"]
+
+
+async def _structured_messages_posted_for(rows: list[AllMessageValues]) -> list[JsonValue]:
+    posted: Final[list[JsonValue]] = []  # mutable-ok: records what the endpoint received
+
+    def record(request_json: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
+        posted.append(request_json["structured_messages"])
+        return {"action": "NONE"}
+
+    await _guardrail_answering(record).apply_guardrail(
+        inputs={"texts": ["hi"], "structured_messages": rows}, request_data={}, input_type="request"
+    )
+    return posted
+
+
+async def _llm_bound_anthropic_messages(
+    guardrail: GenericGuardrailAPI, messages: list[AllAnthropicMessageValues]
+) -> object:
+    data: Final = await AnthropicMessagesHandler().process_input_messages(
+        data={"model": "claude-opus-5-5", "max_tokens": 64, "messages": messages}, guardrail_to_apply=guardrail
+    )
+    return data["messages"]
+
+
+class TestEchoedRowsReachingTheLLM:
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        ("messages", "expected"),
+        [
+            (
+                [{"role": "user", "content": [{"type": "text", "text": f"my ssn is {_SSN}"}, _image_part()]}],
+                [{"role": "user", "content": [{"type": "text", "text": "my ssn is [SSN]"}, _image_part()]}],
+            ),
+            (
+                [
+                    {"role": "system", "content": "You are helpful."},
+                    {"role": "user", "content": [{"type": "text", "text": f"ssn {_SSN}"}, _image_part()]},
+                    {"role": "user", "content": f"again {_SSN}"},
+                ],
+                [
+                    {"role": "system", "content": "You are helpful."},
+                    {"role": "user", "content": [{"type": "text", "text": "ssn [SSN]"}, _image_part()]},
+                    {"role": "user", "content": "again [SSN]"},
+                ],
+            ),
+            (
+                [{"role": "user", "content": f"my ssn is {_SSN}"}],
+                [{"role": "user", "content": "my ssn is [SSN]"}],
+            ),
+        ],
+        ids=["multipart", "multipart_among_string_rows", "string_only"],
+    )
+    async def test_every_row_echoed_applies_the_masked_texts(
+        self, messages: list[AllMessageValues], expected: list[AllMessageValues]
+    ) -> None:
+        guardrail: Final = _guardrail_answering(_echo_every_row_and_mask_texts)
+
+        llm_bound: Final = await _llm_bound_messages(guardrail, messages)
+
+        assert llm_bound == expected, "an unchanged echo of every row must leave the rewrite to texts"
+
+    @pytest.mark.asyncio
+    async def test_every_anthropic_content_block_row_echoed_applies_the_masked_texts(self) -> None:
+        guardrail: Final = _guardrail_answering(_echo_every_row_and_mask_texts)
+
+        llm_bound: Final = await _llm_bound_anthropic_messages(
+            guardrail,
+            [
+                {
+                    "role": "user",
+                    "content": [{"type": "text", "text": f"my ssn is {_SSN}"}, {"type": "text", "text": "ok"}],
+                }
+            ],
+        )
+
+        assert llm_bound == [
+            {"role": "user", "content": [{"type": "text", "text": "my ssn is [SSN]"}, {"type": "text", "text": "ok"}]}
+        ], "an unchanged echo of every content block row must leave the rewrite to texts"
+
+    @pytest.mark.asyncio
+    async def test_an_echoed_multipart_row_is_restored_to_the_callers_row(self) -> None:
+        guardrail: Final = _guardrail_answering(_echo_first_row_and_mask_the_rest)
+
+        llm_bound: Final = await _llm_bound_messages(
+            guardrail,
+            [
+                {"role": "user", "name": "pat", "content": [{"type": "text", "text": "what is this?"}, _image_part()]},
+                {"role": "user", "content": f"my ssn is {_SSN}"},
+            ],
+        )
+
+        assert llm_bound == [
+            {"role": "user", "name": "pat", "content": [{"type": "text", "text": "what is this?"}, _image_part()]},
+            {"role": "user", "content": "my ssn is [SSN]"},
+        ], "the echoed row must keep the caller's keys the request model drops"
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "unvalidated_part", [_guarded_text_part, _pdf_document_part], ids=["guarded_text", "pdf_document"]
+    )
+    async def test_a_row_holding_a_part_the_request_model_rejects_is_masked(
+        self, unvalidated_part: Callable[[], Mapping[str, JsonValue]]
+    ) -> None:
+        guardrail: Final = _guardrail_answering(_mask_every_row_and_text)
+
+        llm_bound: Final = await _llm_bound_messages(
+            guardrail,
+            [
+                {"role": "user", "content": [{"type": "text", "text": f"my ssn is {_SSN}"}, unvalidated_part()]},
+                {"role": "user", "content": f"also {_SSN}"},
+            ],
+        )
+
+        assert llm_bound == [
+            {"role": "user", "content": [{"type": "text", "text": "my ssn is [SSN]"}, unvalidated_part()]},
+            {"role": "user", "content": "also [SSN]"},
+        ], "the guardrail must see the whole row so its masking of it reaches the LLM"
+
+    @pytest.mark.asyncio
+    async def test_an_echoed_row_holding_a_part_the_request_model_rejects_is_restored_to_the_callers_row(
+        self,
+    ) -> None:
+        guardrail: Final = _guardrail_answering(_echo_first_row_and_mask_the_rest)
+
+        llm_bound: Final = await _llm_bound_messages(
+            guardrail,
+            [
+                {"role": "user", "name": "pat", "content": [{"type": "text", "text": "hi"}, _guarded_text_part()]},
+                {"role": "user", "content": f"my ssn is {_SSN}"},
+            ],
+        )
+
+        assert llm_bound == [
+            {"role": "user", "name": "pat", "content": [{"type": "text", "text": "hi"}, _guarded_text_part()]},
+            {"role": "user", "content": "my ssn is [SSN]"},
+        ], "the echoed row must keep the caller's keys the request model drops"
+
+    @pytest.mark.asyncio
+    async def test_rows_the_request_model_accepts_are_posted_as_it_dumps_them(self) -> None:
+        posted: Final = await _structured_messages_posted_for(
+            [
+                {
+                    "role": "user",
+                    "name": "pat",
+                    "content": [
+                        {"type": "text", "text": "hi"},
+                        {**_image_part(), "cache_control": {"type": "ephemeral"}},
+                    ],
+                },
+                {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [
+                        {
+                            "id": "call_1",
+                            "type": "function",
+                            "function": {"name": "lookup", "arguments": "{}"},
+                            "index": 0,
+                        }
+                    ],
+                },
+                {"role": "tool", "tool_call_id": "call_1", "content": "done"},
+            ]
+        )
+
+        assert posted == [
+            [
+                {"role": "user", "content": [{"type": "text", "text": "hi"}, _image_part()]},
+                {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [
+                        {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}
+                    ],
+                },
+                {"role": "tool", "tool_call_id": "call_1", "content": "done"},
+            ]
+        ], "rows the request model dumps in full must reach the guardrail exactly as before"
+
+    @pytest.mark.asyncio
+    async def test_a_row_holding_a_part_the_request_model_rejects_is_posted_with_the_callers_content(self) -> None:
+        posted: Final = await _structured_messages_posted_for(
+            [{"role": "user", "name": "pat", "content": [{"type": "text", "text": "hi"}, _pdf_document_part()]}]
+        )
+
+        assert posted == [[{"role": "user", "content": [{"type": "text", "text": "hi"}, _pdf_document_part()]}]], (
+            "only the content the request model emptied is taken from the caller, the row keys stay as dumped"
+        )
+
+    @pytest.mark.asyncio
+    async def test_a_deeply_nested_part_the_proxy_accepts_is_posted_in_full(self) -> None:
+        deep_document: Final = {"type": "document", "source": _nested(300)}
+
+        posted: Final = await _structured_messages_posted_for(
+            [{"role": "user", "content": [{"type": "text", "text": "hi"}, deep_document]}]
+        )
+
+        assert posted == [[{"role": "user", "content": [{"type": "text", "text": "hi"}, deep_document]}]], (
+            "content nested as deep as the proxy's own JSON parser allows must still reach the guardrail"
+        )
+
+
 class TestImageSupport:
     """Test image handling in guardrail requests"""
 

tests/test_litellm/proxy/guardrails/test_content_utils.py

@@ -2,6 +2,7 @@
 
 from litellm.proxy.guardrails._content_utils import (
     apply_redacted_messages_back,
+    as_json_value,
     build_inspection_messages,
     has_non_string_content,
     is_non_conversational_call_type,
@@ -741,3 +742,14 @@ def test_is_non_conversational_call_type_defaults_to_inspecting_unknown_call_typ
     """A call type this module has never heard of must still be inspected —
     failing closed is the point of the deny-list."""
     assert is_non_conversational_call_type("some_future_call_type") is False
+
+
+def _nested(depth: int) -> dict[str, object]:
+    return {"leaf": "x"} if depth == 0 else {"nested": _nested(depth - 1)}
+
+
+def test_as_json_value_keeps_content_nested_past_the_pydantic_serializer_limit_and_decodes_bytes():
+    assert as_json_value([{"type": "document", "source": _nested(600), "data": b"raw"}, ("a", 1)]) == [
+        {"type": "document", "source": _nested(600), "data": "raw"},
+        ["a", 1],
+    ], "nothing may be truncated to '...' and non-JSON types must become their JSON form"