fix(streaming): count every choice a chunk finishes when an n>1 stream drops before [DONE]

litellm/litellm_core_utils/streaming_handler.py

@@ -109,6 +109,27 @@ def _first_choice_index(chunk: ChatCompletionChunk | ModelResponseStream | Compl
             return 0
 
 
+def _choice_finish_reasons(
+    chunk: ChatCompletionChunk | ModelResponseStream | Completion | None,
+) -> tuple[tuple[int, str | None], ...]:
+    match chunk:
+        case ChatCompletionChunk(choices=choices):
+            return tuple((choice.index, choice.finish_reason) for choice in choices)
+        case ModelResponseStream(choices=choices):
+            return tuple((choice.index, choice.finish_reason) for choice in choices)
+        case Completion(choices=choices):
+            return tuple((choice.index, choice.finish_reason) for choice in choices)
+        case _:
+            return ()
+
+
+def _indices_of_finished_choices(
+    chunk: ChatCompletionChunk | ModelResponseStream | Completion | None,
+) -> frozenset[int]:
+    finished: Final = frozenset(index for index, finish_reason in _choice_finish_reasons(chunk) if finish_reason)
+    return finished or frozenset({_first_choice_index(chunk)})
+
+
 def _stamp_usage_cost(
     usage_chunk: ModelResponse | TextCompletionResponse | ModelResponseStream,
     complete_response: ModelResponse | TextCompletionResponse | ModelResponseStream,
@@ -1526,7 +1547,7 @@ class CustomStreamWrapper:
             print_verbose(f"completion obj content: {completion_obj['content']}")
             if response_obj["is_finished"]:
                 self.received_finish_reason = response_obj["finish_reason"]
-                self._record_finished_choice(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None)
+                self._record_finished_choices(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None)
             if response_obj["usage"] is not None:
                 _text_completion_usage: Final[Usage] = response_obj["usage"]
                 setattr(
@@ -1566,7 +1587,7 @@ class CustomStreamWrapper:
             print_verbose(f"completion obj content: {completion_obj['content']}")
             if response_obj["is_finished"]:
                 self.received_finish_reason = response_obj["finish_reason"]
-                self._record_finished_choice(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None)
+                self._record_finished_choices(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None)
         elif self.custom_llm_provider == "cached_response":
             cached_chunk: Final = cast(ModelResponseStream, chunk)
             cached_choice: Final = cached_chunk.choices[0] if cached_chunk.choices else None
@@ -1617,7 +1638,7 @@ class CustomStreamWrapper:
                         f"{self.custom_llm_provider} raised a streaming error - finish_reason: error, no content string given. Received Chunk={response_obj}"
                     )
                 self.received_finish_reason = response_obj["finish_reason"]
-                self._record_finished_choice(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None)
+                self._record_finished_choices(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None)
             if response_obj.get("original_chunk", None) is not None:
                 if hasattr(response_obj["original_chunk"], "id"):
                     model_response = self.set_model_id(response_obj["original_chunk"].id, model_response)
@@ -2276,8 +2297,10 @@ class CustomStreamWrapper:
                 raise
             raise StopIteration from e
 
-    def _record_finished_choice(self, chunk: ChatCompletionChunk | ModelResponseStream | Completion | None) -> None:
-        self._finished_choice_indices = (self._finished_choice_indices or frozenset()) | {_first_choice_index(chunk)}
+    def _record_finished_choices(self, chunk: ChatCompletionChunk | ModelResponseStream | Completion | None) -> None:
+        self._finished_choice_indices = (self._finished_choice_indices or frozenset()) | _indices_of_finished_choices(
+            chunk
+        )
 
     def _every_choice_finished(self) -> bool:
         if self.received_finish_reason is None:

tests/unit/litellm_core_utils/test_streaming_handler.py

@@ -5284,6 +5284,67 @@ async def test_two_choice_text_completion_stream_dropped_before_done_ends_cleanl
     assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "HiHel"
 
 
+def _chat_choices_finishing_in_one_chunk() -> tuple[ModelResponseStream, ...]:
+    return (
+        _reset_test_chunk(content="Hi", index=0),
+        _reset_test_chunk(content="Hel", index=1),
+        ModelResponseStream(
+            id="chatcmpl-reset-test",
+            created=1783458104,
+            model="stub-model",
+            object="chat.completion.chunk",
+            choices=[
+                StreamingChoices(index=0, delta=Delta(content=None), finish_reason="stop"),
+                StreamingChoices(index=1, delta=Delta(content=None), finish_reason="stop"),
+            ],
+        ),
+    )
+
+
+def _text_choices_finishing_in_one_chunk() -> tuple[Completion, ...]:
+    return (
+        _text_completion_chunk(0, "Hi"),
+        _text_completion_chunk(1, "Hel"),
+        Completion.model_construct(
+            id="cmpl-two-choices",
+            object="text_completion",
+            created=1,
+            model="gpt-3.5-turbo-instruct",
+            choices=[
+                CompletionChoice.model_construct(index=0, text="", finish_reason="stop", logprobs=None),
+                CompletionChoice.model_construct(index=1, text="", finish_reason="stop", logprobs=None),
+            ],
+        ),
+    )
+
+
+@pytest.mark.asyncio
+@_SYNC_AND_ASYNC
+@pytest.mark.parametrize(
+    ("provider", "chunks"),
+    (("openai", _chat_choices_finishing_in_one_chunk), ("text-completion-openai", _text_choices_finishing_in_one_chunk)),
+    ids=("chat", "text_completion"),
+)
+async def test_one_chunk_finishing_both_requested_choices_then_dropped_before_done_ends_cleanly(
+    logging_obj: Logging,
+    sync_mode: bool,
+    provider: str,
+    chunks: Callable[[], Sequence[ModelResponseStream | Completion]],
+):
+    logging_obj.update_environment_variables(litellm_params={}, optional_params={"n": 2})
+    stream: Final = _chunks_then_failure(chunks(), _httpx2_read_error())
+    wrapper: Final = CustomStreamWrapper(
+        completion_stream=stream if sync_mode else _as_async(stream),
+        model="stub-model",
+        custom_llm_provider=provider,
+        logging_obj=logging_obj,
+    )
+
+    collected: Final = await _collect(wrapper, sync_mode)
+
+    assert "".join(chunk.choices[0].delta.content or "" for chunk in collected if chunk.choices) == "HiHel"
+
+
 @pytest.mark.asyncio
 @_SYNC_AND_ASYNC
 @_SDK3_DROPPED_CONNECTIONS