diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6839e11f8c7..36fa093b60d 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/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 diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 4d7c67972f2..3dd0c53d098 100644 --- a/litellm/proxy/common_request_processing.py +++ b/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, diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 94fd915844a..e2d6023be56 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/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 diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 5c27ad8ce77..8eca7da5002 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/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,