scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(proxy): read tag and opt-out metadata the post-call guardrail selection trusts

Browse source
This commit is contained in:
yucheng 2026-11-03 04:43:54 -08:00
parent 5d8d7b811f
commit 68ecee226a
3 changed files with 74 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

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

52
tests/integration/authorization/test_passthrough_auth_guardrails.py
View file

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

16
tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

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