feat(guardrails): release buffered stream chunks after each passing scan

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py

@@ -248,6 +248,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         streaming_buffer_until_moderated: bool | None = None,
         streaming_sampling_rate: int | None = None,
         streaming_end_of_stream_only: bool | None = None,
+        streaming_buffer_release_on_scan: bool | None = None,
         **kwargs,
     ):
         self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
@@ -258,6 +259,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
                         "streaming_buffer_until_moderated": streaming_buffer_until_moderated,
                         "streaming_sampling_rate": streaming_sampling_rate,
                         "streaming_end_of_stream_only": streaming_end_of_stream_only,
+                        "streaming_buffer_release_on_scan": streaming_buffer_release_on_scan,
                     }
                 )
             )
@@ -321,6 +323,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         self.streaming_buffer_until_moderated = streaming_params.streaming_buffer_until_moderated
         self.streaming_sampling_rate = streaming_params.streaming_sampling_rate
         self.streaming_end_of_stream_only = streaming_params.streaming_end_of_stream_only
+        self.streaming_buffer_release_on_scan = streaming_params.streaming_buffer_release_on_scan
 
     def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
         super().update_in_memory_litellm_params(litellm_params)

litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py

@@ -23,6 +23,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
         event_hook=litellm_params.mode,
         default_on=litellm_params.default_on,
         fail_on_error=litellm_params.fail_on_error,
+        streaming_buffer_until_moderated=streaming_params.streaming_buffer_until_moderated,
+        streaming_buffer_release_on_scan=streaming_params.streaming_buffer_release_on_scan,
         streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only,
         streaming_sampling_rate=streaming_params.streaming_sampling_rate,
     )

litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py

@@ -260,6 +260,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
         api_key: str | None = None,
         api_base: str | None = None,
         fail_on_error: bool | None = True,
+        streaming_buffer_until_moderated: bool | None = None,
+        streaming_buffer_release_on_scan: bool | None = None,
         streaming_end_of_stream_only: bool | None = None,
         streaming_sampling_rate: int | None = None,
         async_handler: AsyncHTTPHandler | None = None,
@@ -287,6 +289,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
             CrowdStrikeAIDRGuardrailConfigModelOptionalParams(
                 streaming_end_of_stream_only=streaming_end_of_stream_only,
                 streaming_sampling_rate=streaming_sampling_rate,
+                streaming_buffer_until_moderated=streaming_buffer_until_moderated,
+                streaming_buffer_release_on_scan=streaming_buffer_release_on_scan,
             )
         )
 
@@ -310,6 +314,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
         )
 
     def _set_streaming_params(self, streaming_params: CrowdStrikeAIDRGuardrailConfigModelOptionalParams) -> None:
+        self.streaming_buffer_until_moderated: bool = streaming_params.streaming_buffer_until_moderated or False
+        self.streaming_buffer_release_on_scan: bool = streaming_params.streaming_buffer_release_on_scan or False
         self.streaming_end_of_stream_only: bool = streaming_params.streaming_end_of_stream_only or False
         self.streaming_sampling_rate: int = streaming_params.streaming_sampling_rate or 5
 

litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py

@@ -956,6 +956,7 @@ class UnifiedLLMGuardrails(CustomLogger):
         buffer_until_moderated: bool = _streaming_flag(
             "streaming_buffer_until_moderated", buffer_until_moderated_default
         )
+        release_on_scan: bool = _streaming_flag("streaming_buffer_release_on_scan", False)
 
         if (
             buffer_until_moderated
@@ -972,7 +973,7 @@ class UnifiedLLMGuardrails(CustomLogger):
 
         # Buffering can only moderate the assembled response, so it always
         # defers to end-of-stream.
-        if buffer_until_moderated:
+        if buffer_until_moderated and not release_on_scan:
             end_of_stream_only = True
 
         if guardrail_to_apply is None:
@@ -1026,6 +1027,7 @@ class UnifiedLLMGuardrails(CustomLogger):
         chunk_counter = 0
         responses_so_far: Final[list[object]] = []
         responses_yielded: Final[list[object]] = []
+        withheld_items: Final[list[object]] = []  # mutable-ok: streaming window must be released incrementally
         pending_end_of_stream_items: Final[list[object]] = []
         # Whether any real response chunk has been forwarded to the client.
         # Drives how a block terminates the stream: continue the in-progress
@@ -1069,9 +1071,13 @@ class UnifiedLLMGuardrails(CustomLogger):
                         chunks_yielded = True
                         responses_yielded.append(item)
                         yield item
+                else:
+                    withheld_items.append(item)
                 continue
 
             # Process chunk based on sampling rate
+            if buffer_until_moderated:
+                withheld_items.append(item)
             if chunk_counter % sampling_rate == 0:
                 endpoint_translation = mappings[CallTypes(call_type)]()
                 scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far)
@@ -1081,9 +1087,16 @@ class UnifiedLLMGuardrails(CustomLogger):
                         chunk_counter,
                         guardrail_to_apply.guardrail_name,
                     )
-                    chunks_yielded = True
-                    responses_yielded.append(item)
-                    yield item
+                    if buffer_until_moderated:
+                        for withheld_item in withheld_items:
+                            chunks_yielded = True
+                            responses_yielded.append(withheld_item)
+                            yield withheld_item
+                        withheld_items.clear()
+                    else:
+                        chunks_yielded = True
+                        responses_yielded.append(item)
+                        yield item
                     continue
 
                 verbose_proxy_logger.debug(
@@ -1093,13 +1106,9 @@ class UnifiedLLMGuardrails(CustomLogger):
                     guardrail_to_apply.guardrail_name,
                 )
 
-                # Deep-copy the current chunk before guardrail processing.
-                # process_output_streaming_response modifies responses_so_far
-                # in-place: it puts the combined guardrailed text in the first
-                # chunk and clears all subsequent chunks to "". Without this
-                # copy, yielding processed_items[-1] would yield an empty
-                # string, permanently losing this chunk's content.
-                original_item = copy.deepcopy(item)
+                original_items = (
+                    tuple(copy.deepcopy(withheld_items)) if buffer_until_moderated else (copy.deepcopy(item),)
+                )
 
                 try:
                     await endpoint_translation.process_output_streaming_response(
@@ -1144,13 +1153,16 @@ class UnifiedLLMGuardrails(CustomLogger):
                     return
                 if scan_key is not None:
                     last_scan_key = scan_key
-                chunks_yielded = True
-                responses_yielded.append(original_item)
-                yield original_item
+                for original_item in original_items:
+                    chunks_yielded = True
+                    responses_yielded.append(original_item)
+                    yield original_item
+                withheld_items.clear()
             else:
-                chunks_yielded = True
-                responses_yielded.append(item)
-                yield item
+                if not buffer_until_moderated:
+                    chunks_yielded = True
+                    responses_yielded.append(item)
+                    yield item
 
         # Stream has ended - do final processing with all collected chunks
         if call_type is not None and CallTypes(call_type) in mappings:
@@ -1162,14 +1174,13 @@ class UnifiedLLMGuardrails(CustomLogger):
 
             endpoint_translation = mappings[CallTypes(call_type)]()
 
-            # When buffering, snapshot the original chunks before moderation.
-            # A shallow copy suffices: end-of-stream
-            # process_output_streaming_response builds a separate assembled
-            # response (it does not mutate the individual chunks in place), and
-            # the chunks themselves are replayed verbatim -- so we only need to
-            # preserve the list, not clone every chunk (deepcopy would double
-            # peak memory for large responses).
-            buffered_items: Final = list(responses_so_far) if buffer_until_moderated else None
+            buffered_items: Final = (
+                tuple(copy.deepcopy(withheld_items))
+                if buffer_until_moderated and release_on_scan and not end_of_stream_only
+                else tuple(withheld_items)
+                if buffer_until_moderated
+                else None
+            )
             end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far)
             if _is_redundant_scan(end_scan_key, last_scan_key):
                 verbose_proxy_logger.debug(

litellm/proxy/guardrails/guardrail_initializers.py

@@ -44,6 +44,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
         streaming_buffer_until_moderated=streaming_params.streaming_buffer_until_moderated,
         streaming_sampling_rate=streaming_params.streaming_sampling_rate,
         streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only,
+        streaming_buffer_release_on_scan=streaming_params.streaming_buffer_release_on_scan,
     )
     litellm.logging_callback_manager.add_litellm_callback(_bedrock_callback)
     return _bedrock_callback

litellm/types/guardrails.py

@@ -682,6 +682,12 @@ class BedrockGuardrailStreamingParams(BaseModel):
         "and the scan result lands in guardrail_information; a flagged response still ends the "
         "stream with a block message (disable_exception_on_block=true) or an error frame.",
     )
+    streaming_buffer_release_on_scan: bool = Field(
+        default=False,
+        description="When buffering, scan the accumulated response every streaming_sampling_rate chunks "
+        "and release the withheld chunks once the scan passes, instead of holding everything to end of stream. "
+        "Flagged content is never released. Ignored when streaming_end_of_stream_only is true.",
+    )
 
     @classmethod
     def from_extras(cls, extras: Mapping[str, object] | None) -> "BedrockGuardrailStreamingParams":

litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py

@@ -4,6 +4,14 @@ from .base import GuardrailConfigModel
 
 
 class CrowdStrikeAIDRGuardrailConfigModelOptionalParams(BaseModel):
+    streaming_buffer_until_moderated: bool | None = Field(
+        default=None,
+        description="When True, withhold streamed chunks until moderation passes. Defaults to False when unset.",
+    )
+    streaming_buffer_release_on_scan: bool | None = Field(
+        default=None,
+        description="When buffering, release withheld chunks after each passing scan. Defaults to False when unset.",
+    )
     streaming_end_of_stream_only: bool | None = Field(
         default=None,
         description="If False (default when unset), post_call scans the accumulated streamed response every "

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py

@@ -5603,6 +5603,7 @@ def test_initialize_bedrock_wires_streaming_flags():
             streaming_buffer_until_moderated=False,
             streaming_sampling_rate=3,
             streaming_end_of_stream_only=True,
+            streaming_buffer_release_on_scan=True,
         ),
         {"guardrail_name": "bedrock-streaming"},
     )
@@ -5616,9 +5617,11 @@ def test_initialize_bedrock_wires_streaming_flags():
     assert configured.streaming_buffer_until_moderated is False
     assert configured.streaming_sampling_rate == 3
     assert configured.streaming_end_of_stream_only is True
+    assert configured.streaming_buffer_release_on_scan is True
     assert defaulted.streaming_buffer_until_moderated is True
     assert defaulted.streaming_sampling_rate == 5
     assert defaulted.streaming_end_of_stream_only is False
+    assert defaulted.streaming_buffer_release_on_scan is False
 
 
 def test_initialize_bedrock_rejects_non_positive_sampling_rate():

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py

@@ -1622,10 +1622,23 @@ def test_initialize_guardrail_rejects_unsupported_mode_instead_of_running_other_
 def test_initialize_guardrail_defaults_streaming_params() -> None:
     handler = _initialize_from_config(mode="post_call")
 
+    assert handler.streaming_buffer_until_moderated is False
+    assert handler.streaming_buffer_release_on_scan is False
     assert handler.streaming_end_of_stream_only is False
     assert handler.streaming_sampling_rate == 5
 
 
+def test_initialize_guardrail_forwards_buffer_streaming_params() -> None:
+    handler = _initialize_from_config(
+        mode="post_call",
+        streaming_buffer_until_moderated=True,
+        streaming_buffer_release_on_scan=True,
+    )
+
+    assert handler.streaming_buffer_until_moderated is True
+    assert handler.streaming_buffer_release_on_scan is True
+
+
 @pytest.mark.parametrize(
     "configured",
     [

tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py

@@ -11,7 +11,7 @@ released unchanged after moderation passes.
 """
 
 import json
-from typing import Any, List, Literal, Optional
+from typing import Any, AsyncGenerator, List, Literal, Optional
 
 import pytest
 
@@ -23,7 +23,7 @@ from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
     UnifiedLLMGuardrails,
 )
-from litellm.types.utils import GenericGuardrailAPIInputs
+from litellm.types.utils import Delta, GenericGuardrailAPIInputs, ModelResponseStream, StreamingChoices
 
 BLOCK_MESSAGE = "Blocked by policy: this response was withheld."
 ORIGINAL_MARKER = "ORIGINAL-SECRET-ANSWER"
@@ -60,6 +60,41 @@ class _PassingGuardrail(CustomGuardrail):
         return inputs
 
 
+class _CountingPassingGuardrail(_PassingGuardrail):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.scan_count = 0
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        self.scan_count += 1
+        return inputs
+
+
+class _SecondScanBlockingGuardrail(_CountingPassingGuardrail):
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        self.scan_count += 1
+        if self.scan_count == 2:
+            raise ModifyResponseException(
+                message=BLOCK_MESSAGE,
+                model="gpt-4",
+                request_data=request_data,
+                guardrail_name=self.guardrail_name,
+            )
+        return inputs
+
+
 def _sse_event(event_type: str, data: dict) -> bytes:
     return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
 
@@ -115,6 +150,67 @@ def _decode(chunks: List[Any]) -> str:
     return "".join(c.decode() if isinstance(c, bytes) else str(c) for c in chunks)
 
 
+def _chat_chunk(content: str = "", finish_reason: str | None = None) -> ModelResponseStream:
+    return ModelResponseStream(
+        id="chatcmpl-windowed",
+        created=1724900000,
+        model="gpt-4",
+        choices=[
+            StreamingChoices(
+                index=0,
+                delta=Delta(role="assistant", content=content),
+                finish_reason=finish_reason,
+            )
+        ],
+    )
+
+
+async def _windowed_chat_stream(
+    yielded_count: List[int], collected: List[Any], content_chunks: List[str]
+) -> AsyncGenerator[ModelResponseStream, None]:
+    for content in content_chunks:
+        yielded_count.append(len(collected))
+        yield _chat_chunk(content)
+    yielded_count.append(len(collected))
+    yield _chat_chunk(finish_reason="stop")
+
+
+async def _run_windowed(
+    guardrail: CustomGuardrail,
+    content_chunks: List[str],
+    end_of_stream_only: bool = False,
+) -> tuple[List[Any], List[int]]:
+    guardrail.streaming_buffer_until_moderated = True
+    guardrail.streaming_buffer_release_on_scan = True
+    guardrail.streaming_end_of_stream_only = end_of_stream_only
+    guardrail.streaming_sampling_rate = 2
+    unified = UnifiedLLMGuardrails()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions")
+    request_data = {
+        "messages": [{"role": "user", "content": "hi"}],
+        "guardrail_to_apply": guardrail,
+        "metadata": {"guardrails": [guardrail.guardrail_name]},
+    }
+    collected: List[Any] = []
+    yielded_count: List[int] = []
+    async for chunk in unified.async_post_call_streaming_iterator_hook(
+        user_api_key_dict=user_api_key_dict,
+        response=_windowed_chat_stream(yielded_count, collected, content_chunks),
+        request_data=request_data,
+    ):
+        collected.append(chunk)
+    return collected, yielded_count
+
+
+def _chat_text(chunks: List[Any]) -> str:
+    return "".join(
+        choice.delta.content or ""
+        for chunk in chunks
+        if isinstance(chunk, ModelResponseStream)
+        for choice in chunk.choices
+    )
+
+
 async def _run(guardrail: CustomGuardrail) -> str:
     # Rubrik's real config: end-of-stream-only moderation. Without buffering
     # this releases every chunk before moderation runs (content leaks on
@@ -159,6 +255,45 @@ async def test_buffered_clean_releases_all_content():
     assert BLOCK_MESSAGE not in raw
 
 
+@pytest.mark.asyncio
+async def test_windowed_buffer_releases_after_each_passing_scan():
+    guardrail = _CountingPassingGuardrail(guardrail_name="windowed-pass", event_hook="post_call")
+    content_chunks = ["one ", "two ", "three ", "four ", "five ", "six "]
+
+    collected, yielded_count = await _run_windowed(guardrail, content_chunks)
+
+    assert yielded_count[2] >= 2
+    assert yielded_count == [0, 0, 2, 2, 4, 4, 6]
+    assert _chat_text(collected) == "".join(content_chunks)
+    assert guardrail.scan_count > 1
+
+
+@pytest.mark.asyncio
+async def test_windowed_buffer_drops_blocked_window():
+    guardrail = _SecondScanBlockingGuardrail(guardrail_name="windowed-block", event_hook="post_call")
+    content_chunks = ["one ", "two ", "MARKER ", "four ", "five ", "six "]
+
+    collected, _ = await _run_windowed(guardrail, content_chunks)
+    raw = _decode(collected)
+
+    assert _chat_text(collected) == "one two "
+    assert "MARKER" not in raw
+    assert BLOCK_MESSAGE in raw
+    assert '"error"' not in raw
+
+
+@pytest.mark.asyncio
+async def test_windowed_buffer_with_explicit_end_of_stream_only_stays_fully_buffered():
+    guardrail = _CountingPassingGuardrail(guardrail_name="windowed-eos", event_hook="post_call")
+    content_chunks = ["one ", "two ", "three ", "four ", "five ", "six "]
+
+    collected, yielded_count = await _run_windowed(guardrail, content_chunks, end_of_stream_only=True)
+
+    assert yielded_count == [0, 0, 0, 0, 0, 0, 0]
+    assert _chat_text(collected) == "".join(content_chunks)
+    assert guardrail.scan_count == 1
+
+
 @pytest.mark.asyncio
 async def test_buffered_mode_disabled_for_content_rewriting_guardrail():
     """Buffered replay yields the withheld *original* chunks verbatim, which