scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

fix: ensure role='assistant' in Azure streaming with include_usage

Browse source 
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>
This commit is contained in:
majiayu000 2026-04-14 00:08:14 +08:00
parent 0eae9f101e
commit f46fdbcb1e
No known key found for this signature in database
GPG key ID: 29C2AE46ADC3B14B
2 changed files with 50 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/litellm_core_utils/streaming_handler.py
View file

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

42
tests/test_litellm/litellm_core_utils/test_streaming_handler.py
View file

@ -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'."
)