fix(policy_engine): apply post_call pipeline text rewrites on streams

Buffered streams governed by post_call policy pipelines now deliver text
rewrites back into the stream per surface (chat SSE, responses SSE,
anthropic messages SSE) instead of rejecting the request with a 400
upfront. Rewrites chain across pipeline steps; tool-call rewrites and
translations without stream write-back still withhold the stream.
This commit is contained in:
mateo-berri 2026-08-29 22:12:59 -07:00
parent 99286cd254
commit 673d1743a6
11 changed files with 711 additions and 130 deletions

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, cast
from typing_extensions import assert_never
@ -120,6 +121,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()
@ -931,11 +934,14 @@ 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,
) -> list[Any]:
"""
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).
"""
from litellm.integrations.custom_guardrail import ModifyResponseException
@ -982,6 +988,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
@ -1093,6 +1108,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: list[Any]) -> str:
"""
Parse streaming responses and extract accumulated text content.

View file

@ -1,6 +1,6 @@
from abc import ABC, abstractmethod
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 litellm.integrations.custom_guardrail import (
@ -33,6 +33,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,
@ -113,6 +120,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.
@ -120,6 +128,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

View file

@ -61,6 +61,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.
@ -440,6 +442,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
user_api_key_dict: Any | 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.
@ -454,6 +457,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.
@ -479,6 +486,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(
@ -489,10 +497,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
litellm_logging_obj: "LiteLLMLoggingObj | None",
user_api_key_dict: Any | 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:
@ -501,20 +511,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
@ -591,6 +595,38 @@ 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: object,
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,
)
@staticmethod
def _accumulate_string_content_by_choice_index(
responses_so_far: list["ModelResponseStream"],
@ -922,6 +958,41 @@ 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, ...],
) -> None:
"""Write ended-stream guardrail text rewrites back across the buffered
chunks: each rewritten choice's full text lands in its 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(
(choice_idx, after)
for choice_idx, (before, after) in enumerate(zip(pre_guardrail_texts, post_guardrail_texts))
if before is not None and after is not None and after != before
)
if not changed:
return
await self._apply_guardrail_responses_to_output_streaming(
responses=responses_so_far,
guardrailed_texts=[
after for _choice_idx, after in changed
], # mutable-ok: the callee's signature predates this change and takes lists
task_mappings=[
(choice_idx, None) for choice_idx, _after in changed
], # mutable-ok: the callee's signature predates this change and takes lists
)
async def _apply_guardrail_responses_to_output_streaming(
self,
responses: list["ModelResponseStream"],

View file

@ -28,7 +28,9 @@ Output: response.output is List[GenericResponseOutputItem] where each has:
- text: str
"""
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from itertools import chain, repeat
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Union, cast
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
@ -91,6 +93,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.
@ -482,6 +486,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.
@ -493,7 +498,11 @@ class OpenAIResponsesHandler(BaseTranslation):
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 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 completed response too, so a client reading deltas sees
the rewrite instead of the raw model output.
"""
if not responses_so_far:
return responses_so_far
@ -562,6 +571,19 @@ class OpenAIResponsesHandler(BaseTranslation):
responses=guardrailed_texts,
task_mappings=task_mappings,
)
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
@ -607,6 +629,63 @@ class OpenAIResponsesHandler(BaseTranslation):
)
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.

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
@ -57,28 +59,31 @@ def _tool_call_shape(tool_call: object) -> tuple[object, object]:
return (function.get("name"), function.get("arguments"))
def _rewrote_texts(sent: Sequence[str] | None, returned: Sequence[str] | None) -> bool:
return sent is not None and returned is not None and list(returned) != list(sent)
def _text_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None:
return None if texts is None else tuple(texts)
def _rewrote_tool_calls(sent: Sequence[object] | None, returned: Sequence[object] | None) -> bool:
if sent is None or returned is None:
return False
return [_tool_call_shape(tool_call) for tool_call in returned] != [
_tool_call_shape(tool_call) for tool_call in sent
]
def _tool_call_shapes(tool_calls: Sequence[object] | None) -> tuple[tuple[object, object], ...] | None:
return None if tool_calls is None else tuple(_tool_call_shape(tool_call) for tool_call in tool_calls)
def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool:
return sent is not None and returned is not None and returned != sent
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()
@ -90,13 +95,14 @@ class _StreamRewriteObserver(CustomGuardrail):
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
sent_texts: Final = _text_snapshot(inputs.get("texts"))
sent_tool_shapes: Final = _tool_call_shapes(inputs.get("tool_calls"))
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_texts(inputs.get("texts"), outputs.get("texts"))
or _rewrote_tool_calls(inputs.get("tool_calls"), 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
@ -250,6 +256,41 @@ class PipelineExecutor:
modified_data=working_data if working_data != data else None,
)
@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,
@ -312,16 +353,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(

View file

@ -492,14 +492,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]
@ -516,9 +508,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,),
@ -535,13 +528,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
@ -586,25 +579,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
@ -3485,12 +3459,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:

View file

@ -263,6 +263,71 @@ 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
class TestAnthropicMessagesHandlerInputProcessing:
"""Test input processing preserves litellm_metadata for dynamic guardrails."""

View file

@ -1073,6 +1073,61 @@ 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"
class TestGetStructuredMessages:
"""Test the get_structured_messages method."""

View file

@ -1104,6 +1104,89 @@ 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
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]"
class TestGetStructuredMessages:
"""Test the get_structured_messages method for Responses API handler."""

View file

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

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
@ -1441,7 +1441,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] = {}
@ -1450,24 +1450,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",
@ -1478,25 +1475,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(
@ -1508,13 +1495,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
@ -1618,17 +1603,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)
@ -1638,7 +1616,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)
@ -1655,6 +1633,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",
@ -1761,6 +1814,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