From c9435b5ff3f72c6d9c5ebfa30175498bc5928daf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 18:00:42 -0700 Subject: [PATCH] fix(guardrails): key stream rewrites by choice index and scan delta-only responses buffers Chat streaming write-backs now match chunks by the choice's index field instead of its list position, delivering rewrites to the right choice on n>1 streams; an ended-stream rewrite on a multi-choice buffer fails closed since stream_chunk_builder collapses the choices. The Responses fallback joins output_text.delta events when delivery is expected, so a delta-only buffer is guardrail-checked instead of released raw. --- .../chat/guardrail_translation/handler.py | 50 ++++++---- .../guardrail_translation/handler.py | 31 +++++-- .../test_openai_guardrail_handler.py | 93 +++++++++++++++++++ ...test_openai_responses_guardrail_handler.py | 71 ++++++++++++++ 4 files changed, 222 insertions(+), 23 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 02206a36b0e..182dec81937 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -620,6 +620,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): 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( @@ -747,8 +748,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): @@ -761,7 +762,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 @@ -772,7 +773,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: @@ -973,24 +974,38 @@ class OpenAIChatCompletionsHandler(BaseTranslation): responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place guardrailed_response: "ModelResponse", pre_guardrail_texts: tuple[str | None, ...], + guardrail_name: str, ) -> None: """Write ended-stream guardrail text rewrites back across the buffered - chunks: each rewritten choice's full text lands in its first + 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.""" + 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( - (choice_idx, after) - for choice_idx, (before, after) in enumerate(zip(pre_guardrail_texts, post_guardrail_texts)) + 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=[after for _choice_idx, after in changed], # mutable-ok: callee takes lists - task_mappings=[(choice_idx, None) for choice_idx, _after in changed], # mutable-ok: callee takes lists + 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( @@ -1008,7 +1023,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. """ @@ -1024,9 +1040,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): @@ -1039,7 +1057,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 @@ -1060,7 +1078,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 3c902f1b829..ebd1070a8e8 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -634,14 +634,20 @@ class OpenAIResponsesHandler(BaseTranslation): return responses_so_far # ------------------------------------------------------------------ # - # Fallback: apply guardrail to the accumulated text string. # - # No structured write-back is possible here; guardrails that only # - # need to block/flag (not rewrite) still work correctly, and a # - # rewrite a caller expects delivered fails closed instead. # + # Fallback: apply guardrail to the accumulated text string. When a # + # caller expects rewrites delivered and only output_text.delta events # + # carried the text (a stream cut off before any .done or terminal # + # envelope), the delta text is scanned instead so nothing escapes # + # unchecked. No structured write-back is possible here; guardrails # + # that only need to block/flag (not rewrite) still work correctly, # + # and a rewrite a caller expects delivered fails closed instead. # # ------------------------------------------------------------------ # string_so_far: Final = self.get_streaming_string_so_far(responses_so_far) - if string_so_far: - fallback_inputs: Final = GenericGuardrailAPIInputs(texts=[string_so_far]) + text_to_check: Final = string_so_far or ( + self._delta_text_so_far(responses_so_far) if deliver_ended_stream_rewrites else "" + ) + if text_to_check: + fallback_inputs: Final = GenericGuardrailAPIInputs(texts=[text_to_check]) response_model = ( final_chunk.get("response", {}).get("model") if isinstance(final_chunk.get("response"), dict) else None ) @@ -654,7 +660,7 @@ class OpenAIResponsesHandler(BaseTranslation): 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,): + if deliver_ended_stream_rewrites and fallback_texts and tuple(fallback_texts) != (text_to_check,): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown") @@ -754,6 +760,17 @@ class OpenAIResponsesHandler(BaseTranslation): """ return "".join([response.get("text", "") for response in responses_so_far]) + @staticmethod + def _delta_text_so_far(responses_so_far: Sequence[ResponsesStreamChunk]) -> str: + """Accumulate the text carried by ``response.output_text.delta`` events, + for buffers where no ``.done`` event or terminal envelope repeats it.""" + deltas: Final = ( + response.get("delta") + for response in responses_so_far + if response.get("type") == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA.value + ) + return "".join(delta for delta in deltas if isinstance(delta, str)) + def _has_text_content(self, response: "ResponsesAPIResponse") -> bool: """ Check if response has any text content to process. 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 26442a4a6ed..b177d4a73d4 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 @@ -1128,6 +1128,99 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: 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 dfd6352f9ec..ae3700b1123 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 @@ -1212,6 +1212,77 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: deliver_ended_stream_rewrites=True, ) + @staticmethod + def _recording_guardrail() -> "tuple[CustomGuardrail, List[List[str]]]": + seen: List[List[str]] = [] + + class Recorder(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + seen.append(list(inputs.get("texts", []))) + return inputs + + return Recorder(guardrail_name="recorder"), seen + + @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_fallback_scans_delta_text_when_delivery_expected(self): + handler = OpenAIResponsesHandler() + guardrail, seen = self._recording_guardrail() + 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": "there"}, + ] + + 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 seen == [["hello there"]] + + @pytest.mark.asyncio + async def test_fallback_ignores_delta_text_without_delivery_expected(self): + handler = OpenAIResponsesHandler() + guardrail, seen = self._recording_guardrail() + events = [ + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello world"}, + ] + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + assert result is events + assert seen == [] + @pytest.mark.asyncio async def test_fallback_rewrite_without_delivery_expected_does_not_raise(self): handler = OpenAIResponsesHandler()