feat(guardrails): streaming text transformation in generic_guardrail_api - #32915
feat(guardrails): streaming text transformation in generic_guardrail_api#32915yucheng-berri wants to merge 12 commits into
Conversation
… streaming transform
…call passthrough, n>1 finish_reason)
… text can't bypass the transform
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR adds opt-in
Confidence Score: 4/5The incremental_diff path is opt-in and the default block_only behavior is unchanged; the new findings are edge cases that do not affect well-formed streaming responses. The implementation addresses prior review concerns but two edge-case bugs remain: raw-accumulator corruption for malformed streams without finish_reason, and missing saw_text_content tracking for mixed-chunk-only streams. Open issues from prior threads (double guardrail call, n>1 finish_reason drop) are also unresolved. unified_guardrail.py - the _inspect_full_response_for_block interaction and saw_text_content tracking for mixed chunks deserve a second look.
|
| Filename | Overview |
|---|---|
| litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py | Large new incremental_diff stream implementation. Two new edge-case findings: raw-accumulator corruption via _inspect_full_response_for_block on malformed streams, and saw_text_content never set for mixed-only streams. |
| litellm/llms/openai/chat/guardrail_translation/handler.py | Extracts _process_streaming_block_only and adds _process_streaming_transform. Core logic is correct; has_stream_ended check still inspects only choices[0] (pre-existing). |
| litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py | Adds streaming_transform_mode config param and propagates stream_holdback_chars from guardrail response. |
| litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py | Adds streaming_transform_mode optional param, stream_holdback_chars response field, coerce_stream_holdback_value helper. All new fields are additive and backward-compatible. |
| tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py | 783 lines of new streaming-transform tests covering core scenarios. All tests use sampling_rate=1, which masks behaviors that differ at the default sampling_rate=5. |
Reviews (15): Last reviewed commit: "fix(guardrails): four correctness fixes ..." | Re-trigger Greptile
Greptile SummaryThis PR adds an opt-in
Confidence Score: 3/5The incremental_diff path is safe for text-only and tool-call-only streams, but a protocol violation in the mixed content+tool_calls edge case means the guardrail transformed text is silently lost for SSE-compliant clients when finish_reason rides on the tool passthrough chunk before the synthetic text is emitted. The incremental_diff path works correctly for the common cases but contains a protocol violation in the mixed content+tool_calls edge case that would cause silent data loss for conformant SSE clients. The existing test exercises this code path but only checks the security property (no raw text leak), not the ordering violation. unified_guardrail.py (lines 580-587, tool-call passthrough branch) and test_unified_guardrail.py (mixed-chunk test should assert output ordering relative to finish_reason).
|
| Filename | Overview |
|---|---|
| litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py | Core orchestration for incremental_diff streaming — protocol violation in mixed content+tool_calls edge case; also makes a redundant guardrail API call for tool-call-containing streams. |
| litellm/llms/openai/chat/guardrail_translation/handler.py | Cleanly refactored into block-only and transform paths; accumulates by StreamingChoices.index correctly; does not mutate responses_so_far. |
| litellm/llms/base_llm/guardrail_translation/base_translation.py | Adds StreamTransformSink and stream_transform_sink parameter; backward-compatible, no issues. |
| litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py | Adds streaming_transform_mode, stream_holdback_chars, and coerce_stream_holdback_value; defensive coercion degrades invalid values to 0. |
| litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py | Adds streaming_transform_mode param and surfaces stream_holdback_chars through apply_guardrail return; no issues. |
| litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/init.py | One-line change forwarding streaming_transform_mode to initialize_guardrail; no issues. |
| litellm/types/utils.py | Adds stream_holdback_chars to GenericGuardrailAPIInputs TypedDict with total=False; backward-compatible. |
| tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py | Well-structured mock-only unit tests for response parsing, holdback coercion, and config defaults. |
| tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py | Comprehensive transform-path tests; mixed-chunk test verifies no raw leak but does not assert output ordering relative to finish_reason. |
| litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/example_config.yaml | Documents new streaming options clearly; no issues. |
Reviews (2): Last reviewed commit: "docs(generic_guardrail_api): document st..." | Re-trigger Greptile
PR overviewThis PR adds streaming text transformation support to the generic guardrail API, updating the unified guardrail hook path to process and emit transformed text during streamed responses. One issue remains open around how transformed streaming output is released. With the current behavior, text can be sent to the client before later guardrail processing determines it should be redacted, allowing sensitive content split across chunks to leak during streaming. One earlier issue has already been addressed, but this remaining streaming-buffering problem leaves a meaningful confidentiality risk until fixed. Open issues (1)
Fixed/addressed: 1 · PR risk: 7/10 |
93acde1 to
42227d0
Compare
|
bugbot run |
| mappings=endpoint_guardrail_translation_mappings, | ||
| ): | ||
| yield transformed_item | ||
| return |
There was a problem hiding this comment.
Buffer flag ignored on transform
Low Severity
When streaming_transform_mode is incremental_diff and the route resolves, the hook enters _run_incremental_transform_stream and returns before the legacy loop. streaming_buffer_until_moderated is never applied there, so configured end-of-stream buffering is silently ignored and transformed deltas can reach the client before moderation would have completed on the block-only path.
Reviewed by Cursor Bugbot for commit 42227d0. Configure here.
There was a problem hiding this comment.
streaming_buffer_until_moderated is an all-or-nothing buffer flag: it withholds every chunk until end-of-stream moderation passes, then releases the (clean) raw chunks or emits the block message. That model is inherently incompatible with incremental_diff, which is a stream-and-transform mode — buffering everything until the guardrail says OK would defeat the whole purpose of emitting synthetic deltas as text arrives.
The two config knobs already model this decision:
streaming_buffer_until_moderated: forblock_only, delay every chunk until end of stream.streaming_transform_mode: incremental_diff: stream transformed text as it arrives, no buffering.
Under block_only (the default) both knobs interact as documented. Under incremental_diff the transform flag takes precedence, which the _resolve_transform_call_type guard at line 306-320 makes explicit — we enter _run_incremental_transform_stream and return before the buffer logic. Operators who want moderation-first behavior should stick with block_only; operators who want in-flight rewrites use incremental_diff and accept that the buffer flag is a no-op.
A doc note on this interaction is worth adding — will include it in the follow-up docs.litellm.ai section already tracked in the PR body. Not fixing the code here.
Live e2e — 6/6 scenarios greenReal proxy at Scenarios
What this proves
Test infrastructure (for reproduction)
All 6 scenarios ran end-to-end against a real LLM producing real chunks — no mocks, no synthetic streams. |
6b05b08 to
06ff6b9
Compare
06ff6b9 to
c6af8f5
Compare
2 similar comments
|
bugbot run |
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.
There are 2 total unresolved issues (including 1 from previous review).
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit c47afcd. Configure here.
| # semantic on the streaming_transform_mode. | ||
| if self.streaming_transform_mode == "incremental_diff": | ||
| if guardrail_response.texts is not None: | ||
| return_inputs["texts"] = guardrail_response.texts |
There was a problem hiding this comment.
Non-streaming ignores empty texts
Medium Severity
With streaming_transform_mode set to incremental_diff, _build_guardrail_return_inputs forwards an explicit empty texts list from the guardrail as full suppression. Non-streaming process_output_response still skips applying guardrail text when guardrailed_texts is empty, so the model’s original message content is returned unchanged on post-call paths.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit c47afcd. Configure here.
There was a problem hiding this comment.
This finding is technically accurate but the fix would introduce a different regression. Reasoning:
streaming_transform_modenames its scope: it controls streaming behavior. On the non-streaming path (process_output_response), "empty texts" retains the historical block_only contract of "no change" — that's the shape thousands of deployed guardrails were built against.- The streaming path is where
incremental_diffgenuinely needs the is-not-None semantic, because a rewrite-mode guardrail's only signal for "suppress the whole response" istexts: [](a BLOCK action terminates the stream entirely; GUARDRAIL_INTERVENED with empty texts is the only way to say "emit nothing"). - Non-streaming callers have a proper channel for suppression: BLOCKED action with a blocked_reason. Overloading empty texts to mean "suppress" on non-streaming would silently change behavior for every deployed guardrail returning texts: [] as a no-op — which is exactly the P1 regression I was called out on in the previous Greptile round.
The scope-limited gate (if self.streaming_transform_mode == 'incremental_diff' inside _build_guardrail_return_inputs) is intentional. The asymmetry between streaming and non-streaming here mirrors the asymmetry of the modes themselves: incremental_diff is opt-in, streaming-only, and adds new semantics; non-streaming preserves the historical contract regardless of the flag.
Not fixing on this PR. If we ever want a true response-suppression signal on non-streaming, that belongs in a separate design conversation (probably a new action code, e.g. SUPPRESSED, rather than overloading empty texts).
There was a problem hiding this comment.
Update on this thread: rather than gate the is-not-None semantic on the mode flag (which still crossed into non-streaming callers when the operator sets incremental_diff), reverted _build_guardrail_return_inputs back to its original if guardrail_response.texts: form in 6388c3ee52. Documented as a known limitation on the incremental_diff path.
This PR now touches only the incremental_diff streaming code path — no changes to block_only, non-streaming, or pre_call semantics for any guardrail. A dedicated suppression primitive (probably a new action code, since overloading empty texts is a scope-leak trap) is deferred to a future PR.
|
@greptileai please rescore on the latest HEAD |
Ready-for-review summaryHEAD: 6a950ff (after final merge of What this PR doesOpt-in Adopted from #32084 by @schneidermr. Mirrored on a Bot findings resolved
Live e2e proof (posted earlier: comment)6/6 scenarios green on a real proxy + real Bedrock LLM + real HTTP guardrail server. No mocks.
CI
Docs
Deferred to follow-up (with team sign-off)
Note: Greptile's last score was 3/5 on |
Scope attestationAudited the PR diff for scope leakage. Found one — the Reverted the change in 6388c3ee52. Known limitation (documented in class docstring and example_config.yaml)A guardrail returning What this PR touches, verifiedAll incremental-diff-specific additions live behind
No changes to |
mateo-berri
left a comment
There was a problem hiding this comment.
The main concern is a performance issue in _handle_tool_call_chunk: a guardrail API call fires before every tool-call streaming chunk once text is in flight, meaning a streaming function call with many argument chunks generates proportionally many redundant HTTP round-trips even though subsequent calls produce no new delta
Is this greptile comment legit?
| should_run_final_round = ( | ||
| state["text_pending_flush"] | ||
| or state["deferred_finish_reason_for_text"] | ||
| or (state["saw_text_content"] and not saw_tool_calls) | ||
| ) |
There was a problem hiding this comment.
finish_reason silently dropped for text choices in n>1 streams when they finish before separate tool-call chunks
When n>1 and one choice streams text (finishing with finish_reason="stop" in its own terminal chunk) while a separate subsequent chunk carries another choice's tool_calls, the text choice's finish_reason is permanently lost. The sequence is:
- Choice 0 terminal text chunk →
_record_finish_reasonswritesfinish_reason_per_choice[0]="stop",text_pending_flushstaysTrue - Choice 1 tool-call chunk →
_handle_tool_call_chunkpre-flush fires (is_final=False), emits choice 0's text but not itsfinish_reason, then clearstext_pending_flush - End of loop →
should_run_final_roundevaluatestext_pending_flush=False OR deferred_finish_reason_for_text=False OR (saw_text_content AND NOT saw_tool_calls)=False→False, so the final round never runs andfinish_reason_per_choice[0]is never emitted
The existing test suite covers the case where text and tool-call choices arrive in the same chunk (test_n_gt_1_text_and_tool_call_in_same_chunk_no_text_leak) and text-only n>1 (test_per_choice_finish_reason_when_choices_finish_in_different_chunks), but not the scenario where they arrive in separate chunks.
A minimal fix is to add a fourth condition to should_run_final_round: or bool(finish_reason_per_choice). finish_reason_per_choice has entries only for deferred/unreported text-choice finish reasons (tool-only choices get their finish_reason on the passthrough and are not written there), so this triggers the final round only when needed. _build_transform_chunk with is_final=True would then emit an empty-delta chunk carrying finish_reason="stop" for choice 0.
| ) | ||
| holdback = 0 if is_final else max(0, holdback_per_choice.get(choice_idx, 0)) | ||
| end = max(len(already), len(text) - holdback) | ||
| deltas[choice_idx] = text[len(already) : end] |
There was a problem hiding this comment.
Medium: Late guardrail rewrites can expose already streamed text
With the default zero holdback, this emits the entire currently accepted prefix. An attacker can prompt the model to split sensitive content across processing rounds: the guardrail can pass an incomplete identifier initially and redact it after later characters arrive, but the subsequent startswith failure only aborts after the earlier bytes have reached the client. Retain a safe uncommitted suffix by default or buffer until end of stream for transformations that are not guaranteed to be prefix-stable, rather than relying solely on an optional guardrail-provided holdback.
…itellm_guardrail_streaming_text_transform
…g path Four bug fixes on top of the OSS PR's incremental_diff streaming text transformation, all inside the incremental_diff code paths only. No existing block_only, non-streaming, or pre_call behavior is touched. Fix #1 — Mixed content+tool_call finish_reason ordering _tool_call_passthrough_chunk now takes an optional finish_reason_per_choice map. For a choice carrying both delta.content and delta.tool_calls, finish_reason is stripped from the passthrough and recorded on the map so the final synthetic text chunk delivers it. Without this, SSE-compliant clients stopping at finish_reason drop the guardrailed text — defeating the redaction the whole feature exists for. (Greptile P1 twice, Veria.) Fix #2 — Choice index sort in _process_streaming_transform indices/texts_to_check were derived from dict insertion order. For n>1 streams where choice 1 emits before choice 0, guardrail-returned texts aligned to the input order mapped back to the wrong choice indices on write-back — wrong text goes to wrong choice. Sort raw_by_index.keys() up front so realignment is deterministic. (Bugbot Medium.) Fix #3 — Cross-chunk pre-tool-call text flush With default streaming_sampling_rate=5, text chunks followed by a pure tool-call chunk carrying finish_reason='tool_calls' would emit the passthrough with finish_reason before any transformed text delta had fired. Same failure mode as fix #1 but cross-chunk. Now we flush any accumulated text via _round(is_final=False) BEFORE yielding the tool-call passthrough. (Greptile P1.) Fix #4 — Terminator chunk for deferred finish_reason on empty mutated_text _build_transform_chunk returned None early when mutated_text_per_choice was empty. If a mixed content+tool_call chunk had deferred its finish_reason (via fix #1) and the guardrail then suppressed the text (empty return), the deferred finish_reason was never delivered. Now on is_final=True with empty mutated_text_per_choice, we emit a terminator carrying finish_reason per choice from finish_reason_per_choice. (Bugbot High.) Also normalized Optional[X] → X | None across the OSS PR's added surface via ruff UP045 autofix to keep the strict-rule gate within budget. Pure mechanical typing style change, no semantic effect. Regression tests for all four fixes: - test_mixed_chunk_finish_reason_arrives_after_transformed_text (#1) - test_text_flush_precedes_tool_call_passthrough (#3) - test_final_finish_reason_flushed_when_guardrail_suppresses_text (#4) - test_transform_sends_texts_sorted_by_choice_index (#2) All fixes reachable only when streaming_transform_mode == 'incremental_diff' is configured (via _run_incremental_transform_stream) or when a StreamTransformSink is present (via _process_streaming_transform). Verified scope-clean: no changes to block_only, non-streaming, pre_call, moderation, or sibling guardrails.
2ba08f2 to
527f110
Compare
Rescoped: only four correctness fixes on top of the OSS PRForce-pushed a clean history — single commit What this PR fixes (4 correctness bugs)
Follow-ups (opt-in path only, no correctness impact)
What was included in earlier revisions but dropped
Live e2e proof (still applies)The 6/6 e2e scenarios from comment above were run against a HEAD that included fixes 1-4 (semantically identical to |


Title
feat(guardrails): streaming text transformation in generic_guardrail_api
Relevant issues
Linear ticket
Pre-Submission checklist
Type
New Feature
Changes
Adds an opt-in
streaming_transform_mode: incremental_diffconfig ongeneric_guardrail_apiso an HTTP guardrail can rewrite streamed response text (PII masking, pseudonym reversal, redaction) instead of only blocking on it.The default
block_onlypreserves historical behavior: the guardrail can terminate a stream via aBLOCKEDaction, butGUARDRAIL_INTERVENEDtext rewrites were silently dropped on streams.incremental_diffwithholds the raw upstream chunks, runs the guardrail every N chunks over the accumulated text, and emits the guardrail's rewritten text as synthetic deltas by diffing the mutated accumulated text against what has already been streamed.Companion
stream_holdback_charson the guardrail response contract lets a guardrail server request the framework withhold N trailing chars per choice (word-boundary safety across rounds, e.g. pause on "Thomas W" until the next chunk arrives so the guardrail can decide whether to redact "Thomas Weber").v1 scope:
incremental_diff; other routes silently fall back toblock_only.delta.contenttext is transformed.delta.tool_callschunks pass through withcontentstripped and are inspected for a block decision at end of stream (matchingblock_only), but never text-rewritten.stream_transform_underflow; a guardrail that needs to rewrite recent output must withhold it first viastream_holdback_chars.Included:
StreamTransformSinkonBaseTranslation.process_output_streaming_responsefor handlers that support the streaming text-diff protocolStreamingChoices.index(not enumerate position) son>1streams do not collapsecoerce_stream_holdback_valuefor defensive parsing of malformed guardrail responsesincremental_diffmatchesblock_only's guardrail coverage)delta.content+delta.tool_callsin the same item, andn>1with text on one choice and a tool call on another)Credit
Adopted from #32084 by @schneidermr (PalenaAI). Mirrored onto a
litellm_branch so CircleCI and the internal lint workflow run against the current HEAD. Original PR: #32084Follow-ups (not in this PR)
docs.litellm.ai/docs/adding_provider/generic_guardrail_apineeds a "Streaming text transformations" section covering the diff protocol,stream_holdback_chars, the fail-closed rule, and the v1 scope. Source lives outside this monorepo.build_block_sse_chunksoverride to the OpenAI chat translation handler so a mid-streamstream_transform_underflowHTTPException surfaces as a well-formed SSE error frame +data: [DONE]instead of relying on the outerasync_data_generatorfallback. The current behavior mirrors the pre-existing pattern used by 9 other streaming guardrails on OpenAI-format streams, so this is not a regression, but underincremental_diffthe fail-closed path fires far more routinely than underblock_onlyand would benefit from a cleaner client signal.Note
High Risk
Large changes to proxy streaming guardrail behavior on the hot path for chat completions, including fail-closed suppression and tool-call/mixed-chunk edge cases where incorrect ordering could leak unredacted text.
Overview
Adds opt-in
streaming_transform_mode: incremental_diffongeneric_guardrail_apiso HTTP guardrails can rewrite streamed chat text (masking, redaction) instead of only blocking. Defaultblock_onlyis unchanged: streams still pass raw deltas and only BLOCK stops the stream; GUARDRAIL_INTERVENED text edits were previously dropped on streams.With
incremental_diff, the unified streaming hook withholds raw chunks, runs the guardrail on accumulated stringdelta.content(sampled or end-of-stream), and emits synthetic deltas by diffing guardrailed text against what was already sent. Guardrail responses can includestream_holdback_charsper choice for word-boundary withholding; rewrites that would retract already-sent bytes fail closed withstream_transform_underflow.Implementation spans
StreamTransformSinkonBaseTranslation.process_output_streaming_response, a non-mutating transform path in the OpenAI chat handler (accumulation byStreamingChoices.index), wiring inUnifiedLLMGuardrails(tool-call passthrough with content stripped, end-of-stream block inspection, mixed text+tool SSE ordering), and typed/config plumbing plus extensive regression tests.Reviewed by Cursor Bugbot for commit c47afcd. Bugbot is set up for automated code reviews on this repo. Configure here.