diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index dfe508f3630..3b444519fc6 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/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: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 1528579770d..e5281f452e8 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/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 diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index c0fa3a233b8..85c31a5e2bc 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/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() == () diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index d007efdc93e..1d9a3be5e2f 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/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"]})