diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6ec426d6da5..6b467cf1db7 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -507,6 +507,14 @@ class CustomStreamWrapper: def check_send_stream_usage(self, stream_options: dict | None): return stream_options is not None and stream_options.get("include_usage", False) is True + def _should_forward_continuous_usage(self) -> bool: + options: Final[Mapping[str, object] | None] = getattr(self, "stream_options", None) + return ( + self.send_stream_usage is True + and options is not None + and options.get("continuous_usage_stats", False) is True + ) + def check_is_function_call(self, logging_obj) -> bool: from litellm.litellm_core_utils.prompt_templates.common_utils import ( is_function_call, @@ -1236,6 +1244,8 @@ class CustomStreamWrapper: return self._handle_special_delta_content(model_response) else: if hasattr(model_response, "usage"): + if self._should_forward_continuous_usage() and getattr(model_response, "usage", None) is not None: + return model_response self.chunks.append(model_response) return @@ -1570,6 +1580,11 @@ class CustomStreamWrapper: return _ProviderChunkParsed(response_obj) def chunk_creator(self, chunk: Any): + finish_reason_seen: Final = ( + self.sent_last_chunk is True + or self.received_finish_reason is not None + or self.intermittent_finish_reason is not None + ) if self.custom_llm_provider != "cached_response" and isinstance( chunk, (ChatCompletionChunk, ModelResponseStream) ): @@ -1666,7 +1681,9 @@ class CustomStreamWrapper: model_response.choices[0].delta = plain_delta if plain_delta is not None else Delta() plain_delta_created = plain_delta is not None # rebind-ok: the first chunk builds its delta else: - if self.send_stream_usage is True: + if self.send_stream_usage is True and ( + not self._should_forward_continuous_usage() or not finish_reason_seen + ): model_response.choices = [] return model_response self._record_usage_only_chunk(model_response=model_response) @@ -2045,7 +2062,7 @@ class CustomStreamWrapper: obj_dict = response.model_dump() - if "usage" in obj_dict: + if "usage" in obj_dict and not self._should_forward_continuous_usage(): del obj_dict["usage"] response_hidden_params_for_model = cast( # cast-ok: preserve dynamic mapping behavior @@ -2227,10 +2244,8 @@ class CustomStreamWrapper: # directly to avoid expensive model_copy() per chunk. self.chunks.append(processed_chunk.model_copy()) - # Strip usage from the outgoing chunk so it's not sent twice - # (once in the chunk, once in _hidden_params). obj_dict = processed_chunk.model_dump() - if "usage" in obj_dict: + if "usage" in obj_dict and not self._should_forward_continuous_usage(): del obj_dict["usage"] processed_chunk = self.model_response_creator( chunk=obj_dict, diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 68aabc3c01e..e84dcfca49c 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -523,6 +523,265 @@ def test_strip_sse_data_from_chunk(): assert CustomStreamWrapper.strip_sse_data_from_chunk(None) is None +def _continuous_usage_fixture_chunks() -> tuple[ModelResponseStream, ...]: + return ( + ModelResponseStream( + id="chatcmpl-continuous-usage", + created=1, + model="stream-usage-test", + choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="a"))], + usage=Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), + ), + ModelResponseStream( + id="chatcmpl-continuous-usage", + created=1, + model="stream-usage-test", + choices=[StreamingChoices(index=0, delta=Delta(content="b"), finish_reason="stop")], + usage=Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12), + ), + ModelResponseStream( + id="chatcmpl-continuous-usage", + created=1, + model="stream-usage-test", + choices=[], + usage=Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12), + ), + ) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_continuous_usage_forwards_running_counts(logging_obj: Logging, sync_mode: bool) -> None: + wrapper: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=_continuous_usage_fixture_chunks()), + model="stream-usage-test", + logging_obj=logging_obj, + custom_llm_provider="hosted_vllm", + stream_options={"include_usage": True, "continuous_usage_stats": True}, + ) + collected: Final = tuple(wrapper) if sync_mode else tuple([chunk async for chunk in wrapper]) + content_chunks: Final = tuple(chunk for chunk in collected if chunk.choices and chunk.choices[0].delta.content) + + assert tuple(chunk.choices[0].delta.content for chunk in content_chunks) == ("a", "b") + assert tuple(getattr(chunk, "usage", None) for chunk in content_chunks) == ( + Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), + Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12), + ) + expected_final_usage: Final = Usage( + prompt_tokens=10, + completion_tokens=2, + total_tokens=12, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=2), + ) + assert collected[-1].usage == expected_final_usage + assert tuple(getattr(chunk, "usage", None) for chunk in collected) == ( + Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), + Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12), + None, + expected_final_usage, + ) + + +def _continuous_usage_stream( + logging_obj: Logging, + chunks: tuple[ModelResponseStream, ...], + stream_options: dict[str, bool] | None, +) -> CustomStreamWrapper: + return CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=chunks), + model="stream-usage-test", + logging_obj=logging_obj, + custom_llm_provider="hosted_vllm", + stream_options=stream_options, + ) + + +async def _continuous_usage_collected(wrapper: CustomStreamWrapper, sync_mode: bool) -> tuple[ModelResponseStream, ...]: + return tuple(wrapper) if sync_mode else tuple([chunk async for chunk in wrapper]) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.parametrize( + ("stream_options", "include_final"), + ( + (None, False), + ({"include_usage": True}, True), + ({"include_usage": True, "continuous_usage_stats": False}, True), + ({"include_usage": False, "continuous_usage_stats": True}, False), + ({"continuous_usage_stats": True}, False), + ), +) +@pytest.mark.asyncio +async def test_continuous_usage_default_contract( + logging_obj: Logging, sync_mode: bool, stream_options: dict[str, bool] | None, include_final: bool +) -> None: + wrapper: Final = _continuous_usage_stream(logging_obj, _continuous_usage_fixture_chunks(), stream_options) + collected: Final = await _continuous_usage_collected(wrapper, sync_mode) + content: Final = tuple(chunk for chunk in collected if chunk.choices and chunk.choices[0].delta.content) + assert tuple(getattr(chunk, "usage", None) for chunk in content) == (None, None) + assert getattr(collected[-1], "usage", None) == ( + Usage( + prompt_tokens=10, + completion_tokens=2, + total_tokens=12, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=2), + ) + if include_final + else None + ) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_continuous_usage_preserves_zero_counts(logging_obj: Logging, sync_mode: bool) -> None: + chunks: Final = _continuous_usage_fixture_chunks() + zero: Final = chunks[0].model_copy( + update={ + "choices": [StreamingChoices(index=0, delta=Delta(role="assistant", content=""))], + "usage": Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10), + } + ) + wrapper: Final = _continuous_usage_stream( + logging_obj, (zero, *chunks[1:]), {"include_usage": True, "continuous_usage_stats": True} + ) + collected: Final = await _continuous_usage_collected(wrapper, sync_mode) + assert getattr(collected[0], "usage", None) == Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_continuous_usage_does_not_invent_intermediate_counts(logging_obj: Logging, sync_mode: bool) -> None: + chunks: Final = tuple(chunk.model_copy(update={"usage": None}) for chunk in _continuous_usage_fixture_chunks()[:2]) + wrapper: Final = _continuous_usage_stream( + logging_obj, chunks, {"include_usage": True, "continuous_usage_stats": True} + ) + collected: Final = await _continuous_usage_collected(wrapper, sync_mode) + content: Final = tuple(chunk for chunk in collected if chunk.choices and chunk.choices[0].delta.content) + assert tuple(getattr(chunk, "usage", None) for chunk in content) == (None, None) + assert collected[-1].usage is not None + assert collected[-1].usage.total_tokens == ( + collected[-1].usage.prompt_tokens + collected[-1].usage.completion_tokens + ) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_continuous_usage_only_chunk_is_visible(logging_obj: Logging, sync_mode: bool) -> None: + chunks: Final = _continuous_usage_fixture_chunks() + progress: Final = chunks[-1].model_copy( + update={"usage": Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11)} + ) + wrapper: Final = _continuous_usage_stream( + logging_obj, (chunks[0], progress, *chunks[1:]), {"include_usage": True, "continuous_usage_stats": True} + ) + collected: Final = await _continuous_usage_collected(wrapper, sync_mode) + assert collected[1].choices == [StreamingChoices(index=0, delta=Delta())] + assert getattr(collected[1], "usage", None) == Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11) + 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: + chunks: Final = _continuous_usage_fixture_chunks() + late: Final = chunks[-1].model_copy(update={"usage": Usage(prompt_tokens=10, completion_tokens=3, total_tokens=13)}) + wrapper: Final = _continuous_usage_stream( + logging_obj, (*chunks[:2], late), {"include_usage": True, "continuous_usage_stats": True} + ) + collected: Final = await _continuous_usage_collected(wrapper, sync_mode) + assert len(collected) == 4 + assert collected[-2].choices[0].finish_reason == "stop" + assert collected[-1].usage == Usage( + prompt_tokens=10, + completion_tokens=3, + total_tokens=13, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=3), + ) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_continuous_usage_preserves_details_and_choices(logging_obj: Logging, sync_mode: bool) -> None: + chunks: Final = _continuous_usage_fixture_chunks() + detailed: Final = chunks[0].model_copy( + update={ + "choices": [ + StreamingChoices(index=0, delta=Delta(role="assistant", content="a")), + StreamingChoices(index=1, delta=Delta(content="other")), + ], + "usage": Usage( + prompt_tokens=10, + completion_tokens=1, + total_tokens=11, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=1), + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=4), + ), + } + ) + default_wrapper: Final = _continuous_usage_stream(logging_obj, (detailed, *chunks[1:]), {"include_usage": True}) + default_first: Final = next(default_wrapper) if sync_mode else await anext(default_wrapper) + await default_wrapper.aclose() + wrapper: Final = _continuous_usage_stream( + logging_obj, (detailed, *chunks[1:]), {"include_usage": True, "continuous_usage_stats": True} + ) + collected: Final = await _continuous_usage_collected(wrapper, sync_mode) + assert collected[0].choices == [ + StreamingChoices(index=0, delta=Delta(role="assistant", content="a")), + StreamingChoices(index=1, delta=Delta(content="other")), + ] + assert collected[0].choices == default_first.choices + assert getattr(collected[0], "usage", None) == Usage( + prompt_tokens=10, + completion_tokens=1, + total_tokens=11, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=1), + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=4), + ) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_continuous_usage_is_available_before_completion(logging_obj: Logging, sync_mode: bool) -> None: + wrapper: Final = _continuous_usage_stream( + logging_obj, _continuous_usage_fixture_chunks(), {"include_usage": True, "continuous_usage_stats": True} + ) + first: Final = next(wrapper) if sync_mode else await anext(wrapper) + assert first.choices[0].delta.content == "a" + assert getattr(first, "usage", None) == Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11) + await wrapper.aclose() + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_streaming_handler_with_usage(