fix(hosted_vllm): keep the upstream choice index on token-id-only stream chunks

litellm/litellm_core_utils/streaming_handler.py

@@ -117,13 +117,25 @@ def _prepend_upstream_choice_extras(
 _NO_UPSTREAM_CHOICE_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
 
 
-def _upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]:
+def _upstream_first_choice(upstream_chunk: object) -> BaseModel | None:
     match getattr(upstream_chunk, "choices", None):
         case [BaseModel() as first_choice, *_]:
-            upstream_extras: Final[Mapping[str, object]] = first_choice.model_extra or {}
-            return MappingProxyType({field: upstream_extras[field] for field in fields if field in upstream_extras})
+            return first_choice
         case _:
-            return _NO_UPSTREAM_CHOICE_FIELDS
+            return None
+
+
+def _upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]:
+    first_choice: Final = _upstream_first_choice(upstream_chunk)
+    if first_choice is None:
+        return _NO_UPSTREAM_CHOICE_FIELDS
+    upstream_extras: Final[Mapping[str, object]] = first_choice.model_extra or {}
+    return MappingProxyType({field: upstream_extras[field] for field in fields if field in upstream_extras})
+
+
+def _upstream_choice_index(upstream_chunk: object) -> int:
+    upstream_index: Final[object] = getattr(_upstream_first_choice(upstream_chunk), "index", 0)
+    return upstream_index if isinstance(upstream_index, int) else 0
 
 
 def _non_null_upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]:
@@ -1308,7 +1320,9 @@ class CustomStreamWrapper:
         elif upstream_chunk_fields := _non_null_upstream_choice_fields(
             response_obj.get("original_chunk"), STREAM_CHUNK_CHOICE_FIELDS
         ):
+            upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk"))
             model_response.choices[0].delta = Delta(content=None)  # 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], upstream_chunk_fields)
             return model_response
         else:

tests/unit/litellm_core_utils/test_streaming_handler.py

@@ -6094,6 +6094,25 @@ def test_hosted_vllm_stream_keeps_stop_reason_and_every_token_id(
     assert list(chain.from_iterable(choice.get("token_ids") or () for choice in emitted)) == [16, 11, 17, 20]
 
 
+def test_hosted_vllm_stream_keeps_held_back_token_ids_on_their_own_choice(logging_obj: Logging) -> None:
+    emitted: Final = _hosted_vllm_stream_choices(
+        logging_obj,
+        (
+            {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None},
+            {"index": 1, "delta": {"role": "assistant", "content": ""}, "finish_reason": None},
+            {"index": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]},
+            {"index": 0, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [17]},
+            {"index": 1, "delta": {"content": "2"}, "finish_reason": None, "token_ids": [18]},
+        ),
+    )
+
+    assert [(choice["index"], choice.get("token_ids")) for choice in emitted if choice.get("token_ids")] == [
+        (1, [16]),
+        (0, [17]),
+        (1, [18]),
+    ]
+
+
 def _held_back_stop_string_stream(token_ids: Mapping[str, None]) -> tuple[Mapping[str, object], ...]:
     return (
         {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None},

tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py

@@ -1,4 +1,5 @@
 import json
+from typing import Final
 from unittest.mock import MagicMock
 
 import httpx
@@ -401,7 +402,7 @@ def test_hosted_vllm_custom_tools_use_top_level_input_schema():
 
 
 def _transform_vllm_chat_response(choice_extras: dict[str, object]) -> ModelResponse:
-    payload = {
+    payload: Final = {
         "id": "chatcmpl-1",
         "object": "chat.completion",
         "created": 1,
@@ -417,7 +418,7 @@ def _transform_vllm_chat_response(choice_extras: dict[str, object]) -> ModelResp
         ],
         "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8},
     }
-    raw_response = httpx.Response(
+    raw_response: Final = httpx.Response(
         200, json=payload, request=httpx.Request("POST", "http://vllm.local/v1/chat/completions")
     )
     return HostedVLLMChatConfig().transform_response(
@@ -442,17 +443,17 @@ def _transform_vllm_chat_response(choice_extras: dict[str, object]) -> ModelResp
     ],
 )
 def test_hosted_vllm_keeps_vllm_choice_fields_on_the_choice(choice_extras: dict[str, object]):
-    response = _transform_vllm_chat_response(choice_extras)
+    response: Final = _transform_vllm_chat_response(choice_extras)
 
-    dumped_choice: dict[str, dict[str, object]] = json.loads(response.model_dump_json())["choices"][0]
+    dumped_choice: Final[dict[str, dict[str, object]]] = json.loads(response.model_dump_json())["choices"][0]
     for field, value in choice_extras.items():
         assert dumped_choice[field] == value
         assert dumped_choice["provider_specific_fields"][field] == value
 
 
 def test_hosted_vllm_does_not_invent_vllm_choice_fields():
-    response = _transform_vllm_chat_response({})
+    response: Final = _transform_vllm_chat_response({})
 
-    dumped_choice: dict[str, object] = json.loads(response.model_dump_json())["choices"][0]
+    dumped_choice: Final[dict[str, object]] = json.loads(response.model_dump_json())["choices"][0]
     assert "stop_reason" not in dumped_choice
     assert "token_ids" not in dumped_choice

tests/unit/proxy/proxy_server/test_streaming_helpers.py

@@ -577,7 +577,7 @@ def test_serialize_streaming_chunk_simple_uses_fast_path_bytes():
 
 
 def test_serialize_streaming_chunk_keeps_provider_choice_fields():
-    chunk = _simple_chunk(content="1,2,3,4")
+    chunk: Final = _simple_chunk(content="1,2,3,4")
     chunk.choices[0]["stop_reason"] = ",5"
     chunk.choices[0]["token_ids"] = [20]
 

tests/unit/proxy/test_proxy_server_endpoints_and_startup.py

@@ -13605,12 +13605,12 @@ def test_is_injected_stream_usage_artifact():
 
     assert _is_injected_stream_usage_artifact({"usage": {"prompt_tokens": 1}}) is False
 
-    held_back_token_ids_chunk = ModelResponseStream(
+    held_back_token_ids_chunk: Final = ModelResponseStream(
         model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=[16])]
     )
     assert _is_injected_stream_usage_artifact(held_back_token_ids_chunk) is False
 
-    null_token_ids_chunk = ModelResponseStream(
+    null_token_ids_chunk: Final = ModelResponseStream(
         model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=None)]
     )
     assert _is_injected_stream_usage_artifact(null_token_ids_chunk) is True