diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 9deac7923b6..a19be338edb 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -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 diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index db0f879a08d..21193041c16 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -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])