Merge branch 'litellm_lakera_v2_gemini_contents' into litellm_passthrough_applied_guardrails_header

litellm/integrations/custom_guardrail.py

@@ -5,7 +5,7 @@ import os
 import secrets
 from collections.abc import Iterable, Mapping, Sequence
 from datetime import datetime
-from itertools import chain
+from itertools import chain, repeat
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
 
@@ -219,11 +219,16 @@ def merged_guardrail_entries(*entry_lists: Iterable[object]) -> list[object]:
     any plain name in either order, else its first plain name stays"""
     entries: Final = tuple(chain.from_iterable(entry_lists))
     settings_names: Final = tuple(_settings_entry_names(entry) for entry in entries)
-    named_with_settings: Final = frozenset[str]().union(*settings_names)
+    first_plain_index: Final = _first_indexes(
+        (entry, index) for index, entry in enumerate(entries) if isinstance(entry, str)
+    )
+    first_settings_index: Final = _first_indexes(
+        chain.from_iterable(zip(names, repeat(index)) for index, names in enumerate(settings_names))
+    )
     return [
         entry
         for index, entry in enumerate(entries)
-        if not _superseded_guardrail_entry(entries, settings_names, named_with_settings, index)
+        if not _superseded_guardrail_entry(entry, index, settings_names[index], first_plain_index, first_settings_index)
     ]
 
 
@@ -239,16 +244,20 @@ def _settings_entry_names(entry: object) -> frozenset[str]:
     return frozenset(guardrail_entry_names(entry)) if isinstance(entry, dict) else frozenset()
 
 
+def _first_indexes(pairs: Iterable[tuple[str, int]]) -> Mapping[str, int]:
+    return MappingProxyType(dict(reversed(tuple(pairs))))
+
+
 def _superseded_guardrail_entry(
-    entries: tuple[object, ...],
-    settings_names: tuple[frozenset[str], ...],
-    named_with_settings: frozenset[str],
+    entry: object,
     index: int,
+    settings_names: frozenset[str],
+    first_plain_index: Mapping[str, int],
+    first_settings_index: Mapping[str, int],
 ) -> bool:
-    entry: Final = entries[index]
     if isinstance(entry, str):
-        return entry in named_with_settings or entry in entries[:index]
-    return bool(settings_names[index]) and settings_names[index] <= frozenset[str]().union(*settings_names[:index])
+        return entry in first_settings_index or first_plain_index[entry] < index
+    return bool(settings_names) and all(first_settings_index[name] < index for name in settings_names)
 
 
 def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None = None) -> bool:

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -57,6 +57,7 @@ from litellm.integrations.custom_guardrail import (
     CustomGuardrail,
     guardrail_entry_names,
     guardrail_request_data_with_streaming,
+    merged_guardrail_entries,
     without_server_streaming_classification,
 )
 from litellm.integrations.custom_logger import CustomLogger
@@ -959,19 +960,27 @@ async def _inherit_auth_guardrails(
 
     The endpoint's configured guardrails win. Without any, and with ``inherit`` and a JSON body, the key, team,
     project, policy and request-body guardrails are merged into ``parsed_body["metadata"]`` for the pre-call hook,
-    and returned only when at least one of them runs post_call, so the response is relayed untouched otherwise.
+    with the ones the client named in ``metadata.guardrails`` kept after them, and returned only when at least one
+    of them runs post_call, so the response is relayed untouched otherwise.
     """
     if configured_guardrails or not inherit or not parsed_body:
         return configured_guardrails
     if not isinstance(parsed_body.get("metadata"), dict):
         parsed_body["metadata"] = {}  # rebind-ok: the merge below writes into the live request body's metadata
+    client_named: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails")
     await move_guardrails_to_metadata(
         data=parsed_body,
         _metadata_variable_name="metadata",
         user_api_key_dict=user_api_key_dict,
     )
-    merged: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails")
-    entries: Final = _GUARDRAIL_ENTRIES.validate_python(merged) if isinstance(merged, list) else []
+    metadata: Final = parsed_body["metadata"]
+    merged: Final = (as_str_mapping(metadata) or {}).get("guardrails")
+    entries: Final = merged_guardrail_entries(
+        _GUARDRAIL_ENTRIES.validate_python(merged) if isinstance(merged, list) else (),
+        _GUARDRAIL_ENTRIES.validate_python(client_named) if isinstance(client_named, list) else (),
+    )
+    if entries and isinstance(metadata, dict):
+        metadata["guardrails"] = entries
     return entries if entries and _named_guardrail_runs_post_call(parsed_body, entries) else None
 
 

tests/integration/authorization/test_passthrough_auth_guardrails.py

@@ -365,3 +365,19 @@ def test_key_apply_guardrail_post_call_scans_the_provider_reply(rig: AuthGuardra
 
         assert response.status_code == 200, response.text
         assert [b'"Hi' in request.body for request in rig.generic_post_call_vendor.drain()] == [True]
+
+
+@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE])
+def test_guardrail_the_client_names_in_metadata_still_runs_for_a_guarded_key(rig: AuthGuardrailRig, route: str) -> None:
+    gateway: Final = rig.owned.gateway
+    with gateway.scenario() as scenario:
+        key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]})
+        _drain(rig)
+
+        response: Final = _send(
+            gateway, route, {**_route_body(route, BLOCK_MARKER), "metadata": {"guardrails": [GUARDRAIL]}}, key
+        )
+
+        assert response.status_code == 400, response.text
+        assert [BLOCK_MARKER.encode() in request.body for request in rig.vendor.drain()] == [True]
+        assert rig.upstreams[route].drain() == ()

tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py

@@ -11943,6 +11943,28 @@ class TestInheritAuthGuardrails:
         assert pre_call[0]["metadata"]["guardrails"] == ["key-guard", "body-guard"]
         assert "guardrails" not in pre_call[0]
 
+    @pytest.mark.asyncio
+    async def test_guardrails_the_client_names_in_metadata_still_run_for_a_guarded_key(self):
+        auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})
+
+        pre_call, post_call = await self._run(
+            self._chat_body(metadata={"guardrails": ["client-pre", "client-post"]}),
+            inherit=True,
+            user_api_key_dict=auth,
+            guardrails=(
+                self._guardrail("key-guard", GuardrailEventHooks.pre_call),
+                self._guardrail("client-post", GuardrailEventHooks.post_call),
+            ),
+        )
+
+        assert pre_call[0]["metadata"] == {"guardrails": ["key-guard", "client-pre", "client-post"]}
+        post_call.assert_awaited_once()
+        assert post_call.call_args.kwargs["data"]["metadata"]["guardrails"] == [
+            "key-guard",
+            "client-pre",
+            "client-post",
+        ]
+
     @pytest.mark.asyncio
     async def test_empty_request_body_guardrails_keep_the_key_guardrails(self):
         auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})