mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
Merge remote-tracking branch 'origin/main' into litellm_llm_guard_proxy_call_types
This commit is contained in:
commit
4e5d8004a4
10 changed files with 418 additions and 95 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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]"])])
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue