ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #24354 from Chesars/fix/azure-streaming-role-include-usage

Browse Source 
fix: preserve role='assistant' in Azure streaming with include_usage
This commit is contained in:
Cesar Garcia 2026-03-22 11:19:53 -03:00 committed by GitHub
commit 2132db4f60
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 173 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View File

@ -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:

32
litellm/litellm_core_utils/streaming_handler.py
View File

@ -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
)

7
litellm/main.py
View File

@ -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

144
tests/test_litellm/litellm_core_utils/test_streaming_handler.py
View File

@ -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
])
)