scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 8

Merge pull request #39233 from BerriAI/litellm_post_call_pipeline_stream_rewrite

Browse source 
fix(policy_engine): apply post_call pipeline text rewrites on streams
This commit is contained in:
Mateo Wang 2026-09-01 22:35:55 -07:00 • committed by GitHub
commit bca3bade40
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 1026 additions and 152 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -13,9 +13,10 @@ Pattern Overview:
"""
import json
from collections.abc import Mapping, Sequence
from collections.abc import Iterator, Mapping, Sequence
from copy import deepcopy
from dataclasses import dataclass
from itertools import chain, repeat
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable
from typing_extensions import ReadOnly, TypedDict, assert_never
@ -164,6 +165,8 @@ class AnthropicMessagesHandler(BaseTranslation):
them through guardrail rewrites; downstream provider handling is out of scope.
"""
delivers_ended_stream_text_rewrites = True
def __init__(self):
super().__init__()
self.adapter = LiteLLMAnthropicMessagesAdapter()
@ -1010,11 +1013,15 @@ class AnthropicMessagesHandler(BaseTranslation):
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
deliver_ended_stream_rewrites: bool = False,
) -> Sequence[object]:
"""
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 fails closed instead.
"""
from litellm.integrations.custom_guardrail import ModifyResponseException
@ -1061,6 +1068,15 @@ class AnthropicMessagesHandler(BaseTranslation):
responses_so_far, request_data
)
raise
guardrailed_texts: Final = _guardrailed_inputs.get("texts")
if (
deliver_ended_stream_rewrites
and isinstance(string_so_far, str)
and string_so_far
and guardrailed_texts
and guardrailed_texts[0] != string_so_far
):
self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0])
else:
verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices")
return responses_so_far
@ -1083,6 +1099,11 @@ class AnthropicMessagesHandler(BaseTranslation):
if e.original_response is None:
e.original_response = self._build_streaming_usage_response(responses_so_far, request_data)
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")
return responses_so_far
def _prepare_request_data(
@ -1176,6 +1197,63 @@ class AnthropicMessagesHandler(BaseTranslation):
inputs["model"] = response_model
return inputs
@staticmethod
def _write_ended_stream_text_rewrite(
responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place
rewritten_text: 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. Handles both chunk formats
this stream carries (parsed event dicts and raw SSE bytes)."""
replacements: Final = chain((rewritten_text,), repeat(""))
for idx, item in enumerate(responses_so_far):
if isinstance(item, dict):
delta = item.get("delta")
if item.get("type") == "content_block_delta" and isinstance(delta, dict):
if delta.get("type") == "text_delta":
delta["text"] = next(replacements)
elif isinstance(item, (bytes, bytearray)):
responses_so_far[idx] = ( # rebind-ok: delivers the rewrite into the caller's buffer
AnthropicMessagesHandler._rewrite_sse_text_deltas(bytes(item), replacements)
)
@staticmethod
def _rewrite_sse_text_deltas(sse_bytes: bytes, replacements: "Iterator[str]") -> bytes:
"""Rewrite every ``text_delta`` data line in one SSE chunk with the next
replacement text, leaving all other events and framing byte-identical."""
try:
decoded: Final = sse_bytes.decode("utf-8")
except UnicodeDecodeError:
return sse_bytes
return "\n\n".join(
AnthropicMessagesHandler._rewrite_sse_block(block, replacements) for block in decoded.split("\n\n")
).encode("utf-8")
@staticmethod
def _rewrite_sse_block(block: str, replacements: "Iterator[str]") -> str:
return "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, replacements) for line in block.split("\n"))
@staticmethod
def _rewrite_sse_line(line: str, replacements: "Iterator[str]") -> str:
if not line.startswith("data:"):
return line
try:
data: Final[str | int | float | bool | None | Sequence[object] | Mapping[str, object]] = json.loads(
line[len("data:") :].strip()
)
except json.JSONDecodeError:
return line
if not isinstance(data, dict) or data.get("type") != "content_block_delta":
return line
delta: Final = data.get("delta")
if not isinstance(delta, dict) or delta.get("type") != "text_delta":
return line
return "data: " + json.dumps(
{**data, "delta": {**delta, "text": next(replacements)}} # mutable-ok: json.dumps needs plain dicts
)
def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str:
"""
Parse streaming responses and extract accumulated text content.

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

@ -1,7 +1,7 @@
from abc import ABC, abstractmethod
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Final, Optional
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
if TYPE_CHECKING:
from fastapi import HTTPException
@ -36,6 +36,13 @@ class StreamTransformSink:
class BaseTranslation(ABC):
delivers_ended_stream_text_rewrites: ClassVar[bool] = False
"""Whether ``process_output_streaming_response`` accepts
``deliver_ended_stream_rewrites=True`` and, on an ended (fully buffered)
stream, writes guardrail text rewrites back across ``responses_so_far`` so
a buffered pipeline can release rewritten chunks instead of withholding the
stream. Tool-call rewrites stay undeliverable everywhere."""
@staticmethod
def transform_user_api_key_dict_to_metadata(
user_api_key_dict: Any | None,
@ -141,6 +148,7 @@ class BaseTranslation(ABC):
user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
request_data: dict | None = None,
stream_transform_sink: StreamTransformSink | None = None,
deliver_ended_stream_rewrites: bool = False,
) -> Any:
"""
Process output streaming response with guardrails.
@ -148,6 +156,11 @@ class BaseTranslation(ABC):
Optional to override in subclasses. ``stream_transform_sink`` is the
out-parameter used by handlers that support streaming text
transformations (see ``StreamTransformSink``); base handlers ignore it.
``deliver_ended_stream_rewrites`` is passed True only when the caller
holds the whole buffered stream and the subclass declares
``delivers_ended_stream_text_rewrites``: the handler then writes
guardrail text rewrites back across ``responses_so_far`` instead of
discarding them.
"""
return responses_so_far

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

@ -75,6 +75,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
Methods can be overridden to customize behavior for different message formats.
"""
delivers_ended_stream_text_rewrites = True
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
"""
Convert chat completions request data to OpenAI-spec structured messages.
@ -450,6 +452,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
stream_transform_sink: StreamTransformSink | None = None,
deliver_ended_stream_rewrites: bool = False,
) -> list["ModelResponseStream"]:
"""
Process output streaming responses by applying guardrails to text content.
@ -464,6 +467,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
accumulated text (``responses_so_far`` is left untouched so it stays
a correct raw accumulator across rounds) and the guardrailed text
plus requested holdback are reported per choice on the sink.
deliver_ended_stream_rewrites: When True and the buffered stream has
ended, guardrail text rewrites are written back across
``responses_so_far`` (full rewritten text in each choice's first
content-carrying chunk, the rest blanked) instead of discarded.
Returns:
The (unmodified) list of responses.
@ -489,6 +496,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
deliver_ended_stream_rewrites=deliver_ended_stream_rewrites,
)
async def _process_streaming_block_only(
@ -499,10 +507,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
litellm_logging_obj: "LiteLLMLoggingObj | None",
user_api_key_dict: "UserAPIKeyAuth | None",
request_data: dict | None,
deliver_ended_stream_rewrites: bool = False,
) -> list["ModelResponseStream"]:
"""Block-only streaming path: run the guardrail so an in-flight BLOCK can
terminate the stream. Text rewrites are not propagated to the client here
(see ``_process_streaming_transform`` for the incremental_diff path)."""
(see ``_process_streaming_transform`` for the incremental_diff path) unless
``deliver_ended_stream_rewrites`` opts the ended-stream branch in."""
# check if the stream has ended
has_stream_ended = False
for chunk in responses_so_far:
@ -511,20 +521,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
break
if has_stream_ended:
# convert to model response
model_response: Final = cast(
ModelResponse,
stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj),
)
# run process_output_response
await self.process_output_response(
response=model_response,
await self._process_ended_stream(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail_to_apply,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
deliver_ended_stream_rewrites=deliver_ended_stream_rewrites,
)
return responses_so_far
# Step 0: Check if any response has text content to process
@ -597,6 +601,39 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
return responses_so_far
async def _process_ended_stream(
self,
*,
responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None",
user_api_key_dict: "UserAPIKeyAuth | None",
request_data: dict[str, object] | None, # mutable-ok: same request-payload shape the hooks take
deliver_ended_stream_rewrites: bool,
) -> None:
"""Ended-stream path: rebuild the full response, run the non-streaming
output guardrail against it, and (when opted in) write any text rewrite
back across the buffered chunks."""
model_response: Final = cast(
ModelResponse,
stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj),
)
pre_guardrail_texts: Final = self._string_choice_contents(model_response)
await self.process_output_response(
response=model_response,
guardrail_to_apply=guardrail_to_apply,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
)
if deliver_ended_stream_rewrites:
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_to_apply.guardrail_name or "unknown",
)
def build_stream_error_items(
self,
exc: "HTTPException",
@ -722,8 +759,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
"""
combined_texts: Final[dict[tuple[int, int | None], str]] = {}
for response_idx, response in enumerate(responses_so_far):
for choice_idx, choice in enumerate(response.choices):
for response in responses_so_far:
for choice in response.choices:
if isinstance(choice, litellm.StreamingChoices):
content = choice.delta.content
elif isinstance(choice, litellm.Choices):
@ -736,7 +773,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if isinstance(content, str):
# String content - accumulate for this choice
str_key: tuple[int, int | None] = (choice_idx, None)
str_key: tuple[int, int | None] = (choice.index, None)
if str_key not in combined_texts:
combined_texts[str_key] = ""
combined_texts[str_key] += content
@ -747,7 +784,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
text_str = content_item.get("text")
if text_str:
list_key: tuple[int, int | None] = (
choice_idx,
choice.index,
content_idx,
)
if list_key not in combined_texts:
@ -937,6 +974,51 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if "name" in func_dict:
existing_tool_call.function.name = func_dict["name"]
@staticmethod
def _string_choice_contents(response: "ModelResponse") -> tuple[str | None, ...]:
return tuple(
choice.message.content if isinstance(choice.message.content, str) else None for choice in response.choices
)
async def _write_ended_stream_text_rewrites(
self,
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 fails closed."""
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
)
if not changed:
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: withhold the stream
# 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
)
async def _apply_guardrail_responses_to_output_streaming(
self,
responses: list["ModelResponseStream"],
@ -952,7 +1034,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
Args:
responses: List of ModelResponseStream objects to modify
guardrailed_texts: List of guardrailed text responses (combined from all chunks)
task_mappings: List of tuples (choice_idx, content_idx)
task_mappings: List of tuples (choice_idx, content_idx), where choice_idx
is the choice's ``index`` field, not its position in a chunk's list
Override this method to customize how responses are applied to streaming responses.
"""
@ -968,9 +1051,11 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
# Key: (choice_idx, content_idx), Value: boolean (True if already set)
already_set: Final[dict[tuple[int, int | None], bool]] = {}
# Iterate through all responses and update content
for response_idx, response in enumerate(responses):
for choice_idx_in_response, choice in enumerate(response.choices):
# Iterate through all responses and update content, matching each chunk's
# choice by its index field: on n>1 streams a chunk usually carries one
# choice at list position 0 whose index names the logical choice.
for response in responses:
for choice in response.choices:
if isinstance(choice, litellm.StreamingChoices):
content = choice.delta.content
elif isinstance(choice, litellm.Choices):
@ -983,7 +1068,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if isinstance(content, str):
# String content
str_key: tuple[int, int | None] = (choice_idx_in_response, None)
str_key: tuple[int, int | None] = (choice.index, None)
if str_key in guardrail_map:
if str_key not in already_set:
# First chunk - set the complete guardrailed text
@ -1004,7 +1089,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
for content_idx, content_item in enumerate(content):
if "text" in content_item:
list_key: tuple[int, int | None] = (
choice_idx_in_response,
choice.index,
content_idx,
)
if list_key in guardrail_map:

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

@ -32,6 +32,7 @@ import time
import uuid
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from itertools import chain, repeat
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Union, cast
@ -110,6 +111,15 @@ class ResponsesStreamChunk(TypedDict, total=False):
content_index: ReadOnly[int]
_TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset(
{
ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value,
ResponsesAPIStreamEvents.RESPONSE_FAILED.value,
ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value,
}
)
def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int:
sequence_numbers: Final = (
item.get("sequence_number") if isinstance(item, dict) else getattr(item, "sequence_number", None)
@ -129,6 +139,8 @@ class OpenAIResponsesHandler(BaseTranslation):
Methods can be overridden to customize behavior for different message formats.
"""
delivers_ended_stream_text_rewrites = True
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
"""
Convert Responses API request data to OpenAI-spec structured messages.
@ -520,6 +532,7 @@ class OpenAIResponsesHandler(BaseTranslation):
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
deliver_ended_stream_rewrites: bool = False,
) -> list[Any]:
"""
Process output streaming response by applying guardrails to text content.
@ -528,10 +541,17 @@ class OpenAIResponsesHandler(BaseTranslation):
chunk, apply the guardrail, then write the result back in-place so the
caller sees the modified content (e.g. PII tokens replaced).
For ``response.completed`` events (the normal end-of-stream signal) we
use the same per-item extraction + task-mapping approach as
``process_output_response`` so that unmasking / blocking works correctly
for every output item.
For terminal envelope events (``response.completed``, and equally
``response.incomplete`` / ``response.failed``, whose envelopes carry the
partial output) we use the same per-item extraction + task-mapping
approach as ``process_output_response`` so that unmasking / blocking
works correctly for every output item. With
``deliver_ended_stream_rewrites`` the earlier text-carrying events
(``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 fails closed instead of releasing raw output.
"""
if not responses_so_far:
return responses_so_far
@ -543,14 +563,16 @@ class OpenAIResponsesHandler(BaseTranslation):
return responses_so_far
# ------------------------------------------------------------------ #
# Case 1: response.completed — full response is available in the #
# final chunk; iterate output items, apply guardrail, write back. #
# Case 1: terminal envelope events (completed/incomplete/failed). #
# the accumulated response is available in the final chunk; iterate #
# output items, apply guardrail, write back. Falls through to the #
# string fallback when the envelope yields nothing to check. #
# ------------------------------------------------------------------ #
if final_chunk.get("type") == "response.completed":
if final_chunk.get("type") in _TERMINAL_ENVELOPE_EVENT_TYPES:
response_obj: Final[ResponseOutputEnvelope] = final_chunk.get("response") or {}
if not hasattr(response_obj, "get"):
return responses_so_far
outputs: Final[Sequence[object]] = response_obj.get("output") or []
outputs: Final[Sequence[object]] = (
(response_obj.get("output") or []) if hasattr(response_obj, "get") else []
)
texts_to_check: Final[list[str]] = []
tool_calls_to_check: Final[list[ChatCompletionToolCallChunk]] = []
@ -600,11 +622,25 @@ class OpenAIResponsesHandler(BaseTranslation):
responses=guardrailed_texts,
task_mappings=task_mappings,
)
return responses_so_far
if deliver_ended_stream_rewrites:
rewrites_by_position: Final = MappingProxyType(
{
task_mappings[task_idx]: rewritten
for task_idx, rewritten in enumerate(guardrailed_texts)
if task_idx < len(texts_to_check) and rewritten != texts_to_check[task_idx]
}
)
if rewrites_by_position:
self._sync_stream_events_with_rewrites(
stream_events=responses_so_far[:-1],
rewrites_by_position=rewrites_by_position,
)
return responses_so_far
# ------------------------------------------------------------------ #
# Case 2: response.output_item.done — extract tool calls only. #
# Case 2: response.output_item.done — extract tool calls only, then #
# fall through to the text fallback when a caller expects rewrites #
# delivered, so a buffer truncated here still fails closed on text. #
# ------------------------------------------------------------------ #
if final_chunk.get("type") == "response.output_item.done":
model_response_stream: Final = (
@ -622,12 +658,14 @@ class OpenAIResponsesHandler(BaseTranslation):
input_type="response",
logging_obj=litellm_logging_obj,
)
return responses_so_far
if not deliver_ended_stream_rewrites:
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. #
# need to block/flag (not rewrite) still work correctly, and a #
# rewrite a caller expects delivered fails closed instead. #
# ------------------------------------------------------------------ #
string_so_far: Final = self.get_streaming_string_so_far(responses_so_far)
if string_so_far:
@ -637,26 +675,83 @@ class OpenAIResponsesHandler(BaseTranslation):
)
if response_model:
fallback_inputs["model"] = response_model
await guardrail_to_apply.apply_guardrail(
fallback_outputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=fallback_inputs,
request_data=request_data if request_data is not None else {},
input_type="response",
logging_obj=litellm_logging_obj,
)
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")
return responses_so_far
@staticmethod
def _write_event_field(event: object, field: str, value: str) -> None:
if isinstance(event, dict):
event[field] = value # rebind-ok: delivering the rewrite means editing the buffered event in place
else:
setattr(event, field, value)
def _sync_stream_events_with_rewrites(
self,
stream_events: Sequence[Any],
rewrites_by_position: Mapping[tuple[int, int], str],
) -> None:
"""Sync pre-completion stream events with the rewritten completed
response, keyed by ``(output_index, content_index)``: the first
``output_text.delta`` for a rewritten item carries the full rewritten
text and the rest are blanked, while ``output_text.done``,
``content_part.done``, and ``output_item.done`` events carry the full
rewritten text, so every event a client may read agrees with the
rewritten ``response.completed`` payload."""
delta_replacements: Final = MappingProxyType(
{position: chain((rewritten,), repeat("")) for position, rewritten in rewrites_by_position.items()}
)
for event in stream_events:
if not (isinstance(event, dict) or hasattr(event, "get")):
continue
event_type = event.get("type")
output_index = event.get("output_index")
content_index = event.get("content_index")
if event_type == "response.output_item.done" and isinstance(output_index, int):
self._sync_output_item_done_event(event.get("item"), output_index, rewrites_by_position)
continue
if not isinstance(output_index, int) or not isinstance(content_index, int):
continue
position = (output_index, content_index)
if event_type == "response.output_text.delta" and position in delta_replacements:
self._write_event_field(event, "delta", next(delta_replacements[position]))
elif event_type == "response.output_text.done" and position in rewrites_by_position:
self._write_event_field(event, "text", rewrites_by_position[position])
elif event_type == "response.content_part.done" and position in rewrites_by_position:
part = event.get("part")
if isinstance(part, dict) or hasattr(part, "text"):
self._write_event_field(part, "text", rewrites_by_position[position])
@staticmethod
def _sync_output_item_done_event(
item: object,
output_index: int,
rewrites_by_position: Mapping[tuple[int, int], str],
) -> None:
content: Final = item.get("content") if isinstance(item, dict) else getattr(item, "content", None)
if not isinstance(content, list):
return
for (item_idx, content_idx), rewritten in rewrites_by_position.items():
if item_idx != output_index or content_idx >= len(content):
continue
OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten)
def _check_streaming_has_ended(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> bool:
"""
Check if the streaming has ended.
"""
if not responses_so_far:
return False
terminal_types: Final = {
ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value,
ResponsesAPIStreamEvents.RESPONSE_FAILED.value,
ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value,
}
return responses_so_far[-1].get("type") in terminal_types
return responses_so_far[-1].get("type") in _TERMINAL_ENVELOPE_EVENT_TYPES
def build_stream_error_items(
self,

72
litellm/proxy/policy_engine/pipeline_executor.py
View file

@ -34,6 +34,7 @@ if TYPE_CHECKING:
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
)
from litellm.proxy._types import UserAPIKeyAuth
try:
from fastapi.exceptions import HTTPException
@ -44,7 +45,8 @@ except ImportError:
class UndeliverableStreamRewrite(Exception):
def __init__(self, guardrail_name: str) -> None:
super().__init__(
f"Guardrail '{guardrail_name}' rewrote the streamed response, which streaming pipelines cannot deliver"
f"Guardrail '{guardrail_name}' rewrote the streamed response in a way this endpoint's "
"streaming pipeline cannot deliver"
)
self.guardrail_name: Final = guardrail_name
@ -71,14 +73,17 @@ def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | Non
class _StreamRewriteObserver(CustomGuardrail):
"""Stand-in handed to the endpoint translation in place of a streaming pipeline step's
guardrail. Translations cannot rewrite every buffered chunk consistently, so the gate
withholds the stream whenever the guardrail returned different output than it was given,
which for guardrails like Bedrock's ANONYMIZED action is only known at runtime."""
guardrail. It records whether the guardrail returned different output than it was given,
which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text
rewrites are deliverable on translations that write them back across the buffered chunks
(``delivers_ended_stream_text_rewrites``); tool-call rewrites and text rewrites on any
other translation make the gate withhold the stream."""
def __init__(self, inner: CustomGuardrail) -> None:
super().__init__(guardrail_name=inner.guardrail_name)
self.inner: Final = inner
self.rewrote = False
self.rewrote_texts = False
self.rewrote_tool_calls = False
def structured_messages_cover_full_request(self) -> bool:
return self.inner.structured_messages_cover_full_request()
@ -95,10 +100,9 @@ class _StreamRewriteObserver(CustomGuardrail):
outputs: Final = await self.inner.apply_guardrail(
inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj
)
self.rewrote = (
self.rewrote
or _rewrote(sent_texts, _text_snapshot(outputs.get("texts")))
or _rewrote(sent_tool_shapes, _tool_call_shapes(outputs.get("tool_calls")))
self.rewrote_texts = self.rewrote_texts or _rewrote(sent_texts, _text_snapshot(outputs.get("texts")))
self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote(
sent_tool_shapes, _tool_call_shapes(outputs.get("tool_calls"))
)
return outputs
@ -248,6 +252,41 @@ class PipelineExecutor:
# Ran out of steps without a terminal action → default allow
return _allow_result(step_results=step_results, working_data=working_data, request_data=data)
@staticmethod
async def _run_streaming_step(
step: PipelineStep,
callback: CustomGuardrail,
endpoint_translation: "BaseTranslation",
streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks the translation rewrites in place
hook_input: dict[str, object], # mutable-ok: same request-payload shape as data
user_api_key_dict: "UserAPIKeyAuth | None",
litellm_logging_obj: "LiteLLMLoggingObj | None",
) -> None:
"""Run one streaming post_call step through the endpoint translation, delivering
text rewrites on translations that support ended-stream write-back and raising
``UndeliverableStreamRewrite`` for any rewrite that cannot reach the client."""
observer: Final = _StreamRewriteObserver(callback)
deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites
if deliver_rewrites:
await endpoint_translation.process_output_streaming_response(
responses_so_far=streaming_chunks,
guardrail_to_apply=observer,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=hook_input,
deliver_ended_stream_rewrites=True,
)
else:
await endpoint_translation.process_output_streaming_response(
responses_so_far=streaming_chunks,
guardrail_to_apply=observer,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=hook_input,
)
if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites):
raise UndeliverableStreamRewrite(step.guardrail)
@staticmethod
async def _run_step(
step: PipelineStep,
@ -310,16 +349,15 @@ class PipelineExecutor:
f"Guardrail '{step.guardrail}' does not support streaming pipeline execution",
None,
)
observer: Final = _StreamRewriteObserver(callback)
await endpoint_translation.process_output_streaming_response(
responses_so_far=streaming_chunks,
guardrail_to_apply=observer,
litellm_logging_obj=data.get("litellm_logging_obj"),
await PipelineExecutor._run_streaming_step(
step=step,
callback=callback,
endpoint_translation=endpoint_translation,
streaming_chunks=streaming_chunks,
hook_input=hook_input,
user_api_key_dict=user_api_key_dict,
request_data=hook_input,
litellm_logging_obj=data.get("litellm_logging_obj"),
)
if observer.rewrote:
raise UndeliverableStreamRewrite(step.guardrail)
response = None
elif mode == "post_call":
response = await target.async_post_call_success_hook(

61
litellm/proxy/utils.py
View file

@ -518,14 +518,6 @@ def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool:
return callback is not None and PipelineExecutor.supports_unified_execution(callback)
def _pipeline_step_rewrites_streamed_content(guardrail_name: str) -> bool:
callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name)
if callback is None:
return False
transform_mode: Final = unified_guardrail.resolve_streaming_flag(callback, "streaming_transform_mode", "block_only")
return callback.rewrites_streamed_output() or transform_mode == "incremental_diff"
class _PipelineErrorBody(TypedDict):
message: ReadOnly[str]
type: ReadOnly[str]
@ -542,9 +534,10 @@ def _undeliverable_stream_rewrite_error(policy_name: str, guardrail_name: str) -
"error": {
"message": (
f"Streaming response withheld by policy pipeline '{policy_name}' because guardrail "
f"'{guardrail_name}' rewrote the streamed output, and streaming pipelines cannot deliver "
"rewrites. Retry with stream=false, or drop it from the pipeline steps so guardrails.add "
"applies it to streamed output."
f"'{guardrail_name}' rewrote the streamed output in a way this endpoint's streaming "
"pipeline cannot deliver (a tool-call rewrite, or a text rewrite on a route without "
"stream write-back). Retry with stream=false, or drop it from the pipeline steps so "
"guardrails.add applies it to streamed output."
),
"type": "guardrail_pipeline_error",
"policies": (policy_name,),
@ -561,13 +554,13 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_ap
Background responses skip the post_call hooks entirely, so a pipeline
governing one would silently never execute. Streaming responses execute
pipelines against the buffered stream through the endpoint guardrail
translation of the request route, releasing the buffered chunks on allow.
That needs every step's guardrail to support the unified apply_guardrail
interface and to only allow or block (a step that rewrites streamed
content, via mask_response_content, a MASK action, or
streaming_transform_mode=incremental_diff, would have its rewrite silently
dropped), and needs the route to have a translation at all; anything else
keeps the 400 rather than letting ungoverned output stream through.
translation of the request route, releasing the buffered chunks on allow
(rewritten in place when a guardrail rewrote text and the translation
delivers ended-stream rewrites; a rewrite the translation cannot deliver
fails closed at runtime instead). That needs every step's guardrail to
support the unified apply_guardrail interface, and needs the route to have
a translation at all; anything else keeps the 400 rather than letting
ungoverned output stream through.
"""
is_stream: Final = data.get("stream") is True
is_background: Final = data.get("background") is True
@ -612,25 +605,6 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_ap
}
}
raise HTTPException(status_code=400, detail=unsupported_detail)
rewriting_guardrails: Final = tuple(
guardrail for guardrail in step_guardrails if _pipeline_step_rewrites_streamed_content(guardrail)
)
if rewriting_guardrails:
rewriting_detail: Final[_PipelineErrorDetail] = {
"error": {
"message": (
"Policies with post_call guardrail pipelines cannot govern streaming responses "
"because these pipeline guardrails rewrite streamed content (mask_response_content, "
"a MASK action, or streaming_transform_mode=incremental_diff), which pipeline steps would release "
f"unmodified: {', '.join(rewriting_guardrails)}. Retry with stream=false, or drop "
"them from the pipeline steps so guardrails.add applies them to streamed output."
),
"type": "guardrail_pipeline_error",
"policies": post_call_policies,
"guardrails": rewriting_guardrails,
}
}
raise HTTPException(status_code=400, detail=rewriting_detail)
route: Final = user_api_key_dict.request_route
if not route or resolve_endpoint_translation(user_api_key_dict, None) is not None:
return
@ -3498,12 +3472,13 @@ class ProxyLogging:
pipeline allows it), then runs each pipeline's steps against the
assembled output through the endpoint guardrail translation, the same
machinery flat post_call guardrails use at end of stream. An allow
releases the buffered chunks verbatim; a step whose guardrail rewrote
the output withholds the stream with a 400 instead, since no
translation rewrites every buffered chunk consistently and some
rewrites (Bedrock's ANONYMIZED action, for one) are only decided at
runtime; a block or modify_response terminates with the translation's
block chunks or the raised error.
releases the buffered chunks: verbatim when no guardrail rewrote the
output, rewritten in place when one rewrote text and the translation
delivers ended-stream rewrites (later steps then re-scan the rewritten
chunks, so rewrites chain). A rewrite the translation cannot deliver
(a tool-call rewrite, or a text rewrite on a route without write-back)
withholds the stream with a 400; a block or modify_response terminates
with the translation's block chunks or the raised error.
"""
buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict
async for item in response:

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

@ -263,6 +263,117 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
# Should return the responses unchanged
assert result == responses_so_far
@staticmethod
def _ended_sse_chunks() -> list:
events = [
("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}),
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello "}}),
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "world"}}),
("content_block_stop", {"type": "content_block_stop", "index": 0}),
("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 2}}),
("message_stop", {"type": "message_stop"}),
]
return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events]
@staticmethod
def _masking_guardrail() -> CustomGuardrail:
class MaskWorld(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
return {**inputs, "texts": [text.replace("world", "[MASKED]") for text in inputs.get("texts", [])]}
return MaskWorld(guardrail_name="test")
@staticmethod
def _delta_texts(chunks: list) -> list:
texts = []
for chunk in chunks:
for line in chunk.decode().split("\n"):
if not line.startswith("data:"):
continue
data = json.loads(line[len("data:") :].strip())
if data.get("type") == "content_block_delta":
texts.append(data["delta"]["text"])
return texts
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_writes_text_back_into_sse_chunks(self):
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()
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: message_stop" in raw
assert '"stop_reason": "end_turn"' in raw
@pytest.mark.asyncio
async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self):
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()
original = [bytes(chunk) for chunk in chunks]
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=MagicMock(),
)
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
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()[:-2]
with pytest.raises(UndeliverableStreamRewrite):
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,
)
@pytest.mark.asyncio
async def test_unended_stream_without_rewrite_is_released_with_delivery_expected(self):
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()[:-2]
original = [bytes(chunk) for chunk in chunks]
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=MockPassThroughGuardrail(guardrail_name="test"),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert chunks == original
@pytest.mark.asyncio
async def test_unended_stream_rewrite_without_delivery_expected_does_not_raise(self):
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()[:-2]
original = [bytes(chunk) for chunk in chunks]
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=MagicMock(),
)
assert result is chunks
assert chunks == original
class TestAnthropicMessagesHandlerInputProcessing:
"""Test input processing preserves litellm_metadata for dynamic guardrails."""

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

@ -1073,6 +1073,154 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
# Should return the responses
assert result == responses_so_far
@staticmethod
def _ended_stream_chunks() -> list:
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
return [
ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[StreamingChoices(index=0, delta=Delta(content="Hello"), finish_reason=None)],
),
ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[StreamingChoices(index=0, delta=Delta(content=" world"), finish_reason="stop")],
),
]
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_writes_text_back_into_chunks(self):
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="test")
chunks = self._ended_stream_chunks()
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 chunks[0].choices[0].delta.content == "HELLO WORLD"
assert chunks[1].choices[0].delta.content in (None, "")
assert chunks[1].choices[0].finish_reason == "stop"
@pytest.mark.asyncio
async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self):
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="test")
chunks = self._ended_stream_chunks()
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
)
assert chunks[0].choices[0].delta.content == "Hello"
assert chunks[1].choices[0].delta.content == " world"
assert chunks[1].choices[0].finish_reason == "stop"
@staticmethod
def _two_choice_stream_chunks() -> list:
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
def chunk(index: int, content: str, finish_reason: Optional[str] = None) -> ModelResponseStream:
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, "safe "),
chunk(1, "hello "),
chunk(0, "text", "stop"),
chunk(1, "world", "stop"),
]
@staticmethod
def _world_masking_guardrail() -> CustomGuardrail:
class MaskWorld(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
texts = inputs.get("texts", [])
return {**inputs, "texts": [t.replace("world", "[MASKED]") for t in texts]}
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
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,
)
@pytest.mark.asyncio
async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self):
handler = OpenAIChatCompletionsHandler()
chunks = self._two_choice_stream_chunks()
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=MockPassThroughGuardrail(guardrail_name="test"),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert [c.choices[0].delta.content for c in chunks] == ["safe ", "hello ", "text", "world"]
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrite_lands_on_nonzero_choice_index(self):
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
handler = OpenAIChatCompletionsHandler()
def chunk(content: str, finish_reason: Optional[str]) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[StreamingChoices(index=1, delta=Delta(content=content), finish_reason=finish_reason)],
)
chunks = [chunk("hello ", None), chunk("world", "stop")]
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 chunks[0].choices[0].delta.content == "hello [MASKED]"
assert chunks[1].choices[0].delta.content in (None, "")
class TestGetStructuredMessages:
"""Test the get_structured_messages method."""

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

@ -1122,6 +1122,209 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
output_text = result[-1]["response"]["output"][0]["content"][0]["text"]
assert output_text == original_text
@staticmethod
def _ended_stream_events() -> List[dict]:
content = [{"type": "output_text", "text": "hello world"}]
item = {
"type": "message",
"id": "msg_123",
"status": "completed",
"role": "assistant",
"content": content,
}
return [
{"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"},
{"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"},
{
"type": "response.content_part.done",
"output_index": 0,
"content_index": 0,
"part": {"type": "output_text", "text": "hello world"},
},
{"type": "response.output_item.done", "output_index": 0, "item": {**item, "content": [dict(c) for c in content]}},
{
"type": "response.completed",
"response": {
"id": "resp_123",
"model": "gpt-4o",
"output": [{**item, "content": [dict(c) for c in content]}],
"status": "completed",
},
},
]
@staticmethod
def _masking_guardrail() -> CustomGuardrail:
class MaskWorld(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
texts = inputs.get("texts", [])
return {**inputs, "texts": [t.replace("world", "[MASKED]") for t in texts]}
return MaskWorld(guardrail_name="test-mask")
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_syncs_all_stream_events(self):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()
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]"
assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]"
@pytest.mark.asyncio
@pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"])
async def test_deliver_ended_stream_rewrites_syncs_non_completed_terminals(self, terminal_type):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()
events[-1]["type"] = terminal_type
events[-1]["response"]["status"] = terminal_type.split(".")[-1]
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]"
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
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,
)
@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
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"},
]
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,
)
@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
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,
)
@pytest.mark.asyncio
async def test_output_item_done_last_scans_text_with_delivery_expected(self):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()[:-1]
guardrail = MockRecordingGuardrail(guardrail_name="test")
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is events
assert [inputs.get("texts") for inputs in guardrail.seen_inputs] == [["hello world"]]
@pytest.mark.asyncio
async def test_output_item_done_last_without_delivery_expected_skips_text(self):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()[:-1]
guardrail = MockRecordingGuardrail(guardrail_name="test")
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
)
assert result is events
assert guardrail.seen_inputs == []
@pytest.mark.asyncio
async def test_fallback_rewrite_without_delivery_expected_does_not_raise(self):
handler = OpenAIResponsesHandler()
events = [
{"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"},
]
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
)
assert result is events
@pytest.mark.asyncio
async def test_ended_stream_rewrite_leaves_delta_events_untouched_by_default(self):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
)
assert events[0]["delta"] == "hello "
assert events[1]["delta"] == "world"
assert events[2]["text"] == "hello world"
assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]"
@pytest.mark.asyncio
async def test_failed_stream_scans_delta_text(self):
"""A stream ending in response.failed has text only in delta events; the

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

@ -923,6 +923,8 @@ class _TextReturningGuardrail(CustomGuardrail):
class _TextTranslation:
delivers_ended_stream_text_rewrites = False
def __init__(self):
self.seen_guardrail_names = []

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

@ -24,7 +24,7 @@ from litellm.integrations.custom_guardrail import (
ModifyResponseException,
)
from litellm.integrations.prometheus import PrometheusLogger
from litellm.proxy._types import ProxyException
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
from litellm.proxy.utils import ProxyLogging, _raise_for_streaming_post_call_pipelines
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
@ -1479,7 +1479,7 @@ async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_lacks_uni
("guardrail_config", {"streaming_transform_mode": "incremental_diff"}),
],
)
async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_rewrites_streamed_content(
async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_rewrites_streamed_content(
proxy_logging, make_user_api_key_auth, monkeypatch, rewrite_attribute, value
):
seen: Dict[str, Any] = {}
@ -1488,24 +1488,21 @@ async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_rewrites_
monkeypatch.setattr(litellm, "callbacks", [guardrail])
data = _post_call_pipeline_data(stream=True)
with pytest.raises(HTTPException) as info:
await proxy_logging.pre_call_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
data=data,
call_type="completion",
guardrails_only=True,
)
out = await proxy_logging.pre_call_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
data=data,
call_type="completion",
guardrails_only=True,
)
assert info.value.status_code == 400
assert info.value.detail["error"]["guardrails"] == ("gr-post",)
assert "rewrite streamed content" in info.value.detail["error"]["message"]
assert seen.get("count") is None
assert out is not None
assert out.get("stream") is True
@pytest.mark.asyncio
@pytest.mark.parametrize("action, rejected", [(ContentFilterAction.MASK, True), (ContentFilterAction.BLOCK, False)])
async def test_pre_call_hook_rejects_streaming_only_when_content_filter_step_masks(
proxy_logging, make_user_api_key_auth, monkeypatch, action, rejected
@pytest.mark.parametrize("action", [ContentFilterAction.MASK, ContentFilterAction.BLOCK])
async def test_pre_call_hook_allows_streaming_when_content_filter_step_masks_or_blocks(
proxy_logging, make_user_api_key_auth, monkeypatch, action
):
guardrail = ContentFilterGuardrail(
guardrail_name="gr-post",
@ -1516,25 +1513,15 @@ async def test_pre_call_hook_rejects_streaming_only_when_content_filter_step_mas
data = _post_call_pipeline_data(stream=True)
user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions")
if not rejected:
out = await proxy_logging.pre_call_hook(
user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True
)
assert out is not None and out.get("stream") is True
return
out = await proxy_logging.pre_call_hook(
user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True
)
with pytest.raises(HTTPException) as info:
await proxy_logging.pre_call_hook(
user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True
)
assert info.value.status_code == 400
assert info.value.detail["error"]["guardrails"] == ("gr-post",)
assert "a MASK action" in info.value.detail["error"]["message"]
assert out is not None and out.get("stream") is True
@pytest.mark.asyncio
async def test_pre_call_hook_rejects_streaming_when_content_filter_category_masks(
async def test_pre_call_hook_allows_streaming_when_content_filter_category_masks(
proxy_logging, make_user_api_key_auth, monkeypatch
):
guardrail = ContentFilterGuardrail(
@ -1546,13 +1533,11 @@ async def test_pre_call_hook_rejects_streaming_when_content_filter_category_mask
data = _post_call_pipeline_data(stream=True)
user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions")
with pytest.raises(HTTPException) as info:
await proxy_logging.pre_call_hook(
user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True
)
out = await proxy_logging.pre_call_hook(
user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True
)
assert info.value.status_code == 400
assert info.value.detail["error"]["guardrails"] == ("gr-post",)
assert out is not None and out.get("stream") is True
@pytest.mark.asyncio
@ -1656,17 +1641,10 @@ def _echoed_tool_call_dicts(arguments: str) -> List[Dict[str, Any]]:
@pytest.mark.asyncio
@pytest.mark.parametrize("on_fail, on_error", [("block", None), ("next", "next")])
@pytest.mark.parametrize(
"make_chunks, transform",
[
(_stream_chunks, lambda inputs: {"texts": ["hello [MASKED]"]}),
(_tool_call_stream_chunks, lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')}),
],
ids=["texts", "tool_calls"],
)
async def test_streaming_iterator_hook_pipeline_withholds_runtime_rewrite(
proxy_logging, make_user_api_key_auth, monkeypatch, make_chunks, transform, on_fail, on_error
async def test_streaming_iterator_hook_pipeline_withholds_runtime_tool_call_rewrite(
proxy_logging, make_user_api_key_auth, monkeypatch, on_fail, on_error
):
transform = lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')} # noqa: E731
monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
step = PipelineStep(guardrail="gr-post", on_pass="allow", on_fail=on_fail, on_error=on_error)
@ -1676,7 +1654,7 @@ async def test_streaming_iterator_hook_pipeline_withholds_runtime_rewrite(
async def _drain() -> None:
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
response=_async_chunk_iter(make_chunks()),
response=_async_chunk_iter(_tool_call_stream_chunks()),
request_data=data,
):
delivered.append(item)
@ -1693,6 +1671,81 @@ async def test_streaming_iterator_hook_pipeline_withholds_runtime_rewrite(
assert "stream=false" in error["message"]
@pytest.mark.asyncio
async def test_streaming_iterator_hook_pipeline_delivers_runtime_text_rewrite(
proxy_logging, make_user_api_key_auth, monkeypatch
):
transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731
monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
chunks = _stream_chunks()
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
response=_async_chunk_iter(chunks),
request_data=data,
)
]
assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks]
assert delivered[0].choices[0].delta.content == "hello [MASKED]"
assert delivered[1].choices[0].delta.content in (None, "")
assert delivered[1].choices[0].finish_reason == "stop"
@pytest.mark.asyncio
async def test_streaming_iterator_hook_pipeline_chains_text_rewrites_across_steps(
proxy_logging, make_user_api_key_auth, monkeypatch
):
second_step_saw: Dict[str, Any] = {}
class FirstMask(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
return {**inputs, "texts": [text.replace("world", "[MASKED]") for text in inputs["texts"]]}
class SecondMask(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
second_step_saw["texts"] = list(inputs["texts"])
return {**inputs, "texts": [text.replace("hello", "[GREETING]") for text in inputs["texts"]]}
monkeypatch.setattr(
litellm,
"callbacks",
[
FirstMask(guardrail_name="gr-first", event_hook=GuardrailEventHooks.post_call, default_on=False),
SecondMask(guardrail_name="gr-second", event_hook=GuardrailEventHooks.post_call, default_on=False),
],
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
pipeline = GuardrailPipeline(
mode="post_call",
steps=[
PipelineStep(guardrail="gr-first", on_pass="next", on_fail="block"),
PipelineStep(guardrail="gr-second", on_pass="allow", on_fail="block"),
],
)
data = _post_call_pipeline_data(stream=True)
data["metadata"]["_guardrail_pipelines"] = [("response-governance", pipeline)]
chunks = _stream_chunks()
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
response=_async_chunk_iter(chunks),
request_data=data,
)
]
assert second_step_saw["texts"] == ["hello [MASKED]"]
assert delivered[0].choices[0].delta.content == "[GREETING] [MASKED]"
assert delivered[1].choices[0].delta.content in (None, "")
assert delivered[1].choices[0].finish_reason == "stop"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"make_chunks, transform",
@ -1799,6 +1852,79 @@ async def test_streaming_iterator_hook_pipeline_modify_response_emits_translated
assert not any(item is chunk for item in delivered for chunk in chunks)
@pytest.mark.asyncio
async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthropic_sse(
proxy_logging, make_user_api_key_auth, monkeypatch
):
transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731
monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
chunks = _anthropic_sse_chunks()
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"),
response=_async_chunk_iter(chunks),
request_data=data,
)
]
raw = b"".join(delivered).decode()
assert "hello [MASKED]" in raw
assert "hello world" not in raw
assert raw.count("event: content_block_delta") == 1
for expected_event in ("message_start", "content_block_start", "content_block_stop", "message_delta", "message_stop"):
assert f"event: {expected_event}" in raw
@pytest.mark.asyncio
async def test_pipeline_executor_withholds_text_rewrite_when_translation_lacks_write_back(monkeypatch):
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite
class NoWriteBackTranslation(BaseTranslation):
async def process_input_messages(self, data, guardrail_to_apply, litellm_logging_obj):
return data
async def process_output_response(self, response, guardrail_to_apply, litellm_logging_obj, **kwargs):
return response
async def process_output_streaming_response(
self,
responses_so_far,
guardrail_to_apply,
litellm_logging_obj,
user_api_key_dict=None,
request_data=None,
stream_transform_sink=None,
deliver_ended_stream_rewrites=False,
):
assert deliver_ended_stream_rewrites is False
await guardrail_to_apply.apply_guardrail(
inputs={"texts": ["hello world"]},
request_data=request_data or {},
input_type="response",
)
return responses_so_far
transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731
monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)])
with pytest.raises(UndeliverableStreamRewrite):
await PipelineExecutor.execute_steps(
steps=[PipelineStep(guardrail="gr-post", on_pass="allow", on_fail="block")],
mode="post_call",
data={"metadata": {}},
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
call_type="acompletion",
policy_name="response-governance",
streaming_chunks=_stream_chunks(),
endpoint_translation=NoWriteBackTranslation(),
)
@pytest.mark.asyncio
async def test_streaming_iterator_hook_pipeline_gates_without_iterator_overrides(
proxy_logging, make_user_api_key_auth, monkeypatch