scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

fix(L6): Carry cumulative usage onto the synthesized finish chunk

Browse source 
Signed-off-by: Alberto Perdomo <aperdomo@redhat.com>
This commit is contained in:
Alberto Perdomo 2026-10-26 17:19:57 +00:00
parent 4c97a0471d
commit 92dd8b6792
No known key found for this signature in database
2 changed files with 8 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/litellm_core_utils/streaming_handler.py
View file

@ -2010,6 +2010,8 @@ class CustomStreamWrapper:
if self.stream_options is None: # add usage as hidden param
usage = calculate_total_usage(chunks=self.chunks)
processed_chunk.hidden_params["usage"] = usage
elif self._continuous_usage_stats:
processed_chunk.usage = calculate_total_usage(chunks=self.chunks)
## LOGGING
executor.submit(
self.run_success_logging_and_cache_storage,
@ -2289,6 +2291,8 @@ class CustomStreamWrapper:
dict[str, object], getattr(processed_chunk, HIDDEN_PARAMS_ATTR)
)
processed_chunk_hidden_params["usage"] = usage
elif self._continuous_usage_stats:
processed_chunk.usage = calculate_total_usage(chunks=self.chunks)
# see sync __next__'s sibling branch: deliberately do NOT restore
# here - this chunk is still this call's own data, and restoring
# before returning it would corrupt the caller's own log

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

@ -6774,7 +6774,10 @@ async def test_continuous_usage_stats_finishes_before_the_final_usage_chunk(sync
assert finish_indexes[0] < len(emitted) - 1
terminal_chunk: Final = emitted[-1]
assert not (terminal_chunk.choices and terminal_chunk.choices[0].finish_reason)
assert terminal_chunk.usage.total_tokens == final_usage["total_tokens"]
assert [chunk.usage.total_tokens for chunk in emitted] == [
content_usage["total_tokens"],
*[final_usage["total_tokens"]] * 3,
]
@pytest.mark.parametrize("sync_mode", [True, False])