From 4d596082deecbf414abac3ee33e465717fa1e0e7 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:20:32 +0000 Subject: [PATCH] feat(guardrails): release buffered stream chunks after each passing scan Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/bedrock_guardrails.py | 3 + .../crowdstrike_aidr/__init__.py | 2 + .../crowdstrike_aidr/crowdstrike_aidr.py | 6 + .../unified_guardrail/unified_guardrail.py | 61 ++++---- .../guardrails/guardrail_initializers.py | 1 + litellm/types/guardrails.py | 6 + .../guardrail_hooks/crowdstrike_aidr.py | 8 + .../test_bedrock_guardrails.py | 3 + .../guardrail_hooks/test_crowdstrike_aidr.py | 13 ++ .../test_streaming_buffer_until_moderated.py | 139 +++++++++++++++++- 10 files changed, 215 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 2c407d91a48..db68c7b5c3f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py index c88e6e97a96..59f02817e5f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py +++ b/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, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index 8fed1f906e5..924bbd2bc1a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index ee5cd7c4cb8..f231ac68c6a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/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( diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 7858adeb55d..5dc9af2c081 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/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 diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 92fe41ba717..6346e13f3ba 100644 --- a/litellm/types/guardrails.py +++ b/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": diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py b/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py index 6beca030a3a..df1caab6af6 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py +++ b/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 " diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index d173c5f5c70..dc275c912c8 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/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(): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index 9849ad7ec88..a1aae119d56 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/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", [ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py index 2b163ee5233..89d51866533 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py +++ b/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