From f46fdbcb1efd7f521a95252946286abba2abe664 Mon Sep 17 00:00:00 2001
From: majiayu000 <1835304752@qq.com>
Date: Tue, 14 Apr 2026 00:08:14 +0800
Subject: [PATCH] fix: ensure role='assistant' in Azure streaming with
 include_usage
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When Azure sends an initial prompt_filter_results chunk with empty
choices and stream_options.include_usage=True, the empty-choices
early-return path was returning model_response to the caller. This
caused the outer __anext__/__next__ loop to set sent_first_chunk=True,
so when the real first content chunk arrived, strip_role_from_delta()
stripped its role instead of keeping it — leaving no role='assistant'
in the entire stream.

Fix: only return the empty-choices chunk when it actually carries usage
data (the final usage-only chunk). Skip it otherwise, matching the
behavior of the content-processing path (line 978) which already
returns None for empty choices.

Fixes #24221

Signed-off-by: majiayu000 <1835304752@qq.com>
---
 .../litellm_core_utils/streaming_handler.py   | 14 ++++---
 .../test_streaming_handler.py                 | 42 +++++++++++++++++++
 2 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py
index dee3e2dfb4c..156015b5526 100644
--- a/litellm/litellm_core_utils/streaming_handler.py
+++ b/litellm/litellm_core_utils/streaming_handler.py
@@ -127,9 +127,9 @@ class CustomStreamWrapper:
 
         self.system_fingerprint: Optional[str] = None
         self.received_finish_reason: Optional[str] = None
-        self.intermittent_finish_reason: Optional[
-            str
-        ] = None  # finish reasons that show up mid-stream
+        self.intermittent_finish_reason: Optional[str] = (
+            None  # finish reasons that show up mid-stream
+        )
         self.special_tokens = [
             "<|assistant|>",
             "<|system|>",
@@ -1524,9 +1524,9 @@ class CustomStreamWrapper:
                                                 t.function.arguments = ""
                             _json_delta = delta.model_dump()
                             if "role" not in _json_delta or _json_delta["role"] is None:
-                                _json_delta[
-                                    "role"
-                                ] = "assistant"  # mistral's api returns role as None
+                                _json_delta["role"] = (
+                                    "assistant"  # mistral's api returns role as None
+                                )
                             if "tool_calls" in _json_delta and isinstance(
                                 _json_delta["tool_calls"], list
                             ):
@@ -1563,6 +1563,8 @@ class CustomStreamWrapper:
                     if (
                         self.stream_options is not None
                         and self.stream_options["include_usage"] is True
+                        and hasattr(model_response, "usage")
+                        and model_response.usage is not None
                     ):
                         return model_response
                     return
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 47a77c110b0..0db34d56958 100644
--- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
+++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
@@ -1946,3 +1946,45 @@ def test_gemini_legacy_vertex_tool_calls_finish_reason_with_stop_enum():
         f"Expected 'tool_calls' but got {final.choices[0].finish_reason!r}. "
         "STOP enum was not normalised through map_finish_reason()."
     )
+
+
+def test_azure_empty_choices_with_include_usage_preserves_role():
+    """Regression test for #24221: Azure streaming with include_usage drops role='assistant'.
+
+    When Azure sends an initial prompt_filter_results chunk with empty choices
+    and stream_options.include_usage=True, the empty-choices early-return path
+    must not cause sent_first_chunk to become True. Otherwise the real first
+    content chunk will have its role stripped.
+    """
+    wrapper = CustomStreamWrapper(
+        completion_stream=None,
+        model="azure/gpt-4",
+        logging_obj=MagicMock(),
+        custom_llm_provider="azure",
+        stream_options={"include_usage": True},
+    )
+
+    # Azure's initial empty-choices chunk (prompt_filter_results)
+    empty_chunk = ModelResponseStream(
+        id="chatcmpl-abc123",
+        created=1700000000,
+        model="gpt-4",
+        object="chat.completion.chunk",
+        choices=[],
+        usage=None,
+    )
+
+    assert wrapper.sent_first_chunk is False
+
+    result = wrapper.chunk_creator(chunk=empty_chunk)
+
+    # Empty-choices chunk without usage should be skipped (return None)
+    assert (
+        result is None
+    ), "Empty choices chunk without usage data should not be returned"
+
+    # sent_first_chunk must still be False
+    assert wrapper.sent_first_chunk is False, (
+        "sent_first_chunk should remain False after an empty-choices chunk. "
+        "If it becomes True, the next real content chunk will lose role='assistant'."
+    )