Merge 59c661bd90 into 3f38730d2c

litellm/litellm_core_utils/streaming_handler.py

@@ -27,6 +27,7 @@ from litellm.constants import STREAM_CHUNK_CHOICE_FIELDS, STREAM_FINISH_CHOICE_F
 from litellm.litellm_core_utils.asyncify import asyncify
 from litellm.litellm_core_utils.core_helpers import (
     UPSTREAM_USAGE_COST,
+    as_str_mapping,
     attribute_or_none,
     chunk_usage,
     decoded_json_object,
@@ -81,6 +82,7 @@ TOOL_CALLS_ATTRIBUTE: Final = "tool_calls"
 FUNCTION_CALL_ATTRIBUTE: Final = "function_call"
 
 _SYNC_ITER_EXHAUSTED: Final = object()
+_STREAM_OPTIONS_VALUE: Final = TypeAdapter(object)
 
 _GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__)
 _CompleteStreamingResponse = ModelResponse | TextCompletionResponse | ModelResponseStream | None
@@ -494,10 +496,18 @@ class CustomStreamWrapper:
         self.received_model_id: str | None = None
         self.logging_loop = None
         self.rules = Rules()
-        self.stream_options = stream_options or getattr(logging_obj, "stream_options", None)
+        self.stream_options: object = _STREAM_OPTIONS_VALUE.validate_python(
+            stream_options or getattr(logging_obj, "stream_options", None)
+        )
         self.messages = getattr(logging_obj, "messages", None)
         self.sent_stream_usage = False
         self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False
+        stream_options_fields: Final = as_str_mapping(self.stream_options)
+        self._continuous_usage_stats = (
+            self.send_stream_usage
+            and stream_options_fields is not None
+            and stream_options_fields.get("continuous_usage_stats") is True
+        )
         self.caller_requested_stream_usage: Final = (
             self.check_send_stream_usage(getattr(logging_obj, "stream_options", None))
             and getattr(logging_obj, "stream_usage_injected", False) is not True
@@ -664,8 +674,11 @@ class CustomStreamWrapper:
                     )
         self._restore_consumer_correlation_context()
 
-    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 check_send_stream_usage(self, stream_options: object) -> bool:
+        if stream_options is None:
+            return False
+        get_option: Final[Callable[[str, object], object] | None] = getattr(stream_options, "get", None)
+        return get_option is not None and get_option("include_usage", False) is True
 
     def check_is_function_call(self, logging_obj) -> bool:
         from litellm.litellm_core_utils.prompt_templates.common_utils import (
@@ -1894,7 +1907,7 @@ 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._continuous_usage_stats:
                         model_response.choices = []
                         return model_response
                     self._record_usage_only_chunk(model_response=model_response)
@@ -2377,7 +2390,7 @@ class CustomStreamWrapper:
                     # ModelResponseStream declares `usage` as a field, so
                     # hasattr(response, "usage") is always True — must check
                     # `is not None` to avoid running this path on every chunk.
-                    if getattr(response, "usage", None) is not None:
+                    if getattr(response, "usage", None) is not None and not self._continuous_usage_stats:
                         usage_to_preserve = response.usage
                         if usage_to_preserve:
                             response_hidden_params_for_usage = cast(  # cast-ok: preserve dynamic mapping behavior
@@ -2486,6 +2499,8 @@ class CustomStreamWrapper:
                     self.chunks.append(final_chunk_for_logging)
                 if self.stream_options is None:
                     self._set_final_chunk_usage(processed_chunk)
+                elif self._continuous_usage_stats:
+                    processed_chunk.usage = calculate_total_usage(chunks=self.chunks)
                 ## LOGGING
                 executor.submit(
                     self.run_success_logging_and_cache_storage,
@@ -2572,7 +2587,7 @@ class CustomStreamWrapper:
                         hasattr(processed_chunk, "usage") and getattr(processed_chunk, "usage", None) is not None
                     )
 
-                    if _has_usage:
+                    if _has_usage and not self._continuous_usage_stats:
                         usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk)
                         if usage_free_chunk is None:
                             continue
@@ -2722,6 +2737,8 @@ class CustomStreamWrapper:
                 self.chunks.append(final_chunk_for_logging)
             if self.stream_options is None:
                 self._set_final_chunk_usage(processed_chunk)
+            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

litellm/proxy/common_request_processing.py

@@ -63,6 +63,7 @@ from litellm.litellm_core_utils.bug_report import (
     strip_bug_report_notice,
 )
 from litellm.litellm_core_utils.core_helpers import (
+    as_str_mapping,
     attribute_or_none,
     finite_cost,
     get_or_create_metadata_bucket,
@@ -627,6 +628,7 @@ def stream_usage_tracking_updates(
         return scrub
     always_include: Final = general_settings.get("always_include_stream_usage")
     stream_options: Final = data.get("stream_options")
+    stream_options_fields: Final = as_str_mapping(stream_options)
     if always_include is True:
         if "stream_options" not in data:
             return {**scrub, "stream_options": {"include_usage": True}, STREAM_USAGE_INJECTED_FIELD: True}
@@ -645,7 +647,11 @@ def stream_usage_tracking_updates(
         return scrub
     if not supports_stream_options():
         return scrub
-    merged_stream_options: Final = {**stream_options} if isinstance(stream_options, dict) else {}
+    merged_stream_options: Final = (
+        {key: value for key, value in stream_options_fields.items() if key != "continuous_usage_stats"}
+        if isinstance(stream_options, dict) and stream_options_fields is not None
+        else {}
+    )
     return {
         "stream_options": {**merged_stream_options, "include_usage": True},
         "_litellm_strip_stream_usage": True,

tests/unit/litellm_core_utils/test_streaming_handler.py

@@ -9618,3 +9618,113 @@ def test_chunk_creator_usage_only_chunk_with_dict_usage_keeps_details(
     assert usage.prompt_tokens_details.cached_tokens == 80
     assert usage.completion_tokens_details is not None
     assert usage.completion_tokens_details.reasoning_tokens == 5
+
+
+def _continuous_usage_wrapper(
+    chunks: Sequence[ModelResponseStream],
+    stream_options: Mapping[str, bool],
+) -> CustomStreamWrapper:
+    return CustomStreamWrapper(
+        completion_stream=ModelResponseListIterator(model_responses=chunks),
+        model="hosted_vllm/m",
+        custom_llm_provider="hosted_vllm",
+        stream_options=stream_options,
+        logging_obj=Logging(
+            model="hosted_vllm/m",
+            messages=[{"role": "user", "content": "hi"}],
+            stream=True,
+            call_type="acompletion",
+            start_time=1754900000.0,
+            litellm_call_id="continuous-usage",
+            function_id="test",
+        ),
+    )
+
+
+def _continuous_usage_chunk(choice: Mapping[str, object], usage: Mapping[str, int]) -> ModelResponseStream:
+    return ModelResponseStream(id="c", created=1, model="hosted_vllm/m", choices=[choice], usage=Usage(**usage))
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_continuous_usage_stats_keeps_usage_on_every_chunk(sync_mode: bool) -> None:
+    from litellm.proxy.proxy_server import _serialize_streaming_chunk
+
+    counts: Final = [
+        {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
+        {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12},
+        {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13},
+    ]
+    chunks: Final = [
+        _continuous_usage_chunk({"index": 0, "delta": {"content": "Hi"}, "finish_reason": None}, usage=counts[0]),
+        _continuous_usage_chunk({"index": 0, "delta": {}, "finish_reason": "stop"}, usage=counts[1]),
+        ModelResponseStream(id="c", created=1, model="hosted_vllm/m", choices=[], usage=Usage(**counts[2])),
+    ]
+    wrapper: Final = _continuous_usage_wrapper(
+        chunks, stream_options={"include_usage": True, "continuous_usage_stats": True}
+    )
+
+    emitted: Final = [chunk async for chunk in wrapper] if not sync_mode else list(wrapper)
+
+    assert all(getattr(chunk, "usage", None) is not None for chunk in emitted)
+    assert [chunk.usage.total_tokens for chunk in emitted] == [11, 12, 13]
+    assert litellm.stream_chunk_builder(chunks=emitted).usage.total_tokens == 13
+    serialized: Final = tuple(json.loads(_serialize_streaming_chunk(chunk)) for chunk in emitted)
+    assert [chunk["usage"]["total_tokens"] for chunk in serialized] == [count["total_tokens"] for count in counts]
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_continuous_usage_stats_finishes_before_the_final_usage_chunk(sync_mode: bool) -> None:
+    content_usage: Final = {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}
+    final_usage: Final = {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13}
+    chunks: Final = [
+        _continuous_usage_chunk({"index": 0, "delta": {"content": "Hi"}, "finish_reason": None}, usage=content_usage),
+        _continuous_usage_chunk({"index": 0, "delta": {"content": "!"}, "finish_reason": "length"}, usage=final_usage),
+        ModelResponseStream(id="c", created=1, model="hosted_vllm/m", choices=[], usage=Usage(**final_usage)),
+    ]
+    wrapper: Final = _continuous_usage_wrapper(
+        chunks, stream_options={"include_usage": True, "continuous_usage_stats": True}
+    )
+
+    emitted: Final = [chunk async for chunk in wrapper] if not sync_mode else list(wrapper)
+
+    finish_indexes: Final = [i for i, chunk in enumerate(emitted) if chunk.choices and chunk.choices[0].finish_reason]
+    assert len(finish_indexes) == 1
+    assert emitted[finish_indexes[0]].choices[0].finish_reason == "length"
+    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 [chunk.usage.total_tokens for chunk in emitted] == [
+        content_usage["total_tokens"],
+        *[final_usage["total_tokens"]] * 3,
+    ]
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.parametrize("continuous_usage_stats", [None, False])
+@pytest.mark.asyncio
+async def test_include_usage_without_continuous_usage_stats_still_emits_only_a_final_usage_chunk(
+    sync_mode: bool,
+    continuous_usage_stats: bool | None,
+) -> None:
+    counts: Final = [
+        {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
+        {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13},
+    ]
+    chunks: Final = [
+        _continuous_usage_chunk({"index": 0, "delta": {"content": "Hi"}, "finish_reason": None}, usage=counts[0]),
+        _continuous_usage_chunk({"index": 0, "delta": {}, "finish_reason": "stop"}, usage=counts[1]),
+    ]
+    options: Final = {
+        "include_usage": True,
+        **({"continuous_usage_stats": continuous_usage_stats} if continuous_usage_stats is not None else {}),
+    }
+    wrapper: Final = _continuous_usage_wrapper(chunks, stream_options=options)
+
+    emitted: Final = [chunk async for chunk in wrapper] if not sync_mode else list(wrapper)
+
+    content_chunks: Final = [chunk for chunk in emitted[:-1] if chunk.choices and chunk.choices[0].delta.content]
+    assert content_chunks
+    assert all(getattr(chunk, "usage", None) is None for chunk in content_chunks)
+    assert getattr(emitted[-1], "usage", None) is not None

tests/unit/proxy/test_common_request_processing.py

@@ -9787,6 +9787,23 @@ class TestApplyStreamUsageTracking:
         assert "stream_options" not in data
         assert "_litellm_strip_stream_usage" not in data
 
+    def test_client_continuous_usage_stats_is_not_injected_onto_the_caller(self) -> None:
+        data: Final = {"stream": True, "stream_options": {"continuous_usage_stats": True}}
+
+        _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion")
+
+        assert data["stream_options"] == {"include_usage": True}
+        assert data["_litellm_strip_stream_usage"] is True
+
+    def test_client_requested_continuous_usage_stats_reaches_the_backend_intact(self) -> None:
+        requested: Final = {"include_usage": True, "continuous_usage_stats": True}
+        data: Final = {"stream": True, "stream_options": dict(requested)}
+
+        _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion")
+
+        assert data["stream_options"] == requested
+        assert "_litellm_strip_stream_usage" not in data
+
     def test_client_supplied_strip_marker_is_neutralized(self):
         data = {
             "stream": True,