Merge branch 'litellm_guardrail_choice_isolation' into litellm_guardrail_stream_inspection

# Conflicts:
#	tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py

litellm/integrations/custom_guardrail.py

@@ -256,8 +256,8 @@ class CustomGuardrail(CustomLogger):
         scan_raw_request: bool = False,
         only_scan_new_messages: bool = False,
         timeout: float | None = None,
-        **kwargs,
-    ):
+        **kwargs: object,
+    ) -> None:
         """
         Initialize the CustomGuardrail class
 
@@ -303,7 +303,7 @@ class CustomGuardrail(CustomLogger):
         self.run_in_parallel: bool = run_in_parallel
         self.scan_raw_request: bool = scan_raw_request
         self.only_scan_new_messages: bool = only_scan_new_messages
-        stream_scope_arg: Final[object] = cast(object, kwargs.pop("stream_scope", None))  # cast-ok: config
+        stream_scope_arg: Final[object] = kwargs.pop("stream_scope", None)
         self.apply_stream_scope(stream_scope_arg)
         self.logging_only_scope = None
         self.logging_only_continue_on_input_failure = False

litellm/llms/base_llm/guardrail_translation/base_translation.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
 
@@ -34,6 +34,10 @@ class StreamTransformSink:
     mutated_text_per_choice: dict[int, str] = field(default_factory=dict)
     holdback_per_choice: dict[int, int] = field(default_factory=dict)
 
+    def record(self, texts: Mapping[int, str], holdback: Mapping[int, int]) -> None:
+        self.mutated_text_per_choice = dict(texts)
+        self.holdback_per_choice = dict(holdback)
+
 
 @dataclass(frozen=True, slots=True)
 class StreamingScanKey:

litellm/llms/openai/chat/guardrail_translation/handler.py

@@ -18,9 +18,11 @@ import json
 import time
 import uuid
 from collections.abc import Iterator, Mapping, Sequence
+from itertools import chain
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Union, cast
 
+from pydantic import TypeAdapter
 from typing_extensions import NotRequired, ReadOnly, TypedDict
 
 import litellm
@@ -46,7 +48,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
     stream_item_items,
     unappliable_request_rewrite,
 )
-from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
+from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolParam
 from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
     coerce_stream_holdback_value,
 )
@@ -168,8 +170,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
                 logging_obj=litellm_logging_obj,
             )
 
-            guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
-            guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", [])
+            guardrailed_texts: Final = guardrailed_inputs.get("texts", ())
+            guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", ())
             guardrailed_tools: Final = guardrailed_inputs.get("tools")
             if guardrailed_tools is not None:
                 data["tools"] = (
@@ -258,12 +260,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
     def extract_request_tool_names(self, data: dict) -> list[str]:
         """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name)."""
         names: Final[list[str]] = []
-        for tool in data.get("tools") or []:
+        for tool in data.get("tools") or ():
             if isinstance(tool, dict) and tool.get("type") == "function":
                 fn = tool.get("function")
                 if isinstance(fn, dict) and fn.get("name"):
                     names.append(str(fn["name"]))
-        for fn in data.get("functions") or []:
+        for fn in data.get("functions") or ():
             if isinstance(fn, dict) and fn.get("name"):
                 names.append(str(fn["name"]))
         return names
@@ -387,7 +389,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         guardrail_to_apply: "CustomGuardrail",
         litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
         user_api_key_dict: "UserAPIKeyAuth | None" = None,
-        request_data: dict | None = None,
+        request_data: dict[str, object] | None = None,
     ) -> ModelResponse:
         """
         Process output response by applying guardrails to text content.
@@ -438,9 +440,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
             # create a standalone dict (SDK / direct-call path).
             if request_data is None:
                 request_data = {"response": response}
-            else:
-                if "response" not in request_data:
-                    request_data["response"] = response
+            request_data.setdefault("response", response)
 
             self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
 
@@ -460,7 +460,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
                 logging_obj=litellm_logging_obj,
             )
 
-            guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
+            guardrailed_texts: Final = guardrailed_inputs.get("texts", ())
             returned_tool_calls: Final = guardrailed_inputs.get("tool_calls")
             guardrailed_tool_calls: Final[list[dict[str, object]]] = (
                 cast(list[dict[str, object]], returned_tool_calls)
@@ -494,7 +494,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         guardrail_to_apply: "CustomGuardrail",
         litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
         user_api_key_dict: "UserAPIKeyAuth | None" = None,
-        request_data: dict | None = None,
+        request_data: dict[str, object] | None = None,
         stream_transform_sink: StreamTransformSink | None = None,
         deliver_ended_stream_rewrites: bool = False,
     ) -> list["ModelResponseStream"]:
@@ -550,14 +550,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         guardrail_to_apply: "CustomGuardrail",
         litellm_logging_obj: "LiteLLMLoggingObj | None",
         user_api_key_dict: "UserAPIKeyAuth | None",
-        request_data: dict | None,
+        request_data: dict[str, object] | 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) unless
         ``deliver_ended_stream_rewrites`` opts the ended-stream branch in."""
-        has_stream_ended: Final = self._first_choice_has_finished(responses_so_far)
+        has_stream_ended: Final = deliver_ended_stream_rewrites or self._first_choice_has_finished(responses_so_far)
 
         if has_stream_ended:
             await self._process_ended_stream(
@@ -623,7 +623,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
                 logging_obj=litellm_logging_obj,
             )
 
-            guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
+            guardrailed_texts: Final = guardrailed_inputs.get("texts", ())
 
             # Step 4: Apply guardrailed text back to all streaming chunks
             # For each choice, replace the combined text across all chunks
@@ -656,13 +656,36 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         model_response: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj)
         pre_guardrail_texts: Final = self._string_choice_contents(model_response)
         pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(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,
+        inspection_responses: Final = (
+            tuple(
+                model_response.model_copy(update=MappingProxyType({"choices": [choice]}))
+                for choice in model_response.choices
+            )
+            if pre_guardrail_tool_calls and len(model_response.choices) > 1
+            else (model_response,)
         )
+        inspection_request_data: Final = request_data if request_data is not None else {}
+        try:
+            for inspection_response in inspection_responses:
+                inspection_request_data["response"] = inspection_response
+                inspection_request_data["responses"] = (
+                    [
+                        self._narrowed_to_choice(chunk, inspection_response.choices[0].index)
+                        for chunk in responses_so_far
+                    ]
+                    if len(inspection_response.choices) == 1
+                    else responses_so_far
+                )
+                await self.process_output_response(
+                    response=inspection_response,
+                    guardrail_to_apply=guardrail_to_apply,
+                    litellm_logging_obj=litellm_logging_obj,
+                    user_api_key_dict=user_api_key_dict,
+                    request_data=inspection_request_data,
+                )
+        finally:
+            inspection_request_data["response"] = model_response
+            inspection_request_data["responses"] = responses_so_far
         if not deliver_ended_stream_rewrites:
             return
         await self._write_ended_stream_text_rewrites(
@@ -763,7 +786,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         guardrail_to_apply: "CustomGuardrail",
         litellm_logging_obj: "LiteLLMLoggingObj | None",
         user_api_key_dict: "UserAPIKeyAuth | None",
-        request_data: dict | None,
+        request_data: dict[str, object] | None,
         sink: StreamTransformSink,
     ) -> None:
         """Run the guardrail over the raw accumulated text and report the
@@ -775,8 +798,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         """
         raw_by_index: Final = self._accumulate_string_content_by_choice_index(responses_so_far)
         if not raw_by_index:
-            sink.mutated_text_per_choice = {}
-            sink.holdback_per_choice = {}
+            sink.record(MappingProxyType({}), MappingProxyType({}))
             return
 
         # Fix #2 — sort by StreamingChoices.index so an n>1 stream that emits
@@ -788,18 +810,72 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         indices: Final = sorted(raw_by_index.keys())
         texts_to_check: Final = [raw_by_index[i] for i in indices]
 
-        if request_data is None:
-            request_data = {"responses": responses_so_far}
-        elif "responses" not in request_data:
-            request_data["responses"] = responses_so_far
-        self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
+        resolved_request_data: Final[dict[str, object]] = (
+            request_data if request_data is not None else {"responses": responses_so_far}
+        )
+        if "responses" not in resolved_request_data:
+            resolved_request_data["responses"] = responses_so_far
+        self.merge_user_api_key_metadata_into_request(resolved_request_data, user_api_key_dict)
 
         inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
+        if self._streamed_tool_call_fingerprints(responses_so_far):
+            assembled: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj)
+            resolved_request_data["response"] = assembled
+            if len(assembled.choices) > 1:
+                choice_rounds: Final = tuple(
+                    (
+                        choice,
+                        StreamTransformSink(),
+                        [self._narrowed_to_choice(chunk, choice.index) for chunk in responses_so_far],
+                    )
+                    for choice in assembled.choices
+                )
+                try:
+                    for choice, choice_sink, choice_chunks in choice_rounds:
+                        resolved_request_data["response"] = assembled.model_copy(
+                            update=MappingProxyType({"choices": [choice]})
+                        )
+                        resolved_request_data["responses"] = choice_chunks
+                        await self._process_streaming_transform(
+                            responses_so_far=choice_chunks,
+                            guardrail_to_apply=guardrail_to_apply,
+                            litellm_logging_obj=litellm_logging_obj,
+                            user_api_key_dict=user_api_key_dict,
+                            request_data=resolved_request_data,
+                            sink=choice_sink,
+                        )
+                finally:
+                    resolved_request_data["response"] = assembled
+                    resolved_request_data["responses"] = responses_so_far
+                merged_texts: Final = MappingProxyType(
+                    dict(
+                        chain.from_iterable(
+                            choice_sink.mutated_text_per_choice.items() for _, choice_sink, _ in choice_rounds
+                        )
+                    )
+                )
+                merged_holdback: Final = MappingProxyType(
+                    dict(
+                        chain.from_iterable(
+                            choice_sink.holdback_per_choice.items() for _, choice_sink, _ in choice_rounds
+                        )
+                    )
+                )
+                sink.record(merged_texts, merged_holdback)
+                return
+            tool_calls: Final = chain.from_iterable(choice.message.tool_calls or () for choice in assembled.choices)
+            inputs["tool_calls"] = TypeAdapter(list[ChatCompletionToolCallChunk]).validate_python(
+                tuple(
+                    MappingProxyType({"index": index, **converted})
+                    for index, tool_call in enumerate(tool_calls)
+                    if (converted := self._convert_tool_call_to_dict(tool_call)) is not None
+                )
+            )
         if responses_so_far and getattr(responses_so_far[0], "model", None):
             inputs["model"] = responses_so_far[0].model
         guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
             inputs=inputs,
-            request_data=request_data,
+            request_data=resolved_request_data,
             input_type="response",
             logging_obj=litellm_logging_obj,
         )
@@ -819,13 +895,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
                 len(texts_to_check),
             )
 
-        holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or []
-        sink.mutated_text_per_choice = {
-            idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts)
-        }
-        sink.holdback_per_choice = {
-            indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback)
-        }
+        holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or ()
+        transformed_texts: Final = MappingProxyType(
+            {idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts)}
+        )
+        transformed_holdback: Final = MappingProxyType(
+            {indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback)}
+        )
+        sink.record(transformed_texts, transformed_holdback)
 
     def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
         chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream))
@@ -1145,20 +1222,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
     def _function_tool_call_fragments(
         responses_so_far: Sequence["ModelResponseStream"],
     ) -> tuple[tuple[ChatCompletionDeltaToolCall, ...], ...]:
-        """Group the stream's function tool-call fragments by their tool-call index, in
-        the index order ``stream_chunk_builder`` lists the rebuilt tool calls, keeping
-        only the indices the builder keeps (an id and a name somewhere in the stream)."""
         fragments: Final = tuple(
-            tool_call
+            (choice.index, tool_call)
             for response in responses_so_far
             for choice in response.choices
             for tool_call in choice.delta.tool_calls or ()
             if isinstance(tool_call, ChatCompletionDeltaToolCall)
         )
-        identified: Final = frozenset(fragment.index for fragment in fragments if fragment.id)
-        named: Final = frozenset(fragment.index for fragment in fragments if fragment.function.name)
+        identified: Final = frozenset((choice, fragment.index) for choice, fragment in fragments if fragment.id)
+        named: Final = frozenset((choice, fragment.index) for choice, fragment in fragments if fragment.function.name)
         return tuple(
-            tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named)
+            tuple(fragment for choice, fragment in fragments if (choice, fragment.index) == key)
+            for key in sorted(identified & named)
         )
 
     def _write_ended_stream_tool_call_rewrites(
@@ -1168,28 +1243,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         pre_guardrail_tool_calls: tuple[tuple[str | None, str], ...],
         guardrail_name: str,
     ) -> None:
-        """Write ended-stream guardrail tool-call rewrites back across the buffered
-        chunks: the rewritten name and full arguments land in the tool call's first
-        fragment and the arguments of its later fragments are blanked, mirroring the
-        text write-back. A rewrite on a stream carrying more than one distinct choice
-        index, or whose fragments do not line up with the rebuilt tool calls, is
-        reported as undeliverable, so the pipeline executor discards it and releases
-        the original chunks."""
         post_guardrail_tool_calls: Final = self._function_tool_call_shapes(guardrailed_response)
         if post_guardrail_tool_calls == pre_guardrail_tool_calls:
             return
-        stream_choice_indices: Final = frozenset(
-            choice.index for response in responses_so_far for choice in response.choices
-        )
         fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far)
-        if len(stream_choice_indices) != 1:
-            from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
-
-            raise UndeliverableStreamRewrite(
-                guardrail_name,
-                f"the stream carries {len(stream_choice_indices)} choices and tool-call rewrites are only written "
-                "back on single-choice streams",
-            )
         if len(fragments_by_tool_call) != len(post_guardrail_tool_calls):
             from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
 
@@ -1212,7 +1269,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
     async def _apply_guardrail_responses_to_output_streaming(
         self,
         responses: list["ModelResponseStream"],
-        guardrailed_texts: list[str],
+        guardrailed_texts: Sequence[str],
         task_mappings: list[tuple[int, int | None]],
     ) -> None:
         """
@@ -1230,12 +1287,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         Override this method to customize how responses are applied to streaming responses.
         """
         # Build a mapping of what guardrailed text to use for each (choice_idx, content_idx)
-        guardrail_map: Final[dict[tuple[int, int | None], str]] = {}
-        for task_idx, guardrail_response in enumerate(guardrailed_texts):
-            mapping = task_mappings[task_idx]
-            choice_idx = cast(int, mapping[0])
-            content_idx_optional = cast(int | None, mapping[1])
-            guardrail_map[(choice_idx, content_idx_optional)] = guardrail_response
+        guardrail_map: Final = MappingProxyType(
+            {task_mappings[index]: text for index, text in enumerate(guardrailed_texts)}
+        )
 
         # Track which choices we've already set the guardrailed text for
         # Key: (choice_idx, content_idx), Value: boolean (True if already set)
@@ -1245,34 +1299,20 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         # 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):
-                    content = choice.message.content
-                else:
-                    continue
-
+            for choice, message in (
+                (choice, choice.delta if isinstance(choice, StreamingChoices) else choice.message)
+                for choice in response.choices
+                if isinstance(choice, (StreamingChoices, Choices))
+            ):
+                content = message.content
                 if content is None:
                     continue
 
                 if isinstance(content, str):
-                    # String content
                     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
-                            if isinstance(choice, litellm.StreamingChoices):
-                                choice.delta.content = guardrail_map[str_key]
-                            elif isinstance(choice, litellm.Choices):
-                                choice.message.content = guardrail_map[str_key]
-                            already_set[str_key] = True
-                        else:
-                            # Subsequent chunks - clear the content
-                            if isinstance(choice, litellm.StreamingChoices):
-                                choice.delta.content = ""
-                            elif isinstance(choice, litellm.Choices):
-                                choice.message.content = ""
+                        message.content = guardrail_map[str_key] if str_key not in already_set else ""
+                        already_set[str_key] = True
 
                 elif isinstance(content, list):
                     # List content - handle each content item

tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py

@@ -7,7 +7,7 @@ with guardrail transformations, including tool calls.
 
 import json
 from collections.abc import Mapping
-from typing import Any, Literal, Optional
+from typing import Any, Final, Literal, Optional
 
 import pytest
 
@@ -1395,29 +1395,65 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
         )
 
         assert [
-            (tool_call["id"], tool_call["function"]["arguments"])
-            for tool_call in guardrail.seen_inputs[-1]["tool_calls"]
-        ] == [("call_1", '{"fruit": "persimmon"}'), ("call_2", '{"fruit": "durian"}')]
+            [(tool_call["id"], tool_call["function"]["arguments"]) for tool_call in inputs["tool_calls"]]
+            for inputs in guardrail.seen_inputs
+        ] == [[("call_1", '{"fruit": "persimmon"}')], [("call_2", '{"fruit": "durian"}')]]
 
     @pytest.mark.asyncio
-    async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self):
-        from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
+    @pytest.mark.parametrize("transform", [False, True])
+    async def test_each_choice_refreshes_shared_guardrail_response_context(self, transform: bool) -> None:
+        from fastapi import HTTPException
+        from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink
+        from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
 
-        handler = OpenAIChatCompletionsHandler()
-        chunks = self._two_choice_tool_call_stream_chunks()
+        class ContextGuardrail(CustomGuardrail):
+            async def apply_guardrail(
+                self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object],
+                input_type: Literal["request", "response"], logging_obj: object = None,
+            ) -> GenericGuardrailAPIInputs:
+                response: Final = request_data["response"]
+                assert isinstance(response, ModelResponse)
+                request_data["visited_choices"] = (*request_data.get("visited_choices", ()), response.choices[0].index)
+                if response.choices[0].message.content == "forbidden":
+                    raise HTTPException(status_code=400, detail="later choice blocked")
+                return inputs
 
-        with pytest.raises(UndeliverableStreamRewrite, match="the stream carries 2 choices") as raised:
-            await handler.process_output_streaming_response(
-                responses_so_far=chunks,
-                guardrail_to_apply=MockGuardrail(guardrail_name="test"),
-                litellm_logging_obj=None,
-                deliver_ended_stream_rewrites=True,
+        chunks: Final = [ModelResponseStream(choices=[StreamingChoices(
+            index=index, delta=Delta(content=text, tool_calls=[{
+                "index": 0, "id": f"call_{index}", "type": "function",
+                "function": {"name": "contact", "arguments": "{}"},
+            }]), finish_reason="tool_calls",
+        ) for index, text in enumerate(("allowed", "forbidden"))])]
+        request_data: Final = {"metadata": {"trace": "retained"}}
+        with pytest.raises(HTTPException, match="later choice blocked"):
+            await OpenAIChatCompletionsHandler().process_output_streaming_response(
+                responses_so_far=chunks, guardrail_to_apply=ContextGuardrail(guardrail_name="context"),
+                request_data=request_data, deliver_ended_stream_rewrites=True,
+                stream_transform_sink=StreamTransformSink() if transform else None,
             )
+        assert request_data["visited_choices"] == (0, 1)
+        assert request_data["metadata"]["trace"] == "retained"
+        assert request_data["responses"] is chunks
+        restored: Final = request_data["response"]
+        assert isinstance(restored, ModelResponse)
+        assert tuple(choice.index for choice in restored.choices) == (0, 1)
 
-        assert raised.value.guardrail_name == "test"
-        assert raised.value.reason == (
-            "the stream carries 2 choices and tool-call rewrites are only written back on single-choice streams"
+    @pytest.mark.asyncio
+    async def test_deliver_ended_stream_tool_rewrites_keep_choice_indices(self) -> None:
+        handler: Final = OpenAIChatCompletionsHandler()
+        chunks: Final = self._two_choice_tool_call_stream_chunks()
+        await handler.process_output_streaming_response(
+            responses_so_far=chunks,
+            guardrail_to_apply=MockGuardrail(guardrail_name="test"),
+            litellm_logging_obj=None,
+            deliver_ended_stream_rewrites=True,
         )
+        arguments: Final = tuple(
+            "".join(call.function.arguments for chunk in chunks for choice in chunk.choices
+                    if choice.index == index for call in choice.delta.tool_calls or ())
+            for index in range(2)
+        )
+        assert arguments == ('{"fruit": "PERSIMMON"}', '{"fruit": "DURIAN"}')
 
     @pytest.mark.asyncio
     async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self):