Skip to content

feat(guardrails): streaming text transformation + compliance mode-match fix - #32084

Closed
schneidermr wants to merge 10 commits into
BerriAI:litellm_internal_stagingfrom
PalenaAI:litellm_guardrail_streaming_text_transform
Closed

feat(guardrails): streaming text transformation + compliance mode-match fix#32084
schneidermr wants to merge 10 commits into
BerriAI:litellm_internal_stagingfrom
PalenaAI:litellm_guardrail_streaming_text_transform

Conversation

@schneidermr

@schneidermr schneidermr commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Relevant issues

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

  • I have added meaningful tests
  • My PR passes all CI/CD checks (e.g., lint, format, unit tests)
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

Screenshots / Proof of Fix

Pending a live-proxy run. To exercise it end to end, point a generic_guardrail_api guardrail with streaming_transform_mode: incremental_diff at a real HTTP guardrail service that returns action: GUARDRAIL_INTERVENED with modified texts (and optionally stream_holdback_chars), then stream a /v1/chat/completions request against a real provider and confirm the client receives the transformed deltas whose concatenation equals the guardrailed text

Type

🆕 New Feature
🐛 Bug Fix

Changes

This PR contains two changes. The primary one is the streaming text-transformation feature; the second is a small, self-contained fix to ComplianceChecker that surfaced while testing the feature against a mode-based compliance setup. They are described separately below and can be split into their own commits if a reviewer would prefer that

Streaming text transformation (feature)

LiteLLM's HTTP guardrail protocol is bidirectional: a guardrail service can return action: GUARDRAIL_INTERVENED with a modified texts list, and the non-streaming chat/completions post-call path already writes those texts back before returning to the client. The streaming path did not propagate them. It accumulated the mutated text into the response accumulator but always yielded the raw pre-mutation chunk, so any guardrail whose job is to rewrite text (PII masking, pseudonym reversal, redaction, translation) could not affect a streamed response over HTTP; only the in-process custom_guardrail on_stream_* hooks could

This adds an opt-in streaming_transform_mode with two values. block_only is the default and preserves today's behavior exactly. incremental_diff withholds the raw chunks and instead emits the guardrailed accumulated text as new deltas, computed by diffing the mutated text against what has already been sent to the client

How it works: streaming_transform_mode lives on GenericGuardrailAPIOptionalParams, is read from guardrail_config and optional_params like the sibling streaming_* flags, and is forwarded through initialize_guardrail. On each sampled round the guardrail runs over the raw accumulated text (the accumulator is never mutated in place, so a rewrite guardrail always sees consistent input), keyed by StreamingChoices.index so n > 1 completions stay correct, and the hook emits only the newly guardrailed portion per choice. A new stream_holdback_chars field on the guardrail response (indexed like texts) lets the guardrail withhold a number of trailing chars until the next round, so a pseudonym like "Thomas We" is never emitted before "Weber" fully arrives; the framework carries the holdback forward and forces it to 0 on the end-of-stream flush. The final flush chunk carries each choice's own finish_reason from the last raw chunk to preserve OpenAI wire semantics, and a terminal chunk is guardrailed once (by the flush) rather than twice

The path fails closed rather than silently leaking. If a transform is not a forward extension of what was already streamed (shorter than, or rewriting, the emitted prefix), the framework raises HTTPException(400, stream_transform_underflow) because emitted bytes cannot be retracted; for A2A call types this surfaces as an in-stream JSON-RPC error. Synthetic chunks also drop tool_calls, since v1 does not transform streamed tool calls and passing the raw upstream ones through would bypass the guardrail

The guardrailed text and requested holdback travel from the handler back to the hook through a typed StreamTransformSink out-parameter on process_output_streaming_response, rather than by mutating responses_so_far (which must stay a correct raw accumulator) or changing the method's return type (several existing handler tests assert on it)

Scope: this v1 targets the OpenAI chat completions streaming path with string delta.content only. Non-OpenAI translation handlers (Anthropic, Gemini), list-of-blocks content, and image or tool-call streaming transforms remain block_only and are candidates for follow-up PRs. incremental_diff is gated to the OpenAI chat handler via the resolved request route; any other surface logs a warning and falls back to block_only

ComplianceChecker mode matching (fix)

A guardrail's guardrail_mode can be a str, a list, or a tag-based dict, but the mode matcher compared it to the target mode with ==, so a guardrail configured with mode: [pre_call, post_call] matched no mode and every mode-based compliance check reported NON-COMPLIANT. Mode matching now handles all of those shapes

Tests

Unit tests cover block_only dropping rewrites, incremental_diff emitting transformed deltas whose concatenation equals the full guardrailed text, holdback boundary behavior with no loss or duplication, the fail-closed underflow and prefix-rewrite cases, per-choice finish_reason preservation for n > 1, tool_calls being dropped from synthetic chunks, the guardrail always receiving the raw accumulated text, index-based accumulation for non-zero choice indices, a terminal chunk not being guardrailed twice, the end-of-stream-only single-chunk path, unsupported-route fallback, and stream_holdback_chars parsing (including malformed values degrading to 0) on GenericGuardrailAPIResponse.from_dict. make lint passes with zero net budget additions

Docs: a follow-up should add a "Streaming text transformations" section to the generic guardrail API docs page documenting streaming_transform_mode, stream_holdback_chars, the diff-based emission semantics, and the fail-closed rule


Note

Medium Risk
Touches live streaming response paths and guardrail wire semantics; misconfiguration or guardrail rewrites that shrink/revise already-emitted text can error mid-stream, though defaults preserve prior behavior.

Overview
Adds opt-in streaming_transform_mode: incremental_diff so HTTP guardrails can stream rewritten text (PII masking, redaction, etc.) on OpenAI chat completions: raw chunks are withheld, guardrails run on the raw accumulator via StreamTransformSink, and the client gets synthetic deltas diffed against what was already sent, with optional stream_holdback_chars for word-boundary safety. Default block_only is unchanged. The path fails closed on non-forward transforms (stream_transform_underflow), drops streamed tool_calls from synthetic chunks, and falls back to block-only when the route isn’t OpenAI chat streaming.

Separately fixes ComplianceChecker so guardrails logged with guardrail_mode as a list or tag-based dict count toward mode-based compliance (previously list modes never matched).

Reviewed by Cursor Bugbot for commit 22375d6. Bugbot is set up for automated code reviews on this repo. Configure here.

@schneidermr

Copy link
Copy Markdown
Contributor Author

@greptileai

@codecov

codecov Bot commented Jul 3, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 88.94009% with 24 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...drail_hooks/unified_guardrail/unified_guardrail.py 87.91% 18 Missing ⚠️
.../llms/openai/chat/guardrail_translation/handler.py 87.75% 6 Missing ⚠️

📢 Thoughts on this report? Let us know!

@greptile-apps

greptile-apps Bot commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds opt-in streaming_transform_mode: incremental_diff for the OpenAI chat completions streaming path, enabling HTTP guardrails (PII masking, redaction, etc.) to deliver rewritten text to clients as synthetic deltas diffed against already-emitted bytes. It also claims to fix ComplianceChecker mode matching for list-valued guardrail_mode, but that file is absent from the diff.

  • Streaming transform feature: Raw chunks are withheld; on each sampled round the guardrail runs against the immutable raw accumulator via StreamTransformSink, and only the newly-guardrailed portion is forwarded. The path fails closed on non-forward transforms (stream_transform_underflow), drops tool-call content from synthetic chunks, and falls back to block_only for non-OpenAI-chat routes. Default behavior (block_only) is unchanged.
  • Missing ComplianceChecker fix: litellm/proxy/compliance_checks.py is not in the diff; _get_guardrails_by_mode still uses elif g_mode == mode:, which silently drops guardrails whose guardrail_mode is a list (e.g. [\"pre_call\", \"post_call\"]), causing those guardrails to never count toward mode-based compliance checks.
  • Tests: ~620 new unit tests cover incremental_diff, holdback, underflow, tool-call passthrough, per-choice finish_reason preservation, A2A in-stream errors, and end-of-stream-only mode — all mocked with no network calls.

Confidence Score: 4/5

The streaming transform feature is safe to merge, but the claimed ComplianceChecker mode-matching fix is missing and compliance checks remain broken for list-valued guardrail modes.

The streaming transform implementation is thoroughly tested and defaults to the unchanged block_only path, so it carries little regression risk. However, the PR description explicitly claims to fix ComplianceChecker._get_guardrails_by_mode for list and dict modes — a fix absent from the diff. The equality check elif g_mode == mode: is still in place, meaning a guardrail configured with mode: [pre_call, post_call] is silently excluded from compliance accounting and mode-based compliance checks report NON-COMPLIANT even when those guardrails ran correctly.

litellm/proxy/compliance_checks.py — _get_guardrails_by_mode still uses plain equality and needs to be updated to handle list and dict guardrail_mode values as described in the PR.

Important Files Changed

Filename Overview
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py Adds the full incremental_diff streaming transform path (~420 lines): chunk accumulation, synthetic-delta emission, fail-closed underflow guard, A2A in-stream error handling, and tool-call passthrough with block inspection. Logic is well-structured and unit-tested; default block_only behavior is unchanged.
litellm/llms/openai/chat/guardrail_translation/handler.py Splits process_output_streaming_response into block-only and transform paths; adds _process_streaming_transform and _accumulate_string_content_by_choice_index. StreamTransformSink out-parameter keeps responses_so_far immutable across rounds. Implementation looks correct.
litellm/llms/base_llm/guardrail_translation/base_translation.py Adds StreamTransformSink dataclass and optional stream_transform_sink parameter to process_output_streaming_response; base implementation ignores the sink. Changes are additive and backward-compatible.
litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py Adds streaming_transform_mode field to GenericGuardrailAPIOptionalParams, stream_holdback_chars to GenericGuardrailAPIResponse, and the coerce_stream_holdback_value helper. from_dict parsing is robust (malformed values degrade to 0).
litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py Wires streaming_transform_mode into GenericGuardrailAPI.init and propagates stream_holdback_chars from guardrail response to return_inputs. Changes are additive.
litellm/proxy/compliance_checks.py NOT modified by this PR, but the PR description claims a mode-matching fix was applied here. _get_guardrails_by_mode still uses plain equality (g_mode == mode) which silently drops list-valued guardrail_mode entries, leaving compliance checks reporting NON-COMPLIANT for those guardrails.
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py Adds ~620 lines of unit tests for the streaming transform path covering incremental_diff, holdback, underflow, tool-call passthrough, per-choice finish_reason, A2A error emission, and end-of-stream-only; all use mocks with no real network calls.
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py Adds tests for streaming_transform_mode defaults, GenericGuardrailAPIResponse.from_dict parsing of stream_holdback_chars (including malformed values), and apply_guardrail holdback propagation. All mocked, no network calls.

Reviews (3): Last reviewed commit: "fix(guardrails): strip content from tool..." | Re-trigger Greptile

@schneidermr
schneidermr force-pushed the litellm_guardrail_streaming_text_transform branch from f01dc11 to 6e7aee8 Compare July 3, 2026 22:30
Comment thread litellm/llms/base_llm/guardrail_translation/base_translation.py
Comment thread litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py Outdated
@schneidermr
schneidermr marked this pull request as ready for review July 3, 2026 22:34
Comment thread litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py Outdated
@veria-ai

veria-ai Bot commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

PR overview

All previously flagged issues have been addressed. No open security concerns remain on this pull request.

Security review

No open security issues remain on this pull request.

Fixed/addressed: 3 · PR risk: 0/10

Comment thread litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py Outdated
@schneidermr

Copy link
Copy Markdown
Contributor Author

If you require an HTTP guardrail that transforms the text for the review, you can use this one: https://github.com/PalenaAI/palena-litellm-pseudonymizer

@krrish-berri-2

Copy link
Copy Markdown
Contributor

please address the veria comments.

cc: @yucheng-berri for review

@yucheng-berri

Copy link
Copy Markdown
Contributor

bugbot run

Comment thread litellm/llms/openai/chat/guardrail_translation/handler.py Outdated
Comment thread litellm/llms/openai/chat/guardrail_translation/handler.py
Comment thread litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py Outdated
@yucheng-berri

Copy link
Copy Markdown
Contributor

Hi, thanks for the effort here!

I have a couple of comments:

  1. I believe compliance_checks.py is outside the scope of the current PR description. Could you either update this PR to include that scope or move those changes into a separate PR?
  2. Could you also review and address the Bugbot comments?

@schneidermr schneidermr changed the title feat(guardrails): support streaming text transformation in generic_guardrail_api feat(guardrails): streaming text transformation + compliance mode-match fix Jul 7, 2026
@schneidermr

Copy link
Copy Markdown
Contributor Author

Hi, thanks for the effort here!

I have a couple of comments:

  1. I believe compliance_checks.py is outside the scope of the current PR description. Could you either update this PR to include that scope or move those changes into a separate PR?
  2. Could you also review and address the Bugbot comments?

Hi @yucheng-berri

The compliance_checks (bug)fix is also related to this PR. The PR’s main objective is to implement the missing streaming support for text transformation guardrails. If we utilize such guardrails, it’s highly likely that we’ll apply them to both pre- and post-call processes. However, due to the compliance_checks bug, the compliance validator cannot effectively handle these use cases without this fix. Consequently, the UI will incorrectly mark every such request as non-compliant with GDPR and the EU AI Act.

The compliance_checks function compares an array with a string. It works fine if only one mode is selected. However, if the user adds a guardrail with more than one mode, the compliance_checks string comparison will fail.

@schneidermr

Copy link
Copy Markdown
Contributor Author

@cursoragent review

@cursor

cursor Bot commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

Unable to authenticate your request. Please make sure to connect your GitHub account to Cursor. Go to Cursor

@yucheng-berri

Copy link
Copy Markdown
Contributor

bugbot run

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes using high effort and found 2 potential issues.

Fix All in Cursor

❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.

Reviewed by Cursor Bugbot for commit 22375d6. Configure here.

Comment thread litellm/llms/openai/chat/guardrail_translation/handler.py Outdated
@yucheng-berri

yucheng-berri commented Jul 8, 2026

Copy link
Copy Markdown
Contributor

Thanks for getting back.

I’m not comfortable approving yet though, since the live-proxy run is pending and I don’t see any follow-up proof.

For a streaming guardrail feature, mock tests aren’t enough; we should see an actual proxy run showing the client only receives the transformed text and not the raw sensitive text.

Also, can you add more targeted tests for ComplianceChecker as well. Thanks

@schneidermr

Copy link
Copy Markdown
Contributor Author

@yucheng-berri Thanks for getting back to me!

I’ve taken care of your request about ComplianceChecker targeted tests and also resolved the issue BugBot pointed out.

You’ll find some extra information below that might be helpful as you review it. However, since we’re a bit short on time, if this isn’t quite enough, I’ll need to create a custom build myself with these changes.


Live proxy run — streaming transform, chunk-by-chunk

Ran against a real litellm proxy (Postgres-backed) with a generic_guardrail_api
guardrail configured in streaming_transform_mode: incremental_diff. The upstream
model emits the token Jordan Avery; the guardrail transforms it to
Alice Johnson in-stream. Capturing the raw SSE the client receives:

$ curl -sN localhost:4000/v1/chat/completions \
    -H 'authorization: Bearer ***' -H 'content-type: application/json' \
    -d '{"model":"mock-gpt","stream":true,
         "messages":[{"role":"user","content":"Please book Alice Johnson for the trip."}]}'

# per-chunk delta.content actually delivered to the client:
chunk[0] = 'Alice Johnson ha'
chunk[1] = 's been booked f'
chunk[2] = 'or the trip suc'
chunk[3] = 'cessfully. I wi'
chunk[4] = 'll contact them'
chunk[5] = ' with the detai'
chunk[6] = 'ls.'

reassembled = 'Alice Johnson has been booked for the trip successfully. I will contact them with the details.'

Assertions on the client-visible stream:

  • Raw upstream token Jordan Avery present in ANY chunk? → False
  • Transformed Alice Johnson present? → True

The very first chunk is already 'Alice Johnson ha' — the framework withheld the
differing prefix via stream_holdback_chars until the transform was known, so the
untransformed upstream bytes never reach the client on any chunk boundary. This
is the guarantee the feature exists to provide, verified end-to-end through the
proxy (not a mock).

Fail-closed is also exercised: with the guardrail's detector unreachable, the proxy
returns HTTP 500 (unreachable_fallback: fail_closed) and the request is blocked —
the raw text is never streamed.

Note on direction: here the pre-transform bytes are the provider-side pseudonym
and the transform restores the client's own value; the holdback/no-leak mechanism
is identical for a redaction-style guardrail (client never sees the pre-redaction
sensitive bytes).

@codspeed-hq

codspeed-hq Bot commented Jul 8, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing PalenaAI:litellm_guardrail_streaming_text_transform (b9c7e46) with litellm_internal_staging (11aeeea)1

Open in CodSpeed

Footnotes

  1. No successful run was found on litellm_internal_staging (270406b) during the generation of this report, so 11aeeea was used instead as the comparison base. There might be some changes unrelated to this pull request in this report.

@yucheng-berri

yucheng-berri commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

Hi, thanks for adding the live test results. I have a couple of remaining concerns I’d like to understand:

  1. In the SSE path, mid-stream underflow appears to re-raise after the stream has already started, so the client may not receive a clear fail-closed signal.

  2. Tool-call chunks appear to be yielded raw and skip guardrail inspection, which seems like a security regression compared to block_only.

Could you share your thinking on these cases? I’d like to understand whether these are intended limitations of incremental_diff or issues we should address before considering the feature complete.

Also, it would be great if the compliance_checks change could be split into a separate PR, which will help us to maintain our repo.

@schneidermr
schneidermr force-pushed the litellm_guardrail_streaming_text_transform branch from 3e725cb to d895afa Compare July 9, 2026 21:23
@schneidermr

Copy link
Copy Markdown
Contributor Author

@yucheng-berri Thanks for the careful review. My thinking on each:

1. Mid-stream underflow re-raise

This matches the existing block_only behavior for OpenAI chat rather than being something incremental_diff introduced. The OpenAI chat translation handler has no build_block_sse_chunks override (only the A2A handler does), so a mid-stream fail-closed re-raises there today in both modes; for A2A the underflow already surfaces as an in-stream JSON-RPC error. I agree it is not a clean client signal though. The proper fix is to add a build_block_sse_chunks for the OpenAI chat handler that emits a terminating SSE error frame plus [DONE], which would improve both block_only and incremental_diff. I left it out of this PR because it changes block_only's behavior for existing OpenAI moderation guardrails, but I am happy to add it here or as a fast follow if you want it in scope.

2. Tool-call chunks skipping guardrail inspection

You are right, that was a real gap versus block_only, and I have fixed it. block_only assembles the full response with stream_chunk_builder at end of stream and passes tool_calls to the guardrail, so it inspects them; my incremental path was only sending text. Now the tool-call chunks are still delivered raw (v1 does not transform tool calls), but they are kept in the accumulator and the guardrail's block inspection runs over the full assembled response including the tool calls at end of stream, matching block_only. A guardrail that blocks on a tool call terminates the stream instead of being bypassed. New tests cover both that the tool call is inspected and that a tool-call-blocking guardrail terminates the incremental_diff stream.

One nuance worth flagging: like block_only's sampled mode, a raw tool-call chunk is delivered before the end-of-stream block lands. If you want block-before-delivery for tool calls under streaming_end_of_stream_only, I can withhold the tool-call chunks until the inspection passes; let me know if that should be in scope.

3. Splitting compliance_checks

Will do. I am moving the ComplianceChecker mode-matching fix and its tests into a separate PR so this one is scoped to the streaming feature only. I will link it here once it is up.

@schneidermr

schneidermr commented Jul 9, 2026

Copy link
Copy Markdown
Contributor Author

@yucheng-berri #32676

Comment thread litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py Outdated
@yucheng-berri

Copy link
Copy Markdown
Contributor

@greptileai

@yucheng-berri

Copy link
Copy Markdown
Contributor

merged in #33110, please follow up on this. Thanks!

@schneidermr

Copy link
Copy Markdown
Contributor Author

merged in #33110, please follow up on this. Thanks!

Thank you so much to @yucheng-berri and @mateo-berri

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants