diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 0e01577b20e..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,8 +366,7 @@ class ChunkProcessor: model: Final = ChunkProcessor._get_model_from_chunks(chunks, first_chunk_model) system_fingerprint: Final = chunk.get("system_fingerprint", None) - first_chunk_with_choices: Final = next((c for c in chunks if c.get("choices")), chunk) - role: Final = first_chunk_with_choices["choices"][0]["delta"]["role"] + 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 bacbcbf132b..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 @@ -1476,3 +1478,79 @@ def test_calculate_usage_fills_unknown_split_from_reasoning_estimate( assert usage.completion_tokens == 100 assert usage.completion_tokens_details.reasoning_tokens == expected_reasoning_tokens assert usage.completion_tokens_details.text_tokens == expected_text_tokens + + +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": "gpt-5.4-mini", + "choices": list(choices), + } + 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( + [ModelResponseStream(model="gpt-5.4-mini", choices=[]) for _ in range(2)], + id="all_empty_choices_objects", + ), + ], +) +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": "Hi"}, id="delta_without_role"), pytest.param({}, id="empty_delta")], +) +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"}]), + ] + + 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"