diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c7efbf9dacf..6e091391842 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -815,6 +815,29 @@ def _drop_response_steps_that_did_not_run( _drop_applied_guardrails_in_place(data.get(bucket_key), not_run) +def _add_applied_guardrails_in_place( + data: dict[str, object], # mutable-ok: the names are recorded in the live request body's metadata + names: Sequence[str], +) -> None: + """Adds each name the way add_guardrail_to_applied_guardrails_header does, in time linear in the names""" + from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header + + if not names: + return + metadata: Final = cast("dict[str, object]", get_or_create_metadata_bucket(data)[1]) # cast-ok: str-keyed bucket + if "applied_guardrails" not in metadata: + metadata["applied_guardrails"] = list(dict.fromkeys(names)) + return + applied: Final = metadata["applied_guardrails"] + if not isinstance(applied, list): + for name in names: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=name) + return + entries: Final = cast("list[object]", applied) # cast-ok: entries are only compared with names + present: Final = frozenset(name for name in entries if isinstance(name, str)) + entries.extend(dict.fromkeys(name for name in names if name not in present)) + + def _drop_applied_guardrails_in_place(bucket: object, names: frozenset[str]) -> None: if not isinstance(bucket, dict): return @@ -2831,10 +2854,6 @@ class ProxyLogging: policy pipeline steps that have not run. With ``skip_guardrails`` the pipelines run on another request, so every step is added. """ - from litellm.proxy.common_utils.callback_utils import ( - add_guardrail_to_applied_guardrails_header, - ) - metadata_standard: Final = data.get("metadata") or {} metadata_litellm: Final = data.get("litellm_metadata") or {} @@ -2854,18 +2873,21 @@ class ProxyLogging: if not isinstance(applied_guardrails, list): applied_guardrails = [] - not_run: Final = ( - frozenset[str]() - if skip_guardrails - else _request_steps_that_have_not_run(cast(dict[str, object], data)) # cast-ok: request payload + request: Final = cast("dict[str, object]", data) # cast-ok: request payload + not_run: Final = frozenset[str]() if skip_guardrails else _request_steps_that_have_not_run(request) + already_applied: Final = frozenset( + name + for name in cast("Sequence[object]", applied_guardrails) # cast-ok: checked to be a list above + if isinstance(name, str) + ) + _add_applied_guardrails_in_place( + request, + [ + name + for name in guardrails_in_metadata + if isinstance(name, str) and name not in already_applied and name not in not_run + ], ) - for guardrail_name in guardrails_in_metadata: - if ( - isinstance(guardrail_name, str) - and guardrail_name not in applied_guardrails - and guardrail_name not in not_run - ): - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_name) async def _maybe_execute_pipelines( self, 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 5fd40696224..a46b56294bf 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -11,6 +11,7 @@ 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 @@ -275,54 +276,43 @@ async def test_process_guardrail_callback_enriches_and_reraises_http_exception( # --------------------------------------------------------------------------- -def test_process_guardrail_metadata_calls_header_helper(proxy_logging, monkeypatch): - calls: List[Dict[str, Any]] = [] - - def fake_add(request_data, guardrail_name): - calls.append({"data": request_data, "name": guardrail_name}) - - from litellm.proxy.common_utils import callback_utils - - monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) - data = {"metadata": {"guardrails": ["g1", "g2"]}} +def test_process_guardrail_metadata_adds_each_guardrail_once_in_order(proxy_logging): + data = {"metadata": {"guardrails": ["g1", "g2", "g1"]}} proxy_logging._process_guardrail_metadata(data) - snapshot = { - "call_count": len(calls), - "first_name": calls[0]["name"], - "second_name": calls[1]["name"], - "data_passed_is_input": all(c["data"] is data for c in calls), - } - assert snapshot == { - "call_count": 2, - "first_name": "g1", - "second_name": "g2", - "data_passed_is_input": True, - } + assert data["metadata"]["applied_guardrails"] == ["g1", "g2"] -def test_process_guardrail_metadata_skips_already_applied(proxy_logging, monkeypatch): - calls: List[str] = [] - - def fake_add(request_data, guardrail_name): - calls.append(guardrail_name) - - from litellm.proxy.common_utils import callback_utils - - monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) +def test_process_guardrail_metadata_skips_already_applied(proxy_logging): data = {"metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g1"]}} proxy_logging._process_guardrail_metadata(data) - assert calls == ["g2"] + assert data["metadata"]["applied_guardrails"] == ["g1", "g2"] -def test_process_guardrail_metadata_no_metadata_is_noop(proxy_logging, monkeypatch): - from litellm.proxy.common_utils import callback_utils +def test_process_guardrail_metadata_records_in_the_litellm_metadata_bucket(proxy_logging): + data = {"litellm_metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g2"]}} + proxy_logging._process_guardrail_metadata(data) + assert data == {"litellm_metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g2", "g1"]}} - monkeypatch.setattr( - callback_utils, - "add_guardrail_to_applied_guardrails_header", - MagicMock(side_effect=AssertionError("should not be called")), - ) - proxy_logging._process_guardrail_metadata({}) + +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() + proxy_logging._process_guardrail_metadata(data) + assert time.perf_counter() - started < 2 + assert data["metadata"]["applied_guardrails"] == names + + +def test_process_guardrail_metadata_no_metadata_is_noop(proxy_logging): + data: Dict[str, Any] = {} + proxy_logging._process_guardrail_metadata(data) + assert data == {} + + +def test_process_guardrail_metadata_with_only_settings_entries_leaves_data_unchanged(proxy_logging): + data = {"metadata": {"guardrails": [{"g1": {"extra_body": {}}}]}} + proxy_logging._process_guardrail_metadata(data) + assert data == {"metadata": {"guardrails": [{"g1": {"extra_body": {}}}]}} def test_process_guardrail_metadata_invalid_data_raises(proxy_logging):