diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f7bd86f3a0c..361e576f3e1 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -82,6 +82,8 @@ FUNCTION_CALL_ATTRIBUTE: Final = "function_call" _SYNC_ITER_EXHAUSTED: Final = object() +_STREAMING_CHOICE_KEYS: Final = frozenset({"finish_reason", "index", "delta", "logprobs", "message"}) + _GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__) _CompleteStreamingResponse = ModelResponse | TextCompletionResponse | ModelResponseStream | None _USAGE_COST_HEADER_PROVIDERS: Final[frozenset[str]] = frozenset( @@ -139,11 +141,18 @@ def _entry_len(entries: Mapping[str, object], key: str) -> int: def _prepend_upstream_choice_extras( - choice: StreamingChoices, upstream_choice: ChunkChoice, extra_keys: frozenset[str] + choice: StreamingChoices, upstream_choice: ChunkChoice | StreamingChoices, extra_keys: frozenset[str] ) -> None: if not extra_keys: return - upstream_extras: Final[Mapping[str, object]] = upstream_choice.model_dump(include=set(extra_keys)) + stored_extras: Final = { + key: value for key, value in (upstream_choice.__pydantic_extra__ or {}).items() if key in extra_keys + } + upstream_extras: Final[Mapping[str, object]] = ( + stored_extras + if _holds_only_scalars(stored_extras.items()) + else upstream_choice.model_dump(include=set(extra_keys)) + ) choice.__pydantic_extra__ = {**upstream_extras, **(choice.__pydantic_extra__ or {})} # rebind-ok: reuse choice choice.__pydantic_fields_set__.update(upstream_extras) @@ -439,7 +448,7 @@ class CustomStreamWrapper: self.completion_stream = completion_stream self.sent_first_chunk = False self.sent_last_chunk = False - self._openai_handler_receives_chat_chunks = False + self._openai_handler_receives_chunk_objects = False self._failure_logging: asyncio.Task[None] | None = None self._closed = False self._stream_created_time: float = time.time() @@ -1162,13 +1171,13 @@ class CustomStreamWrapper: def _prebuilt_plain_delta(self, chunk: object) -> Delta | None: if ( - not self._openai_handler_receives_chat_chunks - or not isinstance(chunk, ChatCompletionChunk) + not self._openai_handler_receives_chunk_objects + or not isinstance(chunk, (ChatCompletionChunk, ModelResponseStream)) or not chunk.choices ): return None upstream_delta: Final = cast( # cast-ok: [LIT006] runtime delta may be absent - ChoiceDelta | None, + ChoiceDelta | Delta | None, chunk.choices[0].delta, ) if ( @@ -1296,22 +1305,23 @@ class CustomStreamWrapper: upstream_choice: Final = ( original_chunk_for_reuse.choices[0] if plain_delta_created - and isinstance(original_chunk_for_reuse, ChatCompletionChunk) + and isinstance(original_chunk_for_reuse, (ChatCompletionChunk, ModelResponseStream)) and len(original_chunk_for_reuse.choices) == 1 else None ) upstream_delta: Final = cast( # cast-ok: [LIT006] runtime delta may be absent - ChoiceDelta | None, + ChoiceDelta | Delta | None, getattr(upstream_choice, "delta", None), ) upstream_choice_extra_keys: Final = frozenset( (upstream_choice.__pydantic_extra__ or {}) if upstream_choice is not None else () ) + upstream_logprobs: Final[object] = getattr(upstream_choice, "logprobs", None) can_reuse_choice: Final = ( upstream_choice is not None and upstream_delta is not None and "enhancements" not in upstream_choice_extra_keys - and upstream_choice.logprobs is None + and upstream_logprobs is None and _holds_only_scalars(upstream_delta) ) if can_reuse_choice and upstream_choice is not None: @@ -1320,7 +1330,7 @@ class CustomStreamWrapper: _prepend_upstream_choice_extras( model_response.choices[0], upstream_choice, - upstream_choice_extra_keys - {"message"}, + upstream_choice_extra_keys - _STREAMING_CHOICE_KEYS, ) else: choices: Final = [] @@ -1737,14 +1747,16 @@ class CustomStreamWrapper: self._record_finish_choice_fields(cached_chunk) else: # openai / azure chat model self.uses_openai_chat_stream = True - if isinstance(chunk, ChatCompletionChunk): - self._openai_handler_receives_chat_chunks = True - if any(str(choice.finish_reason) == "abort" for choice in chunk.choices or ()): - raise litellm.InternalServerError( - message="Upstream generation aborted (finish_reason=abort)", - model=chunk.model, - llm_provider=self.custom_llm_provider, - ) + if isinstance(chunk, (ChatCompletionChunk, ModelResponseStream)): + self._openai_handler_receives_chunk_objects = True + if isinstance(chunk, ChatCompletionChunk) and any( + str(choice.finish_reason) == "abort" for choice in chunk.choices or () + ): + raise litellm.InternalServerError( + message="Upstream generation aborted (finish_reason=abort)", + model=chunk.model, + llm_provider=self.custom_llm_provider, + ) if self.custom_llm_provider in [ LlmProviders.AZURE.value, LlmProviders.AZURE_AI.value, diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index df50ae9dc44..3b798b35096 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -7427,6 +7427,121 @@ def test_azure_chunk_creator_streams_choice_content_filter_results_on_later_chun assert ["content_filter_results" in output.choices[0].model_fields_set for output in outputs] == [True] +def _parsed_openai_like_chunk( + chunk_id: str, delta: Mapping[str, object], finish_reason: str | None = None, **choice_extras: object +) -> ModelResponseStream: + parser: Final = OpenAIChatCompletionStreamingHandler(streaming_response=None, sync_stream=True) + return parser.chunk_parser( + { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": 1, + "model": "served-model", + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason, **choice_extras}], + } + ) + + +def _parsed_stream_wrapper(custom_llm_provider: str, logging_obj: Logging) -> CustomStreamWrapper: + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(()), + model="served-model", + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ) + wrapper.chunk_creator(chunk=_parsed_openai_like_chunk("chunk-0", {"role": "assistant", "content": ""})) + return wrapper + + +@pytest.mark.parametrize("custom_llm_provider", ["hosted_vllm", "deepseek", "groq"]) +def test_parsed_openai_like_chunk_creator_builds_one_delta_and_choice_and_no_dump_per_later_chunk( + logging_obj: Logging, custom_llm_provider: str +) -> None: + wrapper: Final = _parsed_stream_wrapper(custom_llm_provider, logging_obj) + later_chunks: Final = ( + _parsed_openai_like_chunk("chunk-1", {"content": "hello"}, stop_reason=None), + _parsed_openai_like_chunk("chunk-2", {"reasoning": "thought"}, stop_reason=None), + _parsed_openai_like_chunk("chunk-3", {"content": " world", "reasoning_content": None}, logprobs=None), + ) + delta_init: Final = Delta.__init__.__code__ + streaming_choices_init: Final = StreamingChoices.__init__.__code__ + model_dump: Final = BaseModel.model_dump.__code__ + + outputs, calls = _chunk_creator_calls( + wrapper, later_chunks, frozenset({delta_init, streaming_choices_init, model_dump}) + ) + + assert sum(call.code is delta_init for call in calls) == len(later_chunks) + assert sum(call.code is streaming_choices_init for call in calls) == len(later_chunks) + assert sum(call.code is model_dump for call in calls) == 0 + assert [_wire_dump(output)["choices"] for output in outputs if output is not None] == [ + [{"index": 0, "delta": {"content": "hello"}}], + [{"index": 0, "delta": {"reasoning_content": "thought"}}], + [{"index": 0, "delta": {"content": " world"}}], + ] + + +def test_parsed_openai_like_chunk_creator_keeps_choice_extras_ahead_of_choice_keys(logging_obj: Logging) -> None: + wrapper: Final = _parsed_stream_wrapper("hosted_vllm", logging_obj) + later_chunks: Final = ( + _parsed_openai_like_chunk("chunk-1", {"content": "hi"}, stop_reason=None, token_ids=[7, 8]), + _parsed_openai_like_chunk("chunk-2", {"content": "!"}, stop_reason=None, matched_stop=2), + ) + + outputs: Final = tuple(wrapper.chunk_creator(chunk=chunk) for chunk in later_chunks) + + choice_dumps: Final = [output.choices[0].model_dump() for output in outputs if output is not None] + assert [list(choice_dump) for choice_dump in choice_dumps] == [ + ["stop_reason", "token_ids", "finish_reason", "index", "delta", "logprobs"], + ["stop_reason", "matched_stop", "finish_reason", "index", "delta", "logprobs"], + ] + assert [choice_dump["token_ids"] for choice_dump in choice_dumps[:1]] == [[7, 8]] + assert [choice_dump["matched_stop"] for choice_dump in choice_dumps[1:]] == [2] + assert [choice_dump["delta"]["content"] for choice_dump in choice_dumps] == ["hi", "!"] + + +def test_parsed_openai_like_stream_still_emits_tool_calls_and_finish_after_reused_chunks( + logging_obj: Logging, +) -> None: + wrapper: Final = _parsed_stream_wrapper("deepseek", logging_obj) + later_chunks: Final = ( + _parsed_openai_like_chunk("chunk-1", {"content": "checking"}), + _parsed_openai_like_chunk( + "chunk-2", + { + "tool_calls": [ + {"index": 0, "id": "call_1", "type": "function", "function": {"name": "f", "arguments": ""}} + ] + }, + ), + _parsed_openai_like_chunk("chunk-3", {"tool_calls": [{"index": 0, "function": {"arguments": "{}"}}]}), + _parsed_openai_like_chunk("chunk-4", {}, finish_reason="tool_calls"), + ) + + outputs: Final = tuple(wrapper.chunk_creator(chunk=chunk) for chunk in later_chunks) + + assert [_wire_dump(output)["choices"] for output in outputs if output is not None] == [ + [{"index": 0, "delta": {"content": "checking"}}], + [ + { + "index": 0, + "delta": { + "tool_calls": [ + {"id": "call_1", "function": {"arguments": "", "name": "f"}, "type": "function", "index": 0} + ] + }, + } + ], + [ + { + "index": 0, + "delta": {"tool_calls": [{"function": {"arguments": "{}"}, "type": "function", "index": 0}]}, + } + ], + [{"finish_reason": "tool_calls", "index": 0, "delta": {}}], + ] + + @pytest.mark.parametrize( "special_delta", [