From 1d7e81cf5d3a29dd4731b3282cf0842aac854ea1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 17:47:06 -0700 Subject: [PATCH] fix(streaming): guard empty choices and missing role when assembling stream chunks --- .../streaming_chunk_builder_utils.py | 23 ++-- .../test_streaming_chunk_builder_utils.py | 129 +++++++----------- 2 files changed, 66 insertions(+), 86 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 81955fe769e..cf9604a0fd5 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -24,6 +24,7 @@ from litellm.types.utils import ( Choices, CompletionTokensDetails, CompletionTokensDetailsWrapper, + Delta, Function, FunctionCall, ModelResponse, @@ -326,6 +327,18 @@ class ChunkProcessor: return chunk_id return "" + @staticmethod + def _get_role_from_chunks(chunks: Sequence["_BaseChunk"]) -> str: + return ChunkProcessor._role_of_choice(next((c["choices"][0] for c in chunks if c.get("choices")), None)) + + @staticmethod + def _role_of_choice(choice: object) -> str: + match choice: + case StreamingChoices(delta=Delta(role=str() as role)) | {"delta": {"role": str() as role}} if role: + return role + case _: + return "assistant" + @staticmethod def _get_model_from_chunks(chunks: Sequence["_BaseChunk"], first_chunk_model: str) -> str: """ @@ -353,15 +366,7 @@ class ChunkProcessor: model: Final = ChunkProcessor._get_model_from_chunks(chunks, first_chunk_model) system_fingerprint: Final = chunk.get("system_fingerprint", None) - # Fall back to None rather than `chunk`: if no chunk carries a non-empty - # `choices` array, indexing [0] on the first chunk raises IndexError. - first_chunk_with_choices = next((c for c in chunks if c.get("choices")), None) - role: str = "assistant" - if first_chunk_with_choices is not None: - _choices = first_chunk_with_choices["choices"] - if len(_choices) > 0: - # `delta` may be absent or omit `role` (e.g. content-only deltas). - role = _choices[0].get("delta", {}).get("role") or "assistant" + role: Final = ChunkProcessor._get_role_from_chunks(chunks) finish_reason = "stop" for chunk in chunks: if "choices" in chunk and len(chunk["choices"]) > 0: diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 2d2451e73f7..626b8a63b20 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1,4 +1,6 @@ import json +from collections.abc import Mapping, Sequence +from typing import Final import pytest @@ -1478,104 +1480,77 @@ def test_calculate_usage_fills_unknown_split_from_reasoning_estimate( assert usage.completion_tokens_details.text_tokens == expected_text_tokens -def _empty_choices_chunk(**extra): - chunk = { - "id": "chatcmpl-empty-choices", +def _openai_chunk( + choices: Sequence[Mapping[str, object]], usage: Mapping[str, int] | None = None +) -> dict[str, object]: + base: Final = { + "id": "chatcmpl-lit6552", "object": "chat.completion.chunk", "created": 1, - "model": "claude-opus-4-8", - "choices": [], + "model": "gpt-5.4-mini", + "choices": list(choices), } - chunk.update(extra) - return chunk + return base if usage is None else {**base, "usage": dict(usage)} @pytest.mark.parametrize( "chunks", [ + pytest.param([_openai_chunk(choices=[]), _openai_chunk(choices=[])], id="all_empty_choices_dicts"), pytest.param( - [_empty_choices_chunk(), _empty_choices_chunk()], - id="all_chunks_have_empty_choices", - ), - pytest.param( - [ - _empty_choices_chunk(usage={"prompt_tokens": 10}), - _empty_choices_chunk(usage={"completion_tokens": 0}), - ], - id="usage_only_chunks", + [ModelResponseStream(model="gpt-5.4-mini", choices=[]) for _ in range(2)], + id="all_empty_choices_objects", ), ], ) -def test_build_base_response_handles_empty_choices(chunks): - """Empty `choices` arrays must not raise IndexError. - - `next((c for c in chunks if c.get("choices")), chunk)` used to fall back to the - first chunk, whose `choices` may be `[]`, so `["choices"][0]` went out of range. - The resulting error is surfaced to the client mid-stream and the request never - reaches SpendLogs. - """ - processor = ChunkProcessor(chunks=list(chunks)) - - response = processor.build_base_response(list(chunks)) +def test_stream_chunk_builder_survives_all_empty_choices(chunks: Sequence[object]) -> None: + response: Final = stream_chunk_builder(chunks=list(chunks)) + assert response is not None assert response.choices[0].message.role == "assistant" + assert response.choices[0].finish_reason == "stop" + + +def test_stream_chunk_builder_keeps_usage_from_usage_only_frames() -> None: + usage_frame: Final = _openai_chunk( + choices=[], usage={"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10} + ) + + response: Final = stream_chunk_builder(chunks=[usage_frame]) + + assert response is not None + assert response.choices[0].message.role == "assistant" + assert response.usage.prompt_tokens == 10 + assert response.usage.total_tokens == 10 @pytest.mark.parametrize( "delta", - [ - pytest.param({"content": "Hello"}, id="delta_without_role"), - pytest.param({}, id="delta_empty_dict"), - ], + [pytest.param({"content": "Hi"}, id="delta_without_role"), pytest.param({}, id="empty_delta")], ) -def test_build_base_response_handles_delta_without_role(delta): - """A `delta` that omits `role` must not raise KeyError.""" - chunks = [ - { - "id": "chatcmpl-no-role", - "object": "chat.completion.chunk", - "created": 1, - "model": "claude-opus-4-8", - "choices": [{"index": 0, "delta": delta, "finish_reason": None}], - } +def test_stream_chunk_builder_defaults_role_when_delta_omits_it(delta: Mapping[str, str]) -> None: + chunks: Final = [ + _openai_chunk(choices=[{"index": 0, "delta": dict(delta), "finish_reason": None}]), + _openai_chunk(choices=[{"index": 0, "delta": {"content": "!"}, "finish_reason": "stop"}]), ] - processor = ChunkProcessor(chunks=list(chunks)) - response = processor.build_base_response(list(chunks)) - - assert response.choices[0].message.role == "assistant" - - -def test_build_base_response_still_reads_role_and_finish_reason(): - """Regression guard: well-formed chunks keep their role and finish_reason.""" - chunks = [ - _empty_choices_chunk(), - { - "id": "chatcmpl-normal", - "object": "chat.completion.chunk", - "created": 1, - "model": "claude-opus-4-8", - "choices": [ - { - "index": 0, - "delta": {"role": "assistant", "content": "Hi"}, - "finish_reason": None, - } - ], - }, - { - "id": "chatcmpl-normal", - "object": "chat.completion.chunk", - "created": 2, - "model": "claude-opus-4-8", - "choices": [ - {"index": 0, "delta": {"content": "!"}, "finish_reason": "stop"} - ], - }, - ] - processor = ChunkProcessor(chunks=list(chunks)) - - response = processor.build_base_response(list(chunks)) + response: Final = stream_chunk_builder(chunks=chunks) + assert response is not None assert response.choices[0].message.role == "assistant" + assert response.choices[0].message.content == delta.get("content", "") + "!" assert response.choices[0].finish_reason == "stop" + + +def test_stream_chunk_builder_reads_role_from_first_frame_with_choices() -> None: + chunks: Final = [ + _openai_chunk(choices=[]), + _openai_chunk(choices=[{"index": 0, "delta": {"role": "user", "content": "Hi"}, "finish_reason": None}]), + _openai_chunk(choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}]), + ] + + response: Final = stream_chunk_builder(chunks=chunks) + + assert response is not None + assert response.choices[0].message.role == "user" + assert response.choices[0].message.content == "Hi"