From 153342dca8139389209c103dd44c87f4e393ece4 Mon Sep 17 00:00:00 2001 From: Alberto Perdomo Date: Tue, 27 Oct 2026 10:17:41 +0000 Subject: [PATCH] fix: Preserve deferred streaming choice metadata --- .../convert_dict_to_response.py | 16 ++--- .../litellm_core_utils/streaming_handler.py | 38 ++++++++-- litellm/proxy/proxy_server.py | 2 +- .../test_streaming_handler.py | 70 ++++++++++++++++++- 4 files changed, 108 insertions(+), 18 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 1af6b68cce3..b486aaec898 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -56,10 +56,10 @@ from litellm.types.utils import Logprobs as TextCompletionLogprobs 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"} +_CHOICES_FIELDS: Final[frozenset[str]] = frozenset(Choices.model_fields.keys()) +_CHOICES_EXPLICIT_KWARGS: Final[frozenset[str]] = _CHOICES_FIELDS | {"enhancements"} _MODEL_RESPONSE_FIELDS: Final[frozenset] = frozenset(ModelResponse.model_fields.keys()) | {"usage"} -_CACHED_CHOICE_RESERVED_KEYS: Final[frozenset] = STREAMING_CHOICE_OPENAI_FIELDS | {"message", "finish_details"} +_CACHED_CHOICE_RESERVED_KEYS: Final[frozenset[str]] = STREAMING_CHOICE_OPENAI_FIELDS | {"message", "finish_details"} def _normalize_images_for_message( @@ -148,8 +148,8 @@ def _split_assembled_content_for_replay(content: str | None) -> list[str]: return _REPLAY_CONTENT_SLICE_RE.findall(content) -def _cached_choice_provider_extras(choice: Mapping[str, object]) -> Mapping[str, object]: - return {key: value for key, value in choice.items() if key not in _CACHED_CHOICE_RESERVED_KEYS} +def _choice_provider_extras(choice: Mapping[str, object], reserved_keys: frozenset[str]) -> Mapping[str, object]: + return {key: value for key, value in choice.items() if key not in reserved_keys} def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None: @@ -260,7 +260,7 @@ async def convert_to_streaming_response_async( index=idx, delta=delta, logprobs=logprobs, - **_cached_choice_provider_extras(choice), + **_choice_provider_extras(choice, _CACHED_CHOICE_RESERVED_KEYS), ) choice_list.append(choice) @@ -360,7 +360,7 @@ def convert_to_streaming_response( delta=delta, logprobs=logprobs, enhancements=enhancements, - **_cached_choice_provider_extras(choice), + **_choice_provider_extras(choice, _CACHED_CHOICE_RESERVED_KEYS), ) choice_list.append(choice) @@ -814,7 +814,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_EXPLICIT_KWARGS}, + **_choice_provider_extras(provider_specific_fields, _CHOICES_EXPLICIT_KWARGS), ) if provider_sent_logprobs and logprobs is None: choice.logprobs = None diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index a19be338edb..7ff029e9ec8 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -33,6 +33,7 @@ from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import OpenAIChatCompletionChunk from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( + STREAMING_CHOICE_OPENAI_FIELDS, CacheCreationTokenDetails, CompletionTokensDetailsWrapper, Delta, @@ -257,6 +258,7 @@ class CustomStreamWrapper: self.system_fingerprint: str | None = None self._provider_response_model: str | None = None self.received_finish_reason: str | None = None + self._deferred_finish_chunk: ModelResponseStream | None = None self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream self.special_tokens = [ "<|assistant|>", @@ -293,7 +295,9 @@ class CustomStreamWrapper: self.received_model_id: str | None = None self.logging_loop = None self.rules = Rules() - self.stream_options = stream_options or getattr(logging_obj, "stream_options", None) + self.stream_options: Mapping[str, object] | None = stream_options or getattr( + logging_obj, "stream_options", None + ) self.messages = getattr(logging_obj, "messages", None) self.sent_stream_usage = False self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False @@ -301,7 +305,7 @@ class CustomStreamWrapper: 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.chunks: list[ModelResponseStream] = [] self._repeated_messages_count = 1 self.is_function_call = self.check_is_function_call(logging_obj=logging_obj) self.created: int | None = None @@ -461,7 +465,7 @@ class CustomStreamWrapper: ) self._restore_consumer_correlation_context() - def check_send_stream_usage(self, stream_options: dict | None): + def check_send_stream_usage(self, stream_options: Mapping[str, object] | None) -> bool: return stream_options is not None and stream_options.get("include_usage", False) is True def check_is_function_call(self, logging_obj) -> bool: @@ -509,9 +513,7 @@ class CustomStreamWrapper: last_content: Final = self.chunks[-1].choices[0].delta.content - if ( - last_content is None or not isinstance(last_content, str) or len(last_content) <= 2 - ): # ignore empty content - https://github.com/BerriAI/litellm/issues/5158#issuecomment-2287156946 + if last_content is None or len(last_content) <= 2: self._repeated_messages_count = 1 return @@ -1070,6 +1072,24 @@ class CustomStreamWrapper: self._optional_combine_thinking_block_in_choices(model_response=model_response) + if self.received_finish_reason is not None: + self._deferred_finish_chunk = ModelResponseStream(choices=model_response.choices).model_copy( + deep=True + ) + return model_response.model_copy( + update={ + "choices": [ + StreamingChoices( + **{ + key: value + for key, value in choice.model_dump().items() + if key in STREAMING_CHOICE_OPENAI_FIELDS + } + ) + for choice in model_response.choices + ] + } + ) return model_response else: return @@ -1773,9 +1793,13 @@ class CustomStreamWrapper: self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) def finish_reason_handler(self): + from litellm.litellm_core_utils.core_helpers import preserve_streaming_choice_provider_extras + 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} ) + if self._deferred_finish_chunk is not None: + preserve_streaming_choice_provider_extras(model_response, self._deferred_finish_chunk) _finish_reason: Final = self.received_finish_reason or self.intermittent_finish_reason if _finish_reason is not None: model_response.choices[0].finish_reason = _finish_reason @@ -2559,7 +2583,7 @@ def _coerce_token_details( return details_type(**(raw if isinstance(raw, dict) else raw.model_dump())) -def calculate_total_usage(chunks: list[ModelResponse]) -> Usage: +def calculate_total_usage(chunks: Sequence[ModelResponse | ModelResponseStream]) -> Usage: """Assume most recent usage chunk has total usage uptil then.""" from litellm.litellm_core_utils.streaming_chunk_builder_utils import ( attach_cache_creation_token_details, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d390f8eb303..03434b8622a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9858,7 +9858,7 @@ def _restamp_streaming_chunk_model( def _streaming_choice_carries_provider_extras(choice: object) -> bool: - extra: Final = getattr(choice, "__pydantic_extra__", None) + extra: Final[Mapping[str, object] | None] = getattr(choice, "__pydantic_extra__", None) if not isinstance(extra, dict): return False return any(key not in STREAMING_CHOICE_OPENAI_FIELDS for key in extra) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 21193041c16..bcecf9bb1b5 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -2,10 +2,12 @@ import asyncio import json import threading import time +from collections.abc import Sequence from typing import Final, NoReturn, Optional from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest +from openai.types.chat import ChatCompletionChunk import litellm from litellm import verbose_logger @@ -6691,11 +6693,15 @@ async def test_leaving_an_async_with_block_closes_the_provider_stream(): assert stream.close_calls == 1 -def _vllm_wrapper(chunks: list[ModelResponseStream], stream_options: dict | None = None) -> CustomStreamWrapper: +def _vllm_wrapper( + chunks: Sequence[ModelResponseStream | ChatCompletionChunk], + stream_options: dict | None = None, + custom_llm_provider: str = "hosted_vllm", +) -> CustomStreamWrapper: return CustomStreamWrapper( completion_stream=ModelResponseListIterator(model_responses=chunks), model="hosted_vllm/m", - custom_llm_provider="hosted_vllm", + custom_llm_provider=custom_llm_provider, stream_options=stream_options, logging_obj=Logging( model="hosted_vllm/m", @@ -6732,6 +6738,66 @@ async def test_vllm_choice_extras_survive_the_terminal_finish_chunk(sync_mode: b assert getattr(finish_choice, "token_ids", None) == [9] +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("source", ["sdk", "http_handler", "cache"]) +@pytest.mark.parametrize("content", ["", "last token"]) +@pytest.mark.parametrize(("stop_reason", "token_ids"), [("eos", (9,)), (0, ())]) +@pytest.mark.asyncio +async def test_vllm_finish_choice_extras_survive_serialized_live_and_cached_streams( + sync_mode: bool, source: str, content: str, stop_reason: str | int, token_ids: tuple[int, ...] +) -> None: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response, + convert_to_streaming_response_async, + ) + from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler + from litellm.proxy.proxy_server import _serialize_streaming_chunk + + terminal_choice: Final = { + "index": 0, + "finish_reason": "stop", + "stop_reason": stop_reason, + "token_ids": list(token_ids), + } + cached_response: Final = { + "id": "c", + "created": 1, + "model": "hosted_vllm/m", + "choices": [{**terminal_choice, "message": {"role": "assistant", "content": content}}], + } + replay_chunks: Final = ( + list(convert_to_streaming_response(cached_response)) + if sync_mode + else [chunk async for chunk in convert_to_streaming_response_async(cached_response)] + ) + parser: Final = OpenAIChatCompletionStreamingHandler(streaming_response=None, sync_stream=True) + live_body: Final = { + "id": "c", + "object": "chat.completion.chunk", + "created": 1, + "model": "hosted_vllm/m", + "choices": [{**terminal_choice, "delta": {"content": content}}], + } + live_chunk: Final = ChatCompletionChunk(**live_body) if source == "sdk" else parser.chunk_parser(live_body) + wrapper: Final = _vllm_wrapper( + replay_chunks if source == "cache" else [live_chunk], + custom_llm_provider="cached_response" if source == "cache" else "hosted_vllm", + ) + + emitted: Final = list(wrapper) if sync_mode else [chunk async for chunk in wrapper] + serialized: Final = [json.loads(_serialize_streaming_chunk(chunk)) for chunk in emitted] + finish_choices: Final = [chunk["choices"][0] for chunk in serialized if chunk["choices"][0].get("finish_reason")] + + assert len(finish_choices) == 1 + assert finish_choices[0]["stop_reason"] == terminal_choice["stop_reason"] + assert finish_choices[0]["token_ids"] == terminal_choice["token_ids"] + assert not finish_choices[0]["delta"].get("content") + assert [ + chunk["choices"][0]["token_ids"] for chunk in serialized if "token_ids" in chunk["choices"][0] + ] == [terminal_choice["token_ids"]] + assert "".join(chunk["choices"][0]["delta"].get("content") or "" for chunk in serialized) == content + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_continuous_usage_stats_keeps_usage_on_every_chunk(sync_mode: bool):