fix(hosted_vllm): keep another choice's final token_ids and stop_reason on their own stream chunk

With n above 1, a finish-only vLLM chunk for a choice other than the one the single finish chunk
carries was dropped or taken as choice 0's finish, losing that choice's last token_ids. It now goes
out as an id-only chunk on its own choice index.

litellm/litellm_core_utils/streaming_handler.py

@@ -166,6 +166,15 @@ def _apply_choice_fields(choice: StreamingChoices, fields: Mapping[str, object])
         setattr(choice, field, value)
 
 
+def _choice_fields_chunk(
+    model_response: ModelResponseStream, upstream_index: int, fields: Mapping[str, object]
+) -> ModelResponseStream:
+    model_response.choices[0].delta = Delta(content="")  # rebind-ok: chunks are built in place here
+    model_response.choices[0].index = upstream_index  # rebind-ok: chunks are built in place here
+    _apply_choice_fields(model_response.choices[0], fields)
+    return model_response
+
+
 def _holds_only_scalars(fields: Iterable[tuple[str, object]]) -> bool:
     return all(value is None or isinstance(value, (str, int, float, bool)) for _, value in fields)
 
@@ -1291,6 +1300,11 @@ class CustomStreamWrapper:
             else:
                 return
         elif self.received_finish_reason is not None:
+            another_choice_fields_chunk: Final = self._another_choice_fields_chunk(
+                model_response, response_obj, chunk_choice_fields
+            )
+            if another_choice_fields_chunk is not None:
+                return another_choice_fields_chunk
             if self.sent_last_chunk is True:
                 # Bedrock returns the guardrail trace in the last chunk - we want to return this here
                 if self.custom_llm_provider == "bedrock" and "trace" in model_response:
@@ -1341,11 +1355,9 @@ class CustomStreamWrapper:
         elif self._has_special_delta_content(model_response):
             return self._handle_special_delta_content(model_response)
         elif chunk_choice_fields:
-            upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk"))
-            model_response.choices[0].delta = Delta(content="")  # rebind-ok: chunks are built in place here
-            model_response.choices[0].index = upstream_index  # rebind-ok: chunks are built in place here
-            _apply_choice_fields(model_response.choices[0], chunk_choice_fields)
-            return model_response
+            return _choice_fields_chunk(
+                model_response, _upstream_choice_index(response_obj.get("original_chunk")), chunk_choice_fields
+            )
         else:
             if hasattr(model_response, "usage"):
                 self.chunks.append(model_response)
@@ -2010,6 +2022,27 @@ class CustomStreamWrapper:
     def _apply_finish_choice_fields(self, choice: StreamingChoices) -> None:
         _apply_choice_fields(choice, self.received_finish_choice_fields.get(choice.index, _NO_UPSTREAM_CHOICE_FIELDS))
 
+    def _another_choice_fields_chunk(
+        self,
+        model_response: ModelResponseStream,
+        response_obj: Mapping[str, object],
+        chunk_choice_fields: Mapping[str, object],
+    ) -> ModelResponseStream | None:
+        upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk"))
+        is_finished: Final = bool(response_obj.get("is_finished"))
+        if is_finished and upstream_index == model_response.choices[0].index and not self.sent_last_chunk:
+            return None
+        finish_fields: Final = (
+            self.received_finish_choice_fields.get(upstream_index, _NO_UPSTREAM_CHOICE_FIELDS)
+            if is_finished
+            else _NO_UPSTREAM_CHOICE_FIELDS
+        )
+        if not chunk_choice_fields and all(value is None for value in finish_fields.values()):
+            return None
+        return _choice_fields_chunk(
+            model_response, upstream_index, MappingProxyType({**finish_fields, **chunk_choice_fields})
+        )
+
     def finish_reason_handler(self):
         model_response: Final = self.model_response_creator(
             hidden_params=None if self.received_model_id is None else {"received_model_id": self.received_model_id}

tests/unit/litellm_core_utils/test_streaming_handler.py

@@ -6341,7 +6341,7 @@ def test_hosted_vllm_stream_with_two_choices_keeps_each_stop_reason_on_its_final
     content_chunk_stop_reasons: Final = frozenset(
         (choice["index"], choice["stop_reason"])
         for choice in emitted
-        if choice["finish_reason"] is None and choice.get("stop_reason") is not None
+        if _DELTA.validate_python(choice["delta"]).get("content") and choice.get("stop_reason") is not None
     )
     finished_with_content: Final = frozenset(
         (index, _TWO_CHOICE_STOP_REASONS[index])
@@ -6351,6 +6351,27 @@ def test_hosted_vllm_stream_with_two_choices_keeps_each_stop_reason_on_its_final
 
     assert content_chunk_stop_reasons == finished_with_content
 
+
+@pytest.mark.parametrize("first_finished", [0, 1])
+@pytest.mark.parametrize(
+    "finish_contents",
+    [("", ""), ("2", ""), ("", "2"), ("2", "2")],
+    ids=["both-finish-only", "content-then-finish-only", "finish-only-then-content", "both-content-then-finish"],
+)
+def test_hosted_vllm_stream_with_two_choices_delivers_every_token_id_and_stop_reason_of_each_choice(
+    logging_obj: Logging, first_finished: int, finish_contents: tuple[str, str]
+) -> None:
+    emitted: Final = _hosted_vllm_stream_choices(logging_obj, _two_choice_stream(first_finished, finish_contents))
+
+    assert {
+        index: _joined_token_ids([choice for choice in emitted if choice["index"] == index])
+        for index in _TWO_CHOICE_TOKEN_IDS
+    } == _TWO_CHOICE_TOKEN_IDS
+    assert frozenset(
+        (choice["index"], choice["stop_reason"]) for choice in emitted if choice.get("stop_reason") is not None
+    ) == frozenset(_TWO_CHOICE_STOP_REASONS.items())
+
+
 _CACHED_VLLM_TOKEN_IDS: Final = (16, 11, 220, 17, 11, 220, 18, 11, 220, 19)