diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index b9ca18c7843..89c8431dfe4 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -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. diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index ba96ab3dc99..4b0cc0fd97c 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -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 diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 54673c77f80..1358cf7c37a 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -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"], diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 7c5d8ac99ad..0f475aa04c8 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -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. diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index acd2c2c973a..1264cdfc14a 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -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( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 6c990222e51..7f3c3aecd87 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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: diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index af3ccd65b11..ba26da50bc8 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -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.""" diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index a29e0be4655..26442a4a6ed 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -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.""" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 447175b09a6..4f95e08cb71 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -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.""" diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 52fd8777a19..908c9f12c9e 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -823,6 +823,8 @@ class _TextReturningGuardrail(CustomGuardrail): class _TextTranslation: + delivers_ended_stream_text_rewrites = False + def __init__(self): self.seen_guardrail_names = [] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 74bd2e483cc..73270ec5671 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -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