fix(policy_engine): keep the per-choice rebuilt response's choices a list so legacy hook rewrites survive the model_dump round-trip

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

@@ -696,11 +696,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
                     ModelResponse,
                     stream_chunk_builder(
                         chunks=[  # mutable-ok: callee takes a list
-                            response.model_copy(
-                                update=MappingProxyType(
-                                    {"choices": tuple(choice for choice in response.choices if choice.index == index)}
-                                )
-                            )
+                            OpenAIChatCompletionsHandler._narrowed_to_choice(response, index)
                             for response in responses_so_far
                         ],
                         logging_obj=litellm_logging_obj,
@@ -710,16 +706,16 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
             for index in choice_indices
         )
         (_, base_response), *_ = rebuilt_by_index
-        return base_response.model_copy(
-            update=MappingProxyType(
-                {
-                    "choices": tuple(
-                        rebuilt.choices[0].model_copy(update=MappingProxyType({"index": index}))
-                        for index, rebuilt in rebuilt_by_index
-                    )
-                }
-            )
-        )
+        stitched_choices: Final = [  # mutable-ok: choices is a List field; a tuple there breaks model_dump round-trips
+            rebuilt.choices[0].model_copy(update=MappingProxyType({"index": index}))
+            for index, rebuilt in rebuilt_by_index
+        ]
+        return base_response.model_copy(update=MappingProxyType({"choices": stitched_choices}))
+
+    @staticmethod
+    def _narrowed_to_choice(response: "ModelResponseStream", index: int) -> "ModelResponseStream":
+        narrowed: Final = [choice for choice in response.choices if choice.index == index]  # mutable-ok: List field
+        return response.model_copy(update=MappingProxyType({"choices": narrowed}))
 
     def build_stream_error_items(
         self,

tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py

@@ -369,6 +369,7 @@ async def test_openai_moderation_guardrail_streaming_safe_content():
             chunk1.choices[0].delta = MagicMock()
             chunk1.choices[0].delta.content = "Hello "
             chunk1.choices[0].finish_reason = None
+            chunk1.choices[0].index = 0
 
             chunk2 = MagicMock()
             chunk2.model = "gpt-4"
@@ -376,6 +377,7 @@ async def test_openai_moderation_guardrail_streaming_safe_content():
             chunk2.choices[0].delta = MagicMock()
             chunk2.choices[0].delta.content = "world"
             chunk2.choices[0].finish_reason = None
+            chunk2.choices[0].index = 0
 
             # Last chunk with finish_reason
             chunk3 = MagicMock()
@@ -384,6 +386,7 @@ async def test_openai_moderation_guardrail_streaming_safe_content():
             chunk3.choices[0].delta = MagicMock()
             chunk3.choices[0].delta.content = "!"
             chunk3.choices[0].finish_reason = "stop"
+            chunk3.choices[0].index = 0
 
             for chunk in [chunk1, chunk2, chunk3]:
                 yield chunk
@@ -480,6 +483,7 @@ async def test_openai_moderation_guardrail_streaming_harmful_content():
             chunk1.choices[0].delta = MagicMock()
             chunk1.choices[0].delta.content = "This is "
             chunk1.choices[0].finish_reason = None
+            chunk1.choices[0].index = 0
 
             # Last chunk - with finish_reason to signal end of stream
             chunk2 = MagicMock()
@@ -488,6 +492,7 @@ async def test_openai_moderation_guardrail_streaming_harmful_content():
             chunk2.choices[0].delta = MagicMock()
             chunk2.choices[0].delta.content = "harmful content"
             chunk2.choices[0].finish_reason = "stop"
+            chunk2.choices[0].index = 0
 
             for chunk in [chunk1, chunk2]:
                 yield chunk

tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py

@@ -42,6 +42,7 @@ async def test_openai_moderation_guardrail_streaming_latency():
                 choice.delta.content = content
                 # Last chunk gets finish_reason
                 choice.finish_reason = "stop" if i == len(chunks_data) - 1 else None
+                choice.index = 0
                 chunk.choices = [choice]
                 yield chunk
 
@@ -122,6 +123,7 @@ async def test_openai_moderation_guardrail_streaming_harmful_content():
                 choice.delta.content = content
                 # Last chunk gets finish_reason
                 choice.finish_reason = "stop" if i == len(chunks_data) - 1 else None
+                choice.index = 0
                 chunk.choices = [choice]
                 yield chunk
 
@@ -224,6 +226,7 @@ async def test_openai_moderation_streaming_end_of_stream_request_data_passthroug
                 choice.delta = MagicMock()
                 choice.delta.content = content
                 choice.finish_reason = "stop" if i == len(chunks_data) - 1 else None
+                choice.index = 0
                 chunk.choices = [choice]
                 yield chunk
 

tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py

@@ -1836,6 +1836,22 @@ def _rewritten_model_response(response: Any) -> litellm.ModelResponse:
     return litellm.ModelResponse(**payload)
 
 
+def _two_choice_stream_chunks() -> List[Any]:
+    return [
+        litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "hello "}, "finish_reason": None}]),
+        litellm.ModelResponseStream(choices=[{"index": 1, "delta": {"content": "bonjour "}, "finish_reason": None}]),
+        litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "world"}, "finish_reason": "stop"}]),
+        litellm.ModelResponseStream(choices=[{"index": 1, "delta": {"content": "monde"}, "finish_reason": "stop"}]),
+    ]
+
+
+def _rewritten_every_choice(response: Any) -> litellm.ModelResponse:
+    payload = response.model_dump()
+    for choice in payload["choices"]:
+        choice["message"]["content"] = "[REWRITTEN] " + choice["message"]["content"]
+    return litellm.ModelResponse(**payload)
+
+
 def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator_only(
     make_user_api_key_auth, monkeypatch, caplog
 ):
@@ -1984,6 +2000,39 @@ async def test_streaming_iterator_hook_runs_legacy_hook_and_delivers_its_rewrite
     assert _warnings(caplog) == []
 
 
+@pytest.mark.asyncio
+async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_every_choice(
+    proxy_logging, make_user_api_key_auth, monkeypatch, caplog
+):
+    seen: Dict[str, Any] = {}
+    guardrail = _legacy_hook_stream_guardrail(seen, rewrite=_rewritten_every_choice)
+    monkeypatch.setattr(litellm, "callbacks", [guardrail])
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    data = _post_call_pipeline_data(stream=True)
+    chunks = _two_choice_stream_chunks()
+    auth = make_user_api_key_auth(request_route="/v1/chat/completions")
+
+    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
+        await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="completion", guardrails_only=True)
+        delivered = [
+            item
+            async for item in proxy_logging.async_post_call_streaming_iterator_hook(
+                user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data
+            )
+        ]
+
+    assert [choice.message.content for choice in seen["response"].choices] == ["hello world", "bonjour monde"]
+    assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks]
+    assert [(item.choices[0].index, item.choices[0].delta.content) for item in delivered] == [
+        (0, "[REWRITTEN] hello world"),
+        (1, "[REWRITTEN] bonjour monde"),
+        (0, ""),
+        (1, ""),
+    ]
+    assert [item.choices[0].finish_reason for item in delivered] == [None, None, "stop", "stop"]
+    assert _warnings(caplog) == []
+
+
 @pytest.mark.asyncio
 async def test_streaming_iterator_hook_releases_stream_untouched_when_legacy_hook_returns_none(
     proxy_logging, make_user_api_key_auth, monkeypatch