diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 433a9ca460f..10947a85489 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/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: diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index b25757af3dc..e901a8d27b7 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/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}, diff --git a/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py index d7baf44bec7..21c9179773c 100644 --- a/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py +++ b/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 diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 458a48fff76..3128f678e09 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/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] diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index c157825a039..7c405b63220 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/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