diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 3b444519fc6..965cfbb9ee8 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, repeat +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args @@ -211,55 +211,32 @@ def guardrail_request_data_with_streaming( } -_GUARDRAIL_ENTRY_WITH_SETTINGS: Final = TypeAdapter(Mapping[str, object]) - - def merged_guardrail_entries(*entry_lists: Iterable[object]) -> list[object]: """Merges guardrail lists so each guardrail is named once: the first ``{name: settings}`` entry naming it wins over 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) - first_plain_index: Final = _first_indexes( - (entry, index) for index, entry in enumerate(entries) if isinstance(entry, str) + named_with_settings: Final = frozenset( + chain.from_iterable(guardrail_entry_names(entry) for entry in entries if isinstance(entry, dict)) ) - 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(entry, index, settings_names[index], first_plain_index, first_settings_index) - ] + seen: Final[set[str]] = set() # mutable-ok: local to this walk + merged: Final[list[object]] = [] # mutable-ok: local to this walk + for entry in entries: + names = guardrail_entry_names(entry) + if (names and seen.issuperset(names)) or (isinstance(entry, str) and entry in named_with_settings): + continue + merged.append(entry) + seen.update(names) + return merged def guardrail_entry_names(entry: object) -> tuple[str, ...]: if isinstance(entry, str): return (entry,) if isinstance(entry, dict): - return tuple(_GUARDRAIL_ENTRY_WITH_SETTINGS.validate_python(entry)) + return tuple(cast("Mapping[str, object]", entry)) # cast-ok: a guardrail settings entry is keyed by name return () -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( - entry: object, - index: int, - settings_names: frozenset[str], - first_plain_index: Mapping[str, int], - first_settings_index: Mapping[str, int], -) -> bool: - if isinstance(entry, str): - 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: if event_type in _REALTIME_STREAMING_HOOKS: return True diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index a554cd13a09..8cf8a2ab40e 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -115,6 +115,7 @@ from litellm.proxy.common_request_processing import ( open_sse_before_first_byte, resolve_litellm_call_id, ) +from litellm.proxy.common_utils.callback_utils import strip_callback_config from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -897,13 +898,52 @@ async def _inherit_auth_guardrails( ) metadata: Final = parsed_body["metadata"] merged: Final = (as_str_mapping(metadata) or {}).get("guardrails") + if not isinstance(merged, list): + return None 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 (), + _GUARDRAIL_ENTRIES.validate_python(merged), _client_named_guardrail_entries(client_named) ) if entries and isinstance(metadata, dict): metadata["guardrails"] = entries - return entries if entries and _named_guardrail_runs_post_call(parsed_body, entries) else None + admin_scoped_body: Final = { + **parsed_body, + "metadata": {**(as_str_mapping(metadata) or {}), **_admin_guardrail_metadata(user_api_key_dict)}, + } + return entries if entries and _named_guardrail_runs_post_call(admin_scoped_body, entries) else None + + +def _client_named_guardrail_entries(client_named: object) -> Sequence[object]: + """The guardrails ``metadata.guardrails`` names: its entries, or the keys of the dict form the guardrails read""" + if isinstance(client_named, list): + return _GUARDRAIL_ENTRIES.validate_python(client_named) + return tuple(as_str_mapping(client_named) or ()) + + +def _admin_guardrail_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object]: + """The key and team metadata ``should_run_guardrail`` reads a global guardrail opt-out from, as the native routes set it""" + return MappingProxyType( + { + "user_api_key_metadata": strip_callback_config(user_api_key_dict.metadata), + "user_api_key_team_metadata": strip_callback_config(user_api_key_dict.team_metadata), + } + ) + + +def _post_call_hook_data( + parsed_body: Mapping[str, object] | None, + guardrails: Sequence[object], + logging_obj: LiteLLMLoggingObj, + admin_metadata: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: the guardrails record their results in this request payload + """The body with what the post-call guardrails read and the kwargs split stripped: ``metadata`` naming the guardrails, + which ``should_run_guardrail`` checks, and the logging object, which unified (apply_guardrail) guardrails resolve a + pass-through response's call type from when the route has no API_ROUTE_TO_CALL_TYPES entry""" + body: Final = parsed_body or {} + return { + **body, + "metadata": {**(as_str_mapping(body.get("metadata")) or {}), "guardrails": guardrails, **admin_metadata}, + "litellm_logging_obj": logging_obj, + } def _named_guardrail_runs_post_call(data: Mapping[str, object], entries: Sequence[object]) -> bool: @@ -1987,26 +2027,18 @@ async def pass_through_request( ) if response.status_code < 400 and response_body is not None and post_call_guardrails: - # Build an enriched data dict: _parsed_body has been stripped of - # `metadata` by both pre_call_hook and _init_kwargs_for_pass_through_endpoint, - # so we re-attach the configured guardrails here so should_run_guardrail - # sees them. - hook_data: Final = dict(_parsed_body or {}) - existing_metadata = hook_data.get("metadata") - if not isinstance(existing_metadata, dict): - existing_metadata = {} - hook_data["metadata"] = { - **existing_metadata, - "guardrails": post_call_guardrails, - } - # The kwargs split popped the logging object off _parsed_body; unified (apply_guardrail) guardrails - # resolve a pass-through response's call type from it when the route has no API_ROUTE_TO_CALL_TYPES entry. - hook_data["litellm_logging_obj"] = logging_obj + hook_data: Final = _post_call_hook_data( + cast(dict[str, object] | None, _parsed_body), # cast-ok: json + post_call_guardrails, + logging_obj, + _admin_guardrail_metadata(user_api_key_dict) if inherit_auth_guardrails else MappingProxyType({}), + ) post_call_guardrail_data = hook_data response_body = await proxy_logging_obj.post_call_success_hook( data=hook_data, user_api_key_dict=user_api_key_dict, response=response_body, + merge_model_level_guardrails=not inherit_auth_guardrails, ) if isinstance(response_body, dict): content = json.dumps(response_body).encode("utf-8") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2fe2c46f2ad..516f5eac4d4 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4343,6 +4343,7 @@ class ProxyLogging: data: dict, response: LLMResponseTypes, user_api_key_dict: UserAPIKeyAuth, + merge_model_level_guardrails: bool = True, ): """ Allow user to modify outgoing data @@ -4352,6 +4353,9 @@ class ProxyLogging: 2. /embeddings 3. /image/generation 4. /files + + ``merge_model_level_guardrails`` is off for a request that never reaches a router deployment, so the guardrails + of a model that merely shares the body's ``model`` name don't run on it. """ from litellm.proxy.proxy_server import llm_router @@ -4369,7 +4373,11 @@ class ProxyLogging: guardrail_callbacks, other_callbacks = _partition_post_call_callbacks() try: # Merge model-level guardrails before checking which guardrails to run - guardrail_data: Final = check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) + guardrail_data: Final = ( + check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) + if merge_model_level_guardrails + else data + ) parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple( callback diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 85c31a5e2bc..6a1f86ee018 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -21,7 +21,9 @@ GUARDRAIL: Final = "passthrough-auth-guardrail" GLOBAL_POST_CALL: Final = "passthrough-global-post-call" KEY_POST_CALL: Final = "passthrough-key-post-call" GENERIC_POST_CALL: Final = "passthrough-generic-post-call" +MODEL_POST_CALL: Final = "passthrough-model-post-call" GUARDED_MODEL: Final = "passthrough-guarded-model" +POST_CALL_GUARDED_MODEL: Final = "passthrough-post-call-guarded-model" BLOCK_MARKER: Final = "block-this-prompt" OPENAI_ROUTE: Final = "/openai/v1/chat/completions" ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" @@ -104,6 +106,7 @@ class AuthGuardrailRig: post_call_vendor: Wire key_post_call_vendor: Wire generic_post_call_vendor: Wire + model_post_call_vendor: Wire upstreams: Mapping[str, Wire] @@ -115,6 +118,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I post_call_vendor: Final = stack.enter_context(wire_server(_lakera_clean)) key_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_flagged)) generic_post_call_vendor: Final = stack.enter_context(wire_server(_verdict)) + model_post_call_vendor: Final = stack.enter_context(wire_server(_verdict)) upstreams: Final = MappingProxyType( { OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)), @@ -149,6 +153,14 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "guardrails": [GUARDRAIL], }, }, + { + "model_name": POST_CALL_GUARDED_MODEL, + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "synthetic-openai-key", + "guardrails": [MODEL_POST_CALL], + }, + }, ], "guardrails": [ { @@ -181,6 +193,16 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "api_key": "synthetic-guardrail-key", }, }, + { + "guardrail_name": MODEL_POST_CALL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "post_call", + "default_on": False, + "api_base": model_post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, { "guardrail_name": GLOBAL_POST_CALL, "litellm_params": { @@ -205,7 +227,13 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I } owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path)) yield AuthGuardrailRig( - owned, vendor, post_call_vendor, key_post_call_vendor, generic_post_call_vendor, upstreams + owned, + vendor, + post_call_vendor, + key_post_call_vendor, + generic_post_call_vendor, + model_post_call_vendor, + upstreams, ) @@ -230,6 +258,7 @@ def _drain(rig: AuthGuardrailRig) -> None: rig.post_call_vendor.drain() rig.key_post_call_vendor.drain() rig.generic_post_call_vendor.drain() + rig.model_post_call_vendor.drain() for upstream in rig.upstreams.values(): upstream.drain() @@ -381,3 +410,77 @@ def test_guardrail_the_client_names_in_metadata_still_runs_for_a_guarded_key(rig 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() == () + + +@pytest.mark.parametrize("settings", [{}, True]) +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_guardrail_the_client_names_in_the_dict_form_still_runs_for_a_guarded_key( + rig: AuthGuardrailRig, route: str, settings: JsonValue +) -> 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: settings}}}, 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() == () + + +def _key_with_opt_out(scenario: Scenario, opt_out: str) -> str: + guarded: Final[dict[str, JsonValue]] = {"guardrails": [GENERIC_POST_CALL]} + match opt_out: + case "key_disable": + return scenario.key(metadata=guarded, disable_global_guardrails=True) + case "key_opted_out": + return scenario.key(metadata={**guarded, "opted_out_global_guardrails": [GLOBAL_POST_CALL]}) + case "team_disable": + return scenario.key(team_id=scenario.team(disable_global_guardrails=True), metadata=guarded) + case "team_opted_out": + return scenario.key( + team_id=scenario.team(metadata={"opted_out_global_guardrails": [GLOBAL_POST_CALL]}), metadata=guarded + ) + case _: + return scenario.key(metadata=guarded) + + +@pytest.mark.parametrize( + ("opt_out", "global_vendor_calls"), + [("key_disable", 0), ("key_opted_out", 0), ("team_disable", 0), ("team_opted_out", 0), ("none", 1)], +) +@pytest.mark.parametrize("route", [OPENAI_ROUTE, "/v1/chat/completions"]) +def test_post_call_pass_honors_a_global_guardrail_opt_out_as_the_native_route_does( + rig: AuthGuardrailRig, route: str, opt_out: str, global_vendor_calls: int +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_opt_out(scenario, opt_out) + _drain(rig) + + response: Final = _send(gateway, route, _route_body(OPENAI_ROUTE, "What is the capital of France?"), key) + + assert response.status_code == 200, response.text + assert len(rig.generic_post_call_vendor.drain()) == 1 + assert len(rig.post_call_vendor.drain()) == global_vendor_calls + + +def test_model_guardrail_does_not_run_on_pass_through_traffic_that_names_the_model(rig: AuthGuardrailRig) -> 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, + OPENAI_ROUTE, + {**_route_body(OPENAI_ROUTE, "What is the capital of France?"), "model": POST_CALL_GUARDED_MODEL}, + key, + ) + + assert response.status_code == 200, response.text + assert len(rig.generic_post_call_vendor.drain()) == 1 + assert rig.model_post_call_vendor.drain() == () diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index c32c1b71f75..302a41a16ff 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -4,6 +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 @@ -3802,6 +3804,24 @@ 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 + + names: Final = [f"g{index}" for index in range(300_000)] + + def _best_cpu_seconds(run: Callable[[], object]) -> float: + def _timed() -> float: + started: Final = time.process_time() + run() + return time.process_time() - started + + return min(_timed() for _ in range(3)) + + assert _best_cpu_seconds(lambda: merged_guardrail_entries(names)) < 25 * _best_cpu_seconds( + lambda: list(dict.fromkeys(names)) + ) + + def test_merged_guardrail_entries_names_each_guardrail_once_preferring_its_first_entry_with_settings(): from litellm.integrations.custom_guardrail import merged_guardrail_entries 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 5700335bae8..b82cb8c7f91 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 @@ -13368,6 +13368,79 @@ class TestInheritAuthGuardrails: merged: Final = _merge_guardrails_with_existing(post_call.call_args.kwargs["data"], ["model-guard"]) assert merged["metadata"]["guardrails"] == ["key-guard", "model-guard"] + @pytest.mark.asyncio + @pytest.mark.parametrize("settings", [{}, True]) + async def test_guardrail_a_client_names_in_the_dict_form_stays_after_the_keys(self, settings: object): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + pre_call, _ = await self._run( + {**self._chat_body(), "metadata": {"guardrails": {"client-guard": settings}}}, + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("key-guard", GuardrailEventHooks.pre_call),), + ) + + assert pre_call[0]["metadata"]["guardrails"] == ["key-guard", "client-guard"] + + @pytest.mark.asyncio + async def test_dict_form_stays_as_sent_for_a_key_without_guardrails(self): + pre_call, post_call = await self._run( + {**self._chat_body(), "metadata": {"guardrails": {"client-guard": {}}}}, + inherit=True, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + guardrails=(self._guardrail("client-guard", GuardrailEventHooks.post_call),), + ) + + assert pre_call[0]["metadata"]["guardrails"] == {"client-guard": {}} + post_call.assert_not_awaited() + + @pytest.mark.asyncio + async def test_post_call_pass_reads_opt_outs_from_the_key_and_team_not_the_body(self): + auth: Final = UserAPIKeyAuth( + api_key="sk-test", + metadata={"guardrails": ["key-guard"], "disable_global_guardrails": True}, + team_metadata={"opted_out_global_guardrails": ["global-post"]}, + ) + forged: Final = {"user_api_key_metadata": {"disable_global_guardrails": False}} + + _, post_call = await self._run( + {**self._chat_body(), "metadata": forged}, + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("key-guard", GuardrailEventHooks.post_call),), + ) + + metadata: Final = post_call.call_args.kwargs["data"]["metadata"] + assert metadata["user_api_key_metadata"]["disable_global_guardrails"] is True + assert metadata["user_api_key_team_metadata"] == {"opted_out_global_guardrails": ["global-post"]} + + @pytest.mark.asyncio + async def test_inherited_post_call_pass_skips_model_level_guardrails(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + _, post_call = await self._run( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("key-guard", GuardrailEventHooks.post_call),), + ) + + assert post_call.call_args.kwargs["merge_model_level_guardrails"] is False + + @pytest.mark.asyncio + async def test_named_global_guardrail_the_key_opted_out_of_does_not_start_the_post_call_pass(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"opted_out_global_guardrails": ["global-post"]}) + + response, _, post_call = await self._run_with_response( + {**self._chat_body(), "metadata": {"guardrails": ["global-post"]}}, + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("global-post", GuardrailEventHooks.post_call, default_on=True),), + ) + + post_call.assert_not_awaited() + assert response.body == b'{\n "id":"chatcmpl-1"\n}' + @pytest.mark.asyncio async def test_pre_call_only_guardrails_relay_the_response_untouched(self): auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py index 7b06905c8aa..7b1ce703a4c 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py @@ -130,6 +130,29 @@ class TestPassthroughPostCallGuardrails: call_kwargs = mock_proxy_logging.post_call_success_hook.call_args assert call_kwargs.kwargs["response"] == _GEMINI_RESPONSE + @patch(_COLLECT, return_value=["rubrik"]) + async def test_configured_endpoint_keeps_model_guardrails_and_reads_no_key_opt_outs( + self, + mock_collect, + ): + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_success_hook = AsyncMock(return_value=_GEMINI_RESPONSE) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) + + with _common_patches(mock_proxy_logging, _make_httpx_response(_GEMINI_RESPONSE)): + await pass_through_request( + request=_make_mock_request(), + target="https://example.com/v1/generateContent", + custom_headers={"Content-Type": "application/json"}, + user_api_key_dict=_make_user_api_key_dict(), + stream=False, + ) + + call_kwargs = mock_proxy_logging.post_call_success_hook.call_args.kwargs + assert call_kwargs["merge_model_level_guardrails"] is True + assert call_kwargs["data"]["metadata"].keys() == {"guardrails"} + @patch(_COLLECT, return_value=[]) async def test_post_call_success_hook_skipped_when_no_guardrails( self, @@ -216,7 +239,7 @@ class TestPassthroughPostCallGuardrails: block path dropping it.""" mock_response = _make_httpx_response(_GEMINI_RESPONSE) - def _block(*, data, user_api_key_dict, response): + def _block(*, data, user_api_key_dict, response, merge_model_level_guardrails): metadata = data.setdefault("metadata", {}) metadata.setdefault("standard_logging_guardrail_information", []).append( {"guardrail_name": "rubrik", "guardrail_status": "guardrail_intervened"} diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py index 53d8948869f..7f6429750c8 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py @@ -59,6 +59,27 @@ async def test_post_call_success_hook_runs_other_callback_and_replaces_response( assert out == new_response +@pytest.mark.asyncio +@pytest.mark.parametrize(("merge", "guardrails_seen"), [(True, ["key-guard", "model-guard"]), (False, ["key-guard"])]) +async def test_post_call_success_hook_merges_model_level_guardrails_unless_told_not_to( + proxy_logging, make_user_api_key_auth, monkeypatch, merge, guardrails_seen +): + router = MagicMock() + router.registry.return_value.deployments.return_value = [{"litellm_params": {"guardrails": ["model-guard"]}}] + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + g = _make_guardrail(should_run=False) + monkeypatch.setattr(litellm, "callbacks", [g]) + + await proxy_logging.post_call_success_hook( + data={"model": "alias", "metadata": {"guardrails": ["key-guard"]}}, + response=MagicMock(), + user_api_key_dict=make_user_api_key_auth(), + merge_model_level_guardrails=merge, + ) + + assert g.should_run_guardrail.call_args.kwargs["data"]["metadata"]["guardrails"] == guardrails_seen + + @pytest.mark.asyncio async def test_post_call_success_hook_guardrail_should_not_run_skipped( proxy_logging, make_user_api_key_auth, monkeypatch