From a17ad50dd3c8893027411187cd697c4424c65627 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 23:05:58 -0700 Subject: [PATCH] fix(responses): run streaming-scoped guardrails on socket turns, keep guardrail fields on blocks, merge guardrail lists by name --- litellm/integrations/custom_guardrail.py | 27 ++++- litellm/proxy/litellm_pre_call_utils.py | 70 ++++-------- .../proxy/response_api_endpoints/endpoints.py | 59 ++++++---- litellm/proxy/utils.py | 4 +- litellm/responses/streaming_iterator.py | 36 ++++-- .../integrations/test_custom_guardrail.py | 10 ++ .../response_api_endpoints/test_endpoints.py | 108 +++++++++++++++++- .../unit/proxy/test_litellm_pre_call_utils.py | 46 ++++++++ .../utils/helpers/test_guardrail_merge.py | 10 +- 9 files changed, 282 insertions(+), 88 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9e2dc902d55..15a29ddfb7e 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -3,12 +3,14 @@ import copy import hashlib import os import secrets -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from datetime import datetime +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args import httpx +from pydantic import TypeAdapter from litellm._internal_context import with_service_target from litellm._logging import verbose_logger @@ -207,6 +209,29 @@ 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 in order, dropping an entry whose guardrails an earlier one already names, so the + first entry for a guardrail, plain or ``{name: settings}``, is the one that stays""" + entries: Final = tuple(chain.from_iterable(entry_lists)) + names: Final = tuple(_guardrail_entry_names(entry) for entry in entries) + return [entry for index, entry in enumerate(entries) if not _named_earlier(names, index)] + + +def _guardrail_entry_names(entry: object) -> frozenset[str]: + if isinstance(entry, str): + return frozenset((entry,)) + if isinstance(entry, dict): + return frozenset(_GUARDRAIL_ENTRY_WITH_SETTINGS.validate_python(entry)) + return frozenset() + + +def _named_earlier(names: tuple[frozenset[str], ...], index: int) -> bool: + return bool(names[index]) and names[index] <= frozenset[str]().union(*names[:index]) + + 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/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b063670c346..0285bc270c7 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -48,7 +48,7 @@ from litellm.constants import ( VECTOR_STORE_CALLER_METADATA_KWARG, X_LITELLM_DISABLE_CALLBACKS, ) -from litellm.integrations.custom_guardrail import without_server_streaming_classification +from litellm.integrations.custom_guardrail import merged_guardrail_entries, without_server_streaming_classification from litellm.litellm_core_utils.core_helpers import is_codex_user_agent from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -3895,30 +3895,14 @@ def add_guardrails_from_key_or_team_metadata( """ from litellm.proxy.utils import premium_user_check - # Initialize guardrails set (avoiding duplicates) - combined_guardrails: Final = set() - - # Add key-level guardrails first - if key_metadata and "guardrails" in key_metadata: - if isinstance(key_metadata["guardrails"], list) and len(key_metadata["guardrails"]) > 0: - premium_user_check() - combined_guardrails.update(key_metadata["guardrails"]) - - # Add team-level guardrails (set automatically handles duplicates) - if team_metadata and "guardrails" in team_metadata: - if isinstance(team_metadata["guardrails"], list) and len(team_metadata["guardrails"]) > 0: - premium_user_check() - combined_guardrails.update(team_metadata["guardrails"]) - - # Add project-level guardrails (set automatically handles duplicates) - if project_metadata and "guardrails" in project_metadata: - if isinstance(project_metadata["guardrails"], list) and len(project_metadata["guardrails"]) > 0: - premium_user_check() - combined_guardrails.update(project_metadata["guardrails"]) - - # Set combined guardrails in metadata as list - if combined_guardrails: - data[metadata_variable_name]["guardrails"] = list(combined_guardrails) + configured_guardrails: Final = tuple( + metadata["guardrails"] + for metadata in (key_metadata, team_metadata, project_metadata) + if metadata and isinstance(metadata.get("guardrails"), list) and len(metadata["guardrails"]) > 0 + ) + if configured_guardrails: + premium_user_check() + data[metadata_variable_name]["guardrails"] = merged_guardrail_entries(*configured_guardrails) _add_guardrails_from_key_or_team_metadata: Final = add_guardrails_from_key_or_team_metadata @@ -4013,14 +3997,10 @@ def _add_guardrails_from_policies_in_metadata( if metadata_variable_name not in data: data[metadata_variable_name] = {} - existing_guardrails = data[metadata_variable_name].get("guardrails", []) - if not isinstance(existing_guardrails, list): - existing_guardrails = [] - - # Combine existing guardrails with policy-resolved guardrails (no duplicates) - combined: Final = set(existing_guardrails) - combined.update(resolved_guardrails) - data[metadata_variable_name]["guardrails"] = list(combined) + existing_guardrails: Final = data[metadata_variable_name].get("guardrails", []) + data[metadata_variable_name]["guardrails"] = merged_guardrail_entries( + existing_guardrails if isinstance(existing_guardrails, list) else (), resolved_guardrails + ) # Store applied policies in metadata for tracking if "applied_policies" not in data[metadata_variable_name]: @@ -4259,16 +4239,7 @@ def _apply_resolved_guardrails_to_metadata( _REQUEST_ENTRIES.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else () ) - selected_names: Final = frozenset( - itertools.chain.from_iterable( - _REQUEST_FIELDS.validate_python(entry).keys() if isinstance(entry, dict) else (entry,) - for entry in existing_guardrails_list - ) - ) - combined: Final = [ - *existing_guardrails_list, - *(name for name in dict.fromkeys(resolved_guardrails) if name not in selected_names), - ] + combined: Final = merged_guardrail_entries(existing_guardrails_list, resolved_guardrails) data[metadata_variable_name]["guardrails"] = combined verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", combined) @@ -4328,12 +4299,13 @@ async def add_guardrails_from_policy_engine( all_tags: Final = _policy_match_tags(data, metadata_variable_name) _team_alias: Final = user_api_key_dict.team_alias _key_alias: Final = user_api_key_dict.key_alias - _model: Final = data.get("model") - context: Final = PolicyMatchContext( - team_alias=_team_alias if isinstance(_team_alias, str) else None, - key_alias=_key_alias if isinstance(_key_alias, str) else None, - model=_model if isinstance(_model, str) else None, - tags=all_tags, + context: Final = PolicyMatchContext.model_validate( + { + "team_alias": _team_alias if isinstance(_team_alias, str) else None, + "key_alias": _key_alias if isinstance(_key_alias, str) else None, + "model": data.get("model"), + "tags": all_tags, + } ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 1f6d3b5bf08..3ebf99d176a 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -36,7 +36,7 @@ from litellm._logging import verbose_proxy_logger from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping from litellm.caching.dual_cache import DualCache from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES -from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.integrations.custom_guardrail import ModifyResponseException, guardrail_request_data_with_streaming from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -89,7 +89,12 @@ from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.proxy.utils import count_request_input_tokens, handle_exception_on_proxy from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction from litellm.responses.litellm_completion_transformation.count_request import MODEL_FILE_ID_MAPPING_FIELD -from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard, spawn_logging_task +from litellm.responses.streaming_iterator import ( + GuardedResponsesWsFrame, + ResponsesWsFrameGuard, + ResponsesWsFrameRejection, + spawn_logging_task, +) from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -2057,8 +2062,8 @@ _CLIENT_METADATA_KEYS: Final = ("metadata", "litellm_metadata") _CLIENT_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config", "policies") _METADATA_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config") _PROXY_ONLY_FRAME_FIELDS: Final = (*_CLIENT_GUARDRAIL_SELECTORS, "include_guardrail_response") -_NESTED_BODY_ONLY_FIELDS: Final = (*_GUARDED_FRAME_FIELDS, *_PROXY_ONLY_FRAME_FIELDS) -_UNFORWARDED_FRAME_FIELDS: Final = (*_NESTED_BODY_ONLY_FIELDS, "litellm_metadata") +_NESTED_BODY_ONLY_FIELDS: Final = (*_GUARDED_FRAME_FIELDS, "guardrails", "policies") +_UNFORWARDED_FRAME_FIELDS: Final = (*_GUARDED_FRAME_FIELDS, *_PROXY_ONLY_FRAME_FIELDS, "litellm_metadata") _OBJECTS: Final = TypeAdapter(list[object]) _TURN_COPIED_CONTAINERS: Final = (list, dict) @@ -2108,26 +2113,27 @@ def _responses_ws_frame_guard( async def guard(frame: str) -> GuardedResponsesWsFrame: frame_obj: Final = _WS_FRAME.validate_json(frame) nested: Final = frame_obj.get("response") - stray: Final = ( - next((key for key in _NESTED_BODY_ONLY_FIELDS if key in frame_obj), None) - if isinstance(nested, dict) - else None - ) + stray: Final = _field_beside_nested_response(frame_obj) if isinstance(nested, dict) else None if stray is not None: raise ValueError( f"response.create sends {stray} at the top level next to response; send it inside response" ) body: Final = _WS_FRAME.validate_python(nested) if isinstance(nested, dict) else frame_obj - frame_data: Final = { - **{ - key: _without_untrusted_metadata(key, value) - for key, value in body.items() - if key != "type" and key not in _PROXY_ONLY_FRAME_FIELDS + frame_data: Final = guardrail_request_data_with_streaming( + { + **{ + key: _without_untrusted_metadata(key, value) + for key, value in body.items() + if key != "type" and key not in _PROXY_ONLY_FRAME_FIELDS + }, + "model": model, + "litellm_call_id": str(uuid.uuid4()), + metadata_key: await _metadata_with_frame_guardrails( + bucket, body, model, metadata_key, user_api_key_dict + ), }, - "model": model, - "litellm_call_id": str(uuid.uuid4()), - metadata_key: await _metadata_with_frame_guardrails(bucket, body, model, metadata_key, user_api_key_dict), - } + is_streaming=True, + ) try: guarded_data: Final = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=frame_data, call_type="aresponses", guardrails_only=True @@ -2137,7 +2143,8 @@ def _responses_ws_frame_guard( ) except Exception as e: await _bill_rejected_frame(rejected_frame_auth, e, frame_data) - raise ValueError(proxy_exception_from_route_error(e).message) from e + proxy_exception: Final = proxy_exception_from_route_error(e) + raise ResponsesWsFrameRejection(proxy_exception.message, proxy_exception.provider_specific_fields) from e if guarded_data.get("model") != model: rerouted: Final = ValueError( f"a guardrail rerouted this turn to {guarded_data.get('model')!r}; an open Responses WebSocket" @@ -2153,7 +2160,10 @@ def _responses_ws_frame_guard( if key not in _UNFORWARDED_FRAME_FIELDS } rewritten: Final = ( - {**frame_obj, "response": {**kept, **guarded_fields}} + { + **{key: value for key, value in frame_obj.items() if key not in _PROXY_ONLY_FRAME_FIELDS}, + "response": {**kept, **guarded_fields}, + } if isinstance(nested, dict) else {**kept, **guarded_fields} ) @@ -2168,6 +2178,15 @@ def _responses_ws_frame_guard( return guard +def _field_beside_nested_response(frame: Mapping[str, object]) -> str | None: + metadata_picks: Final = ( + f"{key}.guardrails" + for key, value in frame.items() + if key in _CLIENT_METADATA_KEYS and isinstance(value, dict) and "guardrails" in value + ) + return next(chain((key for key in _NESTED_BODY_ONLY_FIELDS if key in frame), metadata_picks), None) + + class _FrameGuardrailSelection(TypedDict, total=False): guardrails: ReadOnly[Sequence[str | Mapping[str, Mapping[str, object]]]] guardrail_config: ReadOnly[Mapping[str, object]] diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0ddfc5ce62f..0ce92bc21a1 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -146,6 +146,7 @@ from litellm.exceptions import ( from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, + merged_guardrail_entries, ) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger @@ -9086,8 +9087,7 @@ def _merge_guardrails_with_existing(data: dict, model_level_guardrails: object) if not isinstance(model_level_guardrails, list): model_level_guardrails = [model_level_guardrails] if model_level_guardrails else [] - # Combine existing and model-level guardrails - metadata["guardrails"] = list(set(existing_guardrails + model_level_guardrails)) + metadata["guardrails"] = merged_guardrail_entries(existing_guardrails, model_level_guardrails) return modified_data diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index c1f422c7d7b..4969227f79a 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -2172,6 +2172,29 @@ class GuardedResponsesWsFrame: ResponsesWsFrameGuard: TypeAlias = Callable[[str], Awaitable[GuardedResponsesWsFrame]] +class ResponsesWsFrameRejection(ValueError): + def __init__(self, message: str, provider_specific_fields: Mapping[str, object] | None = None) -> None: + super().__init__(message) + self.provider_specific_fields: Final = provider_specific_fields + + +def _frame_rejection_event(rejection: Exception) -> str: + provider_fields: Final = ( + rejection.provider_specific_fields if isinstance(rejection, ResponsesWsFrameRejection) else None + ) + return json.dumps( + { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": str(rejection), + **({"provider_specific_fields": dict(provider_fields)} if provider_fields else {}), + }, + }, + default=str, + ) + + async def _reject_guarded_turn( websocket: ResponsesClientWebSocket, guarded: GuardedResponsesWsFrame, rejection: Exception ) -> None: @@ -2437,7 +2460,8 @@ class ResponsesWebSocketStreaming: return await self.frame_guard(message) except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event verbose_logger.info("Responses WebSocket: pre_call guardrail rejected a response.create frame: %s", e) - await self._send_error_event("invalid_request_error", str(e)) + with contextlib.suppress(Exception): + await self.websocket.send_text(_frame_rejection_event(e)) return None async def _forward_guarded_frame( @@ -2865,14 +2889,6 @@ class ResponsesWebSocketStreaming: except Exception: # noqa: BLE001, S110 # client may already be gone pass - async def _send_error_event(self, error_type: str, message: str) -> None: - try: - await self.websocket.send_text( - json.dumps({"type": "error", "error": {"type": error_type, "message": message}}) - ) - except Exception: # noqa: BLE001, S110 # client may already be gone - pass - async def client_to_backend(self) -> None: """Forward response.create events from client to backend.""" try: @@ -3286,7 +3302,7 @@ class ManagedResponsesWebSocketHandler: return await self.frame_guard(raw_message) except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event verbose_logger.info("ManagedResponsesWS: pre_call guardrail rejected a response.create frame: %s", e) - await self._send_error(str(e), error_type="invalid_request_error") + await self._send_frame(_frame_rejection_event(e)) return None @staticmethod diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 614ff97553a..de5f6995dcc 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -3786,3 +3786,13 @@ def test_default_on_guardrail() -> None: ) is True ) + + +def test_merged_guardrail_entries_keeps_the_first_entry_per_guardrail_and_any_entry_that_names_a_new_one(): + from litellm.integrations.custom_guardrail import merged_guardrail_entries + + assert merged_guardrail_entries(["a"], [{"a": {}, "b": {}}, {"a": {"x": 1}}], ["b", "c", "c"]) == [ + "a", + {"a": {}, "b": {}}, + "c", + ] diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 63e5315ecd1..5431c633623 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -2002,7 +2002,10 @@ class _FrameGuardrail(CustomGuardrail): self.inputs.append(data["input"]) self.configs.append(cast(dict[str, object], data["litellm_metadata"]).get("guardrail_config")) if self.blocks: - raise HTTPException(status_code=400, detail={"error": f"Content blocked by {self.guardrail_name}"}) + raise HTTPException( + status_code=400, + detail={"error": f"Content blocked by {self.guardrail_name}", "guardrail_name": self.guardrail_name}, + ) return {**data, "input": f"{data['input']} [{self.guardrail_name}]"} if self.rewrites else data @@ -2350,6 +2353,33 @@ class TestResponsesWSFirstFrameModelAuth: ] assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi", "metadata": {}} + @pytest.mark.asyncio + async def test_frame_guard_runs_streaming_scoped_guardrails_and_skips_non_streaming_ones_on_a_socket_turn(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + streaming_only = _FrameGuardrail("streaming-only") + streaming_only.apply_stream_scope("streaming") + streaming_during_call = _FrameGuardrail("streaming-during-call", event_hook=GuardrailEventHooks.during_call) + streaming_during_call.apply_stream_scope("streaming") + non_streaming_only = _FrameGuardrail("non-streaming-only") + non_streaming_only.apply_stream_scope("non_streaming") + frame = { + "type": "response.create", + "input": "hi", + "guardrails": ["streaming-only", "streaming-during-call", "non-streaming-only"], + } + + with _guardrails_on(streaming_only, streaming_during_call, non_streaming_only) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + guarded = await guard(json.dumps(frame)) + + assert streaming_only.inputs == ["hi"] + assert streaming_during_call.inputs == ["hi"] + assert non_streaming_only.inputs == [] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"} + @pytest.mark.asyncio async def test_frame_guard_blocks_a_turn_with_the_message_a_selected_during_call_guardrail_raises(self): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard @@ -2588,7 +2618,9 @@ class TestResponsesWSFirstFrameModelAuth: [ ({"input": "forbidden"}, "input"), ({"guardrails": ["frame-pick"]}, "guardrails"), - ({"include_guardrail_response": True}, "include_guardrail_response"), + ({"policies": ["frame-policy"]}, "policies"), + ({"metadata": {"guardrails": ["frame-pick"]}}, "metadata.guardrails"), + ({"litellm_metadata": {"guardrails": ["frame-pick"]}}, "litellm_metadata.guardrails"), ], ) async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self, top_level, field): @@ -2606,6 +2638,78 @@ class TestResponsesWSFirstFrameModelAuth: proxy_logging_obj.pre_call_hook.assert_not_awaited() + @pytest.mark.asyncio + async def test_frame_guard_forwards_a_nested_turn_without_top_level_fields_that_pick_no_guardrail(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() + ) + + with patch.object(litellm, "callbacks", []): + guarded = await guard( + json.dumps( + { + "type": "response.create", + "include_guardrail_response": True, + "guardrail_config": {"frame-pick": {"threshold": 1}}, + "metadata": {"trace": "t1"}, + "response": {"input": "hi"}, + } + ) + ) + + assert json.loads(guarded.frame) == { + "type": "response.create", + "metadata": {"trace": "t1"}, + "response": {"input": "hi"}, + } + + @pytest.mark.asyncio + async def test_a_blocked_socket_turn_answers_with_the_guardrail_fields_http_returns(self): + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler + + frame_pick = _FrameGuardrail("frame-pick", blocks=True) + websocket = MagicMock(send_text=AsyncMock()) + + with _guardrails_on(frame_pick) as proxy_logging_obj: + handler = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="anthropic/claude-haiku-5-5", + logging_obj=Logging( + model="anthropic/claude-haiku-5-5", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + frame_guard=_responses_ws_frame_guard( + {"model": "claude-managed", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ), + ) + await handler._process_response_create( + json.dumps({"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]}) + ) + + assert json.loads(websocket.send_text.await_args.args[0]) == { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": "Content blocked by frame-pick", + "provider_specific_fields": { + "error": "Content blocked by frame-pick", + "guardrail_name": "frame-pick", + "guardrail_mode": "pre_call", + }, + }, + } + @pytest.mark.asyncio async def test_frame_guard_keeps_the_guardrail_error_when_billing_the_rejected_frame_fails(self): import litellm diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index e82d2d66e3c..2ba12b543de 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -5328,6 +5328,52 @@ async def test_add_guardrails_from_policy_engine_keeps_configured_guardrails_and assert data["metadata"]["guardrails"] == [configured, "audit", "hipaa_audit"] +def test_add_guardrails_from_auth_metadata_lists_each_key_and_team_guardrail_once_keeping_the_key_s_settings(): + from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata + + configured = {"pii_blocker": {"extra_body": {"threshold": 0.9}}} + data = {"model": "gpt-4", "metadata": {}} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + metadata={"guardrails": [configured, "audit"]}, + team_metadata={"guardrails": ["audit", "pii_blocker", "toxic"]}, + ) + + with patch("litellm.proxy.utils.premium_user_check"): + add_guardrails_from_auth_metadata( + user_api_key_dict=user_api_key_dict, data=data, metadata_variable_name="metadata" + ) + + assert data["metadata"]["guardrails"] == [configured, "audit", "toxic"] + + +def test_add_guardrails_from_auth_metadata_keeps_a_selected_guardrail_s_settings_when_a_key_policy_adds_it(): + from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails + + configured = {"pii_blocker": {"extra_body": {"threshold": 0.9}}} + data = {"model": "gpt-4", "metadata": {"guardrails": [configured, "audit"]}} + policy_registry = get_policy_registry() + policy_registry._policies = { + "baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker", "audit", "hipaa_audit"])), + } + policy_registry._initialized = True + + try: + with patch("litellm.proxy.utils.premium_user_check"): + add_guardrails_from_auth_metadata( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", metadata={"policies": ["baseline"]}), + data=data, + metadata_variable_name="metadata", + ) + finally: + policy_registry._policies = {} + policy_registry._initialized = False + + assert data["metadata"]["guardrails"] == [configured, "audit", "hipaa_audit"] + + @pytest.mark.asyncio async def test_add_guardrails_from_policy_engine_applies_inherited_parent_guardrail_when_child_condition_misses(): from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry diff --git a/tests/unit/proxy/utils/helpers/test_guardrail_merge.py b/tests/unit/proxy/utils/helpers/test_guardrail_merge.py index e505517eeff..4ae87dda343 100644 --- a/tests/unit/proxy/utils/helpers/test_guardrail_merge.py +++ b/tests/unit/proxy/utils/helpers/test_guardrail_merge.py @@ -360,7 +360,9 @@ def test_merge_guardrails_with_existing_creates_metadata_when_missing(): } -def test_merge_guardrails_with_existing_raises_on_unhashable_guardrail(): - data = {"metadata": {"guardrails": [{"unhashable": True}]}} - with pytest.raises(TypeError): - _merge_guardrails_with_existing(data, ["g1"]) +def test_merge_guardrails_with_existing_keeps_a_selected_guardrail_s_settings_over_the_model_s_plain_entry(): + configured = {"pii-redact": {"extra_body": {"threshold": 0.9}}} + + result = _merge_guardrails_with_existing({"metadata": {"guardrails": [configured]}}, ["pii-redact", "toxic-filter"]) + + assert result["metadata"]["guardrails"] == [configured, "toxic-filter"]