diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index c7d12e5cf3a..fc276dd7fe6 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, Protocol, cast, overload, runtime_checkable from typing_extensions import ReadOnly, TypedDict, assert_never @@ -164,6 +165,8 @@ class AnthropicMessagesHandler(BaseTranslation): them through guardrail rewrites; downstream provider handling is out of scope. """ + delivers_ended_stream_text_rewrites = True + def __init__(self): super().__init__() self.adapter = LiteLLMAnthropicMessagesAdapter() @@ -1010,11 +1013,15 @@ class AnthropicMessagesHandler(BaseTranslation): litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, + deliver_ended_stream_rewrites: bool = False, ) -> Sequence[object]: """ Process output streaming response by applying guardrails to text content. Get the string so far, check the apply guardrail to the string so far, and return the list of responses so far. + With ``deliver_ended_stream_rewrites``, an ended stream whose guardrail rewrote the text gets the rewrite + written back across the buffered chunks (full rewritten text in the first ``text_delta``, the rest blanked); + a rewrite on a stream that never reported a ``stop_reason`` has no write-back and fails closed instead. """ from litellm.integrations.custom_guardrail import ModifyResponseException @@ -1061,6 +1068,15 @@ class AnthropicMessagesHandler(BaseTranslation): responses_so_far, request_data ) raise + guardrailed_texts: Final = _guardrailed_inputs.get("texts") + if ( + deliver_ended_stream_rewrites + and isinstance(string_so_far, str) + and string_so_far + and guardrailed_texts + and guardrailed_texts[0] != string_so_far + ): + self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0]) else: verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices") return responses_so_far @@ -1083,6 +1099,11 @@ class AnthropicMessagesHandler(BaseTranslation): if e.original_response is None: e.original_response = self._build_streaming_usage_response(responses_so_far, request_data) raise + unended_texts: Final = _guardrailed_inputs.get("texts") + if deliver_ended_stream_rewrites and unended_texts and tuple(unended_texts) != (string_so_far,): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown") return responses_so_far def _prepare_request_data( @@ -1176,6 +1197,63 @@ class AnthropicMessagesHandler(BaseTranslation): inputs["model"] = response_model return inputs + @staticmethod + def _write_ended_stream_text_rewrite( + responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place + rewritten_text: str, + ) -> None: + """Deliver an ended-stream guardrail text rewrite by rewriting the + buffered chunks in place: the first ``text_delta`` carries the full + rewritten text and every later one is blanked, leaving the surrounding + message and content-block framing untouched. Handles both chunk formats + this stream carries (parsed event dicts and raw SSE bytes).""" + replacements: Final = chain((rewritten_text,), repeat("")) + for idx, item in enumerate(responses_so_far): + if isinstance(item, dict): + delta = item.get("delta") + if item.get("type") == "content_block_delta" and isinstance(delta, dict): + if delta.get("type") == "text_delta": + delta["text"] = next(replacements) + elif isinstance(item, (bytes, bytearray)): + responses_so_far[idx] = ( # rebind-ok: delivers the rewrite into the caller's buffer + AnthropicMessagesHandler._rewrite_sse_text_deltas(bytes(item), replacements) + ) + + @staticmethod + def _rewrite_sse_text_deltas(sse_bytes: bytes, replacements: "Iterator[str]") -> bytes: + """Rewrite every ``text_delta`` data line in one SSE chunk with the next + replacement text, leaving all other events and framing byte-identical.""" + try: + decoded: Final = sse_bytes.decode("utf-8") + except UnicodeDecodeError: + return sse_bytes + return "\n\n".join( + AnthropicMessagesHandler._rewrite_sse_block(block, replacements) for block in decoded.split("\n\n") + ).encode("utf-8") + + @staticmethod + def _rewrite_sse_block(block: str, replacements: "Iterator[str]") -> str: + return "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, replacements) for line in block.split("\n")) + + @staticmethod + def _rewrite_sse_line(line: str, replacements: "Iterator[str]") -> str: + if not line.startswith("data:"): + return line + try: + data: Final[str | int | float | bool | None | Sequence[object] | Mapping[str, object]] = json.loads( + line[len("data:") :].strip() + ) + except json.JSONDecodeError: + return line + if not isinstance(data, dict) or data.get("type") != "content_block_delta": + return line + delta: Final = data.get("delta") + if not isinstance(delta, dict) or delta.get("type") != "text_delta": + return line + return "data: " + json.dumps( + {**data, "delta": {**delta, "text": next(replacements)}} # mutable-ok: json.dumps needs plain dicts + ) + def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str: """ Parse streaming responses and extract accumulated text content. diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 220fcedb0f8..591d26ae027 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from collections.abc import Sequence from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional if TYPE_CHECKING: from fastapi import HTTPException @@ -36,6 +36,13 @@ class StreamTransformSink: class BaseTranslation(ABC): + delivers_ended_stream_text_rewrites: ClassVar[bool] = False + """Whether ``process_output_streaming_response`` accepts + ``deliver_ended_stream_rewrites=True`` and, on an ended (fully buffered) + stream, writes guardrail text rewrites back across ``responses_so_far`` so + a buffered pipeline can release rewritten chunks instead of withholding the + stream. Tool-call rewrites stay undeliverable everywhere.""" + @staticmethod def transform_user_api_key_dict_to_metadata( user_api_key_dict: Any | None, @@ -141,6 +148,7 @@ class BaseTranslation(ABC): user_api_key_dict: Optional["UserAPIKeyAuth"] = None, request_data: dict | None = None, stream_transform_sink: StreamTransformSink | None = None, + deliver_ended_stream_rewrites: bool = False, ) -> Any: """ Process output streaming response with guardrails. @@ -148,6 +156,11 @@ class BaseTranslation(ABC): Optional to override in subclasses. ``stream_transform_sink`` is the out-parameter used by handlers that support streaming text transformations (see ``StreamTransformSink``); base handlers ignore it. + ``deliver_ended_stream_rewrites`` is passed True only when the caller + holds the whole buffered stream and the subclass declares + ``delivers_ended_stream_text_rewrites``: the handler then writes + guardrail text rewrites back across ``responses_so_far`` instead of + discarding them. """ return responses_so_far diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 96a5ed663fc..8d9f6b88aa5 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -75,6 +75,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ + delivers_ended_stream_text_rewrites = True + def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ Convert chat completions request data to OpenAI-spec structured messages. @@ -450,6 +452,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, stream_transform_sink: StreamTransformSink | None = None, + deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: """ Process output streaming responses by applying guardrails to text content. @@ -464,6 +467,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): accumulated text (``responses_so_far`` is left untouched so it stays a correct raw accumulator across rounds) and the guardrailed text plus requested holdback are reported per choice on the sink. + deliver_ended_stream_rewrites: When True and the buffered stream has + ended, guardrail text rewrites are written back across + ``responses_so_far`` (full rewritten text in each choice's first + content-carrying chunk, the rest blanked) instead of discarded. Returns: The (unmodified) list of responses. @@ -489,6 +496,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): litellm_logging_obj=litellm_logging_obj, user_api_key_dict=user_api_key_dict, request_data=request_data, + deliver_ended_stream_rewrites=deliver_ended_stream_rewrites, ) async def _process_streaming_block_only( @@ -499,10 +507,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): litellm_logging_obj: "LiteLLMLoggingObj | None", user_api_key_dict: "UserAPIKeyAuth | None", request_data: dict | None, + deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: """Block-only streaming path: run the guardrail so an in-flight BLOCK can terminate the stream. Text rewrites are not propagated to the client here - (see ``_process_streaming_transform`` for the incremental_diff path).""" + (see ``_process_streaming_transform`` for the incremental_diff path) unless + ``deliver_ended_stream_rewrites`` opts the ended-stream branch in.""" # check if the stream has ended has_stream_ended = False for chunk in responses_so_far: @@ -511,20 +521,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation): break if has_stream_ended: - # convert to model response - model_response: Final = cast( - ModelResponse, - stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj), - ) - # run process_output_response - await self.process_output_response( - response=model_response, + await self._process_ended_stream( + responses_so_far=responses_so_far, guardrail_to_apply=guardrail_to_apply, litellm_logging_obj=litellm_logging_obj, user_api_key_dict=user_api_key_dict, request_data=request_data, + deliver_ended_stream_rewrites=deliver_ended_stream_rewrites, ) - return responses_so_far # Step 0: Check if any response has text content to process @@ -597,6 +601,39 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return responses_so_far + async def _process_ended_stream( + self, + *, + responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: "LiteLLMLoggingObj | None", + user_api_key_dict: "UserAPIKeyAuth | None", + request_data: dict[str, object] | None, # mutable-ok: same request-payload shape the hooks take + deliver_ended_stream_rewrites: bool, + ) -> None: + """Ended-stream path: rebuild the full response, run the non-streaming + output guardrail against it, and (when opted in) write any text rewrite + back across the buffered chunks.""" + model_response: Final = cast( + ModelResponse, + stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj), + ) + pre_guardrail_texts: Final = self._string_choice_contents(model_response) + await self.process_output_response( + response=model_response, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) + if deliver_ended_stream_rewrites: + await self._write_ended_stream_text_rewrites( + responses_so_far=responses_so_far, + guardrailed_response=model_response, + pre_guardrail_texts=pre_guardrail_texts, + guardrail_name=guardrail_to_apply.guardrail_name or "unknown", + ) + def build_stream_error_items( self, exc: "HTTPException", @@ -722,8 +759,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): """ combined_texts: Final[dict[tuple[int, int | None], str]] = {} - for response_idx, response in enumerate(responses_so_far): - for choice_idx, choice in enumerate(response.choices): + for response in responses_so_far: + for choice in response.choices: if isinstance(choice, litellm.StreamingChoices): content = choice.delta.content elif isinstance(choice, litellm.Choices): @@ -736,7 +773,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if isinstance(content, str): # String content - accumulate for this choice - str_key: tuple[int, int | None] = (choice_idx, None) + str_key: tuple[int, int | None] = (choice.index, None) if str_key not in combined_texts: combined_texts[str_key] = "" combined_texts[str_key] += content @@ -747,7 +784,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): text_str = content_item.get("text") if text_str: list_key: tuple[int, int | None] = ( - choice_idx, + choice.index, content_idx, ) if list_key not in combined_texts: @@ -937,6 +974,51 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if "name" in func_dict: existing_tool_call.function.name = func_dict["name"] + @staticmethod + def _string_choice_contents(response: "ModelResponse") -> tuple[str | None, ...]: + return tuple( + choice.message.content if isinstance(choice.message.content, str) else None for choice in response.choices + ) + + async def _write_ended_stream_text_rewrites( + self, + responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place + guardrailed_response: "ModelResponse", + pre_guardrail_texts: tuple[str | None, ...], + guardrail_name: str, + ) -> None: + """Write ended-stream guardrail text rewrites back across the buffered + chunks: the full rewritten text lands in the choice's first + content-carrying chunk and the rest are blanked, the same shape the + in-flight write-back uses. Chunks carrying only finish_reason or usage + stay untouched. A rewrite on a stream carrying more than one distinct + choice index fails closed.""" + post_guardrail_texts: Final = self._string_choice_contents(guardrailed_response) + changed: Final = tuple( + after + for before, after in zip(pre_guardrail_texts, post_guardrail_texts) + if before is not None and after is not None and after != before + ) + if not changed: + return + stream_choice_indices: Final = frozenset( + choice.index for response in responses_so_far for choice in response.choices + ) + if len(stream_choice_indices) != 1: + # stream_chunk_builder collapses every choice into one index-0 + # choice, so a rewrite of the rebuilt response cannot be attributed + # back to a single choice on an n>1 stream: withhold the stream + # rather than deliver the rewrite on the wrong choice + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_name) + target_choice_index: Final = next(iter(stream_choice_indices)) + await self._apply_guardrail_responses_to_output_streaming( + responses=responses_so_far, + guardrailed_texts=list(changed), # mutable-ok: callee takes lists + task_mappings=[(target_choice_index, None) for _ in changed], # mutable-ok: callee takes lists + ) + async def _apply_guardrail_responses_to_output_streaming( self, responses: list["ModelResponseStream"], @@ -952,7 +1034,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Args: responses: List of ModelResponseStream objects to modify guardrailed_texts: List of guardrailed text responses (combined from all chunks) - task_mappings: List of tuples (choice_idx, content_idx) + task_mappings: List of tuples (choice_idx, content_idx), where choice_idx + is the choice's ``index`` field, not its position in a chunk's list Override this method to customize how responses are applied to streaming responses. """ @@ -968,9 +1051,11 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Key: (choice_idx, content_idx), Value: boolean (True if already set) already_set: Final[dict[tuple[int, int | None], bool]] = {} - # Iterate through all responses and update content - for response_idx, response in enumerate(responses): - for choice_idx_in_response, choice in enumerate(response.choices): + # Iterate through all responses and update content, matching each chunk's + # choice by its index field: on n>1 streams a chunk usually carries one + # choice at list position 0 whose index names the logical choice. + for response in responses: + for choice in response.choices: if isinstance(choice, litellm.StreamingChoices): content = choice.delta.content elif isinstance(choice, litellm.Choices): @@ -983,7 +1068,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if isinstance(content, str): # String content - str_key: tuple[int, int | None] = (choice_idx_in_response, None) + str_key: tuple[int, int | None] = (choice.index, None) if str_key in guardrail_map: if str_key not in already_set: # First chunk - set the complete guardrailed text @@ -1004,7 +1089,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): for content_idx, content_item in enumerate(content): if "text" in content_item: list_key: tuple[int, int | None] = ( - choice_idx_in_response, + choice.index, content_idx, ) if list_key in guardrail_map: diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 1530c154e93..724a0a1d2f0 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -32,6 +32,7 @@ import time import uuid from collections.abc import Mapping, Sequence from dataclasses import dataclass +from itertools import chain, repeat from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast @@ -110,6 +111,15 @@ class ResponsesStreamChunk(TypedDict, total=False): content_index: ReadOnly[int] +_TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset( + { + ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, + ResponsesAPIStreamEvents.RESPONSE_FAILED.value, + ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, + } +) + + def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int: sequence_numbers: Final = ( item.get("sequence_number") if isinstance(item, dict) else getattr(item, "sequence_number", None) @@ -129,6 +139,8 @@ class OpenAIResponsesHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ + delivers_ended_stream_text_rewrites = True + def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ Convert Responses API request data to OpenAI-spec structured messages. @@ -520,6 +532,7 @@ class OpenAIResponsesHandler(BaseTranslation): litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, + deliver_ended_stream_rewrites: bool = False, ) -> list[Any]: """ Process output streaming response by applying guardrails to text content. @@ -528,10 +541,17 @@ class OpenAIResponsesHandler(BaseTranslation): chunk, apply the guardrail, then write the result back in-place so the caller sees the modified content (e.g. PII tokens replaced). - For ``response.completed`` events (the normal end-of-stream signal) we - use the same per-item extraction + task-mapping approach as - ``process_output_response`` so that unmasking / blocking works correctly - for every output item. + For terminal envelope events (``response.completed``, and equally + ``response.incomplete`` / ``response.failed``, whose envelopes carry the + partial output) we use the same per-item extraction + task-mapping + approach as ``process_output_response`` so that unmasking / blocking + works correctly for every output item. With + ``deliver_ended_stream_rewrites`` the earlier text-carrying events + (``response.output_text.delta`` / ``.done``, + ``response.content_part.done``, ``response.output_item.done``) are synced + to the rewritten envelope too, so a client reading deltas sees the + rewrite instead of the raw model output; a rewrite observed where no + write-back is possible fails closed instead of releasing raw output. """ if not responses_so_far: return responses_so_far @@ -543,14 +563,16 @@ class OpenAIResponsesHandler(BaseTranslation): return responses_so_far # ------------------------------------------------------------------ # - # Case 1: response.completed — full response is available in the # - # final chunk; iterate output items, apply guardrail, write back. # + # Case 1: terminal envelope events (completed/incomplete/failed). # + # the accumulated response is available in the final chunk; iterate # + # output items, apply guardrail, write back. Falls through to the # + # string fallback when the envelope yields nothing to check. # # ------------------------------------------------------------------ # - if final_chunk.get("type") == "response.completed": + if final_chunk.get("type") in _TERMINAL_ENVELOPE_EVENT_TYPES: response_obj: Final[ResponseOutputEnvelope] = final_chunk.get("response") or {} - if not hasattr(response_obj, "get"): - return responses_so_far - outputs: Final[Sequence[object]] = response_obj.get("output") or [] + outputs: Final[Sequence[object]] = ( + (response_obj.get("output") or []) if hasattr(response_obj, "get") else [] + ) texts_to_check: Final[list[str]] = [] tool_calls_to_check: Final[list[ChatCompletionToolCallChunk]] = [] @@ -600,11 +622,25 @@ class OpenAIResponsesHandler(BaseTranslation): responses=guardrailed_texts, task_mappings=task_mappings, ) - - return responses_so_far + if deliver_ended_stream_rewrites: + rewrites_by_position: Final = MappingProxyType( + { + task_mappings[task_idx]: rewritten + for task_idx, rewritten in enumerate(guardrailed_texts) + if task_idx < len(texts_to_check) and rewritten != texts_to_check[task_idx] + } + ) + if rewrites_by_position: + self._sync_stream_events_with_rewrites( + stream_events=responses_so_far[:-1], + rewrites_by_position=rewrites_by_position, + ) + return responses_so_far # ------------------------------------------------------------------ # - # Case 2: response.output_item.done — extract tool calls only. # + # Case 2: response.output_item.done — extract tool calls only, then # + # fall through to the text fallback when a caller expects rewrites # + # delivered, so a buffer truncated here still fails closed on text. # # ------------------------------------------------------------------ # if final_chunk.get("type") == "response.output_item.done": model_response_stream: Final = ( @@ -622,12 +658,14 @@ class OpenAIResponsesHandler(BaseTranslation): input_type="response", logging_obj=litellm_logging_obj, ) - return responses_so_far + if not deliver_ended_stream_rewrites: + return responses_so_far # ------------------------------------------------------------------ # # Fallback: apply guardrail to the accumulated text string. # # No structured write-back is possible here; guardrails that only # - # need to block/flag (not rewrite) still work correctly. # + # need to block/flag (not rewrite) still work correctly, and a # + # rewrite a caller expects delivered fails closed instead. # # ------------------------------------------------------------------ # string_so_far: Final = self.get_streaming_string_so_far(responses_so_far) if string_so_far: @@ -637,26 +675,83 @@ class OpenAIResponsesHandler(BaseTranslation): ) if response_model: fallback_inputs["model"] = response_model - await guardrail_to_apply.apply_guardrail( + fallback_outputs: Final = await guardrail_to_apply.apply_guardrail( inputs=fallback_inputs, request_data=request_data if request_data is not None else {}, input_type="response", logging_obj=litellm_logging_obj, ) + fallback_texts: Final = fallback_outputs.get("texts") + if deliver_ended_stream_rewrites and fallback_texts and tuple(fallback_texts) != (string_so_far,): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown") return responses_so_far + @staticmethod + def _write_event_field(event: object, field: str, value: str) -> None: + if isinstance(event, dict): + event[field] = value # rebind-ok: delivering the rewrite means editing the buffered event in place + else: + setattr(event, field, value) + + def _sync_stream_events_with_rewrites( + self, + stream_events: Sequence[Any], + rewrites_by_position: Mapping[tuple[int, int], str], + ) -> None: + """Sync pre-completion stream events with the rewritten completed + response, keyed by ``(output_index, content_index)``: the first + ``output_text.delta`` for a rewritten item carries the full rewritten + text and the rest are blanked, while ``output_text.done``, + ``content_part.done``, and ``output_item.done`` events carry the full + rewritten text, so every event a client may read agrees with the + rewritten ``response.completed`` payload.""" + delta_replacements: Final = MappingProxyType( + {position: chain((rewritten,), repeat("")) for position, rewritten in rewrites_by_position.items()} + ) + for event in stream_events: + if not (isinstance(event, dict) or hasattr(event, "get")): + continue + event_type = event.get("type") + output_index = event.get("output_index") + content_index = event.get("content_index") + if event_type == "response.output_item.done" and isinstance(output_index, int): + self._sync_output_item_done_event(event.get("item"), output_index, rewrites_by_position) + continue + if not isinstance(output_index, int) or not isinstance(content_index, int): + continue + position = (output_index, content_index) + if event_type == "response.output_text.delta" and position in delta_replacements: + self._write_event_field(event, "delta", next(delta_replacements[position])) + elif event_type == "response.output_text.done" and position in rewrites_by_position: + self._write_event_field(event, "text", rewrites_by_position[position]) + elif event_type == "response.content_part.done" and position in rewrites_by_position: + part = event.get("part") + if isinstance(part, dict) or hasattr(part, "text"): + self._write_event_field(part, "text", rewrites_by_position[position]) + + @staticmethod + def _sync_output_item_done_event( + item: object, + output_index: int, + rewrites_by_position: Mapping[tuple[int, int], str], + ) -> None: + content: Final = item.get("content") if isinstance(item, dict) else getattr(item, "content", None) + if not isinstance(content, list): + return + for (item_idx, content_idx), rewritten in rewrites_by_position.items(): + if item_idx != output_index or content_idx >= len(content): + continue + OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten) + def _check_streaming_has_ended(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> bool: """ Check if the streaming has ended. """ if not responses_so_far: return False - terminal_types: Final = { - ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, - ResponsesAPIStreamEvents.RESPONSE_FAILED.value, - ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, - } - return responses_so_far[-1].get("type") in terminal_types + return responses_so_far[-1].get("type") in _TERMINAL_ENVELOPE_EVENT_TYPES def build_stream_error_items( self, diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 069fbd46185..160938b723c 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 @@ -71,14 +73,17 @@ def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | Non class _StreamRewriteObserver(CustomGuardrail): """Stand-in handed to the endpoint translation in place of a streaming pipeline step's - guardrail. Translations cannot rewrite every buffered chunk consistently, so the gate - withholds the stream whenever the guardrail returned different output than it was given, - which for guardrails like Bedrock's ANONYMIZED action is only known at runtime.""" + guardrail. It records whether the guardrail returned different output than it was given, + which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text + rewrites are deliverable on translations that write them back across the buffered chunks + (``delivers_ended_stream_text_rewrites``); tool-call rewrites and text rewrites on any + other translation make the gate withhold the stream.""" def __init__(self, inner: CustomGuardrail) -> None: super().__init__(guardrail_name=inner.guardrail_name) self.inner: Final = inner - self.rewrote = False + self.rewrote_texts = False + self.rewrote_tool_calls = False def structured_messages_cover_full_request(self) -> bool: return self.inner.structured_messages_cover_full_request() @@ -95,10 +100,9 @@ class _StreamRewriteObserver(CustomGuardrail): outputs: Final = await self.inner.apply_guardrail( inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj ) - self.rewrote = ( - self.rewrote - or _rewrote(sent_texts, _text_snapshot(outputs.get("texts"))) - or _rewrote(sent_tool_shapes, _tool_call_shapes(outputs.get("tool_calls"))) + self.rewrote_texts = self.rewrote_texts or _rewrote(sent_texts, _text_snapshot(outputs.get("texts"))) + self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote( + sent_tool_shapes, _tool_call_shapes(outputs.get("tool_calls")) ) return outputs @@ -248,6 +252,41 @@ class PipelineExecutor: # Ran out of steps without a terminal action → default allow return _allow_result(step_results=step_results, working_data=working_data, request_data=data) + @staticmethod + async def _run_streaming_step( + step: PipelineStep, + callback: CustomGuardrail, + endpoint_translation: "BaseTranslation", + streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks the translation rewrites in place + hook_input: dict[str, object], # mutable-ok: same request-payload shape as data + user_api_key_dict: "UserAPIKeyAuth | None", + litellm_logging_obj: "LiteLLMLoggingObj | None", + ) -> None: + """Run one streaming post_call step through the endpoint translation, delivering + text rewrites on translations that support ended-stream write-back and raising + ``UndeliverableStreamRewrite`` for any rewrite that cannot reach the client.""" + observer: Final = _StreamRewriteObserver(callback) + deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites + if deliver_rewrites: + await endpoint_translation.process_output_streaming_response( + responses_so_far=streaming_chunks, + guardrail_to_apply=observer, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=hook_input, + deliver_ended_stream_rewrites=True, + ) + else: + await endpoint_translation.process_output_streaming_response( + responses_so_far=streaming_chunks, + guardrail_to_apply=observer, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=hook_input, + ) + if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites): + raise UndeliverableStreamRewrite(step.guardrail) + @staticmethod async def _run_step( step: PipelineStep, @@ -310,16 +349,15 @@ class PipelineExecutor: f"Guardrail '{step.guardrail}' does not support streaming pipeline execution", None, ) - observer: Final = _StreamRewriteObserver(callback) - await endpoint_translation.process_output_streaming_response( - responses_so_far=streaming_chunks, - guardrail_to_apply=observer, - litellm_logging_obj=data.get("litellm_logging_obj"), + await PipelineExecutor._run_streaming_step( + step=step, + callback=callback, + endpoint_translation=endpoint_translation, + streaming_chunks=streaming_chunks, + hook_input=hook_input, user_api_key_dict=user_api_key_dict, - request_data=hook_input, + litellm_logging_obj=data.get("litellm_logging_obj"), ) - if observer.rewrote: - raise UndeliverableStreamRewrite(step.guardrail) response = None elif mode == "post_call": response = await target.async_post_call_success_hook( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 3314ccefa65..63dc5a03804 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -518,14 +518,6 @@ def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool: return callback is not None and PipelineExecutor.supports_unified_execution(callback) -def _pipeline_step_rewrites_streamed_content(guardrail_name: str) -> bool: - callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) - if callback is None: - return False - transform_mode: Final = unified_guardrail.resolve_streaming_flag(callback, "streaming_transform_mode", "block_only") - return callback.rewrites_streamed_output() or transform_mode == "incremental_diff" - - class _PipelineErrorBody(TypedDict): message: ReadOnly[str] type: ReadOnly[str] @@ -542,9 +534,10 @@ def _undeliverable_stream_rewrite_error(policy_name: str, guardrail_name: str) - "error": { "message": ( f"Streaming response withheld by policy pipeline '{policy_name}' because guardrail " - f"'{guardrail_name}' rewrote the streamed output, and streaming pipelines cannot deliver " - "rewrites. Retry with stream=false, or drop it from the pipeline steps so guardrails.add " - "applies it to streamed output." + f"'{guardrail_name}' rewrote the streamed output in a way this endpoint's streaming " + "pipeline cannot deliver (a tool-call rewrite, or a text rewrite on a route without " + "stream write-back). Retry with stream=false, or drop it from the pipeline steps so " + "guardrails.add applies it to streamed output." ), "type": "guardrail_pipeline_error", "policies": (policy_name,), @@ -561,13 +554,13 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_ap Background responses skip the post_call hooks entirely, so a pipeline governing one would silently never execute. Streaming responses execute pipelines against the buffered stream through the endpoint guardrail - translation of the request route, releasing the buffered chunks on allow. - That needs every step's guardrail to support the unified apply_guardrail - interface and to only allow or block (a step that rewrites streamed - content, via mask_response_content, a MASK action, or - streaming_transform_mode=incremental_diff, would have its rewrite silently - dropped), and needs the route to have a translation at all; anything else - keeps the 400 rather than letting ungoverned output stream through. + translation of the request route, releasing the buffered chunks on allow + (rewritten in place when a guardrail rewrote text and the translation + delivers ended-stream rewrites; a rewrite the translation cannot deliver + fails closed at runtime instead). That needs every step's guardrail to + support the unified apply_guardrail interface, and needs the route to have + a translation at all; anything else keeps the 400 rather than letting + ungoverned output stream through. """ is_stream: Final = data.get("stream") is True is_background: Final = data.get("background") is True @@ -612,25 +605,6 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_ap } } raise HTTPException(status_code=400, detail=unsupported_detail) - rewriting_guardrails: Final = tuple( - guardrail for guardrail in step_guardrails if _pipeline_step_rewrites_streamed_content(guardrail) - ) - if rewriting_guardrails: - rewriting_detail: Final[_PipelineErrorDetail] = { - "error": { - "message": ( - "Policies with post_call guardrail pipelines cannot govern streaming responses " - "because these pipeline guardrails rewrite streamed content (mask_response_content, " - "a MASK action, or streaming_transform_mode=incremental_diff), which pipeline steps would release " - f"unmodified: {', '.join(rewriting_guardrails)}. Retry with stream=false, or drop " - "them from the pipeline steps so guardrails.add applies them to streamed output." - ), - "type": "guardrail_pipeline_error", - "policies": post_call_policies, - "guardrails": rewriting_guardrails, - } - } - raise HTTPException(status_code=400, detail=rewriting_detail) route: Final = user_api_key_dict.request_route if not route or resolve_endpoint_translation(user_api_key_dict, None) is not None: return @@ -3498,12 +3472,13 @@ class ProxyLogging: pipeline allows it), then runs each pipeline's steps against the assembled output through the endpoint guardrail translation, the same machinery flat post_call guardrails use at end of stream. An allow - releases the buffered chunks verbatim; a step whose guardrail rewrote - the output withholds the stream with a 400 instead, since no - translation rewrites every buffered chunk consistently and some - rewrites (Bedrock's ANONYMIZED action, for one) are only decided at - runtime; a block or modify_response terminates with the translation's - block chunks or the raised error. + releases the buffered chunks: verbatim when no guardrail rewrote the + output, rewritten in place when one rewrote text and the translation + delivers ended-stream rewrites (later steps then re-scan the rewritten + chunks, so rewrites chain). A rewrite the translation cannot deliver + (a tool-call rewrite, or a text rewrite on a route without write-back) + withholds the stream with a 400; a block or modify_response terminates + with the translation's block chunks or the raised error. """ buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict async for item in response: 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 0fe7730e91e..274c351ebc7 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,117 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: # Should return the responses unchanged assert result == responses_so_far + @staticmethod + def _ended_sse_chunks() -> list: + events = [ + ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello "}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "world"}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ("message_stop", {"type": "message_stop"}), + ] + return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] + + @staticmethod + def _masking_guardrail() -> CustomGuardrail: + class MaskWorld(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, "texts": [text.replace("world", "[MASKED]") for text in inputs.get("texts", [])]} + + return MaskWorld(guardrail_name="test") + + @staticmethod + def _delta_texts(chunks: list) -> list: + texts = [] + for chunk in chunks: + for line in chunk.decode().split("\n"): + if not line.startswith("data:"): + continue + data = json.loads(line[len("data:") :].strip()) + if data.get("type") == "content_block_delta": + texts.append(data["delta"]["text"]) + return texts + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_text_back_into_sse_chunks(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_sse_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert self._delta_texts(chunks) == ["hello [MASKED]", ""] + raw = b"".join(chunks).decode() + assert "event: message_start" in raw and "event: message_stop" in raw + assert '"stop_reason": "end_turn"' in raw + + @pytest.mark.asyncio + async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_sse_chunks() + original = [bytes(chunk) for chunk in chunks] + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=MagicMock(), + ) + + assert chunks == original + + @pytest.mark.asyncio + async def test_unended_stream_rewrite_with_delivery_expected_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = AnthropicMessagesHandler() + chunks = self._ended_sse_chunks()[:-2] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_unended_stream_without_rewrite_is_released_with_delivery_expected(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_sse_chunks()[:-2] + original = [bytes(chunk) for chunk in chunks] + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockPassThroughGuardrail(guardrail_name="test"), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert chunks == original + + @pytest.mark.asyncio + async def test_unended_stream_rewrite_without_delivery_expected_does_not_raise(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_sse_chunks()[:-2] + original = [bytes(chunk) for chunk in chunks] + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=MagicMock(), + ) + + assert result is chunks + assert chunks == original + class TestAnthropicMessagesHandlerInputProcessing: """Test input processing preserves litellm_metadata for dynamic guardrails.""" 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 7dd6065063a..af22390f80c 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,154 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: # Should return the responses assert result == responses_so_far + @staticmethod + def _ended_stream_chunks() -> list: + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + return [ + ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[StreamingChoices(index=0, delta=Delta(content="Hello"), finish_reason=None)], + ), + ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[StreamingChoices(index=0, delta=Delta(content=" world"), finish_reason="stop")], + ), + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_text_back_into_chunks(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="test") + chunks = self._ended_stream_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert chunks[0].choices[0].delta.content == "HELLO WORLD" + assert chunks[1].choices[0].delta.content in (None, "") + assert chunks[1].choices[0].finish_reason == "stop" + + @pytest.mark.asyncio + async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="test") + chunks = self._ended_stream_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + assert chunks[0].choices[0].delta.content == "Hello" + assert chunks[1].choices[0].delta.content == " world" + assert chunks[1].choices[0].finish_reason == "stop" + + @staticmethod + def _two_choice_stream_chunks() -> list: + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + def chunk(index: int, content: str, finish_reason: Optional[str] = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)], + ) + + return [ + chunk(0, "safe "), + chunk(1, "hello "), + chunk(0, "text", "stop"), + chunk(1, "world", "stop"), + ] + + @staticmethod + def _world_masking_guardrail() -> CustomGuardrail: + class MaskWorld(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + texts = inputs.get("texts", []) + return {**inputs, "texts": [t.replace("world", "[MASKED]") for t in texts]} + + return MaskWorld(guardrail_name="test-mask") + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrite_on_multi_choice_stream_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIChatCompletionsHandler() + chunks = self._two_choice_stream_chunks() + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._world_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self): + handler = OpenAIChatCompletionsHandler() + chunks = self._two_choice_stream_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockPassThroughGuardrail(guardrail_name="test"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert [c.choices[0].delta.content for c in chunks] == ["safe ", "hello ", "text", "world"] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrite_lands_on_nonzero_choice_index(self): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + handler = OpenAIChatCompletionsHandler() + + def chunk(content: str, finish_reason: Optional[str]) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[StreamingChoices(index=1, delta=Delta(content=content), finish_reason=finish_reason)], + ) + + chunks = [chunk("hello ", None), chunk("world", "stop")] + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._world_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert chunks[0].choices[0].delta.content == "hello [MASKED]" + assert chunks[1].choices[0].delta.content in (None, "") + class TestGetStructuredMessages: """Test the get_structured_messages method.""" 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 315b6948bd8..bd07e924d12 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 @@ -1122,6 +1122,209 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: output_text = result[-1]["response"]["output"][0]["content"][0]["text"] assert output_text == original_text + @staticmethod + def _ended_stream_events() -> List[dict]: + content = [{"type": "output_text", "text": "hello world"}] + item = { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": content, + } + return [ + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "}, + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "world"}, + {"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"}, + { + "type": "response.content_part.done", + "output_index": 0, + "content_index": 0, + "part": {"type": "output_text", "text": "hello world"}, + }, + {"type": "response.output_item.done", "output_index": 0, "item": {**item, "content": [dict(c) for c in content]}}, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "model": "gpt-4o", + "output": [{**item, "content": [dict(c) for c in content]}], + "status": "completed", + }, + }, + ] + + @staticmethod + def _masking_guardrail() -> CustomGuardrail: + class MaskWorld(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + texts = inputs.get("texts", []) + return {**inputs, "texts": [t.replace("world", "[MASKED]") for t in texts]} + + return MaskWorld(guardrail_name="test-mask") + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_all_stream_events(self): + handler = OpenAIResponsesHandler() + events = self._ended_stream_events() + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert events[0]["delta"] == "hello [MASKED]" + assert events[1]["delta"] == "" + assert events[2]["text"] == "hello [MASKED]" + assert events[3]["part"]["text"] == "hello [MASKED]" + assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]" + assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]" + + @pytest.mark.asyncio + @pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"]) + async def test_deliver_ended_stream_rewrites_syncs_non_completed_terminals(self, terminal_type): + handler = OpenAIResponsesHandler() + events = self._ended_stream_events() + events[-1]["type"] = terminal_type + events[-1]["response"]["status"] = terminal_type.split(".")[-1] + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert events[0]["delta"] == "hello [MASKED]" + assert events[1]["delta"] == "" + assert events[2]["text"] == "hello [MASKED]" + assert events[3]["part"]["text"] == "hello [MASKED]" + assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]" + assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]" + + @pytest.mark.asyncio + async def test_fallback_rewrite_with_delivery_expected_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = [ + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "}, + {"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"}, + ] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_fallback_delta_only_rewrite_with_delivery_expected_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = [ + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "}, + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "world"}, + ] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_output_item_done_last_rewrite_with_delivery_expected_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = self._ended_stream_events()[:-1] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_output_item_done_last_scans_text_with_delivery_expected(self): + handler = OpenAIResponsesHandler() + events = self._ended_stream_events()[:-1] + guardrail = MockRecordingGuardrail(guardrail_name="test") + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert [inputs.get("texts") for inputs in guardrail.seen_inputs] == [["hello world"]] + + @pytest.mark.asyncio + async def test_output_item_done_last_without_delivery_expected_skips_text(self): + handler = OpenAIResponsesHandler() + events = self._ended_stream_events()[:-1] + guardrail = MockRecordingGuardrail(guardrail_name="test") + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + assert result is events + assert guardrail.seen_inputs == [] + + @pytest.mark.asyncio + async def test_fallback_rewrite_without_delivery_expected_does_not_raise(self): + handler = OpenAIResponsesHandler() + events = [ + {"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"}, + ] + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + ) + + assert result is events + + @pytest.mark.asyncio + async def test_ended_stream_rewrite_leaves_delta_events_untouched_by_default(self): + handler = OpenAIResponsesHandler() + events = self._ended_stream_events() + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + ) + + assert events[0]["delta"] == "hello " + assert events[1]["delta"] == "world" + assert events[2]["text"] == "hello world" + assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]" + @pytest.mark.asyncio async def test_failed_stream_scans_delta_text(self): """A stream ending in response.failed has text only in delta events; the 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 1152e67e730..d399b4f01cb 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -923,6 +923,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 d9b3578c966..d56936747aa 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 @@ -1479,7 +1479,7 @@ async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_lacks_uni ("guardrail_config", {"streaming_transform_mode": "incremental_diff"}), ], ) -async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_rewrites_streamed_content( +async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_rewrites_streamed_content( proxy_logging, make_user_api_key_auth, monkeypatch, rewrite_attribute, value ): seen: Dict[str, Any] = {} @@ -1488,24 +1488,21 @@ async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_rewrites_ monkeypatch.setattr(litellm, "callbacks", [guardrail]) data = _post_call_pipeline_data(stream=True) - with pytest.raises(HTTPException) as info: - await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), - data=data, - call_type="completion", - guardrails_only=True, - ) + out = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + data=data, + call_type="completion", + guardrails_only=True, + ) - assert info.value.status_code == 400 - assert info.value.detail["error"]["guardrails"] == ("gr-post",) - assert "rewrite streamed content" in info.value.detail["error"]["message"] - assert seen.get("count") is None + assert out is not None + assert out.get("stream") is True @pytest.mark.asyncio -@pytest.mark.parametrize("action, rejected", [(ContentFilterAction.MASK, True), (ContentFilterAction.BLOCK, False)]) -async def test_pre_call_hook_rejects_streaming_only_when_content_filter_step_masks( - proxy_logging, make_user_api_key_auth, monkeypatch, action, rejected +@pytest.mark.parametrize("action", [ContentFilterAction.MASK, ContentFilterAction.BLOCK]) +async def test_pre_call_hook_allows_streaming_when_content_filter_step_masks_or_blocks( + proxy_logging, make_user_api_key_auth, monkeypatch, action ): guardrail = ContentFilterGuardrail( guardrail_name="gr-post", @@ -1516,25 +1513,15 @@ async def test_pre_call_hook_rejects_streaming_only_when_content_filter_step_mas data = _post_call_pipeline_data(stream=True) user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions") - if not rejected: - out = await proxy_logging.pre_call_hook( - user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True - ) - assert out is not None and out.get("stream") is True - return + out = await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True + ) - with pytest.raises(HTTPException) as info: - await proxy_logging.pre_call_hook( - user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True - ) - - assert info.value.status_code == 400 - assert info.value.detail["error"]["guardrails"] == ("gr-post",) - assert "a MASK action" in info.value.detail["error"]["message"] + assert out is not None and out.get("stream") is True @pytest.mark.asyncio -async def test_pre_call_hook_rejects_streaming_when_content_filter_category_masks( +async def test_pre_call_hook_allows_streaming_when_content_filter_category_masks( proxy_logging, make_user_api_key_auth, monkeypatch ): guardrail = ContentFilterGuardrail( @@ -1546,13 +1533,11 @@ async def test_pre_call_hook_rejects_streaming_when_content_filter_category_mask data = _post_call_pipeline_data(stream=True) user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions") - with pytest.raises(HTTPException) as info: - await proxy_logging.pre_call_hook( - user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True - ) + out = await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True + ) - assert info.value.status_code == 400 - assert info.value.detail["error"]["guardrails"] == ("gr-post",) + assert out is not None and out.get("stream") is True @pytest.mark.asyncio @@ -1656,17 +1641,10 @@ def _echoed_tool_call_dicts(arguments: str) -> List[Dict[str, Any]]: @pytest.mark.asyncio @pytest.mark.parametrize("on_fail, on_error", [("block", None), ("next", "next")]) -@pytest.mark.parametrize( - "make_chunks, transform", - [ - (_stream_chunks, lambda inputs: {"texts": ["hello [MASKED]"]}), - (_tool_call_stream_chunks, lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')}), - ], - ids=["texts", "tool_calls"], -) -async def test_streaming_iterator_hook_pipeline_withholds_runtime_rewrite( - proxy_logging, make_user_api_key_auth, monkeypatch, make_chunks, transform, on_fail, on_error +async def test_streaming_iterator_hook_pipeline_withholds_runtime_tool_call_rewrite( + proxy_logging, make_user_api_key_auth, monkeypatch, on_fail, on_error ): + transform = lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')} # noqa: E731 monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) step = PipelineStep(guardrail="gr-post", on_pass="allow", on_fail=on_fail, on_error=on_error) @@ -1676,7 +1654,7 @@ async def test_streaming_iterator_hook_pipeline_withholds_runtime_rewrite( async def _drain() -> None: async for item in proxy_logging.async_post_call_streaming_iterator_hook( user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), - response=_async_chunk_iter(make_chunks()), + response=_async_chunk_iter(_tool_call_stream_chunks()), request_data=data, ): delivered.append(item) @@ -1693,6 +1671,81 @@ async def test_streaming_iterator_hook_pipeline_withholds_runtime_rewrite( assert "stream=false" in error["message"] +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_delivers_runtime_text_rewrite( + proxy_logging, make_user_api_key_auth, monkeypatch +): + transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731 + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + assert delivered[0].choices[0].delta.content == "hello [MASKED]" + assert delivered[1].choices[0].delta.content in (None, "") + assert delivered[1].choices[0].finish_reason == "stop" + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_chains_text_rewrites_across_steps( + proxy_logging, make_user_api_key_auth, monkeypatch +): + second_step_saw: Dict[str, Any] = {} + + class FirstMask(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, "texts": [text.replace("world", "[MASKED]") for text in inputs["texts"]]} + + class SecondMask(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + second_step_saw["texts"] = list(inputs["texts"]) + return {**inputs, "texts": [text.replace("hello", "[GREETING]") for text in inputs["texts"]]} + + monkeypatch.setattr( + litellm, + "callbacks", + [ + FirstMask(guardrail_name="gr-first", event_hook=GuardrailEventHooks.post_call, default_on=False), + SecondMask(guardrail_name="gr-second", event_hook=GuardrailEventHooks.post_call, default_on=False), + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + pipeline = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="gr-first", on_pass="next", on_fail="block"), + PipelineStep(guardrail="gr-second", on_pass="allow", on_fail="block"), + ], + ) + data = _post_call_pipeline_data(stream=True) + data["metadata"]["_guardrail_pipelines"] = [("response-governance", pipeline)] + chunks = _stream_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + assert second_step_saw["texts"] == ["hello [MASKED]"] + assert delivered[0].choices[0].delta.content == "[GREETING] [MASKED]" + assert delivered[1].choices[0].delta.content in (None, "") + assert delivered[1].choices[0].finish_reason == "stop" + + @pytest.mark.asyncio @pytest.mark.parametrize( "make_chunks, transform", @@ -1799,6 +1852,79 @@ async def test_streaming_iterator_hook_pipeline_modify_response_emits_translated assert not any(item is chunk for item in delivered for chunk in chunks) +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthropic_sse( + proxy_logging, make_user_api_key_auth, monkeypatch +): + transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731 + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = _anthropic_sse_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + raw = b"".join(delivered).decode() + assert "hello [MASKED]" in raw + assert "hello world" not in raw + assert raw.count("event: content_block_delta") == 1 + for expected_event in ("message_start", "content_block_start", "content_block_stop", "message_delta", "message_stop"): + assert f"event: {expected_event}" in raw + + +@pytest.mark.asyncio +async def test_pipeline_executor_withholds_text_rewrite_when_translation_lacks_write_back(monkeypatch): + from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation + from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite + + class NoWriteBackTranslation(BaseTranslation): + async def process_input_messages(self, data, guardrail_to_apply, litellm_logging_obj): + return data + + async def process_output_response(self, response, guardrail_to_apply, litellm_logging_obj, **kwargs): + return response + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj, + user_api_key_dict=None, + request_data=None, + stream_transform_sink=None, + deliver_ended_stream_rewrites=False, + ): + assert deliver_ended_stream_rewrites is False + await guardrail_to_apply.apply_guardrail( + inputs={"texts": ["hello world"]}, + request_data=request_data or {}, + input_type="response", + ) + return responses_so_far + + transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731 + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + + with pytest.raises(UndeliverableStreamRewrite): + await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="gr-post", on_pass="allow", on_fail="block")], + mode="post_call", + data={"metadata": {}}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + call_type="acompletion", + policy_name="response-governance", + streaming_chunks=_stream_chunks(), + endpoint_translation=NoWriteBackTranslation(), + ) + + @pytest.mark.asyncio async def test_streaming_iterator_hook_pipeline_gates_without_iterator_overrides( proxy_logging, make_user_api_key_auth, monkeypatch