From fe159ac944b6fb8917b3347f1d1b865bc3056afc Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 13:32:52 -0800 Subject: [PATCH] fix(proxy): run policy pipelines in the pass-through post-call pass and read reply integers outside provider replies A policy's post_call pipeline now runs on built-in pass-through responses, with the pipeline state request setup moved out of the body carried into the post-call data, so the pipeline's on_fail decides instead of the step running as a plain guardrail. A non-provider JSON response is read whole, integers included, and tool-call integers in a provider reply are read, so a number in what the model or API wrote still reaches the guardrail while a reply's timestamps and usage counts stay out. --- .../guardrail_translation/README.md | 2 +- .../guardrail_translation/handler.py | 28 +++++++-- .../pass_through_endpoints.py | 51 ++++++++++++--- .../test_passthrough_auth_guardrails.py | 63 +++++++++++++++++++ .../integrations/test_custom_guardrail.py | 10 +-- .../guardrail_translation/test_handler.py | 31 +++++++++ .../test_pass_through_endpoints.py | 37 +++++++++++ 7 files changed, 200 insertions(+), 22 deletions(-) diff --git a/litellm/llms/pass_through/guardrail_translation/README.md b/litellm/llms/pass_through/guardrail_translation/README.md index 79d96e3049e..8a3501dd972 100644 --- a/litellm/llms/pass_through/guardrail_translation/README.md +++ b/litellm/llms/pass_through/guardrail_translation/README.md @@ -24,7 +24,7 @@ litellm/proxy/pass_through_endpoints/ The `PassThroughEndpointHandler` enables guardrails to run on passthrough endpoint requests by: 1. **Field Targeting**: Extracts specific fields from the request/response using JSONPath expressions configured in `request_fields` / `response_fields` -2. **Full Payload Fallback**: If no field targeting is configured, processes the entire request payload, or the text fields of the response +2. **Full Payload Fallback**: If no field targeting is configured, processes the entire request payload. A response is read whole, integers included, except a provider's reply, where the ids, timestamps and usage counts around the text are left out 3. **Config Access**: Uses `get_passthrough_guardrails_config()` / `set_passthrough_guardrails_config()` helpers to access the passthrough guardrails configuration stored in request metadata ## Example Config diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 3c712a213d4..80d03527acb 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -31,7 +31,7 @@ if TYPE_CHECKING: # The fields of a provider's JSON reply that name or encode something rather than carry its text, by the field they sit -# under. Anything outside these positions, such as a tool call's arguments, is the model's own output and is read whole. +# under, and the fields that hold what the model wrote, which are read whole. _REPLY_TOP_LEVEL_FIELDS: Final = frozenset( ( "id", @@ -57,13 +57,27 @@ _REPLY_NESTED_FIELDS: Final = MappingProxyType( "tool_calls": frozenset(("id", "type")), } ) -_REPLY_BODY_FIELDS: Final = frozenset(("candidates", "choices", "content", "data", "output", "predictions")) +_REPLY_MODEL_OUTPUT_FIELDS: Final = frozenset(("args", "arguments", "input", "parameters")) +_REPLY_MARKER_FIELDS: Final = frozenset(("candidates", "choices", "object", "output", "predictions")) +_IMAGE_REPLY_FIELDS: Final = frozenset(("created", "data")) +_NO_FIELDS: Final[frozenset[str]] = frozenset() + + +def _is_provider_reply(response: Mapping[str, object]) -> bool: + return bool(_REPLY_MARKER_FIELDS & response.keys()) or _IMAGE_REPLY_FIELDS.issubset(response) + + +def _reply_field_position(name: str, skipped: frozenset[str] | None) -> frozenset[str] | None: + if skipped is None or name in _REPLY_MODEL_OUTPUT_FIELDS: + return None + return _REPLY_NESTED_FIELDS.get(name, _NO_FIELDS) def _response_text(response: Mapping[str, object]) -> str: - """The strings a JSON response carries, without the numbers around them or a provider reply's ids and encoded data, - which a guardrail would otherwise read as phone or card numbers.""" - top_level: Final = _REPLY_TOP_LEVEL_FIELDS if _REPLY_BODY_FIELDS & response.keys() else None + """The text a guardrail reads from a JSON response. In a provider's reply that is its strings without the ids and + encoded data around them, so a timestamp or a usage count doesn't read as a phone or card number. What the model + wrote into a tool call, and any other API's JSON, is read whole, integers included. Floats are scores and vectors.""" + top_level: Final = _REPLY_TOP_LEVEL_FIELDS if _is_provider_reply(response) else None pending: Final[list[tuple[object, frozenset[str] | None]]] = [(response, top_level)] # mutable-ok: stack texts: Final[list[str]] = [] # mutable-ok: local to this walk while pending: @@ -72,11 +86,13 @@ def _response_text(response: Mapping[str, object]) -> str: items = as_object_list(value) if isinstance(value, str): texts.append(value) + elif skipped is None and isinstance(value, int) and not isinstance(value, bool): + texts.append(str(value)) elif fields is not None: pending.extend( reversed( [ - (item, None if skipped is None else _REPLY_NESTED_FIELDS.get(name)) + (item, _reply_field_position(name, skipped)) for name, item in fields.items() if skipped is None or name not in skipped ] diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 2ec3fa50fdd..f377756b75a 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -158,7 +158,7 @@ from litellm.proxy.route_priority import ( configured_pass_through_routes_first, shadowed_by_builtin_claim, ) -from litellm.proxy.utils import normalize_route_for_root_path +from litellm.proxy.utils import normalize_route_for_root_path, pipeline_managed_guardrail_names from litellm.repositories.team_repository import TeamRepository from litellm.secret_managers.main import get_secret_str from litellm.types import utils as types_utils @@ -204,6 +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") _KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState)) _PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",)) @@ -896,7 +897,7 @@ async def _inherit_auth_guardrails( The endpoint's configured guardrails win. Without any, and with ``inherit`` and a JSON body, the key, team, project, policy and request-body guardrails are merged into ``parsed_body["metadata"]`` for the pre-call hook, with the ones the client named in ``metadata.guardrails`` kept after them, and returned only when at least one - of them runs post_call, so the response is relayed untouched otherwise. + of them runs post_call. """ if configured_guardrails or not inherit or not parsed_body: return configured_guardrails @@ -941,21 +942,42 @@ 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""" + 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} + + +def _starts_post_call_pass(guardrails: Sequence[object] | None, policy_pipelines: Mapping[str, object]) -> bool: + return bool(guardrails) or bool(pipeline_managed_guardrail_names({"metadata": dict(policy_pipelines)}, "post_call")) + + def _post_call_hook_data( parsed_body: Mapping[str, object] | None, - guardrails: Sequence[object], + guardrails: Sequence[object] | None, logging_obj: LiteLLMLoggingObj, user_api_key_dict: UserAPIKeyAuth, + policy_pipelines: Mapping[str, object], inherited: bool, ) -> dict[str, object]: # mutable-ok: the guardrails record their results in this request payload """The body with what the post-call guardrails read and the kwargs split stripped: ``metadata`` naming the guardrails, - which ``should_run_guardrail`` checks, and the logging object, which unified (apply_guardrail) guardrails resolve a - pass-through response's call type from when the route has no API_ROUTE_TO_CALL_TYPES entry""" + which ``should_run_guardrail`` checks, the policy pipelines to run, and the logging object, which unified + (apply_guardrail) guardrails resolve a pass-through response's call type from when the route has no + API_ROUTE_TO_CALL_TYPES entry""" body: Final = parsed_body or {} admin_metadata: Final = _admin_guardrail_metadata(user_api_key_dict) if inherited else MappingProxyType({}) return { **body, - "metadata": {**(as_str_mapping(body.get("metadata")) or {}), "guardrails": guardrails, **admin_metadata}, + "metadata": { + **(as_str_mapping(body.get("metadata")) or {}), + **policy_pipelines, + "guardrails": guardrails or (), + **admin_metadata, + }, "litellm_logging_obj": logging_obj, } @@ -2086,10 +2108,23 @@ async def pass_through_request( logging_obj=logging_obj, ) - if response.status_code < 400 and response_body is not None and post_call_guardrails: + policy_pipelines: Final = _policy_pipeline_metadata( + cast(dict[str, object], kwargs), # cast-ok: the logging fields init_kwargs built above + inherit_auth_guardrails, + ) + if ( + response.status_code < 400 + and response_body is not None + and _starts_post_call_pass(post_call_guardrails, policy_pipelines) + ): json_body: Final = cast(dict[str, object] | None, _parsed_body) # cast-ok: json hook_data: Final = _post_call_hook_data( - json_body, post_call_guardrails, logging_obj, user_api_key_dict, inherit_auth_guardrails + json_body, + post_call_guardrails, + logging_obj, + user_api_key_dict, + policy_pipelines, + inherit_auth_guardrails, ) post_call_guardrail_data = hook_data response_body = await _run_post_call_guardrails_keeping_the_charge( diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 02030cfb5cb..d292fe41059 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -44,6 +44,9 @@ ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent" CUSTOM_ROUTE: Final = "/passthrough-auth-custom" PHONE_POST_CALL: Final = "passthrough-auth-phone-post-call" +PIPELINE_BLOCKS: Final = "passthrough-pipeline-blocks-the-response" +PIPELINE_ALLOWS: Final = "passthrough-pipeline-allows-the-response" +PIPELINE_BEFORE_THE_CALL: Final = "passthrough-pipeline-before-the-call" _CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" _JSON: Final = TypeAdapter(dict[str, JsonValue]) _JSON_LIST: Final = TypeAdapter(list[JsonValue]) @@ -61,6 +64,13 @@ def _route_body(route: str, prompt: str) -> dict[str, JsonValue]: return {"model": "gpt-5.6", "messages": [{"role": "user", "content": prompt}]} +def _pipeline_policy(mode: str, guardrail: str, on_fail: str) -> dict[str, JsonValue]: + return { + "guardrails": {"add": [guardrail]}, + "pipeline": {"mode": mode, "steps": [{"guardrail": guardrail, "on_pass": "allow", "on_fail": on_fail}]}, + } + + def _verdict(request: Request) -> Reply: action: Final = "BLOCKED" if BLOCK_MARKER.encode() in request.body else "NONE" return Reply(body=json.dumps({"action": action, "blocked_reason": "synthetic block"}).encode()) @@ -167,6 +177,15 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I } ], }, + "policies": { + PIPELINE_BLOCKS: _pipeline_policy("post_call", KEY_POST_CALL, "block"), + PIPELINE_ALLOWS: _pipeline_policy("post_call", KEY_POST_CALL, "allow"), + PIPELINE_BEFORE_THE_CALL: _pipeline_policy("pre_call", GUARDRAIL, "block"), + }, + "policy_attachments": [ + {"policy": policy, "tags": [policy]} + for policy in (PIPELINE_BLOCKS, PIPELINE_ALLOWS, PIPELINE_BEFORE_THE_CALL) + ], "model_list": [ *_JSON_LIST.validate_python(config.get("model_list") or []), { @@ -567,6 +586,50 @@ def _charged_rows(gateway: Gateway, key: str) -> tuple[Mapping[str, object], ... ) +def _tagged_body(route: str, tag: str) -> dict[str, JsonValue]: + return {**_route_body(route, "What is the capital of France?"), "metadata": {"tags": [tag]}} + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_post_call_policy_pipeline_blocks_a_flagged_pass_through_response(rig: AuthGuardrailRig, route: str) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key() + _drain(rig) + + response: Final = _send(gateway, route, _tagged_body(route, PIPELINE_BLOCKS), key) + + assert response.status_code == 400, response.text + assert len(rig.key_post_call_vendor.drain()) == 1 + assert len(rig.upstreams[route].drain()) == 1 + assert [row["status"] for row in _charged_rows(gateway, key)] == ["failure"] + + +def test_post_call_policy_pipeline_decides_what_a_flagged_step_does(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key() + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, _tagged_body(OPENAI_ROUTE, PIPELINE_ALLOWS), key) + + assert response.status_code == 200, response.text + assert len(rig.key_post_call_vendor.drain()) == 1 + + +def test_pre_call_policy_pipeline_relays_the_provider_bytes(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key() + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, _tagged_body(OPENAI_ROUTE, PIPELINE_BEFORE_THE_CALL), key) + + assert response.status_code == 200, response.text + assert response.content == _OPENAI_BODY + assert len(rig.vendor.drain()) == 1 + + @pytest.mark.parametrize("route", [ANTHROPIC_ROUTE, GEMINI_ROUTE]) def test_response_a_post_call_guardrail_blocks_still_charges_the_provider_call( rig: AuthGuardrailRig, route: str diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index e6a11f92a0c..1e93db413b1 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -3802,16 +3802,12 @@ def test_default_on_guardrail() -> None: ) -def test_merged_guardrail_entries_dedupes_an_all_plain_list_without_walking_it(): - from litellm.integrations import custom_guardrail +def test_merged_guardrail_entries_dedupes_an_all_plain_list(): + from litellm.integrations.custom_guardrail import merged_guardrail_entries names: Final = [f"g{index % 500}" for index in range(1000)] - with patch.object(custom_guardrail, "guardrail_entry_names") as entry_names: - merged: Final = custom_guardrail.merged_guardrail_entries(names) - - assert merged == [f"g{index}" for index in range(500)] - entry_names.assert_not_called() + assert merged_guardrail_entries(names) == [f"g{index}" for index in range(500)] def test_merged_guardrail_entries_reads_each_entry_of_a_mixed_list_a_fixed_number_of_times(): diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index 62f49b0047b..b712f6c471c 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -276,6 +276,37 @@ async def test_post_call_scan_reads_a_field_named_like_the_envelope_in_what_the_ assert await _scanned_response_texts(response) != () +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response, expected", + [ + ({"result": {"credit_card": 4111111111111111, "score": 0.4111111111111111, "ok": True}}, "4111111111111111"), + ( + {"candidates": [{"content": {"parts": [{"functionCall": {"args": {"card": 4111111111111111}}}]}}]}, + "4111111111111111", + ), + ({"choices": [{"message": {"tool_calls": [{"function": {"arguments": {"n": 5551234567}}}]}}]}, "5551234567"), + ], +) +async def test_post_call_scan_reads_an_integer_in_what_the_model_or_api_wrote(response, expected): + assert await _scanned_response_texts(response) == (expected,) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response", + [ + {"created": 1793600123, "data": [{"url": "https://x.example/a.png"}]}, + {"id": "batch_1", "object": "batch", "created_at": 1793600123, "request_counts": {"total": 5551234567}}, + {"predictions": [{"embeddings": {"values": [0.0123456789]}, "total": 5551234567}]}, + ], +) +async def test_post_call_scan_leaves_out_the_numbers_of_a_provider_reply(response): + assert all( + "1793600123" not in text and "5551234567" not in text for text in await _scanned_response_texts(response) + ) + + @pytest.mark.asyncio async def test_post_call_scan_skips_a_reply_with_no_text(): response: Final = {"object": "list", "data": [{"object": "embedding", "embedding": [0.4111111111111111]}]} 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 0667a6d47d6..f9a3f04e9e7 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 @@ -85,6 +85,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id from litellm.proxy.utils import InternalUsageCache, ProxyLogging, _merge_guardrails_with_existing, hash_token from litellm.types import utils as types_utils from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep from litellm.types.utils import CallTypes from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, @@ -13382,6 +13383,42 @@ class TestInheritAuthGuardrails: assert pre_call[0]["metadata"]["guardrails"] == ["key-guard", "client-guard"] + @staticmethod + def _body_with_policy_pipeline(mode: str) -> dict[str, object]: + pipeline: Final = GuardrailPipeline(mode=mode, steps=[PipelineStep(guardrail="step-guard")]) + return { + **TestInheritAuthGuardrails._chat_body(), + "metadata": { + "_guardrail_pipelines": [("policy", pipeline)], + "_pipeline_managed_guardrails": {"step-guard"}, + }, + } + + @pytest.mark.asyncio + async def test_post_call_policy_pipeline_starts_the_post_call_pass_with_its_state(self): + _, post_call = await self._run( + self._body_with_policy_pipeline("post_call"), + inherit=True, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + post_call.assert_awaited_once() + metadata: Final = post_call.await_args.kwargs["data"]["metadata"] + assert metadata["_pipeline_managed_guardrails"] == {"step-guard"} + assert [policy for policy, _ in metadata["_guardrail_pipelines"]] == ["policy"] + + @pytest.mark.asyncio + @pytest.mark.parametrize(("mode", "inherit"), [("pre_call", True), ("post_call", False)]) + async def test_policy_pipeline_leaves_the_response_alone_without_a_post_call_run(self, mode: str, inherit: bool): + response, _, post_call = await self._run_with_response( + self._body_with_policy_pipeline(mode), + inherit=inherit, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + post_call.assert_not_awaited() + assert response.body == b'{\n "id":"chatcmpl-1"\n}' + @pytest.mark.asyncio async def test_dict_form_stays_as_sent_for_a_key_without_guardrails(self): pre_call, post_call = await self._run(