diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 00c6ce97e78..625a98de132 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -670,6 +670,17 @@ def _assign_attributes(target: object, values: Iterable[tuple[str, object]]) -> setattr(target, 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 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 8b94c34942c..b64aeeb7271 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 @@ -25,6 +25,7 @@ from litellm.types.llms.openai import ( OpenAIModerationResponse, ) from litellm.types.utils import ( + STREAMING_CHOICE_OPENAI_FIELDS, ChatCompletionDeltaToolCall, ChatCompletionMessageCustomToolCall, ChatCompletionMessageToolCall, @@ -55,8 +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_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[str]] = STREAMING_CHOICE_OPENAI_FIELDS | {"message", "finish_details"} def _normalize_images_for_message( @@ -145,6 +148,10 @@ def _split_assembled_content_for_replay(content: str | None) -> list[str]: return _REPLAY_CONTENT_SLICE_RE.findall(content) +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: # Rebuild the delta as content-only so every accumulate-able field (role, # tool_calls, reasoning_content, thinking_blocks, audio, images, @@ -158,6 +165,12 @@ def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None: del choice.enhancements +def _clear_cached_choice_provider_extras(choice: StreamingChoices) -> None: + for key in tuple(choice.__pydantic_extra__ or {}): + if key not in STREAMING_CHOICE_OPENAI_FIELDS: + delattr(choice, key) + + def _first_replay_delta( message: Mapping[str, object], reasoning_items: Sequence[ChatCompletionReasoningItem] | None ) -> Delta: @@ -242,7 +255,13 @@ async def convert_to_streaming_response_async( logprobs = choice.get("logprobs", None) - choice = StreamingChoices(finish_reason=finish_reason, index=idx, delta=delta, logprobs=logprobs) + choice = StreamingChoices( + finish_reason=finish_reason, + index=idx, + delta=delta, + logprobs=logprobs, + **_choice_provider_extras(_FIELD_MAP.validate_python(choice), _CACHED_CHOICE_RESERVED_KEYS), + ) choice_list.append(choice) model_response_object.choices = choice_list @@ -291,6 +310,8 @@ async def convert_to_streaming_response_async( slice_chunk.choices[0].delta.content = piece if i > 0: _clear_later_replay_slice_metadata(slice_chunk.choices[0]) + if i != last_idx: + _clear_cached_choice_provider_extras(slice_chunk.choices[0]) slice_chunk.choices[0].finish_reason = original_finish_reason if i == last_idx else None if i == last_idx and original_usage is not None: setattr(slice_chunk, "usage", original_usage) @@ -332,6 +353,7 @@ def convert_to_streaming_response( delta=delta, logprobs=logprobs, enhancements=enhancements, + **_choice_provider_extras(_FIELD_MAP.validate_python(choice), _CACHED_CHOICE_RESERVED_KEYS), ) choice_list.append(choice) @@ -376,6 +398,8 @@ def convert_to_streaming_response( slice_chunk.choices[0].delta.content = piece if i > 0: _clear_later_replay_slice_metadata(slice_chunk.choices[0]) + if i != last_idx: + _clear_cached_choice_provider_extras(slice_chunk.choices[0]) slice_chunk.choices[0].finish_reason = original_finish_reason if i == last_idx else None if i == last_idx and original_usage is not None: setattr(slice_chunk, "usage", original_usage) @@ -834,6 +858,9 @@ def convert_to_model_response_object( logprobs=logprobs, enhancements=enhancements, provider_specific_fields=provider_specific_fields, + **_choice_provider_extras( + _FIELD_MAP.validate_python(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 b64e146ca08..7b6a577c91b 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -17,14 +17,14 @@ import httpx from openai.types.chat import ChatCompletionChunk from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice from openai.types.chat.chat_completion_chunk import ChoiceDelta -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import NotRequired, TypedDict import litellm from litellm import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.asyncify import asyncify -from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost +from litellm.litellm_core_utils.core_helpers import as_str_mapping, chunk_usage, provider_reported_usage_cost from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR, get_hidden_params, @@ -40,6 +40,7 @@ from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( + STREAMING_CHOICE_OPENAI_FIELDS, CacheCreationTokenDetails, CompletionTokensDetailsWrapper, Delta, @@ -70,6 +71,7 @@ TOOL_CALLS_ATTRIBUTE: Final = "tool_calls" FUNCTION_CALL_ATTRIBUTE: Final = "function_call" _SYNC_ITER_EXHAUSTED: Final = object() +_STREAM_OPTIONS_VALUE: Final = TypeAdapter(object) _GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__) _CompleteStreamingResponse = ModelResponse | TextCompletionResponse | ModelResponseStream | None @@ -337,6 +339,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|>", @@ -373,10 +376,18 @@ 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: object = _STREAM_OPTIONS_VALUE.validate_python( + 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 + stream_options_fields: Final = as_str_mapping(self.stream_options) + self._continuous_usage_stats = ( + isinstance(self.stream_options, dict) + and stream_options_fields is not None + and stream_options_fields.get("continuous_usage_stats") is True + ) and self.send_stream_usage self.caller_requested_stream_usage: Final = ( self.check_send_stream_usage(getattr(logging_obj, "stream_options", None)) and getattr(logging_obj, "stream_usage_injected", False) is not True @@ -543,8 +554,11 @@ class CustomStreamWrapper: ) self._restore_consumer_correlation_context() - def check_send_stream_usage(self, stream_options: dict | None): - return stream_options is not None and stream_options.get("include_usage", False) is True + def check_send_stream_usage(self, stream_options: object) -> bool: + if stream_options is None: + return False + get_option: Final[Callable[[str, object], object] | None] = getattr(stream_options, "get", None) + return get_option is not None and get_option("include_usage", False) is True def check_is_function_call(self, logging_obj) -> bool: from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -591,9 +605,7 @@ class CustomStreamWrapper: last_content: Final = self.chunks[-1].choices[0].delta.content - if ( - 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 @@ -1129,6 +1141,7 @@ class CustomStreamWrapper: plain_delta_created: bool = False, ): from litellm.litellm_core_utils.core_helpers import ( + preserve_streaming_choice_provider_extras, preserve_upstream_non_openai_attributes, ) @@ -1227,6 +1240,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 @@ -1261,6 +1292,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 @@ -1694,7 +1729,7 @@ class CustomStreamWrapper: model_response.choices[0].delta = plain_delta if plain_delta is not None else Delta() plain_delta_created = plain_delta is not None # rebind-ok: the first chunk builds its delta else: - if self.send_stream_usage is True: + if self.send_stream_usage is True and not self._continuous_usage_stats: model_response.choices = [] return model_response self._record_usage_only_chunk(model_response=model_response) @@ -1912,9 +1947,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 @@ -2071,7 +2110,7 @@ class CustomStreamWrapper: # ModelResponseStream declares `usage` as a field, so # hasattr(response, "usage") is always True — must check # `is not None` to avoid running this path on every chunk. - if getattr(response, "usage", None) is not None: + if getattr(response, "usage", None) is not None and not self._continuous_usage_stats: usage_to_preserve = response.usage if usage_to_preserve: response_hidden_params_for_usage = cast( # cast-ok: preserve dynamic mapping behavior @@ -2173,6 +2212,8 @@ class CustomStreamWrapper: self.chunks.append(processed_chunk) if self.stream_options is None: self._set_final_chunk_usage(processed_chunk) + elif self._continuous_usage_stats: + processed_chunk.usage = calculate_total_usage(chunks=self.chunks) ## LOGGING executor.submit( self.run_success_logging_and_cache_storage, @@ -2259,7 +2300,7 @@ class CustomStreamWrapper: hasattr(processed_chunk, "usage") and getattr(processed_chunk, "usage", None) is not None ) - if _has_usage: + if _has_usage and not self._continuous_usage_stats: usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk) if usage_free_chunk is None: continue @@ -2406,6 +2447,8 @@ class CustomStreamWrapper: self.chunks.append(processed_chunk) if self.stream_options is None: self._set_final_chunk_usage(processed_chunk) + elif self._continuous_usage_stats: + processed_chunk.usage = calculate_total_usage(chunks=self.chunks) # see sync __next__'s sibling branch: deliberately do NOT restore # here - this chunk is still this call's own data, and restoring # before returning it would corrupt the caller's own log diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 24a36a7f2dc..6f8cce85651 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -59,6 +59,7 @@ from litellm.litellm_core_utils.bug_report import ( strip_bug_report_notice, ) from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, get_or_create_metadata_bucket, independent_snapshot, is_expected_client_error, @@ -591,6 +592,7 @@ def stream_usage_tracking_updates( return scrub always_include: Final = general_settings.get("always_include_stream_usage") stream_options: Final = data.get("stream_options") + stream_options_fields: Final = as_str_mapping(stream_options) if always_include is True: if "stream_options" not in data: return {**scrub, "stream_options": {"include_usage": True}, STREAM_USAGE_INJECTED_FIELD: True} @@ -609,7 +611,11 @@ def stream_usage_tracking_updates( return scrub if not supports_stream_options(): return scrub - merged_stream_options: Final = {**stream_options} if isinstance(stream_options, dict) else {} + merged_stream_options: Final = ( + {key: value for key, value in stream_options_fields.items() if key != "continuous_usage_stats"} + if isinstance(stream_options, dict) and stream_options_fields is not None + else {} + ) return { "stream_options": {**merged_stream_options, "include_usage": True}, "_litellm_strip_stream_usage": True, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 04f84f1f7b5..e78819165cc 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -192,6 +192,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, @@ -10056,6 +10057,13 @@ def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) - return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending +def _streaming_choice_carries_provider_extras(choice: object) -> bool: + 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) + + def _fast_serialize_simple_model_response_stream( chunk: ModelResponseStream, ) -> bytes | None: @@ -10078,6 +10086,8 @@ def _fast_serialize_simple_model_response_stream( choice: Final = choices[0] if getattr(choice, "logprobs", None) is not None or getattr(choice, "enhancements", None) is not None: return None + if _streaming_choice_carries_provider_extras(chunk.choices[0]): + return None delta: Final[object] = getattr(choice, "delta", None) if delta is None: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 055c18f2e36..eeabb044838 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2098,6 +2098,9 @@ class UnreportedUsage(Usage): pass +STREAMING_CHOICE_OPENAI_FIELDS: Final = frozenset({"index", "delta", "finish_reason", "logprobs", "enhancements"}) + + class StreamingChoices(OpenAIObject): def __init__( self, diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index 0c601a6b2c8..7faf0268a60 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/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} diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py index 348d0105a6c..ab7c81b52f6 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py @@ -460,3 +460,18 @@ def test_sync_non_enumerated_delta_fields_only_on_first_chunk(): assert getattr(chunks[0].choices[0].delta, "annotations", None) == annotations for c in chunks[1:]: assert getattr(c.choices[0].delta, "annotations", None) is None + + +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_keeps_cached_choice_provider_extras_on_the_finish_chunk(replay): + cached = _async_payload() + cached["choices"][0]["stop_reason"] = "eos" + cached["choices"][0]["token_ids"] = [9] + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + + assert len(chunks) > 1 + assert [getattr(c.choices[0], "stop_reason", None) for c in chunks] == [None] * (len(chunks) - 1) + ["eos"] + assert [getattr(c.choices[0], "token_ids", None) for c in chunks] == [None] * (len(chunks) - 1) + [[9]] + assert chunks[-1].choices[0].finish_reason == "stop" diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 6d700bae514..2e007c11b65 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -15,6 +15,7 @@ import httpx import pytest from openai.types.chat import ChatCompletionChunk from pydantic import BaseModel, TypeAdapter + import litellm from litellm import verbose_logger from litellm._logging import session_id_var, trace_id_var @@ -7846,6 +7847,179 @@ async def test_leaving_an_async_with_block_closes_the_provider_stream(): assert stream.close_calls == 1 +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=custom_llm_provider, + 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.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): + 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) + + assert all(getattr(chunk, "usage", None) is not None for chunk in emitted) + assert [chunk.usage.total_tokens for chunk in emitted] == [11, 12, 13] + assert litellm.stream_chunk_builder(chunks=emitted).usage.total_tokens == 13 + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_continuous_usage_stats_finishes_before_the_final_usage_chunk(sync_mode: bool): + content_usage: Final = {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11} + final_usage: Final = {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13} + chunks: Final = [ + _vllm_chunk({"index": 0, "delta": {"content": "Hi"}, "finish_reason": None}, usage=content_usage), + _vllm_chunk({"index": 0, "delta": {"content": "!"}, "finish_reason": "length"}, usage=final_usage), + ModelResponseStream(id="c", created=1, model="hosted_vllm/m", choices=[], usage=Usage(**final_usage)), + ] + 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) + + finish_indexes: Final = [i for i, chunk in enumerate(emitted) if chunk.choices and chunk.choices[0].finish_reason] + assert len(finish_indexes) == 1 + assert emitted[finish_indexes[0]].choices[0].finish_reason == "length" + assert finish_indexes[0] < len(emitted) - 1 + terminal_chunk: Final = emitted[-1] + assert not (terminal_chunk.choices and terminal_chunk.choices[0].finish_reason) + assert [chunk.usage.total_tokens for chunk in emitted] == [ + content_usage["total_tokens"], + *[final_usage["total_tokens"]] * 3, + ] + + +@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 + def _final_chunk_with_dict_usage(choices) -> ModelResponseStream: """A stream chunk whose ``usage`` is a plain dict. diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 9a3516d20d6..099fa647ad6 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -544,6 +544,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 diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 52816dc50a8..4db1f9f3038 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -9371,6 +9371,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,