From 0eb21e4b403d067b5db773e348767c5966be516e Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 26 Oct 2026 17:33:00 -0700 Subject: [PATCH] fix(guardrails): name a streamed pipeline step whose guardrail records its own results A decision_model step on a held stream ran but was left out of the header and the spend log, where main listed it --- .../proxy/policy_engine/pipeline_executor.py | 3 +- .../test_guardrail_pipeline_applied_header.py | 56 +++++++++++++++++++ .../policy_engine/test_pipeline_executor.py | 13 +++++ 3 files changed, 70 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index c5213f4c138..c8d5e4c3b48 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -573,8 +573,7 @@ class PipelineExecutor: _release_original_chunks(step.guardrail, undeliverable.reason, streaming_chunks, originals) return finally: - if not callback.records_own_guardrail_information: - _name_step_that_ran(hook_input, metadata_key, step.guardrail) + _name_step_that_ran(hook_input, metadata_key, step.guardrail) discard_reason: Final = observer.discard_reason(rewrite_undeliverable_reason) if discard_reason is not None: _release_original_chunks(step.guardrail, discard_reason, streaming_chunks, originals) diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py index e0fb2905a06..58e46141cac 100644 --- a/tests/integration/observability/test_guardrail_pipeline_applied_header.py +++ b/tests/integration/observability/test_guardrail_pipeline_applied_header.py @@ -36,6 +36,8 @@ _RESPONSE_BLOCKER: Final = "response-blocking-step" _LIVE_FIRST: Final = "live-first-step" _LIVE_SECOND: Final = "live-second-step" _UNGATED_STEP: Final = "ungated-step" +_JUDGE: Final = "judge-step" +_JUDGE_MODEL: Final = "pipeline-judge" _STOPS_AFTER_FIRST: Final = "stops-after-first" _RUNS_EVERY_STEP: Final = "runs-every-step" _BLOCKS_AT_FIRST: Final = "blocks-at-first" @@ -46,6 +48,7 @@ _RESPONSE_MOVES_PAST_FIRST: Final = "response-moves-past-first" _SCANS_AFTER_RELEASE: Final = "scans-after-release" _CANNOT_HOLD_STREAM: Final = "cannot-hold-stream" _ERRORS_PAST_FIRST: Final = "errors-past-first" +_JUDGES_THEN_CHECKS: Final = "judges-then-checks" _ENDPOINTS: Final = ("/v1/chat/completions", "/v1/responses", "/v1/messages") @@ -154,6 +157,16 @@ def _completion(request: Request) -> Reply: def _vendor(request: Request) -> Reply: + if request.target == f"/{_JUDGE}/v1/decisions": + return Reply( + body=json.dumps( + { + "model": "gpt-6-luna", + "answers": [{"type": "predicate", "name": "unsafe_output", "probability": 0.01}], + "usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}, + } + ).encode() + ) if request.target.startswith(f"/{_PROMPT_SHIELD}{_PROMPT_SHIELD_PATH}?"): return Reply( body=json.dumps({"userPromptAnalysis": {"attackDetected": False}, "documentsAnalysis": []}).encode() @@ -198,6 +211,17 @@ def _policy(mode: str, steps: tuple[tuple[str, str, str], ...], *, on_error: str def _config(directory: Path, vendor: Wire) -> Path: configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) configuration["litellm_settings"] = {**configuration.get("litellm_settings", {}), "callbacks": ["otel"]} + configuration["model_list"] = [ + *configuration.get("model_list", []), + { + "model_name": _JUDGE_MODEL, + "litellm_params": { + "model": "openai/gpt-6-luna", + "api_base": f"{vendor.url}/{_JUDGE}/v1", + "api_key": "synthetic-judge-key", + }, + }, + ] configuration["guardrails"] = [ *( _guardrail(name, "pre_call", vendor) @@ -218,6 +242,23 @@ def _config(directory: Path, vendor: Wire) -> Path: "api_key": "synthetic-guardrail-key", }, }, + { + "guardrail_name": _JUDGE, + "litellm_params": { + "guardrail": "decision_model", + "mode": "post_call", + "default_on": False, + "decision_model": _JUDGE_MODEL, + "checks": [ + { + "name": "unsafe_output", + "instructions": "Does the response say anything unsafe?", + "action": "block", + "threshold": 0.5, + } + ], + }, + }, { "guardrail_name": _PROMPT_SHIELD, "litellm_params": { @@ -253,6 +294,7 @@ def _config(directory: Path, vendor: Wire) -> Path: _CANNOT_HOLD_STREAM: _policy( "post_call", ((_UNGATED_STEP, "next", "block"), (_RESPONSE_SECOND, "next", "block")) ), + _JUDGES_THEN_CHECKS: _policy("post_call", ((_JUDGE, "next", "block"), (_RESPONSE_SECOND, "next", "block"))), _ERRORS_PAST_FIRST: _policy( "pre_call", ((_REQUEST_ERRORING, "next", "block"), (_REQUEST_SECOND, "next", "block")), on_error="next" ), @@ -516,6 +558,20 @@ def test_the_spend_log_and_the_trace_name_only_the_steps_that_ran( ] * len(spans), spans +@pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"]) +def test_a_response_pipeline_names_a_judge_step_that_records_its_own_results(rig: _Rig, stream: bool) -> None: + team, response = _ask_as_new_team(rig, _JUDGES_THEN_CHECKS, stream=stream) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_JUDGE, _RESPONSE_SECOND) + assert sorted(response.headers.get(_HEADER, "").split(",")) == [_JUDGE, _RESPONSE_SECOND], dict(response.headers) + rows: Final = eventually( + lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)), + lambda found: len(found) == 1, + seconds=70, + ) + assert sorted(object_value(rows[0]["metadata"]).get("applied_guardrails") or []) == [_JUDGE, _RESPONSE_SECOND] + + def _agent_card(name: str, url: str) -> dict[str, object]: return { "protocolVersion": "0.3", diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index e981b7bdeb0..068102aac08 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/unit/proxy/policy_engine/test_pipeline_executor.py @@ -1434,6 +1434,19 @@ async def test_streaming_step_records_the_guardrail_in_the_applied_guardrails_he assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] +class _SelfRecordingGuardrail(_TextReturningGuardrail): + records_own_guardrail_information = True + + +@pytest.mark.asyncio +async def test_streaming_step_whose_guardrail_records_its_own_results_is_named(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [_SelfRecordingGuardrail(["hello world"])]) + + result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) + + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + + class _LitellmMetadataOpeningGuardrail(CustomGuardrail): """Opens ``litellm_metadata`` on the request while scanning, as the A2A translation does."""