diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 121bf0ed2c8..6ae21b02aff 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -204,7 +204,7 @@ _blocked_response_cost_logging: Final = PassThroughEndpointLogging(log_dispatch= _METADATA_KEYS: Final = frozenset(("litellm_metadata", "metadata")) _DECODED_JSON: Final = TypeAdapter(object) _GUARDRAIL_ENTRIES: Final = TypeAdapter(list[object]) -_POLICY_PIPELINE_FIELDS: Final = ("_guardrail_pipelines", "_pipeline_managed_guardrails") +_POST_CALL_SELECTION_FIELDS: Final = ("_guardrail_pipelines", "_pipeline_managed_guardrails", "tags") _KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState)) _PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",)) @@ -930,9 +930,10 @@ async def _inherit_auth_guardrails( ) if entries and isinstance(metadata, dict): metadata["guardrails"] = entries + admin_metadata: Final = _admin_guardrail_metadata(user_api_key_dict) admin_scoped_body: Final = { **parsed_body, - "metadata": {**(as_str_mapping(metadata) or {}), **_admin_guardrail_metadata(user_api_key_dict)}, + **{key: {**(as_str_mapping(parsed_body.get(key)) or {}), **admin_metadata} for key in sorted(_METADATA_KEYS)}, } return entries if entries and _named_guardrail_runs_post_call(admin_scoped_body, entries) else None @@ -954,14 +955,14 @@ def _admin_guardrail_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, ) -def _policy_pipeline_metadata(kwargs: Mapping[str, object], inherited: bool) -> Mapping[str, object]: - """The policy pipelines resolved for a request that inherits auth guardrails, from the logging metadata that - request setup moved the body's ``metadata`` into""" +def _post_call_selection_metadata(kwargs: Mapping[str, object], inherited: bool) -> Mapping[str, object]: + """The policy pipelines resolved for a request that inherits auth guardrails, and the tags a tag-scoped guardrail + mode reads, from the logging metadata that request setup moved the body's ``metadata`` into""" if not inherited: return MappingProxyType({}) litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} metadata: Final = as_str_mapping(litellm_params.get("metadata")) or {} - return {field: metadata[field] for field in _POLICY_PIPELINE_FIELDS if field in metadata} + return {field: metadata[field] for field in _POST_CALL_SELECTION_FIELDS if field in metadata} def _starts_post_call_pass(guardrails: Sequence[object] | None, policy_pipelines: Mapping[str, object]) -> bool: @@ -2141,7 +2142,7 @@ async def pass_through_request( logging_obj=logging_obj, ) - policy_pipelines: Final = _policy_pipeline_metadata( + policy_pipelines: Final = _post_call_selection_metadata( cast(dict[str, object], kwargs), # cast-ok: the logging fields init_kwargs built above inherit_auth_guardrails, ) diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 690b018d8f3..d73e3db7d4c 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -11,7 +11,7 @@ import httpx import pytest import yaml from integration._support.anthropic_sse import message_json -from integration._support.client import Gateway, Scenario, eventually +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually from integration._support.openai_wire import chat_reply from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.safeguards_rig import key_hash @@ -23,6 +23,8 @@ GLOBAL_POST_CALL: Final = "passthrough-global-post-call" KEY_POST_CALL: Final = "passthrough-key-post-call" GENERIC_POST_CALL: Final = "passthrough-generic-post-call" BLOCKING_POST_CALL: Final = "passthrough-post-call-blocks-every-reply" +TAG_SCOPED_POST_CALL: Final = "passthrough-tag-scoped-post-call" +POST_CALL_TAG: Final = "scan-the-reply" MODEL_POST_CALL: Final = "passthrough-model-post-call" GUARDED_MODEL: Final = "passthrough-guarded-model" POST_CALL_GUARDED_MODEL: Final = "passthrough-post-call-guarded-model" @@ -148,7 +150,10 @@ _MONITORED_ANTHROPIC_BODY: Final = json.dumps( _UNPRICED_ANTHROPIC_BODY: Final = json.dumps( - _JSON.validate_json(message_json("msg_passthrough_auth_unpriced", UNPRICED_MODEL, "Hi")) + { + **JSON_OBJECT.validate_json(message_json("msg_passthrough_auth_unpriced", UNPRICED_MODEL, "Hi")), + "usage": {"input_tokens": 300, "output_tokens": 200}, + } ).encode() @@ -289,6 +294,16 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "api_key": "synthetic-guardrail-key", }, }, + { + "guardrail_name": TAG_SCOPED_POST_CALL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": {"tags": {POST_CALL_TAG: "post_call"}, "default": "logging_only"}, + "default_on": False, + "api_base": blocking_post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, { "guardrail_name": MONITOR_POST_CALL, "litellm_params": { @@ -787,6 +802,39 @@ def test_a_blocked_response_with_no_cost_result_keeps_its_usage_on_the_failure_r assert rows[0]["api_base"], rows +def test_a_tpm_limit_counts_blocked_responses_with_no_cost_result(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [BLOCKING_POST_CALL]}, tpm_limit=600) + _drain(rig) + body: Final = {**_route_body(ANTHROPIC_ROUTE, "Hello"), "model": UNPRICED_MODEL} + + limited: Final = eventually( + lambda: _send(gateway, ANTHROPIC_ROUTE, body, key), + lambda response: response.status_code != 400, + seconds=70, + ) + + assert limited.status_code == 429, limited.text + + +def test_a_tag_scoped_post_call_guardrail_scans_a_reply_for_a_request_with_its_tag(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [TAG_SCOPED_POST_CALL]}) + _drain(rig) + + response: Final = _send( + gateway, + OPENAI_ROUTE, + {**_route_body(OPENAI_ROUTE, "Hello"), "metadata": {"tags": [POST_CALL_TAG]}}, + key, + ) + + assert response.status_code == 400, response.text + assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 + + def test_custom_endpoint_writes_no_charge_for_a_response_its_post_call_guardrail_blocks(rig: AuthGuardrailRig) -> None: gateway: Final = rig.owned.gateway with gateway.scenario() as scenario: 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 1fd7ee4d70c..3d89e1f21e3 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 @@ -13517,6 +13517,22 @@ class TestInheritAuthGuardrails: 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_a_body_opt_out_in_litellm_metadata_does_not_stop_the_post_call_pass(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["global-post"]}) + + _, post_call = await self._run( + { + **self._chat_body(), + "litellm_metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}}, + }, + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("global-post", GuardrailEventHooks.post_call, default_on=True),), + ) + + post_call.assert_awaited_once() + @pytest.mark.asyncio async def test_inherited_post_call_pass_runs_only_the_guardrails_the_request_selected(self): auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})