diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 817367b8844..3c712a213d4 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/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) diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 302a41a16ff..e6a11f92a0c 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/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(): diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index d2a439557fe..62f49b0047b 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/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]}]} diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index a46b56294bf..cc84ff9f247 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/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