diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index cc1f6fef019..42274b8645b 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/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 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 d2b0b9124f9..d5836f55db8 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,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 diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 973aeb13b8a..11e6d255e9f 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/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) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index b4f20df9e25..5838e893d62 100644 --- a/litellm/proxy/common_request_processing.py +++ b/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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d6eabf23451..d390f8eb303 100644 --- a/litellm/proxy/proxy_server.py +++ b/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( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a97e35a9021..30d7faa0082 100644 --- a/litellm/types/utils.py +++ b/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, 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/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 93f700b1c2a..a0eafbd33a2 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/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 diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index f53e1931782..be88ea81339 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/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 diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 1d197602543..33165ff4b30 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/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,