fix(L6): Preserve vLLM choice fields and honor continuous_usage_stats

Signed-off-by: Alberto Perdomo <aperdomo@redhat.com>

litellm/litellm_core_utils/core_helpers.py

@@ -564,6 +564,17 @@ def preserve_upstream_non_openai_attributes(
             setattr(model_response, key, value)
 
 
+def preserve_streaming_choice_provider_extras(
+    model_response: "ModelResponseStream", original_chunk: "ModelResponseStream"
+) -> None:
+    from litellm.types.utils import STREAMING_CHOICE_OPENAI_FIELDS
+
+    for outgoing, upstream in zip(model_response.choices, original_chunk.choices):
+        for key, value in (upstream.__pydantic_extra__ or {}).items():
+            if key not in STREAMING_CHOICE_OPENAI_FIELDS:
+                setattr(outgoing, key, value)
+
+
 def safe_deep_copy(data):
     """
     Safe Deep Copy

litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py

@@ -56,6 +56,7 @@ from .get_headers import get_response_headers
 
 _MESSAGE_FIELDS: Final[frozenset] = frozenset(Message.model_fields.keys())
 _CHOICES_FIELDS: Final[frozenset] = frozenset(Choices.model_fields.keys())
+_CHOICES_EXPLICIT_KWARGS: Final[frozenset] = _CHOICES_FIELDS | {"enhancements"}
 _MODEL_RESPONSE_FIELDS: Final[frozenset] = frozenset(ModelResponse.model_fields.keys()) | {"usage"}
 
 
@@ -790,7 +791,7 @@ def convert_to_model_response_object(
                     logprobs=logprobs,
                     enhancements=enhancements,
                     provider_specific_fields=provider_specific_fields,
-                    **{f: v for f, v in provider_specific_fields.items() if f not in _CHOICES_FIELDS},
+                    **{f: v for f, v in provider_specific_fields.items() if f not in _CHOICES_EXPLICIT_KWARGS},
                 )
                 if provider_sent_logprobs and logprobs is None:
                     choice.logprobs = None

litellm/litellm_core_utils/streaming_handler.py

@@ -299,7 +299,7 @@ class CustomStreamWrapper:
         self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False
         self._continuous_usage_stats = (
             isinstance(self.stream_options, dict) and self.stream_options.get("continuous_usage_stats") is True
-        )
+        ) and self.send_stream_usage
         self.tool_call = False
         self.chunks: list = []  # keep track of the returned chunks - used for calculating the input/output tokens for stream options
         self._repeated_messages_count = 1
@@ -1006,6 +1006,7 @@ class CustomStreamWrapper:
         response_obj: dict[str, Any],
     ):
         from litellm.litellm_core_utils.core_helpers import (
+            preserve_streaming_choice_provider_extras,
             preserve_upstream_non_openai_attributes,
         )
 
@@ -1103,6 +1104,10 @@ class CustomStreamWrapper:
                     model_response=model_response,
                     original_chunk=_original_chunk,
                 )
+                preserve_streaming_choice_provider_extras(
+                    model_response=model_response,
+                    original_chunk=_original_chunk,
+                )
 
             if _is_delta_empty:
                 model_response.choices[0].delta = Delta(content=None)  # ensure empty delta chunk returned
@@ -1550,6 +1555,8 @@ class CustomStreamWrapper:
                             model_response.choices[0].delta = Delta()
                 else:
                     if self.send_stream_usage is True:
+                        if self._continuous_usage_stats:
+                            self.sent_stream_usage = True
                         model_response.choices = []
                         return model_response
                     self._record_usage_only_chunk(model_response=model_response)

litellm/proxy/common_request_processing.py

@@ -552,11 +552,6 @@ def stream_usage_tracking_updates(
         return scrub
     if not supports_stream_options():
         return scrub
-    # The proxy injects include_usage for its own cost tracking, so
-    # caller-supplied continuous_usage_stats must not ride along and leak
-    # per-chunk usage to a caller who never asked for usage. A caller that
-    # wants cumulative usage sends include_usage too and takes the early return
-    # above, keeping both flags intact for the backend.
     merged_stream_options: Final = (
         {k: v for k, v in stream_options.items() if k != "continuous_usage_stats"}
         if isinstance(stream_options, dict)

litellm/proxy/proxy_server.py

@@ -171,6 +171,7 @@ from litellm.router_utils.routing_groups import parse_routing_groups
 from litellm.types.caching import RedisPipelineIncrementOperation
 from litellm.types.utils import (
     PRICING_OVERRIDES_KEY,
+    STREAMING_CHOICE_OPENAI_FIELDS,
     CountContextManagement,
     ModelResponse,
     ModelResponseStream,
@@ -9856,14 +9857,11 @@ def _restamp_streaming_chunk_model(
     return chunk, model_mismatch_logged
 
 
-_OPENAI_CHUNK_CHOICE_FIELDS: Final = frozenset({"index", "delta", "finish_reason", "logprobs", "enhancements"})
-
-
 def _streaming_choice_carries_provider_extras(choice: object) -> bool:
     extra: Final = getattr(choice, "__pydantic_extra__", None)
     if not isinstance(extra, dict):
         return False
-    return any(key not in _OPENAI_CHUNK_CHOICE_FIELDS for key in extra)
+    return any(key not in STREAMING_CHOICE_OPENAI_FIELDS for key in extra)
 
 
 def _fast_serialize_simple_model_response_stream(

litellm/types/utils.py

@@ -2043,6 +2043,9 @@ class Usage(SafeAttributeModel, CompletionUsage):
         setattr(self, key, value)
 
 
+STREAMING_CHOICE_OPENAI_FIELDS: Final = frozenset({"index", "delta", "finish_reason", "logprobs", "enhancements"})
+
+
 class StreamingChoices(OpenAIObject):
     def __init__(
         self,

tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py

@@ -2968,3 +2968,64 @@ def test_convert_to_model_response_object_function_output():
     assert result.usage.completion_tokens_details == CompletionTokensDetailsWrapper(reasoning_tokens=0)
 
     assert result._hidden_params is not None
+
+
+def test_convert_to_model_response_object_keeps_provider_choice_level_extras_at_choice_level():
+    response_object = {
+        "id": "chatcmpl-vllm-extras",
+        "object": "chat.completion",
+        "created": 1,
+        "model": "HostedVLLM/Qwen2.5-0.5B-Instruct",
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "Hi"},
+                "finish_reason": "stop",
+                "stop_reason": "eos",
+                "token_ids": [11, 22, 33],
+            }
+        ],
+        "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
+    }
+
+    result = convert_to_model_response_object(
+        model_response_object=ModelResponse(),
+        response_object=response_object,
+        response_type="completion",
+    )
+
+    choice: Final = result.choices[0]
+    assert getattr(choice, "stop_reason", None) == "eos"
+    assert getattr(choice, "token_ids", None) == [11, 22, 33]
+    assert choice.provider_specific_fields is not None
+    assert choice.provider_specific_fields["stop_reason"] == "eos"
+    assert choice.provider_specific_fields["token_ids"] == [11, 22, 33]
+    dumped: Final = json.loads(choice.model_dump_json())
+    assert dumped["stop_reason"] == "eos"
+    assert dumped["token_ids"] == [11, 22, 33]
+
+
+def test_convert_to_model_response_object_accepts_a_choice_enhancements_key():
+    response_object = {
+        "id": "chatcmpl-enhancements",
+        "object": "chat.completion",
+        "created": 1,
+        "model": "azure/gpt-4o",
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "Hi"},
+                "finish_reason": "stop",
+                "enhancements": {"grounding": True},
+            }
+        ],
+        "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
+    }
+
+    result = convert_to_model_response_object(
+        model_response_object=ModelResponse(),
+        response_object=response_object,
+        response_type="completion",
+    )
+
+    assert result.choices[0].enhancements == {"grounding": True}

tests/unit/litellm_core_utils/test_streaming_handler.py

@@ -6689,3 +6689,88 @@ async def test_leaving_an_async_with_block_closes_the_provider_stream():
     async with _wrapper_over(stream):
         pass
     assert stream.close_calls == 1
+
+
+def _vllm_wrapper(chunks: list[ModelResponseStream], stream_options: dict | None = None) -> CustomStreamWrapper:
+    return CustomStreamWrapper(
+        completion_stream=ModelResponseListIterator(model_responses=chunks),
+        model="hosted_vllm/m",
+        custom_llm_provider="hosted_vllm",
+        stream_options=stream_options,
+        logging_obj=Logging(
+            model="hosted_vllm/m",
+            messages=[{"role": "user", "content": "hi"}],
+            stream=True,
+            call_type="acompletion",
+            start_time=time.time(),
+            litellm_call_id="vllm-choice",
+            function_id="test",
+        ),
+    )
+
+
+def _vllm_chunk(choice: dict, usage: dict | None = None) -> ModelResponseStream:
+    body: Final = {"id": "c", "created": 1, "model": "hosted_vllm/m", "choices": [choice]}
+    if usage is not None:
+        body["usage"] = usage
+    return ModelResponseStream(**body)
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_vllm_choice_extras_survive_the_terminal_finish_chunk(sync_mode: bool):
+    chunks: Final = [
+        _vllm_chunk({"index": 0, "delta": {"content": "Hi"}, "finish_reason": None}),
+        _vllm_chunk({"index": 0, "delta": {}, "finish_reason": "stop", "stop_reason": "eos", "token_ids": [9]}),
+    ]
+    wrapper: Final = _vllm_wrapper(chunks)
+
+    emitted: Final = [chunk async for chunk in wrapper] if not sync_mode else list(wrapper)
+
+    finish_choice: Final = emitted[-1].choices[0]
+    assert getattr(finish_choice, "stop_reason", None) == "eos"
+    assert getattr(finish_choice, "token_ids", None) == [9]
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_continuous_usage_stats_keeps_usage_on_every_chunk(sync_mode: bool):
+    counts: Final = [
+        {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
+        {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12},
+        {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13},
+    ]
+    chunks: Final = [
+        _vllm_chunk({"index": 0, "delta": {"content": "Hi"}, "finish_reason": None}, usage=counts[0]),
+        _vllm_chunk({"index": 0, "delta": {}, "finish_reason": "stop"}, usage=counts[1]),
+        ModelResponseStream(id="c", created=1, model="hosted_vllm/m", choices=[], usage=Usage(**counts[2])),
+    ]
+    wrapper: Final = _vllm_wrapper(chunks, stream_options={"include_usage": True, "continuous_usage_stats": True})
+
+    emitted: Final = [chunk async for chunk in wrapper] if not sync_mode else list(wrapper)
+
+    usage_chunks: Final = [chunk for chunk in emitted if getattr(chunk, "usage", None) is not None]
+    assert [chunk.usage.total_tokens for chunk in usage_chunks] == [11, 12, 13]
+    assert emitted[-1].choices == []
+    assert litellm.stream_chunk_builder(chunks=emitted).usage.total_tokens == 13
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_include_usage_without_continuous_usage_stats_still_emits_only_a_final_usage_chunk(sync_mode: bool):
+    counts: Final = [
+        {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
+        {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13},
+    ]
+    chunks: Final = [
+        _vllm_chunk({"index": 0, "delta": {"content": "Hi"}, "finish_reason": None}, usage=counts[0]),
+        _vllm_chunk({"index": 0, "delta": {}, "finish_reason": "stop"}, usage=counts[1]),
+    ]
+    wrapper: Final = _vllm_wrapper(chunks, stream_options={"include_usage": True})
+
+    emitted: Final = [chunk async for chunk in wrapper] if not sync_mode else list(wrapper)
+
+    content_chunks: Final = [chunk for chunk in emitted[:-1] if chunk.choices and chunk.choices[0].delta.content]
+    assert content_chunks
+    assert all(getattr(chunk, "usage", None) is None for chunk in content_chunks)
+    assert getattr(emitted[-1], "usage", None) is not None

tests/unit/proxy/proxy_server/test_streaming_helpers.py

@@ -541,6 +541,19 @@ def test_fast_serialize_simple_model_response_stream_returns_bytes_payload():
     }
 
 
+def test_fast_serialize_falls_back_when_choice_carries_provider_extras():
+    chunk = _simple_chunk()
+    chunk.choices[0].stop_reason = "eos"
+    chunk.choices[0].token_ids = [1, 2, 3]
+
+    assert _fast_serialize_simple_model_response_stream(chunk) is None
+
+    serialized: Final = _serialize_streaming_chunk(chunk)
+    payload: Final = json.loads(serialized)
+    assert payload["choices"][0]["stop_reason"] == "eos"
+    assert payload["choices"][0]["token_ids"] == [1, 2, 3]
+
+
 def test_fast_serialize_simple_model_response_stream_with_usage_returns_none_invalid():
     """Fast path bails (returns None) when ``usage`` is populated — the slow
     path is required to preserve usage fields. Returning None here is the

tests/unit/proxy/test_common_request_processing.py

@@ -8434,6 +8434,23 @@ class TestApplyStreamUsageTracking:
         assert "stream_options" not in data
         assert "_litellm_strip_stream_usage" not in data
 
+    def test_client_continuous_usage_stats_is_not_injected_onto_the_caller(self):
+        data = {"stream": True, "stream_options": {"continuous_usage_stats": True}}
+
+        _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion")
+
+        assert data["stream_options"] == {"include_usage": True}
+        assert data["_litellm_strip_stream_usage"] is True
+
+    def test_client_requested_continuous_usage_stats_reaches_the_backend_intact(self):
+        requested: Final = {"include_usage": True, "continuous_usage_stats": True}
+        data = {"stream": True, "stream_options": dict(requested)}
+
+        _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion")
+
+        assert data["stream_options"] == requested
+        assert "_litellm_strip_stream_usage" not in data
+
     def test_client_supplied_strip_marker_is_neutralized(self):
         data = {
             "stream": True,