scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-10 08:28:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 152

fix(responses): run streaming-scoped guardrails on socket turns, keep guardrail fields on blocks, merge guardrail lists by name

Browse source
This commit is contained in:
mateo-berri 2026-10-28 23:05:58 -07:00
parent ad55e6eaaa
commit a17ad50dd3
9 changed files with 282 additions and 88 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
litellm/integrations/custom_guardrail.py
View file

@ -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

70
litellm/proxy/litellm_pre_call_utils.py
View file

@ -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(

59
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -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]]

4
litellm/proxy/utils.py
View file

@ -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

36
litellm/responses/streaming_iterator.py
View file

@ -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

10
tests/unit/integrations/test_custom_guardrail.py
View file

@ -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",
]

108
tests/unit/proxy/response_api_endpoints/test_endpoints.py
View file

@ -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

46
tests/unit/proxy/test_litellm_pre_call_utils.py
View file

@ -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

10
tests/unit/proxy/utils/helpers/test_guardrail_merge.py
View file

@ -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"]