scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 134

Merge pull request #41425 from BerriAI/litellm_streaming_buffer_release_on_scan

Browse source 
feat(guardrails): release buffered stream chunks after each passing scan
This commit is contained in:
Yassin Kortam 2026-09-16 14:37:56 -07:00 • committed by GitHub
commit cd08c65002
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 656 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -1561,10 +1561,12 @@ class AnthropicMessagesHandler(BaseTranslation):
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
stream_ended: Final = self._check_streaming_has_ended(responses_so_far)
tool_use_fingerprints: Final = self._streamed_tool_use_fingerprints(responses_so_far)
return StreamingScanKey(
texts=(self.get_streaming_string_so_far(responses_so_far),),
tool_calls=self._streamed_tool_use_fingerprints(responses_so_far) if stream_ended else (),
tool_calls=tool_use_fingerprints if stream_ended else (),
stream_ended=stream_ended,
tool_calls_in_flight=bool(tool_use_fingerprints) and not stream_ended,
)
@classmethod

6
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -40,11 +40,15 @@ class StreamingScanKey:
"""What a streaming guardrail round would hand to ``apply_guardrail``. Two keys
compare equal when the round would scan the same content again; ``stream_ended``
stays out of the comparison and only says whether the handler is on its
end-of-stream path, where an empty payload is still scanned today."""
end-of-stream path, where an empty payload is still scanned today.
``tool_calls_in_flight`` also stays out of the comparison: it flags that tool
calls have streamed which this round cannot scan yet, so a buffered window
holding them must stay withheld until the end-of-stream scan covers them."""
texts: tuple[str, ...]
tool_calls: tuple[str, ...] = ()
stream_ended: bool = field(default=False, compare=False)
tool_calls_in_flight: bool = field(default=False, compare=False)
@property
def has_nothing_to_scan(self) -> bool:

12
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -792,10 +792,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream))
stream_ended: Final = self._first_choice_has_finished(responses_so_far)
tool_call_fingerprints: Final = self._streamed_tool_call_fingerprints(responses_so_far)
return StreamingScanKey(
texts=tuple(self._combine_streaming_texts(chunks).values()),
tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (),
tool_calls=tool_call_fingerprints if stream_ended else (),
stream_ended=stream_ended,
tool_calls_in_flight=bool(tool_call_fingerprints) and not stream_ended,
)
@staticmethod
@ -804,7 +806,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
stream_item_fingerprint(tool_call)
for chunk in responses_so_far
for choice in _stream_chunk_choices(chunk)
for tool_call in stream_item_items(stream_item_field(choice, "delta"), "tool_calls")
for tool_call in _streamed_delta_tool_calls(stream_item_field(choice, "delta"))
)
@staticmethod
@ -1342,6 +1344,12 @@ def _stream_chunk_choices(item: object) -> Sequence[object]:
return ()
def _streamed_delta_tool_calls(delta: object) -> tuple[object, ...]:
function_call: Final = stream_item_field(delta, "function_call")
legacy: Final = () if function_call is None else (function_call,)
return stream_item_items(delta, "tool_calls") + legacy
def _blocked_stream_identity(
exc: "ModifyResponseException", responses_so_far: Sequence[object]
) -> tuple[str, int, str]:

15
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -1175,11 +1175,22 @@ class OpenAIResponsesHandler(BaseTranslation):
last_event_type: Final = stream_item_field(last_event, "type")
if last_event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value:
return None
if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value:
if last_event_type in _TERMINAL_ENVELOPE_EVENT_TYPES:
return self._completed_response_scan_key(stream_item_field(last_event, "response"))
return StreamingScanKey(
texts=(self.get_streaming_string_so_far(responses_so_far),),
stream_ended=self._check_streaming_has_ended(responses_so_far),
tool_calls_in_flight=self._has_streamed_tool_call_events(responses_so_far),
)
@staticmethod
def _has_streamed_tool_call_events(responses_so_far: Sequence[object]) -> bool:
return any(
stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES
or (
stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES
and stream_item_field(stream_item_field(event, "item"), "type") in _TOOL_CALL_ITEM_TYPES
)
for event in responses_so_far
)
@staticmethod

6
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -37,6 +37,7 @@ from litellm.types.guardrails import (
ApplyGuardrailResponse,
BaseLitellmParams,
BedrockGuardrailConfigModel,
BedrockGuardrailStreamingParams,
Guardrail,
GuardrailEventHooks,
GuardrailInfoResponse,
@ -1959,7 +1960,10 @@ async def get_provider_specific_params():
```
"""
# Get fields from the models
bedrock_fields: Final = _get_fields_from_model(BedrockGuardrailConfigModel)
bedrock_fields: Final = {
**_get_fields_from_model(BedrockGuardrailConfigModel),
**_get_fields_from_model(BedrockGuardrailStreamingParams),
}
presidio_fields: Final = _get_fields_from_model(PresidioPresidioConfigModelUserInterface)
lakera_v2_fields: Final = _get_fields_from_model(LakeraV2GuardrailConfigModel)
tool_permission_fields: Final = _get_fields_from_model(ToolPermissionGuardrailConfigModel)

9
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -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,13 +323,18 @@ 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)
self._set_streaming_params(BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra))
def _streams_incrementally(self) -> bool:
return not self.streaming_buffer_until_moderated and not self.mask_response_content
if self.mask_response_content:
return False
if not self.streaming_buffer_until_moderated:
return True
return self.streaming_buffer_release_on_scan and not self.streaming_end_of_stream_only
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:

2
litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py
View file

@ -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,
)

6
litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py
View file

@ -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

77
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -956,6 +956,7 @@ class UnifiedLLMGuardrails(CustomLogger):
buffer_until_moderated: bool = _streaming_flag(
"streaming_buffer_until_moderated", buffer_until_moderated_default
)
release_on_scan: Final[bool] = _streaming_flag("streaming_buffer_release_on_scan", False)
if (
buffer_until_moderated
@ -970,9 +971,7 @@ class UnifiedLLMGuardrails(CustomLogger):
)
buffer_until_moderated = False
# 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,12 +1025,14 @@ 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
# message (True) vs emit a standalone block message (False, buffered).
chunks_yielded = False
last_scan_key: StreamingScanKey | None = None # rebind-ok: replaced after every scan round
tool_calls_in_flight = False # rebind-ok: tracks the latest scan key's unscanned tool calls
async for item in response:
chunk_counter += 1
@ -1069,21 +1070,37 @@ 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)
if scan_key is not None:
tool_calls_in_flight = scan_key.tool_calls_in_flight
hold_window = buffer_until_moderated and (scan_key is None or tool_calls_in_flight)
if _is_redundant_scan(scan_key, last_scan_key):
verbose_proxy_logger.debug(
"Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round",
chunk_counter,
guardrail_to_apply.guardrail_name,
)
chunks_yielded = True
responses_yielded.append(item)
yield item
if buffer_until_moderated:
if hold_window:
continue
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 +1110,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 +1157,24 @@ 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
if hold_window:
verbose_proxy_logger.debug(
"Holding %s buffered chunks for guardrail %s: this round could not scan the whole window",
len(withheld_items),
guardrail_to_apply.guardrail_name,
)
withheld_items[:] = original_items
continue
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 +1186,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(

1
litellm/proxy/guardrails/guardrail_initializers.py
View file

@ -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

6
litellm/types/guardrails.py
View file

@ -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":

8
litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py
View file

@ -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 "

3
tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py
View file

@ -2482,7 +2482,10 @@ class TestAnthropicMessagesHandlerStreamingScanKey:
open_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use])
ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use, self._stop("tool_use")])
assert open_key == StreamingScanKey(texts=("hi",))
assert open_key.tool_calls_in_flight is True
assert handler.get_streaming_scan_key([self._text_delta("hi")]).tool_calls_in_flight is False
assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
assert ended_key.tool_calls_in_flight is False
assert ended_key != open_key

25
tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py
View file

@ -2206,10 +2206,35 @@ class TestStreamingScanKey:
[self._chunk("hi"), tool_chunk, self._chunk(None, finish_reason="stop")]
)
assert open_key == StreamingScanKey(texts=("hi",))
assert open_key.tool_calls_in_flight is True
assert handler.get_streaming_scan_key([self._chunk("hi")]).tool_calls_in_flight is False
assert ended_key.texts == ("hi",)
assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
assert ended_key.tool_calls_in_flight is False
assert ended_key != open_key
def test_legacy_function_call_delta_is_held_like_a_tool_call(self):
from litellm.types.utils import Delta, FunctionCall, ModelResponseStream, StreamingChoices
handler = OpenAIChatCompletionsHandler()
function_chunk = ModelResponseStream(
choices=[
StreamingChoices(
index=0,
delta=Delta(content=None, function_call=FunctionCall(name="run_shell", arguments='{"cmd": "rm"}')),
finish_reason=None,
)
]
)
open_key = handler.get_streaming_scan_key([self._chunk("hi"), function_chunk])
ended_key = handler.get_streaming_scan_key(
[self._chunk("hi"), function_chunk, self._chunk(None, finish_reason="function_call")]
)
assert open_key.tool_calls_in_flight is True
assert open_key.tool_calls == ()
assert len(ended_key.tool_calls) == 1 and "run_shell" in ended_key.tool_calls[0]
assert ended_key.tool_calls_in_flight is False
def test_text_after_the_first_choice_finishes_still_changes_the_key(self):
handler = OpenAIChatCompletionsHandler()
first_done = [self._chunk("a", index=0), self._chunk("b", finish_reason="stop", index=0)]

39
tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
View file

@ -3211,3 +3211,42 @@ class TestOpenAIResponsesHandlerStreamingScanKey:
def test_output_item_done_round_is_never_deduped(self):
done = {"type": "response.output_item.done", "sequence_number": 1, "item": {"type": "function_call"}}
assert OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hi"), done]) is None
@pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"])
def test_non_completed_terminal_envelopes_key_their_output_items(self, terminal_type):
handler = OpenAIResponsesHandler()
arguments_delta = {
"type": "response.function_call_arguments.delta",
"sequence_number": 1,
"item_id": "fc_1",
"delta": '{"city":',
}
function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": '{"city":'}
terminal = {"type": terminal_type, "sequence_number": 2, "response": {"id": "resp_1", "output": [function_call]}}
mid_stream_key = handler.get_streaming_scan_key([arguments_delta])
ended_key = handler.get_streaming_scan_key([arguments_delta, terminal])
assert ended_key.stream_ended is True
assert ended_key.tool_calls_in_flight is False
assert len(ended_key.tool_calls) == 1
assert ended_key != mid_stream_key
def test_streamed_tool_call_events_flag_tool_calls_in_flight_until_the_stream_ends(self):
handler = OpenAIResponsesHandler()
added = {
"type": "response.output_item.added",
"sequence_number": 1,
"item": {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "get_weather"},
}
arguments_delta = {
"type": "response.function_call_arguments.delta",
"sequence_number": 2,
"item_id": "fc_1",
"delta": '{"city":',
}
function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "{}"}
assert handler.get_streaming_scan_key([self._delta(0, "hi")]).tool_calls_in_flight is False
assert handler.get_streaming_scan_key([self._delta(0, "hi"), added]).tool_calls_in_flight is True
assert handler.get_streaming_scan_key([self._delta(0, "hi"), arguments_delta]).tool_calls_in_flight is True
ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), added, self._completed(3, [function_call])])
assert ended_key.tool_calls_in_flight is False
assert len(ended_key.tool_calls) == 1

41
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -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():
@ -5721,6 +5724,44 @@ async def test_buffered_default_hook_scans_before_any_chunk():
assert len([e for e in events if e != "scan"]) >= 1
@pytest.mark.asyncio
async def test_buffered_release_on_scan_hook_releases_each_window_after_its_scan():
guardrail = BedrockGuardrail(
guardrail_name="bedrock-release-on-scan",
guardrailIdentifier="test-id",
guardrailVersion="DRAFT",
event_hook=GuardrailEventHooks.post_call,
default_on=True,
streaming_buffer_release_on_scan=True,
streaming_sampling_rate=1,
)
assert guardrail._streams_incrementally() is True
events = await _run_streaming_hook_recording_order(guardrail)
assert events == ["scan", ("chunk", "Hello"), "scan", ("chunk", " world"), ("chunk", "")]
@pytest.mark.asyncio
async def test_buffered_release_on_scan_defers_to_end_of_stream_only():
guardrail = BedrockGuardrail(
guardrail_name="bedrock-release-on-scan-end-only",
guardrailIdentifier="test-id",
guardrailVersion="DRAFT",
event_hook=GuardrailEventHooks.post_call,
default_on=True,
streaming_buffer_release_on_scan=True,
streaming_end_of_stream_only=True,
streaming_sampling_rate=1,
)
assert guardrail._streams_incrementally() is False
events = await _run_streaming_hook_recording_order(guardrail)
assert events.count("scan") == 1
assert events[0] == "scan"
@pytest.mark.asyncio
async def test_masking_keeps_buffered_path_even_when_unbuffered_configured():
guardrail = BedrockGuardrail(

13
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py
View file

@ -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",
[

404
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py
View file

@ -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
@ -19,14 +19,25 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
ModifyResponseException,
)
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
_is_redundant_scan,
)
from litellm.types.utils import (
ChatCompletionDeltaToolCall,
Delta,
Function,
FunctionCall,
GenericGuardrailAPIInputs,
ModelResponseStream,
StreamingChoices,
)
from litellm.types.utils import GenericGuardrailAPIInputs
BLOCK_MESSAGE = "Blocked by policy: this response was withheld."
ORIGINAL_MARKER = "ORIGINAL-SECRET-ANSWER"
TOOL_ARGUMENTS_MARKER = "TOOL-ARGS-SECRET"
class _BlockingGuardrail(CustomGuardrail):
@ -60,6 +71,85 @@ 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 _ToolCallRecordingGuardrail(_CountingPassingGuardrail):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.tool_call_scan_indexes: List[int] = []
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 inputs.get("tool_calls"):
self.tool_call_scan_indexes.append(self.scan_count)
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
class _MarkerBlockingGuardrail(_CountingPassingGuardrail):
"""Blocks as soon as the inspected input field (texts or tool_calls) carries the marker."""
def __init__(self, *args, marker: str, field: Literal["texts", "tool_calls"] = "texts", **kwargs):
super().__init__(*args, **kwargs)
self.marker = marker
self.field = field
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.marker in json.dumps(inputs.get(self.field, [])):
raise ModifyResponseException(
message=BLOCK_MESSAGE,
model="gpt-4o",
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 +205,212 @@ 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,
)
],
)
def _tool_call_chunk(
arguments: str, finish_reason: str | None = None, legacy_function_call: bool = False
) -> ModelResponseStream:
delta = (
Delta(role="assistant", content=None, function_call=FunctionCall(name="run_shell", arguments=arguments))
if legacy_function_call
else Delta(
role="assistant",
content=None,
tool_calls=[
ChatCompletionDeltaToolCall(
id="call_1",
type="function",
index=0,
function=Function(name="run_shell", arguments=arguments),
)
],
)
)
return ModelResponseStream(
id="chatcmpl-windowed",
created=1724900000,
model="gpt-4",
choices=[StreamingChoices(index=0, delta=delta, finish_reason=finish_reason)],
)
async def _windowed_chat_stream(
yielded_count: List[int],
collected: List[Any],
content_chunks: List[str],
tool_argument_chunks: List[str] | None = None,
legacy_function_call: bool = False,
) -> AsyncGenerator[ModelResponseStream, None]:
for content in content_chunks:
yielded_count.append(len(collected))
yield _chat_chunk(content)
for arguments in tool_argument_chunks or []:
yielded_count.append(len(collected))
yield _tool_call_chunk(arguments, legacy_function_call=legacy_function_call)
yielded_count.append(len(collected))
yield _chat_chunk(finish_reason="tool_calls" if tool_argument_chunks else "stop")
def _tool_argument_text(chunks: List[Any]) -> str:
return "".join(
tool_call.function.arguments or ""
for chunk in chunks
if isinstance(chunk, ModelResponseStream)
for choice in chunk.choices
for tool_call in choice.delta.tool_calls or []
)
def _function_call_argument_text(chunks: list[Any]) -> str:
return "".join(
choice.delta.function_call.arguments or ""
for chunk in chunks
if isinstance(chunk, ModelResponseStream)
for choice in chunk.choices
if choice.delta.function_call is not None
)
async def _run_windowed(
guardrail: CustomGuardrail,
content_chunks: List[str],
end_of_stream_only: bool = False,
tool_argument_chunks: List[str] | None = None,
legacy_function_call: 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, tool_argument_chunks, legacy_function_call
),
request_data=request_data,
):
collected.append(chunk)
return collected, yielded_count
def _responses_message_stream_events(text_chunks: List[str]) -> List[dict]:
message = {"type": "message", "id": "msg_1", "status": "completed", "role": "assistant"}
content = [{"type": "output_text", "text": "".join(text_chunks), "annotations": []}]
return [
{"type": "response.output_item.added", "output_index": 0, "item": {**message, "content": []}},
*(
{
"type": "response.output_text.delta",
"item_id": "msg_1",
"output_index": 0,
"content_index": 0,
"delta": text,
}
for text in text_chunks
),
{"type": "response.output_item.done", "output_index": 0, "item": {**message, "content": content}},
{
"type": "response.completed",
"response": {
"id": "resp_1",
"model": "gpt-4o",
"status": "completed",
"output": [{**message, "content": content}],
},
},
]
def _responses_truncated_function_call_events(text: str, argument_chunks: List[str]) -> List[dict]:
message = {"type": "message", "id": "msg_1", "status": "completed", "role": "assistant"}
content = [{"type": "output_text", "text": text, "annotations": []}]
function_call = {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "run_shell"}
return [
{"type": "response.output_item.added", "output_index": 0, "item": {**message, "content": []}},
{
"type": "response.output_text.delta",
"item_id": "msg_1",
"output_index": 0,
"content_index": 0,
"delta": text,
},
{"type": "response.output_item.added", "output_index": 1, "item": {**function_call, "arguments": ""}},
*(
{"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": arguments}
for arguments in argument_chunks
),
{
"type": "response.incomplete",
"response": {
"id": "resp_1",
"model": "gpt-4o",
"status": "incomplete",
"output": [
{**message, "content": content},
{**function_call, "arguments": "".join(argument_chunks), "status": "incomplete"},
],
},
},
]
async def _replay(events: List[dict]) -> AsyncGenerator[dict, None]:
for event in events:
yield event
async def _run_windowed_responses(guardrail: CustomGuardrail, events: List[dict]) -> str:
guardrail.streaming_buffer_until_moderated = True
guardrail.streaming_buffer_release_on_scan = True
guardrail.streaming_sampling_rate = 2
unified = UnifiedLLMGuardrails()
user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/responses")
request_data = {
"input": "hi",
"guardrail_to_apply": guardrail,
"metadata": {"guardrails": [guardrail.guardrail_name]},
}
collected: List[Any] = []
async for chunk in unified.async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
response=_replay(events),
request_data=request_data,
):
collected.append(chunk)
return json.dumps([chunk if isinstance(chunk, dict) else str(chunk) for chunk in collected])
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 +455,110 @@ 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_holds_tool_call_windows_until_end_of_stream_scan():
guardrail = _ToolCallRecordingGuardrail(guardrail_name="windowed-tools", event_hook="post_call")
content_chunks = ["one ", "two ", "three "]
tool_argument_chunks = ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}']
collected, yielded_count = await _run_windowed(guardrail, content_chunks, tool_argument_chunks=tool_argument_chunks)
assert yielded_count == [0, 0, 2, 2, 2, 2, 2]
assert _chat_text(collected) == "".join(content_chunks)
assert _tool_argument_text(collected) == "".join(tool_argument_chunks)
assert guardrail.tool_call_scan_indexes == [guardrail.scan_count]
@pytest.mark.asyncio
async def test_windowed_buffer_holds_legacy_function_call_windows_until_end_of_stream():
guardrail = _PassingGuardrail(guardrail_name="windowed-functions", event_hook="post_call")
content_chunks = ["one ", "two ", "three "]
function_argument_chunks = ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}']
collected, yielded_count = await _run_windowed(
guardrail, content_chunks, tool_argument_chunks=function_argument_chunks, legacy_function_call=True
)
assert yielded_count == [0, 0, 2, 2, 2, 2, 2]
assert _chat_text(collected) == "".join(content_chunks)
assert _function_call_argument_text(collected) == "".join(function_argument_chunks)
def test_tool_call_only_scan_key_is_not_skipped_as_empty():
assert _is_redundant_scan(StreamingScanKey(texts=("",)), None) is True
assert _is_redundant_scan(StreamingScanKey(texts=("",), tool_calls=("run_shell:{}",)), None) is False
@pytest.mark.asyncio
async def test_windowed_responses_output_item_done_round_keeps_text_window_withheld():
guardrail = _MarkerBlockingGuardrail(
guardrail_name="windowed-responses", event_hook="post_call", marker=ORIGINAL_MARKER
)
events = _responses_message_stream_events(["one ", f"{ORIGINAL_MARKER} "])
raw = await _run_windowed_responses(guardrail, events)
assert ORIGINAL_MARKER not in raw, f"unscanned window leaked: {raw!r}"
assert BLOCK_MESSAGE in raw
assert guardrail.scan_count >= 1
@pytest.mark.asyncio
async def test_windowed_responses_incomplete_stream_scans_tool_call_before_release():
guardrail = _MarkerBlockingGuardrail(
guardrail_name="windowed-responses-tools",
event_hook="post_call",
marker=TOOL_ARGUMENTS_MARKER,
field="tool_calls",
)
events = _responses_truncated_function_call_events("hi ", ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}'])
raw = await _run_windowed_responses(guardrail, events)
assert '"hi "' in raw
assert TOOL_ARGUMENTS_MARKER not in raw, f"unscanned tool call leaked: {raw!r}"
assert BLOCK_MESSAGE 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

16
tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py
View file

@ -682,6 +682,22 @@ async def test_provider_specific_params_includes_embedding_toggle():
assert field["default_value"] is False
@pytest.mark.asyncio
async def test_provider_specific_params_exposes_bedrock_streaming_flags():
from litellm.proxy.guardrails.guardrail_endpoints import get_provider_specific_params
provider_params = await get_provider_specific_params()
bedrock = provider_params["bedrock"]
assert "guardrailIdentifier" in bedrock
assert "guardrailVersion" in bedrock
assert bedrock["streaming_buffer_release_on_scan"]["type"] == "boolean"
assert bedrock["streaming_buffer_release_on_scan"]["default_value"] is False
assert bedrock["streaming_buffer_until_moderated"]["default_value"] is True
assert bedrock["streaming_end_of_stream_only"]["type"] == "boolean"
assert bedrock["streaming_sampling_rate"]["type"] == "number"
@pytest.mark.asyncio
async def test_provider_specific_params_includes_hide_secrets():
"""hide-secrets lives in the enterprise package so it is not in