Merge branch 'litellm_passthrough_auth_guardrails' into litellm_lakera_v2_gemini_contents

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,

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):