diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 1e7cdc44c70..9c2d8feb364 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -4220,7 +4220,7 @@ def _policy_match_tags(request_data: Mapping[str, object], metadata_variable_nam async def add_guardrails_from_policy_engine( - data: dict, + data: dict[str, object], # mutable-ok: pops the request's policies and records the resolved guardrails in place metadata_variable_name: str, user_api_key_dict: UserAPIKeyAuth, ) -> None: @@ -4263,10 +4263,11 @@ 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=data.get("model"), + model=_model if isinstance(_model, str) else None, tags=all_tags, ) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 97bbf893e3e..0e78f7afbd9 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import copy import json import uuid from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence @@ -75,6 +76,7 @@ from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decry from litellm.proxy.litellm_pre_call_utils import ( _CLIENT_PRICING_METADATA_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list _UNTRUSTED_METADATA_CONTROL_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list + add_guardrails_from_policy_engine, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( request_tags_from_request_data, @@ -2017,8 +2019,12 @@ _PII_TOKENS: Final = TypeAdapter(dict[str, str]) _GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" _GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools") _CLIENT_METADATA_KEYS: Final = ("metadata", "litellm_metadata") -# the guardrail list is merged from key, team and the admission frame; a later frame cannot shrink it _CLIENT_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config", "policies") +_PROXY_ONLY_FRAME_FIELDS: Final = (*_CLIENT_GUARDRAIL_SELECTORS, "include_guardrail_response") +_NESTED_BODY_ONLY_FIELDS: Final = (*_GUARDED_FRAME_FIELDS, *_PROXY_ONLY_FRAME_FIELDS) +_GUARDRAIL_NAMES: Final = TypeAdapter(list[str]) +_OBJECTS: Final = TypeAdapter(list[object]) +_TURN_COPIED_CONTAINERS: Final = (list, dict) class _PreCallGuardrailWalk(Protocol): @@ -2057,21 +2063,21 @@ 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") - if isinstance(nested, dict) and any(key in frame_obj for key in _GUARDED_FRAME_FIELDS): + if isinstance(nested, dict) and any(key in frame_obj for key in _NESTED_BODY_ONLY_FIELDS): raise ValueError( - "response.create carries input, instructions or tools both at the top level and under response" + "response.create carries input, instructions, tools or a guardrail field both at the top level " + "and under 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 _CLIENT_GUARDRAIL_SELECTORS + if key != "type" and key not in _PROXY_ONLY_FRAME_FIELDS }, "model": model, "litellm_call_id": str(uuid.uuid4()), - # shallow: the bucket carries proxy-owned objects (the OTEL parent span) that cannot be deep-copied - metadata_key: {**bucket}, + metadata_key: await _metadata_with_frame_guardrails(bucket, body, model, metadata_key, user_api_key_dict), } try: guarded_data: Final = await proxy_logging_obj.pre_call_hook( @@ -2089,7 +2095,11 @@ def _responses_ws_frame_guard( raise rerouted # a guardrail may delete a field (the Responses translation pops ``instructions``), so start without them guarded_fields: Final = {key: guarded_data[key] for key in _GUARDED_FRAME_FIELDS if key in guarded_data} - kept: Final = {key: value for key, value in body.items() if key not in _GUARDED_FRAME_FIELDS} + kept: Final = { + key: value + for key, value in body.items() + if key not in _GUARDED_FRAME_FIELDS and key not in _PROXY_ONLY_FRAME_FIELDS + } rewritten: Final = ( {**frame_obj, "response": {**kept, **guarded_fields}} if isinstance(nested, dict) @@ -2106,6 +2116,71 @@ def _responses_ws_frame_guard( return guard +async def _metadata_with_frame_guardrails( + bucket: Mapping[str, object], + body: Mapping[str, object], + model: object, + metadata_key: str, + user_api_key_dict: UserAPIKeyAuth, +) -> Mapping[str, object]: + """A frame selects guardrails the way an HTTP ``/v1/responses`` body does, adding to the key, team and policy + guardrails the socket was admitted with, so a frame can never shrink that list""" + metadata: Final = ( + await _metadata_with_frame_policies(bucket, body["policies"], model, metadata_key, user_api_key_dict) + if "policies" in body + else bucket + ) + existing_guardrails: Final = metadata.get("guardrails") + existing_config: Final = metadata.get("guardrail_config") + return { + **metadata, + **( + { + "guardrails": [ + *(_OBJECTS.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else ()), + *_GUARDRAIL_NAMES.validate_python(body["guardrails"]), + ] + } + if "guardrails" in body + else {} + ), + **( + { + "guardrail_config": { + **(_WS_FRAME.validate_python(existing_config) if isinstance(existing_config, dict) else {}), + **_WS_FRAME.validate_python(body["guardrail_config"]), + } + } + if "guardrail_config" in body + else {} + ), + } + + +async def _metadata_with_frame_policies( + bucket: Mapping[str, object], + policies: object, + model: object, + metadata_key: str, + user_api_key_dict: UserAPIKeyAuth, +) -> Mapping[str, object]: + selection: Final = { + "model": model, + "policies": policies, + metadata_key: {key: _turn_copy(value) for key, value in bucket.items()}, + } + await add_guardrails_from_policy_engine( + data=selection, metadata_variable_name=metadata_key, user_api_key_dict=user_api_key_dict + ) + return _WS_FRAME.validate_python(selection[metadata_key]) + + +def _turn_copy(value: object) -> object: + """Shallow, since the bucket carries proxy-owned objects (the OTEL parent span) that cannot be copied, and the + policy engine writes its lists in place""" + return copy.copy(value) if type(value) in _TURN_COPIED_CONTAINERS else value + + def _without_untrusted_metadata(key: str, value: object) -> object: """Like the HTTP path, drop the proxy-owned ``user_api_key_*`` and guardrail control keys a client may forge.""" if key not in _CLIENT_METADATA_KEYS or not isinstance(value, dict): diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 985aee061bc..1c0a04bdf90 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -2226,6 +2226,112 @@ class TestResponsesWSFirstFrameModelAuth: assert "guardrails" not in (scanned.get("metadata") or {}) assert scanned["litellm_metadata"]["guardrails"] == ["key-required"] + @pytest.mark.asyncio + @pytest.mark.parametrize("nested", [False, True]) + async def test_frame_guard_runs_the_guardrails_a_frame_selects_on_top_of_the_admitted_ones(self, nested: bool): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + scans: list[dict] = [] + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + scans.append(data) + return data + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + admitted = { + "model": "gpt-4o-mini", + "litellm_metadata": { + "guardrails": ["key-required"], + "guardrail_config": {"key-required": {"threshold": 1}}, + }, + } + guard = _responses_ws_frame_guard(admitted, MagicMock(), proxy_logging_obj, MagicMock()) + selection = { + "guardrails": ["frame-pick"], + "guardrail_config": {"frame-pick": {"threshold": 2}}, + "include_guardrail_response": True, + } + body = {"input": "hi", "store": False, **selection} + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + selected = await guard(json.dumps(frame)) + plain = await guard(json.dumps({"type": "response.create", "input": "again"})) + + assert scans[0]["litellm_metadata"]["guardrails"] == ["key-required", "frame-pick"] + assert scans[0]["litellm_metadata"]["guardrail_config"] == { + "key-required": {"threshold": 1}, + "frame-pick": {"threshold": 2}, + } + assert not set(selection) & set(scans[0]) + unselected = {"input": "hi", "store": False} + assert json.loads(selected.frame) == ( + {"type": "response.create", "response": unselected} if nested else {"type": "response.create", **unselected} + ) + assert scans[1]["litellm_metadata"]["guardrails"] == ["key-required"] + assert scans[1]["litellm_metadata"]["guardrail_config"] == {"key-required": {"threshold": 1}} + assert json.loads(plain.frame) == {"type": "response.create", "input": "again"} + assert admitted["litellm_metadata"] == { + "guardrails": ["key-required"], + "guardrail_config": {"key-required": {"threshold": 1}}, + } + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "selection", [{"guardrails": "frame-pick"}, {"guardrails": [{"name": "frame-pick"}]}, {"guardrail_config": []}] + ) + async def test_frame_guard_refuses_a_frame_whose_guardrail_selection_is_malformed(self, selection): + from pydantic import ValidationError + + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + MagicMock(), + proxy_logging_obj, + MagicMock(), + ) + + with pytest.raises(ValidationError): + await guard(json.dumps({"type": "response.create", "input": "hi", **selection})) + + proxy_logging_obj.pre_call_hook.assert_not_awaited() + + @pytest.mark.asyncio + async def test_frame_guard_runs_the_guardrails_of_the_policies_a_frame_selects(self): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + UserAPIKeyAuth(api_key="hashed"), + proxy_logging_obj, + MagicMock(), + ) + policy_registry = get_policy_registry() + policy_registry._policies = {} + policy_registry._policies_by_id = {} + policy_registry._initialized = False + try: + policy_registry.add_policy("frame-policy", Policy(guardrails=PolicyGuardrails(add=["policy-pick"]))) + + guarded = await guard(json.dumps({"type": "response.create", "input": "hi", "policies": ["frame-policy"]})) + finally: + policy_registry._policies = {} + policy_registry._policies_by_id = {} + policy_registry._initialized = False + + scanned = proxy_logging_obj.pre_call_hook.await_args.kwargs["data"] + assert scanned["litellm_metadata"]["guardrails"] == ["key-required", "policy-pick"] + assert "policies" not in scanned + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"} + @pytest.mark.asyncio async def test_frame_guard_copies_an_admitted_bucket_that_holds_the_otel_parent_span(self): import threading @@ -2367,7 +2473,10 @@ class TestResponsesWSFirstFrameModelAuth: assert dict(guarded.pii_tokens) == {"": "alice@example.com"} @pytest.mark.asyncio - async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self): + @pytest.mark.parametrize( + "top_level", [{"input": "forbidden"}, {"guardrails": ["frame-pick"]}, {"include_guardrail_response": True}] + ) + async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self, top_level): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard proxy_logging_obj = MagicMock() @@ -2377,7 +2486,7 @@ class TestResponsesWSFirstFrameModelAuth: ) with patch.object(litellm, "callbacks", []), pytest.raises(ValueError, match="both at the top level"): - await guard(json.dumps({"type": "response.create", "input": "forbidden", "response": {}})) + await guard(json.dumps({"type": "response.create", **top_level, "response": {}})) proxy_logging_obj.pre_call_hook.assert_not_awaited()