fix(streaming): preserve usage from trailing empty-choices chunks

When backends like vLLM send usage in a separate SSE chunk with an empty
choices array after finish_reason, chunk_creator() was returning None
and the usage data was lost.

Now, when such a chunk carries usage data, it is returned so __next__()
can accumulate it in self.chunks for calculate_total_usage(). The
downstream logic strips usage and filters the empty response, so no
extra chunk is yielded to the caller.

Fixes #25389
This commit is contained in:
lawrence3699 2026-04-09 19:09:32 +10:00
parent 97f722f558
commit a59dbc5908
No known key found for this signature in database
2 changed files with 92 additions and 0 deletions

View file

@ -1561,6 +1561,12 @@ class CustomStreamWrapper:
and self.stream_options["include_usage"] is True
):
return model_response
# Still return model_response when it carries usage data
# so trailing usage chunks (e.g. from vLLM) get accumulated
# in self.chunks for calculate_total_usage(). The __next__
# loop will strip the usage field and skip empty responses.
if getattr(model_response, "usage", None) is not None:
return model_response
return
## CHECK FOR TOOL USE

View file

@ -1543,6 +1543,92 @@ def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj):
f"Expected completion_tokens=135 from provider, got {hidden_usage.completion_tokens}"
)
def test_usage_chunk_empty_choices_vllm_pattern(logging_obj):
"""
Test that usage data from a trailing chunk with an empty choices array
(as sent by vLLM) is preserved in _hidden_params.
Reproduces issue #25389: vLLM sends usage in a separate SSE chunk after
finish_reason with choices=[]. Without the fix, chunk_creator() returns
None for this chunk and the usage is lost.
"""
# Simulate vLLM's streaming pattern:
# 1) content chunk
# 2) finish_reason chunk
# 3) usage-only chunk with empty choices array
chunks = [
ModelResponseStream(
id="chatcmpl-abc",
object="chat.completion.chunk",
created=1000000,
model="vllm/qwen",
choices=[
StreamingChoices(
index=0,
delta=Delta(role="assistant", content="Hello"),
finish_reason=None,
)
],
),
ModelResponseStream(
id="chatcmpl-abc",
object="chat.completion.chunk",
created=1000000,
model="vllm/qwen",
choices=[
StreamingChoices(
index=0,
delta=Delta(content=""),
finish_reason="stop",
)
],
),
# vLLM sends usage in a chunk with an EMPTY choices array
ModelResponseStream(
id="chatcmpl-abc",
object="chat.completion.chunk",
created=1000000,
model="vllm/qwen",
choices=[],
usage=Usage(
prompt_tokens=26,
completion_tokens=242,
total_tokens=268,
),
),
]
wrapper = CustomStreamWrapper(
completion_stream=ModelResponseListIterator(model_responses=chunks),
model="vllm/qwen",
logging_obj=logging_obj,
custom_llm_provider="openai",
stream_options=None,
)
collected = []
for chunk in wrapper:
collected.append(chunk)
assert len(collected) > 0, "Expected at least one chunk"
# The empty-choices usage chunk must NOT leak as a visible chunk
for c in collected:
assert len(c.choices) > 0, (
"Empty-choices usage chunk should not be yielded to the caller"
)
last_chunk = collected[-1]
hidden_usage = last_chunk._hidden_params.get("usage")
assert hidden_usage is not None, "Expected usage in _hidden_params"
assert hidden_usage.prompt_tokens == 26, (
f"Expected prompt_tokens=26, got {hidden_usage.prompt_tokens}"
)
assert hidden_usage.completion_tokens == 242, (
f"Expected completion_tokens=242, got {hidden_usage.completion_tokens}"
)
@pytest.mark.asyncio
async def test_custom_stream_wrapper_aclose():
"""Test that aclose() delegates to the underlying completion_stream's aclose()"""