diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 1935372e5d..1e13fb2dcf 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -119,7 +119,10 @@ class ChunkProcessor: model = ChunkProcessor._get_model_from_chunks(chunks, first_chunk_model) system_fingerprint = chunk.get("system_fingerprint", None) - role = chunk["choices"][0]["delta"]["role"] + first_chunk_with_choices = next( + (c for c in chunks if c.get("choices")), chunk + ) + role = first_chunk_with_choices["choices"][0]["delta"]["role"] finish_reason = "stop" for chunk in chunks: if "choices" in chunk and len(chunk["choices"]) > 0: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 67e4fadf63..d9de363d8c 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -831,6 +831,11 @@ class CustomStreamWrapper: "annotations" in model_response.choices[0].delta and model_response.choices[0].delta.annotations is not None ) + or ( + not self.sent_first_chunk + and hasattr(model_response.choices[0].delta, "role") + and model_response.choices[0].delta.role is not None + ) ): return True else: @@ -1556,6 +1561,7 @@ class CustomStreamWrapper: self.stream_options is not None and self.stream_options["include_usage"] is True ): + model_response.choices = [] return model_response return ## CHECK FOR TOOL USE @@ -1855,11 +1861,12 @@ class CustomStreamWrapper: response, cache_hit, ) # log response - choice = response.choices[0] - if isinstance(choice, StreamingChoices): - self.response_uptil_now += choice.delta.get("content", "") or "" - else: - self.response_uptil_now += "" + if response.choices: + choice = response.choices[0] + if isinstance(choice, StreamingChoices): + self.response_uptil_now += choice.delta.get("content", "") or "" + else: + self.response_uptil_now += "" self.rules.post_call_rules( input=self.response_uptil_now, model=self.model ) @@ -1867,7 +1874,7 @@ class CustomStreamWrapper: self.chunks.append(response) # Add mcp_list_tools to first chunk if present - if not self.sent_first_chunk: + if not self.sent_first_chunk and response.choices: response = self._add_mcp_list_tools_to_first_chunk(response) self.sent_first_chunk = True @@ -2035,16 +2042,17 @@ class CustomStreamWrapper: completion_start_time=datetime.datetime.now() ) - choice = processed_chunk.choices[0] - if isinstance(choice, StreamingChoices): - self.response_uptil_now += choice.delta.get("content", "") or "" - else: - self.response_uptil_now += "" + if processed_chunk.choices: + choice = processed_chunk.choices[0] + if isinstance(choice, StreamingChoices): + self.response_uptil_now += choice.delta.get("content", "") or "" + else: + self.response_uptil_now += "" self.rules.post_call_rules( input=self.response_uptil_now, model=self.model ) # Add mcp_list_tools to first chunk if present - if not self.sent_first_chunk: + if not self.sent_first_chunk and processed_chunk.choices: processed_chunk = self._add_mcp_list_tools_to_first_chunk( processed_chunk ) diff --git a/litellm/main.py b/litellm/main.py index eace9c630b..aaffc15384 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7370,8 +7370,11 @@ def stream_chunk_builder( # noqa: PLR0915 if len(chunks) == 0: return None ## Route to the text completion logic - if isinstance( - chunks[0]["choices"][0], litellm.utils.TextChoices + first_chunk_with_choices = next( + (c for c in chunks if c["choices"]), None + ) + if first_chunk_with_choices is not None and isinstance( + first_chunk_with_choices["choices"][0], litellm.utils.TextChoices ): # route to the text completion logic return stream_chunk_builder_text_completion( chunks=chunks, messages=messages diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 20e064ef8f..12bf51d3dd 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1826,3 +1826,147 @@ async def test_custom_stream_wrapper_anext_exhaustion_raises_stop_async_iteratio pass # expected clean termination except RuntimeError as e: pytest.fail(f"PEP 479 regression: StopIteration leaked as RuntimeError: {e}") + + +# Azure streaming chunks that reproduce issue #24221: +# Azure sends an initial chunk with prompt_filter_results and choices=[], +# then a chunk with role='assistant' and content='', then content chunks. +# With stream_options.include_usage=True, the empty-choices chunk was +# forwarded with an inflated default choice, consuming the sent_first_chunk +# flag and causing strip_role_from_delta to strip the role from the real +# first chunk. +_AZURE_CHUNKS_WITH_PROMPT_FILTER = [ + # Chunk 1: prompt_filter_results, no choices (Azure-specific) + ModelResponseStream( + id="chatcmpl-abc123", + created=1742056047, + model=None, + object="chat.completion.chunk", + choices=[], + usage=None, + ), + # Chunk 2: first real chunk with role='assistant' and empty content + ModelResponseStream( + id="chatcmpl-abc123", + created=1742056047, + model=None, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="", role="assistant"), + ) + ], + usage=None, + ), + # Chunk 3: content + ModelResponseStream( + id="chatcmpl-abc123", + created=1742056047, + model=None, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="Hello!"), + ) + ], + usage=None, + ), + # Chunk 4: finish_reason + ModelResponseStream( + id="chatcmpl-abc123", + created=1742056047, + model=None, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(), + ) + ], + usage=None, + ), + # Chunk 5: final usage chunk, no choices + ModelResponseStream( + id="chatcmpl-abc123", + created=1742056047, + model=None, + object="chat.completion.chunk", + choices=[], + usage=Usage( + completion_tokens=10, + prompt_tokens=20, + total_tokens=30, + ), + ), +] + + +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_azure_streaming_role_preserved_with_include_usage(sync_mode: bool): + """ + Regression test for https://github.com/BerriAI/litellm/issues/24221 + + Azure sends an initial chunk with choices=[] (prompt_filter_results) + before the first content chunk. With stream_options.include_usage=True, + this chunk was forwarded with an inflated default choice, which: + 1. Consumed the sent_first_chunk flag + 2. Caused strip_role_from_delta to strip role from the real first chunk + + The fix ensures: + - Chunks with choices=[] are forwarded faithfully (no inflated choices) + - sent_first_chunk is only marked for chunks with real choices + - Chunks with role in delta are not discarded as empty + """ + completion_stream = ModelResponseListIterator( + model_responses=_AZURE_CHUNKS_WITH_PROMPT_FILTER + ) + + response = CustomStreamWrapper( + completion_stream=completion_stream, + model="azure/gpt-5-nano", + custom_llm_provider="azure", + logging_obj=Logging( + model="azure/gpt-5-nano", + messages=[{"role": "user", "content": "Hey"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="12345", + function_id="1245", + ), + stream_options={"include_usage": True}, + ) + + chunks = [] + if sync_mode: + for chunk in response: + chunks.append(chunk) + else: + async for chunk in response: + chunks.append(chunk) + + # The prompt_filter chunk should be forwarded with choices=[] + assert len(chunks[0].choices) == 0, ( + f"Expected prompt_filter chunk with choices=[], got {len(chunks[0].choices)} choices" + ) + + # At least one chunk must have role='assistant' in its delta + has_role = any( + len(c.choices) > 0 + and getattr(c.choices[0].delta, "role", None) == "assistant" + for c in chunks + ) + assert has_role, ( + "No chunk contained role='assistant' in delta (issue #24221). " + "Chunk deltas: " + + str([ + c.choices[0].delta if c.choices else "no choices" + for c in chunks + ]) + )