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

Merge pull request #41933 from BerriAI/litellm_deliver_multi_choice_stream_rewrites

Browse source 
fix(policy_engine): deliver guardrail text rewrites on multi-choice, unfinished, and envelope-less streams
This commit is contained in:
Mateo Wang 2026-09-19 00:22:02 -07:00 • committed by GitHub
commit 5f1268c056
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 418 additions and 95 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1253,10 +1253,9 @@ class AnthropicMessagesHandler(BaseTranslation):
Process output streaming response by applying guardrails to text content.
Get the string so far, check the apply guardrail to the string so far, and return the list of responses so far.
With ``deliver_ended_stream_rewrites``, an ended stream whose guardrail rewrote the text gets the rewrite
written back across the buffered chunks (full rewritten text in the first ``text_delta``, the rest blanked);
a rewrite on a stream that never reported a ``stop_reason`` has no write-back and is reported as
undeliverable, so the pipeline executor discards it and releases the original chunks.
With ``deliver_ended_stream_rewrites``, a stream whose guardrail rewrote the text gets the rewrite
written back across the buffered chunks (full rewritten text in the first ``text_delta``, the rest blanked),
whether or not the stream ever reported a ``stop_reason``.
"""
from litellm.integrations.custom_guardrail import ModifyResponseException
@ -1312,7 +1311,11 @@ class AnthropicMessagesHandler(BaseTranslation):
and guardrailed_texts
and guardrailed_texts[0] != string_so_far
):
self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0])
self._write_ended_stream_text_rewrite(
responses_so_far,
guardrailed_texts[0],
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
if deliver_ended_stream_rewrites:
returned_tool_calls: Final = _guardrailed_inputs.get("tool_calls")
self._write_ended_stream_tool_call_rewrites(
@ -1354,9 +1357,11 @@ class AnthropicMessagesHandler(BaseTranslation):
raise
unended_texts: Final = _guardrailed_inputs.get("texts")
if deliver_ended_stream_rewrites and unended_texts and tuple(unended_texts) != (string_so_far,):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown")
self._write_ended_stream_text_rewrite(
responses_so_far,
unended_texts[0],
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
return responses_so_far
def _prepare_request_data(
@ -1450,26 +1455,40 @@ class AnthropicMessagesHandler(BaseTranslation):
inputs["model"] = response_model
return inputs
@staticmethod
@classmethod
def _write_ended_stream_text_rewrite(
cls,
responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place
rewritten_text: str,
guardrail_name: str,
) -> None:
"""Deliver an ended-stream guardrail text rewrite by rewriting the
buffered chunks in place: the first ``text_delta`` carries the full
rewritten text and every later one is blanked, leaving the surrounding
message and content-block framing untouched."""
message and content-block framing untouched. A buffer with no
``text_delta`` has nowhere to carry the rewrite, so the pipeline
executor discards it and releases the original chunks."""
def is_text_delta(event: Mapping[str, object]) -> bool:
delta: Final = event.get("delta")
return (
event.get("type") == "content_block_delta"
and isinstance(delta, Mapping)
and delta.get("type") == "text_delta"
)
if not any(is_text_delta(event) for item in responses_so_far for event in cls._iter_sse_events(item)):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_name)
replacements: Final = chain((rewritten_text,), repeat(""))
def rewrite_text_delta(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
delta: Final = event.get("delta")
if event.get("type") != "content_block_delta" or not isinstance(delta, Mapping):
return None
if delta.get("type") != "text_delta":
if not is_text_delta(event):
return None
return _SSEFieldRewrite("delta", "text", next(replacements))
AnthropicMessagesHandler._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta)
cls._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta)
@classmethod
def _write_ended_stream_tool_call_rewrites(

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

@ -18,6 +18,7 @@ import json
import time
import uuid
from collections.abc import Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Union, cast
from typing_extensions import NotRequired, ReadOnly, TypedDict
@ -651,10 +652,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
"""Ended-stream path: rebuild the full response, run the non-streaming
output guardrail against it, and (when opted in) write any text or
tool-call rewrite back across the buffered chunks."""
model_response: Final = cast(
ModelResponse,
stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj),
)
model_response: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj)
pre_guardrail_texts: Final = self._string_choice_contents(model_response)
pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response)
await self.process_output_response(
@ -666,20 +664,59 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
)
if not deliver_ended_stream_rewrites:
return
guardrail_name: Final = guardrail_to_apply.guardrail_name or "unknown"
await self._write_ended_stream_text_rewrites(
responses_so_far=responses_so_far,
guardrailed_response=model_response,
pre_guardrail_texts=pre_guardrail_texts,
guardrail_name=guardrail_name,
)
self._write_ended_stream_tool_call_rewrites(
responses_so_far=responses_so_far,
guardrailed_response=model_response,
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
guardrail_name=guardrail_name,
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
@staticmethod
def _rebuild_ended_stream_per_choice(
responses_so_far: Sequence["ModelResponseStream"],
litellm_logging_obj: "LiteLLMLoggingObj | None",
) -> "ModelResponse":
"""``stream_chunk_builder`` folds every choice of a stream into one index-0
choice, so the stream is rebuilt one choice index at a time (every chunk
kept, its choices narrowed to that index, so usage-only chunks still
count) and the rebuilt choices are stitched into one response, each
carrying the index the stream gave it."""
choice_indices: Final = tuple(
sorted(frozenset(choice.index for response in responses_so_far for choice in response.choices))
)
rebuilt_by_index: Final = tuple(
(
index,
cast(
ModelResponse,
stream_chunk_builder(
chunks=[ # mutable-ok: callee takes a list
OpenAIChatCompletionsHandler._narrowed_to_choice(response, index)
for response in responses_so_far
],
logging_obj=litellm_logging_obj,
),
),
)
for index in choice_indices
)
(_, base_response), *_ = rebuilt_by_index
stitched_choices: Final = [ # mutable-ok: choices is a List field; a tuple there breaks model_dump round-trips
rebuilt.choices[0].model_copy(update=MappingProxyType({"index": index}))
for index, rebuilt in rebuilt_by_index
]
return base_response.model_copy(update=MappingProxyType({"choices": stitched_choices}))
@staticmethod
def _narrowed_to_choice(response: "ModelResponseStream", index: int) -> "ModelResponseStream":
narrowed: Final = [choice for choice in response.choices if choice.index == index] # mutable-ok: List field
return response.model_copy(update=MappingProxyType({"choices": narrowed}))
def build_stream_error_items(
self,
exc: "HTTPException",
@ -1058,39 +1095,28 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place
guardrailed_response: "ModelResponse",
pre_guardrail_texts: tuple[str | None, ...],
guardrail_name: str,
) -> None:
"""Write ended-stream guardrail text rewrites back across the buffered
chunks: the full rewritten text lands in the choice's first
content-carrying chunk and the rest are blanked, the same shape the
in-flight write-back uses. Chunks carrying only finish_reason or usage
stay untouched. A rewrite on a stream carrying more than one distinct
choice index is reported as undeliverable, so the pipeline executor
discards it and releases the original chunks."""
chunks, one rewrite per rebuilt choice index: the full rewritten text
lands in that choice's first content-carrying chunk and the rest are
blanked, the same shape the in-flight write-back uses. Chunks carrying
only finish_reason or usage stay untouched."""
post_guardrail_texts: Final = self._string_choice_contents(guardrailed_response)
changed: Final = tuple(
after
for before, after in zip(pre_guardrail_texts, post_guardrail_texts)
if before is not None and after is not None and after != before
rewrites_by_choice: Final = MappingProxyType(
{
choice.index: after
for choice, before, after in zip(
guardrailed_response.choices, pre_guardrail_texts, post_guardrail_texts
)
if before is not None and after is not None and after != before
}
)
if not changed:
if not rewrites_by_choice:
return
stream_choice_indices: Final = frozenset(
choice.index for response in responses_so_far for choice in response.choices
)
if len(stream_choice_indices) != 1:
# stream_chunk_builder collapses every choice into one index-0
# choice, so a rewrite of the rebuilt response cannot be attributed
# back to a single choice on an n>1 stream: report it undeliverable
# rather than deliver the rewrite on the wrong choice
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_name)
target_choice_index: Final = next(iter(stream_choice_indices))
await self._apply_guardrail_responses_to_output_streaming(
responses=responses_so_far,
guardrailed_texts=list(changed), # mutable-ok: callee takes lists
task_mappings=[(target_choice_index, None) for _ in changed], # mutable-ok: callee takes lists
guardrailed_texts=list(rewrites_by_choice.values()), # mutable-ok: callee takes lists
task_mappings=[(index, None) for index in rewrites_by_choice], # mutable-ok: callee takes lists
)
@staticmethod

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

@ -209,6 +209,7 @@ _TOOL_CALL_PAYLOAD_EVENT_TYPES: Final = _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES | f
_TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS
)
_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"})
_OUTPUT_TEXT_EVENT_TYPES: Final = frozenset({"response.output_text.delta", "response.output_text.done"})
_PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
{"function_call_output": "output", "message": "content"}
)
@ -832,9 +833,10 @@ class OpenAIResponsesHandler(BaseTranslation):
(``response.output_text.delta`` / ``.done``,
``response.content_part.done``, ``response.output_item.done``) are synced
to the rewritten envelope too, so a client reading deltas sees the
rewrite instead of the raw model output; a rewrite observed where no
write-back is possible is reported as undeliverable, so the pipeline
executor discards it and releases the original events.
rewrite instead of the raw model output; a stream with no envelope
gets its rewrite spread over the buffered text events, and a rewrite
observed where no write-back is possible is reported as undeliverable,
so the pipeline executor discards it and releases the original events.
"""
if not responses_so_far:
return responses_so_far
@ -958,10 +960,9 @@ class OpenAIResponsesHandler(BaseTranslation):
return responses_so_far
# ------------------------------------------------------------------ #
# Fallback: apply guardrail to the accumulated text string. #
# No structured write-back is possible here; guardrails that only #
# need to block/flag (not rewrite) still work correctly, and a #
# rewrite a caller expects delivered is reported undeliverable. #
# Fallback: apply guardrail to the accumulated text string. With no #
# envelope to rewrite, a rewrite a caller expects delivered is spread #
# over the buffered text events instead. #
# ------------------------------------------------------------------ #
string_so_far: Final = self.get_streaming_string_so_far(responses_so_far)
if string_so_far:
@ -979,11 +980,54 @@ class OpenAIResponsesHandler(BaseTranslation):
)
fallback_texts: Final = fallback_outputs.get("texts")
if deliver_ended_stream_rewrites and fallback_texts and tuple(fallback_texts) != (string_so_far,):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown")
self._spread_text_rewrite_over_stream_events(
stream_events=responses_so_far,
rewritten_text=fallback_texts[0],
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
return responses_so_far
def _spread_text_rewrite_over_stream_events(
self,
stream_events: Sequence[Any],
rewritten_text: str,
guardrail_name: str,
) -> None:
"""Deliver a text rewrite on a stream with no completed envelope by
spreading it over the text parts the guardrail scanned, in stream
order: the whole rewrite on the first part and every later part
blanked, through the same sync the envelope path uses. A scanned
event the sync cannot place (one that is not an ``output_text`` delta
or done, or lacks integer ``output_index`` / ``content_index``) makes
the rewrite undeliverable, so the pipeline executor discards it and
releases the original events."""
scanned_events: Final = tuple(
event
for event in stream_events
if isinstance(stream_item_field(event, "text"), str) or isinstance(stream_item_field(event, "delta"), str)
)
scanned_positions: Final = tuple(
dict.fromkeys(
(stream_item_field(event, "output_index"), stream_item_field(event, "content_index"))
for event in scanned_events
)
)
placeable_positions: Final = tuple(
(output_index, content_index)
for output_index, content_index in scanned_positions
if isinstance(output_index, int) and isinstance(content_index, int)
)
if len(placeable_positions) != len(scanned_positions) or any(
stream_item_field(event, "type") not in _OUTPUT_TEXT_EVENT_TYPES for event in scanned_events
):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_name)
self._sync_stream_events_with_rewrites(
stream_events=stream_events,
rewrites_by_position=MappingProxyType(dict(zip(placeable_positions, chain((rewritten_text,), repeat(""))))),
)
@staticmethod
def _write_event_field(event: object, field: str, value: str) -> None:
if isinstance(event, dict):

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

@ -445,20 +445,45 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
assert chunks == original
@pytest.mark.asyncio
async def test_unended_stream_rewrite_with_delivery_expected_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
async def test_unended_stream_rewrite_with_delivery_expected_lands_in_the_buffered_deltas(self):
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()[:-2]
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert self._delta_texts(chunks) == ["hello [MASKED]", ""]
raw = b"".join(chunks).decode()
assert "event: message_start" in raw and "event: content_block_stop" in raw
assert "event: message_stop" not in raw
@pytest.mark.asyncio
async def test_unended_stream_rewrite_with_no_text_delta_to_carry_it_fails_open(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
class FillEmpty(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
return {**inputs, "texts": ["[INJECTED]" for _ in inputs.get("texts", [])]}
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()[:2]
original = [bytes(chunk) for chunk in chunks]
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._masking_guardrail(),
guardrail_to_apply=FillEmpty(guardrail_name="test"),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
assert chunks == original
@pytest.mark.asyncio
async def test_unended_stream_without_rewrite_is_released_with_delivery_expected(self):
handler = AnthropicMessagesHandler()

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

@ -1262,19 +1262,81 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
return MaskWorld(guardrail_name="test-mask")
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrite_on_multi_choice_stream_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
async def test_deliver_ended_stream_rewrite_lands_on_the_rewritten_choice_only(self):
handler = OpenAIChatCompletionsHandler()
chunks = self._two_choice_stream_chunks()
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._world_masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._world_masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert [(c.choices[0].index, c.choices[0].delta.content) for c in chunks] == [
(0, "safe "),
(1, "hello [MASKED]"),
(0, "text"),
(1, ""),
]
assert [c.choices[0].finish_reason for c in chunks] == [None, None, "stop", "stop"]
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_each_choice_with_its_own_text(self):
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="test")
chunks = self._two_choice_stream_chunks()
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert guardrail.last_inputs["texts"] == ["safe text", "hello world"]
assert [(c.choices[0].index, c.choices[0].delta.content) for c in chunks] == [
(0, "SAFE TEXT"),
(1, "HELLO WORLD"),
(0, ""),
(1, ""),
]
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_every_choice_when_a_usage_only_chunk_closes_the_stream(self):
from litellm.types.utils import ModelResponseStream, Usage
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="test")
usage_chunk = ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[],
usage=Usage(prompt_tokens=5, completion_tokens=7, total_tokens=12),
)
chunks = [*self._two_choice_stream_chunks(), usage_chunk]
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert guardrail.last_inputs["texts"] == ["safe text", "hello world"]
assert [(c.choices[0].index, c.choices[0].delta.content) for c in chunks[:4]] == [
(0, "SAFE TEXT"),
(1, "HELLO WORLD"),
(0, ""),
(1, ""),
]
assert [c.choices[0].finish_reason for c in chunks[:4]] == [None, None, "stop", "stop"]
assert chunks[4].choices == []
assert chunks[4].usage.completion_tokens == 7
@staticmethod
def _two_choice_tool_call_stream_chunks() -> list:

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

@ -1747,33 +1747,82 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]"
@pytest.mark.asyncio
async def test_fallback_rewrite_with_delivery_expected_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
async def test_fallback_rewrite_with_delivery_expected_lands_in_the_delta_and_done_events(self):
handler = OpenAIResponsesHandler()
events = [
{"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "},
{"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"},
]
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is events
assert events[0]["delta"] == "hello [MASKED]"
assert events[1]["text"] == "hello [MASKED]"
@pytest.mark.asyncio
async def test_fallback_delta_only_rewrite_with_delivery_expected_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
async def test_fallback_delta_only_rewrite_with_delivery_expected_spreads_over_the_deltas(self):
handler = OpenAIResponsesHandler()
events = [
{"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "},
{"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "world"},
]
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is events
assert [event["delta"] for event in events] == ["hello [MASKED]", ""]
@pytest.mark.asyncio
async def test_fallback_rewrite_across_parts_lands_whole_on_the_first_part(self):
handler = OpenAIResponsesHandler()
events = [
{"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "},
{"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello "},
{"type": "response.output_text.delta", "output_index": 1, "content_index": 0, "delta": "wor"},
{"type": "response.output_text.delta", "output_index": 1, "content_index": 0, "delta": "ld"},
]
guardrail = MockRecordingGuardrail(guardrail_name="test")
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert [inputs.get("texts") for inputs in guardrail.seen_inputs] == [["hello world"]]
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert events[0]["delta"] == "hello [MASKED]"
assert events[1]["text"] == "hello [MASKED]"
assert [event["delta"] for event in events[2:]] == ["", ""]
@pytest.mark.asyncio
async def test_fallback_rewrite_over_an_unplaceable_scanned_event_fails_open(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = OpenAIResponsesHandler()
events = [
{"type": "response.reasoning_summary_text.delta", "output_index": 0, "summary_index": 0, "delta": "hello "},
{"type": "response.output_text.delta", "output_index": 1, "content_index": 0, "delta": "world"},
]
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=events,
@ -1781,21 +1830,26 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert [event["delta"] for event in events] == ["hello ", "world"]
@pytest.mark.asyncio
async def test_output_item_done_last_rewrite_with_delivery_expected_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
async def test_output_item_done_last_rewrite_with_delivery_expected_syncs_every_text_event(self):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()[:-1]
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is events
assert events[0]["delta"] == "hello [MASKED]"
assert events[1]["delta"] == ""
assert events[2]["text"] == "hello [MASKED]"
assert events[3]["part"]["text"] == "hello [MASKED]"
assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]"
@pytest.mark.asyncio
async def test_output_item_done_last_scans_text_with_delivery_expected(self):

5
tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py
View file

@ -369,6 +369,7 @@ async def test_openai_moderation_guardrail_streaming_safe_content():
chunk1.choices[0].delta = MagicMock()
chunk1.choices[0].delta.content = "Hello "
chunk1.choices[0].finish_reason = None
chunk1.choices[0].index = 0
chunk2 = MagicMock()
chunk2.model = "gpt-4"
@ -376,6 +377,7 @@ async def test_openai_moderation_guardrail_streaming_safe_content():
chunk2.choices[0].delta = MagicMock()
chunk2.choices[0].delta.content = "world"
chunk2.choices[0].finish_reason = None
chunk2.choices[0].index = 0
# Last chunk with finish_reason
chunk3 = MagicMock()
@ -384,6 +386,7 @@ async def test_openai_moderation_guardrail_streaming_safe_content():
chunk3.choices[0].delta = MagicMock()
chunk3.choices[0].delta.content = "!"
chunk3.choices[0].finish_reason = "stop"
chunk3.choices[0].index = 0
for chunk in [chunk1, chunk2, chunk3]:
yield chunk
@ -480,6 +483,7 @@ async def test_openai_moderation_guardrail_streaming_harmful_content():
chunk1.choices[0].delta = MagicMock()
chunk1.choices[0].delta.content = "This is "
chunk1.choices[0].finish_reason = None
chunk1.choices[0].index = 0
# Last chunk - with finish_reason to signal end of stream
chunk2 = MagicMock()
@ -488,6 +492,7 @@ async def test_openai_moderation_guardrail_streaming_harmful_content():
chunk2.choices[0].delta = MagicMock()
chunk2.choices[0].delta.content = "harmful content"
chunk2.choices[0].finish_reason = "stop"
chunk2.choices[0].index = 0
for chunk in [chunk1, chunk2]:
yield chunk

3
tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py
View file

@ -42,6 +42,7 @@ async def test_openai_moderation_guardrail_streaming_latency():
choice.delta.content = content
# Last chunk gets finish_reason
choice.finish_reason = "stop" if i == len(chunks_data) - 1 else None
choice.index = 0
chunk.choices = [choice]
yield chunk
@ -122,6 +123,7 @@ async def test_openai_moderation_guardrail_streaming_harmful_content():
choice.delta.content = content
# Last chunk gets finish_reason
choice.finish_reason = "stop" if i == len(chunks_data) - 1 else None
choice.index = 0
chunk.choices = [choice]
yield chunk
@ -224,6 +226,7 @@ async def test_openai_moderation_streaming_end_of_stream_request_data_passthroug
choice.delta = MagicMock()
choice.delta.content = content
choice.finish_reason = "stop" if i == len(chunks_data) - 1 else None
choice.index = 0
chunk.choices = [choice]
yield chunk

36
tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py
View file

@ -1318,6 +1318,42 @@ async def test_streaming_step_records_guardrail_information_once_on_block(monkey
assert _recorded_guardrail_statuses(result) == ["guardrail_intervened"]
def _two_choice_chat_chunks():
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
def chunk(index, content, finish_reason=None):
return ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)],
)
return [chunk(0, "pers"), chunk(1, "pers"), chunk(0, "immon", "stop"), chunk(1, "immon", "stop")]
@pytest.mark.asyncio
async def test_streaming_step_delivers_text_rewrites_on_every_choice_of_a_chat_stream(monkeypatch, caplog):
from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler
monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["[MASKED]", "[MASKED]"])])
chunks = _two_choice_chat_chunks()
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(OpenAIChatCompletionsHandler(), chunks)
assert result.terminal_action == "allow"
assert not any("discarded" in record.getMessage() for record in caplog.records)
assert [(c.choices[0].index, c.choices[0].delta.content) for c in chunks] == [
(0, "[MASKED]"),
(1, "[MASKED]"),
(0, ""),
(1, ""),
]
assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
@pytest.mark.asyncio
async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewrite(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])])

49
tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py
View file

@ -1836,6 +1836,22 @@ def _rewritten_model_response(response: Any) -> litellm.ModelResponse:
return litellm.ModelResponse(**payload)
def _two_choice_stream_chunks() -> List[Any]:
return [
litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "hello "}, "finish_reason": None}]),
litellm.ModelResponseStream(choices=[{"index": 1, "delta": {"content": "bonjour "}, "finish_reason": None}]),
litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "world"}, "finish_reason": "stop"}]),
litellm.ModelResponseStream(choices=[{"index": 1, "delta": {"content": "monde"}, "finish_reason": "stop"}]),
]
def _rewritten_every_choice(response: Any) -> litellm.ModelResponse:
payload = response.model_dump()
for choice in payload["choices"]:
choice["message"]["content"] = "[REWRITTEN] " + choice["message"]["content"]
return litellm.ModelResponse(**payload)
def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator_only(
make_user_api_key_auth, monkeypatch, caplog
):
@ -1984,6 +2000,39 @@ async def test_streaming_iterator_hook_runs_legacy_hook_and_delivers_its_rewrite
assert _warnings(caplog) == []
@pytest.mark.asyncio
async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_every_choice(
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
):
seen: Dict[str, Any] = {}
guardrail = _legacy_hook_stream_guardrail(seen, rewrite=_rewritten_every_choice)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
chunks = _two_choice_stream_chunks()
auth = make_user_api_key_auth(request_route="/v1/chat/completions")
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="completion", guardrails_only=True)
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data
)
]
assert [choice.message.content for choice in seen["response"].choices] == ["hello world", "bonjour monde"]
assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks]
assert [(item.choices[0].index, item.choices[0].delta.content) for item in delivered] == [
(0, "[REWRITTEN] hello world"),
(1, "[REWRITTEN] bonjour monde"),
(0, ""),
(1, ""),
]
assert [item.choices[0].finish_reason for item in delivered] == [None, None, "stop", "stop"]
assert _warnings(caplog) == []
@pytest.mark.asyncio
async def test_streaming_iterator_hook_releases_stream_untouched_when_legacy_hook_returns_none(
proxy_logging, make_user_api_key_auth, monkeypatch