scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-26 06:03:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

fix(streaming): retain usage updates with empty choices

Browse source
This commit is contained in:
AkilanChithra 2026-10-12 13:53:33 -05:00
parent f9b9251f78
commit 8130a0572e
2 changed files with 33 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/litellm_core_utils/streaming_handler.py
View file

@ -1082,6 +1082,8 @@ class CustomStreamWrapper:
return self._handle_special_delta_content(model_response)
else:
if hasattr(model_response, "usage"):
if self._should_forward_continuous_usage() and model_response.usage is not None:
return model_response
self.chunks.append(model_response)
return

31
tests/unit/litellm_core_utils/test_streaming_handler.py
View file

@ -579,6 +579,37 @@ async def test_continuous_usage_only_chunk_is_visible(logging_obj: Logging, sync
assert len(collected) == 5
@pytest.mark.parametrize("sync_mode", (True, False))
@pytest.mark.parametrize("content", (None, ""))
@pytest.mark.parametrize("completion_tokens", (0, 1))
@pytest.mark.parametrize("continuous", (True, False))
@pytest.mark.asyncio
async def test_continuous_usage_empty_choice_is_visible(
logging_obj: Logging, sync_mode: bool, content: str | None, completion_tokens: int, continuous: bool
) -> None:
chunks: Final = _continuous_usage_fixture_chunks()
progress: Final = chunks[0].model_copy(
update={
"choices": [StreamingChoices(index=0, delta=Delta(content=content))],
"usage": Usage(prompt_tokens=10, completion_tokens=completion_tokens, total_tokens=10 + completion_tokens),
}
)
wrapper: Final = _continuous_usage_stream(
logging_obj, (chunks[0], progress, *chunks[1:]), {"include_usage": True, "continuous_usage_stats": continuous}
)
collected: Final = await _continuous_usage_collected(wrapper, sync_mode)
assert len(collected) == (5 if continuous else 4)
if continuous:
assert collected[1].choices == [StreamingChoices(index=0, delta=Delta(content=content))]
assert getattr(collected[1], "usage", None) == Usage(
prompt_tokens=10, completion_tokens=completion_tokens, total_tokens=10 + completion_tokens
)
else:
assert collected[1].choices[0].delta.content == "b"
assert getattr(collected[1], "usage", None) is None
assert collected[-1].usage.total_tokens == 12
@pytest.mark.parametrize("sync_mode", (True, False))
@pytest.mark.asyncio
async def test_continuous_usage_keeps_final_trailer_contract(logging_obj: Logging, sync_mode: bool) -> None: