diff --git a/litellm/constants.py b/litellm/constants.py index 25770d960bb..075b7d08b90 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -686,6 +686,8 @@ STREAM_SSE_DONE_STRING: Final[str] = "[DONE]" STREAM_SSE_DATA_PREFIX: Final[str] = "data: " STREAM_SSE_KEEPALIVE_PING_CHUNK: Final[str] = 'event: ping\ndata: {"type": "ping"}\n\n' STREAM_SSE_KEEPALIVE_PING_BYTES: Final[bytes] = STREAM_SSE_KEEPALIVE_PING_CHUNK.encode("utf-8") +STREAM_FINISH_CHOICE_FIELDS: Final = ("stop_reason",) +STREAM_CHUNK_CHOICE_FIELDS: Final = ("token_ids",) ### SPEND TRACKING ### DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND: Final = get_env_float( "DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400 diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 2348d8e57d6..ffe047f7bee 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -23,6 +23,7 @@ from typing_extensions import NotRequired, TypedDict import litellm from litellm import verbose_logger from litellm._uuid import uuid +from litellm.constants import STREAM_CHUNK_CHOICE_FIELDS, STREAM_FINISH_CHOICE_FIELDS 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.hidden_params import ( @@ -113,6 +114,41 @@ def _prepend_upstream_choice_extras( choice.__pydantic_fields_set__.update(upstream_extras) +_NO_UPSTREAM_CHOICE_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _upstream_first_choice(upstream_chunk: object) -> BaseModel | None: + match getattr(upstream_chunk, "choices", None): + case [BaseModel() as first_choice, *_]: + return first_choice + case _: + 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]: + return MappingProxyType( + {field: value for field, value in _upstream_choice_fields(upstream_chunk, fields).items() if value is not None} + ) + + +def _apply_choice_fields(choice: StreamingChoices, fields: Mapping[str, object]) -> None: + for field, value in fields.items(): + setattr(choice, field, value) + + def _holds_only_scalars(fields: Iterable[tuple[str, object]]) -> bool: return all(value is None or isinstance(value, (str, int, float, bool)) for _, value in fields) @@ -337,6 +373,7 @@ class CustomStreamWrapper: self.system_fingerprint: str | None = None self._provider_response_model: str | None = None self.received_finish_reason: str | None = None + self.received_finish_choice_fields: Mapping[str, object] = _NO_UPSTREAM_CHOICE_FIELDS self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream self.special_tokens = [ "<|assistant|>", @@ -1268,12 +1305,26 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = map_finish_reason( finish_reason=self.received_finish_reason ) # ensure consistent output to openai + _apply_choice_fields(model_response.choices[0], self.received_finish_choice_fields) + if response_obj.get("is_finished"): + _apply_choice_fields( + model_response.choices[0], + _non_null_upstream_choice_fields(_original_chunk, STREAM_CHUNK_CHOICE_FIELDS), + ) self.sent_last_chunk = True return model_response elif self._has_special_delta_content(model_response): return self._handle_special_delta_content(model_response) + 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: if hasattr(model_response, "usage"): self.chunks.append(model_response) @@ -1497,6 +1548,7 @@ class CustomStreamWrapper: print_verbose(f"completion obj content: {completion_obj['content']}") if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] + self.received_finish_choice_fields = _upstream_choice_fields(chunk, STREAM_FINISH_CHOICE_FIELDS) if response_obj["usage"] is not None: _text_completion_usage: Final[Usage] = response_obj["usage"] setattr( @@ -1586,6 +1638,7 @@ class CustomStreamWrapper: f"{self.custom_llm_provider} raised a streaming error - finish_reason: error, no content string given. Received Chunk={response_obj}" ) self.received_finish_reason = response_obj["finish_reason"] + self.received_finish_choice_fields = _upstream_choice_fields(chunk, STREAM_FINISH_CHOICE_FIELDS) if response_obj.get("original_chunk", None) is not None: if hasattr(response_obj["original_chunk"], "id"): model_response = self.set_model_id(response_obj["original_chunk"].id, model_response) @@ -1931,6 +1984,7 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = _finish_reason else: model_response.choices[0].finish_reason = "stop" + _apply_choice_fields(model_response.choices[0], self.received_finish_choice_fields) ## if tool use if ( diff --git a/litellm/llms/hosted_vllm/chat/transformation.py b/litellm/llms/hosted_vllm/chat/transformation.py index 63e6cb2b5da..4a1635f38d1 100644 --- a/litellm/llms/hosted_vllm/chat/transformation.py +++ b/litellm/llms/hosted_vllm/chat/transformation.py @@ -3,8 +3,11 @@ Translate from OpenAI's `/v1/chat/completions` to VLLM's `/v1/chat/completions` """ import json -from collections.abc import Coroutine -from typing import Final, Literal, cast, overload +from collections.abc import Coroutine, Mapping +from typing import TYPE_CHECKING, Final, Literal, cast, overload + +import httpx +from pydantic import JsonValue from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_image_mime_type_from_url, @@ -22,10 +25,24 @@ from litellm.types.llms.openai import ( ChatCompletionVideoObject, ChatCompletionVideoUrlObject, ) +from litellm.types.utils import Choices, ModelResponse from ....utils import remove_additional_properties, remove_strict_from_schema from ...openai.chat.gpt_transformation import OpenAIGPTConfig +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer + +_VLLM_CHOICE_FIELDS: Final = ("stop_reason", "token_ids") + + +def _vllm_choice_fields(choice: Choices) -> Mapping[str, object]: + provider_fields: Final[Mapping[str, object]] = ( + choice.provider_specific_fields or {} if "provider_specific_fields" in choice else {} + ) + return {field: provider_fields[field] for field in _VLLM_CHOICE_FIELDS if field in provider_fields} + class HostedVLLMChatConfig(OpenAIGPTConfig): def _convert_custom_tools_to_function_tools(self, tools: list[dict[str, object]]) -> list[dict[str, object]]: @@ -106,6 +123,38 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): return super().map_openai_params(non_default_params, optional_params, model, drop_params) + def transform_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: "LiteLLMLoggingObj", + request_data: dict[str, JsonValue], # mutable-ok: BaseConfig signature + messages: list[AllMessageValues], # mutable-ok: BaseConfig signature + optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature + litellm_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature + encoding: "Tokenizer | None", + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ModelResponse: + response: Final = super().transform_response( # pyright: ignore[reportUnknownMemberType] # OpenAIGPTConfig.transform_response takes bare dicts + model=model, + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data=request_data, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + encoding=encoding, + api_key=api_key, + json_mode=json_mode, + ) + for choice in response.choices: + for field, value in _vllm_choice_fields(choice).items(): + choice[field] = value + return response + def _get_openai_compatible_provider_info( self, api_base: str | None, api_key: str | None ) -> tuple[str | None, str | None]: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b263412ce0f..c6120b05ba1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10065,6 +10065,9 @@ def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) - return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending +_STANDARD_STREAM_CHOICE_FIELDS: Final = frozenset({"finish_reason", "index", "delta", "logprobs", "enhancements"}) + + def _fast_serialize_simple_model_response_stream( chunk: ModelResponseStream, ) -> bytes | None: @@ -10087,6 +10090,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 not isinstance(choice, BaseModel) or (choice.model_extra or {}).keys() - _STANDARD_STREAM_CHOICE_FIELDS: + return None delta: Final[object] = getattr(choice, "delta", None) if delta is None: @@ -10165,6 +10170,11 @@ def _is_empty_streaming_choice(choice: StreamingChoices) -> bool: return False if getattr(choice, "logprobs", None) is not None: return False + provider_choice_values: Final = ( + value for field, value in (choice.model_extra or {}).items() if field not in _STANDARD_STREAM_CHOICE_FIELDS + ) + if not _all_values_none(provider_choice_values): + return False delta: Final = choice.delta if hasattr(choice, "delta") else None if delta is None: return True diff --git a/litellm/text_completions/streaming_iterator.py b/litellm/text_completions/streaming_iterator.py index ec1b96fd432..0e27b90480f 100644 --- a/litellm/text_completions/streaming_iterator.py +++ b/litellm/text_completions/streaming_iterator.py @@ -1,5 +1,6 @@ from typing import Final +from litellm.constants import STREAM_FINISH_CHOICE_FIELDS from litellm.types.utils import Choices, ModelResponse, TextChoices, TextCompletionResponse @@ -34,11 +35,15 @@ class TextCompletionStreamWrapper: raise Exception if chunk["choices"]: - delta: Final = chunk["choices"][0]["delta"] + chat_choice: Final = chunk["choices"][0] + delta: Final = chat_choice["delta"] text_choices["text"] = delta["content"] text_choices["reasoning_content"] = delta.get("reasoning_content") - text_choices["index"] = chunk["choices"][0]["index"] - text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"] + text_choices["index"] = chat_choice["index"] + text_choices["finish_reason"] = chat_choice["finish_reason"] + for field in STREAM_FINISH_CHOICE_FIELDS: + if field in chat_choice: + text_choices[field] = chat_choice[field] response["choices"] = [text_choices] else: response["choices"] = [] diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 79978ecfe8e..5697b78ebe2 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -8,6 +8,7 @@ import threading import time from collections.abc import Mapping, Sequence from dataclasses import dataclass +from itertools import chain from types import CodeType, FrameType, MappingProxyType from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -30,6 +31,7 @@ from litellm.litellm_core_utils.streaming_handler import ( _ProviderChunkParsed, ) from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response from litellm.types.utils import ( @@ -6158,6 +6160,89 @@ def test_chunk_creator_keeps_upstream_index_on_single_choice_chunk(logging_obj: assert result.choices[0].delta.content == "hello" +def _hosted_vllm_stream_choices( + logging_obj: Logging, upstream_choices: Sequence[Mapping[str, object]] +) -> tuple[dict[str, object], ...]: + sse_lines: Final = ( + *( + "data: " + json.dumps({"id": "chatcmpl-1", "object": "chat.completion.chunk", "model": "m", "choices": [c]}) + for c in upstream_choices + ), + "data: [DONE]", + ) + completion_stream: Final = HostedVLLMChatConfig().get_model_response_iterator( + streaming_response=iter(sse_lines), sync_stream=True, json_mode=False + ) + return tuple( + chunk.choices[0].model_dump() + for chunk in CustomStreamWrapper( + completion_stream=completion_stream, model="m", custom_llm_provider="hosted_vllm", logging_obj=logging_obj + ) + ) + + +@pytest.mark.parametrize("stop_reason", [",5", 151645]) +@pytest.mark.parametrize("final_content", [",3", ""], ids=["content-then-finish", "finish-only"]) +def test_hosted_vllm_stream_keeps_stop_reason_and_every_token_id( + logging_obj: Logging, final_content: str, stop_reason: str | int +) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]}, + {"index": 0, "delta": {"content": "1,2"}, "finish_reason": None, "token_ids": [11, 17]}, + { + "index": 0, + "delta": {"content": final_content}, + "finish_reason": "stop", + "stop_reason": stop_reason, + "token_ids": [20], + }, + ), + ) + + assert (emitted[-1]["finish_reason"], emitted[-1].get("stop_reason")) == ("stop", stop_reason) + 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}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, **token_ids}, + {"index": 0, "delta": {"content": "1,2"}, "finish_reason": None, **token_ids}, + {"index": 0, "delta": {"content": ""}, "finish_reason": "stop", "stop_reason": ",3", **token_ids}, + ) + + +def test_hosted_vllm_stream_with_null_token_ids_emits_the_same_chunks_as_without_them(logging_obj: Logging) -> None: + with_null_ids: Final = _hosted_vllm_stream_choices(logging_obj, _held_back_stop_string_stream({"token_ids": None})) + without_ids: Final = _hosted_vllm_stream_choices(logging_obj, _held_back_stop_string_stream({})) + + assert [(c["delta"], c["finish_reason"]) for c in with_null_ids] == [ + (c["delta"], c["finish_reason"]) for c in without_ids + ] + + def _later_chunk_outputs( custom_llm_provider: str, later_chunks: Sequence[ChatCompletionChunk], logging_obj: Logging ) -> tuple[ModelResponseStream, ...]: 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 c792d5dffcd..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,14 +1,16 @@ import json -from unittest.mock import MagicMock, patch +from typing import Final +from unittest.mock import MagicMock +import httpx import pytest - from litellm.constants import ( DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, ) from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig +from litellm.types.utils import ModelResponse def test_hosted_vllm_chat_transformation_file_url(): @@ -397,3 +399,61 @@ def test_hosted_vllm_custom_tools_use_top_level_input_schema(): assert tools[0]["function"]["name"] == "search" assert tools[0]["function"]["description"] == "Search docs" assert tools[0]["function"]["parameters"] == input_schema + + +def _transform_vllm_chat_response(choice_extras: dict[str, object]) -> ModelResponse: + payload: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "glm", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "logprobs": None, + "finish_reason": "stop", + **choice_extras, + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, + } + raw_response: Final = httpx.Response( + 200, json=payload, request=httpx.Request("POST", "http://vllm.local/v1/chat/completions") + ) + return HostedVLLMChatConfig().transform_response( + model="hosted_vllm/glm", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +@pytest.mark.parametrize( + "choice_extras", + [ + {"stop_reason": ", 5", "token_ids": [16, 11, 220]}, + {"stop_reason": 128009}, + {"stop_reason": None}, + ], +) +def test_hosted_vllm_keeps_vllm_choice_fields_on_the_choice(choice_extras: dict[str, object]): + response: Final = _transform_vllm_chat_response(choice_extras) + + 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: Final = _transform_vllm_chat_response({}) + + 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 9a3516d20d6..3128f678e09 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -576,6 +576,16 @@ def test_serialize_streaming_chunk_simple_uses_fast_path_bytes(): } +def test_serialize_streaming_chunk_keeps_provider_choice_fields(): + chunk: Final = _simple_chunk(content="1,2,3,4") + chunk.choices[0]["stop_reason"] = ",5" + chunk.choices[0]["token_ids"] = [20] + + assert json.loads(_serialize_streaming_chunk(chunk))["choices"] == [ + {"index": 0, "delta": {"role": "assistant", "content": "1,2,3,4"}, "stop_reason": ",5", "token_ids": [20]} + ] + + def test_serialize_streaming_chunk_invalid_input_raises_attribute_error(): """The helper is typed as ``BaseModel`` — handing it a plain dict trips the attribute-access path (no ``model_dump_json``).""" 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 55d1587669a..7c405b63220 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -13582,7 +13582,7 @@ def _stream_usage_generator_chunks(): def test_is_injected_stream_usage_artifact(): from litellm.proxy.proxy_server import _is_injected_stream_usage_artifact - from litellm.types.utils import ModelResponseStream, Usage + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage content_chunk, finish_chunk, empty_choices_usage_chunk = _stream_usage_test_chunks() assert _is_injected_stream_usage_artifact(empty_choices_usage_chunk) is True @@ -13605,6 +13605,16 @@ def test_is_injected_stream_usage_artifact(): assert _is_injected_stream_usage_artifact({"usage": {"prompt_tokens": 1}}) is False + 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: 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 + async def _collect_async_data_generator_frames(request_data: dict) -> list: from litellm.proxy.proxy_server import async_data_generator diff --git a/tests/unit/text_completions/test_streaming_iterator.py b/tests/unit/text_completions/test_streaming_iterator.py index b7edffc0753..37d079eab26 100644 --- a/tests/unit/text_completions/test_streaming_iterator.py +++ b/tests/unit/text_completions/test_streaming_iterator.py @@ -2,8 +2,11 @@ from collections.abc import AsyncIterator, Iterator from typing import Final import pytest +from openai.types.completion import Completion from litellm.exceptions import APIConnectionError +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper from litellm.types.utils import Choices, Delta, ModelResponseStream, StreamingChoices, Usage @@ -85,6 +88,44 @@ async def test_text_completion_stream_usage_is_opt_in( ] +def _completion_chunk(text: str, finish_reason: str | None, **choice_extras: object) -> Completion: + return Completion.construct( + id="cmpl-1", + created=1, + model="m", + object="text_completion", + choices=[{"index": 0, "text": text, "finish_reason": finish_reason, "logprobs": None, **choice_extras}], + ) + + +@pytest.mark.parametrize("final_text", [", 5, 6", ""], ids=["text-then-finish", "finish-only"]) +def test_text_completion_stream_keeps_stop_reason_on_the_finish_chunk(final_text: str) -> None: + logging_obj: Final = Logging( + model="m", + messages=[{"role": "user", "content": "1, 2, 3, 4"}], + stream=True, + call_type="text_completion", + start_time=0.0, + litellm_call_id="call-1", + function_id="function-1", + ) + upstream: Final = (_completion_chunk(", 5", None), _completion_chunk(final_text, "stop", stop_reason=", 7")) + chat_stream: Final = CustomStreamWrapper( + completion_stream=iter(upstream), + model="m", + custom_llm_provider="text-completion-openai", + logging_obj=logging_obj, + ) + + responses: Final = list(TextCompletionStreamWrapper(completion_stream=chat_stream, model="m")) + + assert responses[-1].choices[0].model_dump(exclude_none=True) == { + "finish_reason": "stop", + "index": 0, + "stop_reason": ", 7", + } + + def test_text_completion_conversion_rejects_non_streaming_choices() -> None: chunk: Final = Choices() wrapper: Final = TextCompletionStreamWrapper(completion_stream=None, model="test-model")