fix(proxy): skip reply metadata by position, not by field name at any depth

A tool call's arguments are the model's own output, so a field named id or status in them stays in the post-call scan, as do those fields on a custom endpoint's own JSON. Timing asserts in the merge and metadata tests are replaced with deterministic checks.

litellm/llms/pass_through/guardrail_translation/handler.py

@@ -7,6 +7,7 @@ to extract specific fields for guardrail processing.
 """
 
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping, Sequence
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, Optional, Protocol, runtime_checkable
 
 from pydantic import TypeAdapter
@@ -29,44 +30,60 @@ if TYPE_CHECKING:
     from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
 
 
-# Strings in a provider's JSON response that name or encode something rather than carry the reply's text.
-_NON_TEXT_RESPONSE_FIELDS: Final = frozenset(
+# The fields of a provider's JSON reply that name or encode something rather than carry its text, by the field they sit
+# under. Anything outside these positions, such as a tool call's arguments, is the model's own output and is read whole.
+_REPLY_TOP_LEVEL_FIELDS: Final = frozenset(
     (
-        "b64_json",
-        "bytesBase64Encoded",
-        "embedding",
-        "finishReason",
-        "finish_reason",
         "id",
         "model",
         "modelVersion",
         "object",
         "responseId",
-        "role",
         "service_tier",
         "status",
         "system_fingerprint",
-        "thoughtSignature",
-        "type",
     )
 )
+_REPLY_NESTED_FIELDS: Final = MappingProxyType(
+    {
+        "candidates": frozenset(("finishReason",)),
+        "choices": frozenset(("finish_reason",)),
+        "content": frozenset(("role", "type")),
+        "data": frozenset(("b64_json", "embedding", "object")),
+        "message": frozenset(("role",)),
+        "output": frozenset(("id", "role", "status", "type")),
+        "parts": frozenset(("thoughtSignature",)),
+        "predictions": frozenset(("bytesBase64Encoded",)),
+        "tool_calls": frozenset(("id", "type")),
+    }
+)
+_REPLY_BODY_FIELDS: Final = frozenset(("candidates", "choices", "content", "data", "output", "predictions"))
 
 
 def _response_text(response: Mapping[str, object]) -> str:
-    """The text a JSON response carries, without its timestamps, usage counts, ids and vectors, which a guardrail
-    would otherwise read as phone or card numbers."""
-    pending: Final[list[object]] = [response]  # mutable-ok: local to this walk
+    """The strings a JSON response carries, without the numbers around them or a provider reply's ids and encoded data,
+    which a guardrail would otherwise read as phone or card numbers."""
+    top_level: Final = _REPLY_TOP_LEVEL_FIELDS if _REPLY_BODY_FIELDS & response.keys() else None
+    pending: Final[list[tuple[object, frozenset[str] | None]]] = [(response, top_level)]  # mutable-ok: stack
     texts: Final[list[str]] = []  # mutable-ok: local to this walk
     while pending:
-        value = pending.pop()
+        value, skipped = pending.pop()
         fields = as_str_mapping(value)
         items = as_object_list(value)
         if isinstance(value, str):
             texts.append(value)
         elif fields is not None:
-            pending.extend(reversed([item for name, item in fields.items() if name not in _NON_TEXT_RESPONSE_FIELDS]))
+            pending.extend(
+                reversed(
+                    [
+                        (item, None if skipped is None else _REPLY_NESTED_FIELDS.get(name))
+                        for name, item in fields.items()
+                        if skipped is None or name not in skipped
+                    ]
+                )
+            )
         elif items is not None:
-            pending.extend(reversed(items))
+            pending.extend((item, skipped) for item in reversed(items))
     return "\n".join(text for text in texts if text)
 
 

tests/unit/integrations/test_custom_guardrail.py

@@ -4,10 +4,8 @@ import datetime as dt
 import json
 import pickle
 import threading
-import time
-from collections.abc import Callable
 from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, patch
 
 import pytest
 
@@ -3804,22 +3802,32 @@ def test_default_on_guardrail() -> None:
     )
 
 
-def test_merged_guardrail_entries_stays_close_to_a_plain_dedup_on_a_large_list():
-    from litellm.integrations.custom_guardrail import merged_guardrail_entries
+def test_merged_guardrail_entries_dedupes_an_all_plain_list_without_walking_it():
+    from litellm.integrations import custom_guardrail
 
-    names: Final = [f"g{index}" for index in range(300_000)]
+    names: Final = [f"g{index % 500}" for index in range(1000)]
 
-    def _best_cpu_seconds(run: Callable[[], object]) -> float:
-        def _timed() -> float:
-            started: Final = time.process_time()
-            run()
-            return time.process_time() - started
+    with patch.object(custom_guardrail, "guardrail_entry_names") as entry_names:
+        merged: Final = custom_guardrail.merged_guardrail_entries(names)
 
-        return min(_timed() for _ in range(3))
+    assert merged == [f"g{index}" for index in range(500)]
+    entry_names.assert_not_called()
 
-    assert _best_cpu_seconds(lambda: merged_guardrail_entries(names)) < 25 * _best_cpu_seconds(
-        lambda: list(dict.fromkeys(names))
-    )
+
+def test_merged_guardrail_entries_reads_each_entry_of_a_mixed_list_a_fixed_number_of_times():
+    from litellm.integrations import custom_guardrail
+
+    entries: Final = [*(f"g{index}" for index in range(1000)), {"g1": {"threshold": 1}}]
+
+    with patch.object(
+        custom_guardrail, "guardrail_entry_names", wraps=custom_guardrail.guardrail_entry_names
+    ) as entry_names:
+        merged: Final = custom_guardrail.merged_guardrail_entries(entries)
+
+    assert "g1" not in merged
+    assert merged[-1] == {"g1": {"threshold": 1}}
+    assert len(merged) == 1000
+    assert entry_names.call_count <= 2 * len(entries)
 
 
 def test_merged_guardrail_entries_names_each_guardrail_once_preferring_its_first_entry_with_settings():

tests/unit/llms/pass_through/guardrail_translation/test_handler.py

@@ -261,6 +261,21 @@ async def test_post_call_scan_still_reads_every_string_a_reply_carries():
     assert await _scanned_response_texts(response) == ('call 555-123-4567\n{"a": "b"}\ncard 4111 1111 1111 1111',)
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "response",
+    [
+        {"candidates": [{"content": {"parts": [{"functionCall": {"args": {"id": "alice@example.com"}}}]}}]},
+        {"content": [{"input": {"id": "alice@example.com", "message": {"role": "alice@example.com"}}}]},
+        {"choices": [{"message": {"tool_calls": [{"function": {"nested": {"type": "alice@example.com"}}}]}}]},
+        {"id": "alice@example.com", "result": {"status": "alice@example.com"}},
+    ],
+)
+async def test_post_call_scan_reads_a_field_named_like_the_envelope_in_what_the_model_or_api_wrote(response):
+    assert all("alice@example.com" in text for text in await _scanned_response_texts(response))
+    assert await _scanned_response_texts(response) != ()
+
+
 @pytest.mark.asyncio
 async def test_post_call_scan_skips_a_reply_with_no_text():
     response: Final = {"object": "list", "data": [{"object": "embedding", "embedding": [0.4111111111111111]}]}

tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py

@@ -11,7 +11,6 @@ from __future__ import annotations
 
 import asyncio
 import json
-import time
 from collections.abc import AsyncGenerator, Iterator
 from contextlib import AbstractContextManager, nullcontext
 from copy import deepcopy
@@ -294,12 +293,14 @@ def test_process_guardrail_metadata_records_in_the_litellm_metadata_bucket(proxy
     assert data == {"litellm_metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g2", "g1"]}}
 
 
-def test_process_guardrail_metadata_handles_many_guardrail_names_quickly(proxy_logging):
-    names = [f"guard-{i}" for i in range(50_000)]
-    data = {"metadata": {"guardrails": names, "applied_guardrails": []}}
-    started = time.perf_counter()
+def test_process_guardrail_metadata_does_not_search_the_applied_list_per_name(proxy_logging):
+    class _NoMembershipChecks(list):
+        def __contains__(self, item):
+            raise AssertionError("searched the applied list")
+
+    names = [f"guard-{i}" for i in range(100)]
+    data = {"metadata": {"guardrails": names, "applied_guardrails": _NoMembershipChecks()}}
     proxy_logging._process_guardrail_metadata(data)
-    assert time.perf_counter() - started < 2
     assert data["metadata"]["applied_guardrails"] == names