From a56a27deb63ee89be84f7ac7f3b3f62a9eee0163 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Mon, 9 Nov 2026 18:45:00 -0800 Subject: [PATCH] fix(prometheus): record litellm_overhead_latency_metric for streams, stamp plus the proxy's first-chunk work (#49920) * fix(prometheus): record litellm_overhead_latency_metric for streams and count the proxy's per-chunk work in it * fix(proxy): sum the per-chunk stream work through Logging.add_streaming_overhead so the lint gates hold * fix(proxy): count the per-chunk work of chat completions streams too, and drop a comment that repeated the code * fix(prometheus): keep the handoff overhead stamp when a /v1/responses stream finishes The Responses streaming iterator calls update_response_metadata again at stream end with include_overhead=False, which replaced the timing stamped at handoff with one that had no overhead, so litellm_overhead_latency_metric still had no sample for streamed /v1/responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): type the parameters of the per-chunk overhead tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(prometheus): add only the proxy's first-chunk work to a stream's overhead, not a sum over every chunk A 1 ms hook on each of 500 chunks delays the client by 1 ms per chunk, not 500 ms, so the overhead metric keeps the stamp plus the work on the first chunk and stops summing later ones Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): drive the first-chunk overhead tests with a fake clock instead of real sleeps Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): a keepalive ping is not a stream's first chunk, and test the overhead with a real ProxyLogging Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 15 ++- .../llm_response_utils/response_metadata.py | 23 +++- litellm/proxy/common_request_processing.py | 19 ++- litellm/proxy/proxy_server.py | 4 + .../test_response_metadata.py | 52 ++++++++ .../test_litellm_logging.py | 39 ++++++ .../proxy/test_common_request_processing.py | 80 +++++++++++++ ...test_proxy_server_endpoints_and_startup.py | 113 ++++++++++++++++-- .../unit/responses/test_streaming_iterator.py | 5 + 9 files changed, 338 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e9342a765bc..d88c8876ae0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -906,6 +906,7 @@ class Logging(LiteLLMLoggingBaseClass): # Timing for results that cannot carry ``_hidden_params`` (plain-dict /v1/messages # responses and the bridge stream wrappers); see ``update_response_metadata``. self.response_timing_metrics: Mapping[str, float] = {} + self.first_chunk_overhead_ms: float | None = None # Passthrough endpoint guardrails config for field targeting self.passthrough_guardrails_config: dict[str, object] | None = None @@ -983,6 +984,11 @@ class Logging(LiteLLMLoggingBaseClass): """Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``.""" self.response_timing_metrics = dict(timing_metrics) + def set_first_chunk_overhead(self, seconds: float) -> None: + """The proxy's own work on the first chunk of a stream, the delay every later chunk inherits.""" + if self.first_chunk_overhead_ms is None: + self.first_chunk_overhead_ms = seconds * 1000 + def add_dynamic_callback(self, callback: CustomLogger) -> None: self.dynamic_input_callbacks = self._with_dynamic_callback(self.dynamic_input_callbacks, callback) self.dynamic_success_callbacks = self._with_dynamic_callback(self.dynamic_success_callbacks, callback) @@ -7984,8 +7990,15 @@ def get_standard_logging_object_payload( if clean_hidden_params["litellm_overhead_time_ms"] is None and status == "success": # /v1/messages dict results and the bridge stream wrappers keep it on the logging object; # failure payloads stay None like every response type that carries its own _hidden_params - timing_metrics: Final = getattr(logging_obj, "response_timing_metrics", None) or {} + timing_metrics: Final[Mapping[str, float]] = getattr(logging_obj, "response_timing_metrics", None) or {} clean_hidden_params["litellm_overhead_time_ms"] = timing_metrics.get("litellm_overhead_time_ms") + # the overhead above was stamped when the stream object came back, before any chunk was + # read; the proxy's work on the first chunk (hook, cost injection, serialization) is the + # delay the client sees on every chunk, so only that one is added, never a sum over chunks + first_chunk_overhead_ms: Final = getattr(logging_obj, "first_chunk_overhead_ms", None) + stamped_overhead_ms: Final = clean_hidden_params["litellm_overhead_time_ms"] + if status == "success" and stamped_overhead_ms is not None and isinstance(first_chunk_overhead_ms, float): + clean_hidden_params["litellm_overhead_time_ms"] = round(stamped_overhead_ms + first_chunk_overhead_ms, 4) model_cost_information: Final = StandardLoggingPayloadSetup.get_model_cost_information( base_model=base_model, diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index 051a8fdf256..9d45426ca95 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py @@ -299,7 +299,7 @@ def update_response_metadata( # cannot hold ``_hidden_params``: keep only the timing on the logging object (no cost # recompute) so the proxy headers and the standard logging payload can still read it. logging_obj.set_response_timing_metrics( - response_timing_metrics(start_time, end_time, logging_obj, include_overhead) + _logging_obj_timing_metrics(start_time, end_time, logging_obj, include_overhead) ) return @@ -307,3 +307,24 @@ def update_response_metadata( metadata.set_timing_metrics(start_time, end_time, logging_obj, include_overhead) metadata.set_hidden_params(logging_obj, model, kwargs, cache_hit=cache_hit) metadata.apply() + if logging_obj.stream is True: + # the response assembled from the chunks never inherits the stream wrapper's + # ``_hidden_params``, so the logging payload reads the timing from the logging object + logging_obj.set_response_timing_metrics( + _logging_obj_timing_metrics(start_time, end_time, logging_obj, include_overhead) + ) + + +def _logging_obj_timing_metrics( + start_time: datetime.datetime, + end_time: datetime.datetime, + logging_obj: LiteLLMLoggingObject, + include_overhead: bool, +) -> Mapping[str, float]: + """A stream read to completion keeps the overhead stamped when the stream was handed back.""" + timing_metrics: Final = response_timing_metrics(start_time, end_time, logging_obj, include_overhead) + stamped: Final[Mapping[str, float]] = getattr(logging_obj, "response_timing_metrics", None) or {} + stamped_overhead_ms: Final = stamped.get("litellm_overhead_time_ms") + if include_overhead or stamped_overhead_ms is None: + return timing_metrics + return {**timing_metrics, "litellm_overhead_time_ms": stamped_overhead_ms} diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index eb356491014..9e2b49ef692 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3,6 +3,7 @@ import contextlib import json import logging import math +import time from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone @@ -4865,6 +4866,15 @@ class ProxyBaseLLMRequestProcessing: _finalize_streaming_generator_cleanup = finalize_streaming_generator_cleanup + @staticmethod + def record_first_chunk_overhead( + logging_obj: LiteLLMLoggingObj | None, chunk_received_at: float, is_provider_output: bool = True + ) -> None: + """The proxy's own work on a chunk; the logging object keeps the first one for the stamped overhead. + A keepalive ping carries no provider output, so it is never that first chunk.""" + if logging_obj is not None and is_provider_output: + logging_obj.set_first_chunk_overhead(time.perf_counter() - chunk_received_at) + @staticmethod async def async_streaming_data_generator( response: object, @@ -4922,12 +4932,12 @@ class ProxyBaseLLMRequestProcessing: response=response, request_data=request_data, ) - release_stream_failure: Final = _ClientStreamFailureReport.hold_while_streaming( - stream_logging_obj if isinstance(stream_logging_obj, LiteLLMLoggingObj) else None - ) + timing_sink: Final = stream_logging_obj if isinstance(stream_logging_obj, LiteLLMLoggingObj) else None + release_stream_failure: Final = _ClientStreamFailureReport.hold_while_streaming(timing_sink) try: str_so_far = "" async for chunk in guarded_stream: + chunk_received_at = time.perf_counter() # ``.format(chunk)`` was previously evaluated for every chunk # regardless of log level; gate it behind the level check. if debug_enabled: @@ -4968,6 +4978,9 @@ class ProxyBaseLLMRequestProcessing: delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES serialized = serialize_chunk(chunk) recent_tail = advance_sse_tail(recent_tail, serialized) + ProxyBaseLLMRequestProcessing.record_first_chunk_overhead( + timing_sink, chunk_received_at, is_provider_output=chunk != STREAM_SSE_KEEPALIVE_PING_BYTES + ) yield serialized held_tail: Final = flush_tail() if flush_tail is not None else b"" if held_tail: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 55540820651..3bc7d5a1b7d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10911,6 +10911,8 @@ async def async_data_generator( # ``get_response_string`` work per chunk on every deployment that # happened to ship a streaming-iterator override (the default). needs_per_chunk_hook: Final = proxy_logging_obj.needs_per_chunk_streaming_hook() + stream_logging_obj: Final[object] = request_data.get("litellm_logging_obj") # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # legacy request data is narrowed below + timing_sink: Final = stream_logging_obj if isinstance(stream_logging_obj, LiteLLMLoggingObj) else None is_raw_sse_stream: Final = bool(request_data.get("_litellm_raw_sse_stream")) strip_stream_usage: Final = bool(request_data.get("_litellm_strip_stream_usage")) raw_sse_buffer = "" @@ -10946,6 +10948,7 @@ async def async_data_generator( yield ": ping\n\n" continue chunk = item + chunk_received_at = time.perf_counter() if needs_per_chunk_hook: ### CALL HOOKS ### - modify outgoing data chunk, _str_so_far = await _apply_streaming_chunk_hooks( @@ -11031,6 +11034,7 @@ async def async_data_generator( error_message = chunk break + ProxyBaseLLMRequestProcessing.record_first_chunk_overhead(timing_sink, chunk_received_at) if not raw_passthrough: try: if error_state is not None: diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py b/tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py index 6d9d8c7997a..4c9b4e55c76 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py @@ -207,6 +207,58 @@ class TestDictResultsSkipMetadataUpdate: logging_obj.set_response_timing_metrics.assert_not_called() assert result._hidden_params["litellm_overhead_time_ms"] == 100.0 + def test_update_response_metadata_keeps_stream_wrapper_timing_on_logging_obj(self): + """A stream wrapper carries _hidden_params, but the response assembled from its chunks does + not inherit them, so the logging payload needs the stamped timing on the logging object.""" + result = ModelResponse() + logging_obj = MagicMock() + logging_obj.stream = True + logging_obj.model_call_details = {"llm_api_duration_ms": 900.0} + logging_obj.caching_details = None + logging_obj.response_cost_calculator = MagicMock(return_value=0.001) + logging_obj.litellm_call_id = "test-call-id" + + update_response_metadata( + result=result, + logging_obj=logging_obj, + model="gpt-4", + kwargs={}, + start_time=datetime.datetime(2025, 1, 1, 0, 0, 0), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1), + ) + + logging_obj.set_response_timing_metrics.assert_called_once_with( + {"_response_ms": 1000.0, "litellm_overhead_time_ms": 100.0} + ) + assert result._hidden_params["litellm_overhead_time_ms"] == 100.0 + + def test_update_response_metadata_keeps_the_handoff_overhead_when_a_stream_is_read_to_completion(self): + """A Responses API stream stamps its overhead when the stream is handed back and calls again at + stream end with include_overhead=False; that second call must not wipe the stamp.""" + result = ModelResponse() + logging_obj = MagicMock() + logging_obj.stream = True + logging_obj.response_timing_metrics = {"_response_ms": 300.0, "litellm_overhead_time_ms": 100.0} + logging_obj.model_call_details = {"llm_api_duration_ms": 200.0} + logging_obj.caching_details = None + logging_obj.response_cost_calculator = MagicMock(return_value=0.001) + logging_obj.litellm_call_id = "test-call-id" + + update_response_metadata( + result=result, + logging_obj=logging_obj, + model="gpt-4", + kwargs={}, + start_time=datetime.datetime(2025, 1, 1, 0, 0, 0), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 10), + include_overhead=False, + ) + + logging_obj.set_response_timing_metrics.assert_called_once_with( + {"_response_ms": 10000.0, "litellm_overhead_time_ms": 100.0} + ) + assert "litellm_overhead_time_ms" not in result._hidden_params + def test_update_response_metadata_omits_overhead_for_completed_stream(self): """LIT-5466: the Responses streaming iterator finishes the whole stream before updating metadata, and the provider call it recorded stopped at the first byte.""" diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 2f5edbd6350..e8f90dc0146 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -10130,6 +10130,45 @@ def test_get_standard_logging_object_payload_reads_overhead_from_logging_obj_for assert payload["hidden_params"]["litellm_overhead_time_ms"] == 100.0 +@pytest.mark.parametrize( + ("stamped_overhead_ms", "first_chunk_overhead_ms", "expected_overhead_ms"), + ( + (100.0, 25.5, 125.5), + (None, 25.5, None), + ), +) +def test_get_standard_logging_object_payload_adds_first_chunk_work_to_overhead( + logging_obj, stamped_overhead_ms, first_chunk_overhead_ms, expected_overhead_ms +): + """The overhead stamped when the stream object came back misses the hook and cost injection on + the first chunk, the delay every chunk the client sees inherits, so that one chunk's work is + added on top. With no stamped overhead the first chunk work alone is not reported.""" + import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + logging_obj.set_response_timing_metrics( + {"_response_ms": 1000.0} + if stamped_overhead_ms is None + else {"_response_ms": 1000.0, "litellm_overhead_time_ms": stamped_overhead_ms} + ) + logging_obj.first_chunk_overhead_ms = first_chunk_overhead_ms + now = datetime.datetime.now() + payload = get_standard_logging_object_payload( + kwargs={"litellm_call_id": "call-1", "model": "gpt-4o", "messages": []}, + init_response_obj={"id": "msg_1", "type": "message", "role": "assistant", "content": []}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["hidden_params"]["litellm_overhead_time_ms"] == expected_overhead_ms + + def test_get_standard_logging_object_payload_survives_logging_obj_without_timing_metrics(logging_obj): """The payload is built inside a blanket except that returns None, so a logging object without the timing carrier (custom subclasses, older pickles) must not silently drop every spend log.""" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 901dbb6ea97..23023560a41 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -2,6 +2,7 @@ import asyncio import contextvars import copy import datetime +import time import gzip import json from collections.abc import AsyncIterable, Awaitable, Coroutine, Mapping @@ -5447,6 +5448,29 @@ class TestHandleLLMApiExceptionFramingHeaders: assert proxy_exc.headers["x-request-id"] == "abc-123" +class _FakeClock: + """Hand-advanced stand-in for time.perf_counter, so the hook's cost is exact and no real time passes.""" + + def __init__(self) -> None: + self.now = 1000.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +class _TimeWithFakeClock: + """The module's `time` with perf_counter swapped for the fake clock, everything else untouched.""" + + def __init__(self, clock: _FakeClock) -> None: + self.perf_counter = clock + + def __getattr__(self, name: str) -> object: + return getattr(time, name) + + class TestAsyncStreamingDataGeneratorFastPath: """Fast/slow path branching in async_streaming_data_generator.""" @@ -5516,6 +5540,62 @@ class TestAsyncStreamingDataGeneratorFastPath: ProxyLogging._callback_capabilities_cache.clear() + @pytest.mark.asyncio + async def test_slow_path_records_first_chunk_work_on_the_logging_object(self, monkeypatch: pytest.MonkeyPatch): + """The overhead header is stamped before any chunk is read, so the time the proxy spends in + the hook on the first chunk is recorded on the logging object for the success callback to add. + Later chunks are not summed, a 20 ms hook on 3 chunks delays the client by 20 ms, not 60, and a + keepalive ping that arrives before any provider output is not the first chunk.""" + import litellm.proxy.common_request_processing as common_request_processing_module + + clock: Final = _FakeClock() + monkeypatch.setattr(common_request_processing_module, "time", _TimeWithFakeClock(clock)) + + class _SlowStreamingCb(CustomLogger): + async def async_post_call_streaming_hook( + self, user_api_key_dict: UserAPIKeyAuth, response: object + ) -> object: + clock.advance(0.5 if response == STREAM_SSE_KEEPALIVE_PING_BYTES else 0.02) + return response + + monkeypatch.setattr(litellm, "callbacks", [_SlowStreamingCb()]) + ProxyLogging._callback_capabilities_cache.clear() + + proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) + logging_obj = Logging( + model="claude-x", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id="call-stream", + function_id="fn", + ) + chunks = [ + STREAM_SSE_KEEPALIVE_PING_BYTES, + *( + ModelResponseStream(choices=[StreamingChoices(delta=Delta(content=word))]) + for word in ("Hello", " world", " more") + ), + ] + + out = [ + c + async for c in ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + response=self._aiter(chunks), + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), + request_data={"model": "claude-x", "litellm_logging_obj": logging_obj}, + proxy_logging_obj=proxy_logging_obj, + serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk, + serialize_error=lambda e: "data: error\n\n", + ) + ] + + assert len(out) == len(chunks) + assert logging_obj.first_chunk_overhead_ms == pytest.approx(20.0) + + ProxyLogging._callback_capabilities_cache.clear() + async def test_opt_in_header_injects_cost_while_global_flag_is_off(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) monkeypatch.setattr(litellm, "callbacks", []) diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 1688dab6b5f..9bbdc7165dd 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -11,7 +11,7 @@ import subprocess import sys import time import types -from collections.abc import Callable, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence from datetime import datetime, timedelta, timezone from pathlib import Path from types import MappingProxyType @@ -961,7 +961,9 @@ async def test_committed_spend_twins_outlive_the_slowest_spend_write_to_the_db( finally: scheduler.shutdown(wait=False) - slowest_db_write: Final = SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS + proxy_batch_write_at + _BATCH_WRITE_JITTER_MAX_SECONDS + slowest_db_write: Final = ( + SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS + proxy_batch_write_at + _BATCH_WRITE_JITTER_MAX_SECONDS + ) assert spend.committed_ttl is not None assert spend.committed_ttl >= slowest_db_write @@ -2003,11 +2005,11 @@ async def test_proxy_startup_after_a_failed_migration_does_not_reuse_the_closed_ async with proxy_startup_event(FastAPI()) as state: assert state["database"] is None + @pytest.mark.asyncio async def test_proxy_startup_names_the_config_file_that_set_the_unsafe_key(monkeypatch, tmp_path): from fastapi import FastAPI - from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError from litellm.proxy.proxy_server import proxy_startup_event @@ -5711,15 +5713,27 @@ def _content_policy_router(**router_kwargs: object) -> litellm.Router: model_list=[ { "model_name": "primary", - "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "fake", "mock_response": "primary answered"}, + "litellm_params": { + "model": "openai/gpt-5.4-mini", + "api_key": "fake", + "mock_response": "primary answered", + }, }, { "model_name": "secondary", - "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "fake", "mock_response": "secondary answered"}, + "litellm_params": { + "model": "openai/gpt-5.4-mini", + "api_key": "fake", + "mock_response": "secondary answered", + }, }, { "model_name": "backup", - "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "fake", "mock_response": "backup answered"}, + "litellm_params": { + "model": "openai/gpt-5.4-nano", + "api_key": "fake", + "mock_response": "backup answered", + }, }, ], num_retries=0, @@ -5799,7 +5813,9 @@ async def test_db_content_policy_fallbacks_layer_over_the_litellm_settings_rules "invalid_rule", [{"secondary": "backup"}, ["backup"], dict[str, list[str]](), {"secondary": ["backup"], "primary": ["secondary"]}], ) -async def test_an_invalid_db_content_policy_rule_skips_only_itself(monkeypatch: pytest.MonkeyPatch, invalid_rule: object): +async def test_an_invalid_db_content_policy_rule_skips_only_itself( + monkeypatch: pytest.MonkeyPatch, invalid_rule: object +): monkeypatch.setattr(litellm, "content_policy_fallbacks", [{"primary": ["backup"]}]) router: Final = _content_policy_router() @@ -7555,6 +7571,89 @@ async def test_async_data_generator_cleanup_on_early_exit(): mock_response.aclose.assert_awaited_once() +class _FakeClock: + """Hand-advanced stand-in for time.perf_counter, so the hook's cost is exact and no real time passes.""" + + def __init__(self) -> None: + self.now = 1000.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +class _TimeWithFakeClock: + """The module's `time` with perf_counter swapped for the fake clock, everything else untouched.""" + + def __init__(self, clock: _FakeClock) -> None: + self.perf_counter = clock + + def __getattr__(self, name: str) -> object: + return getattr(time, name) + + +@pytest.mark.asyncio +async def test_async_data_generator_records_first_chunk_work_on_the_logging_object(monkeypatch: pytest.MonkeyPatch): + """The overhead header is stamped before any chunk is read, so the time the proxy spends in the + hook on the first chunk of a chat stream is recorded on the logging object for the success callback. + Later chunks are not summed, a 20 ms hook on 3 chunks delays the client by 20 ms, not 60.""" + import litellm.proxy.common_request_processing as common_request_processing_module + import litellm.proxy.proxy_server as proxy_server_module + from litellm.integrations.custom_logger import CustomLogger + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import async_data_generator + from litellm.proxy.utils import ProxyLogging + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + clock: Final = _FakeClock() + monkeypatch.setattr(proxy_server_module, "time", _TimeWithFakeClock(clock)) + monkeypatch.setattr(common_request_processing_module, "time", _TimeWithFakeClock(clock)) + + class _SlowStreamingCb(CustomLogger): + async def async_post_call_streaming_hook(self, user_api_key_dict: UserAPIKeyAuth, response: object) -> object: + clock.advance(0.02) + return response + + monkeypatch.setattr(litellm, "callbacks", [_SlowStreamingCb()]) + ProxyLogging._callback_capabilities_cache.clear() + + logging_obj = Logging( + model="gpt-5.4-mini", + messages=[{"role": "user", "content": "test"}], + stream=True, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="call-stream", + function_id="fn", + ) + request_data = { + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "test"}], + "litellm_logging_obj": logging_obj, + } + chunks = [ + ModelResponseStream(choices=[StreamingChoices(delta=Delta(content=word))]) + for word in ("Hello", " world", " more") + ] + + async def provider_stream() -> AsyncIterator[object]: + for chunk in chunks: + yield chunk + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=MagicMock())): + yielded = [ + data async for data in async_data_generator(provider_stream(), MagicMock(spec=UserAPIKeyAuth), request_data) + ] + + ProxyLogging._callback_capabilities_cache.clear() + assert len(yielded) == len(chunks) + 1 + assert yielded[-1] == "data: [DONE]\n\n" + assert logging_obj.first_chunk_overhead_ms == pytest.approx(20.0) + + def _guarded_chain_logging(chain): from litellm.proxy.utils import ProxyLogging diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index b1859ddf33c..3fb9edc9459 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -745,8 +745,10 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead(): """LIT-5466: the provider call is timed to first byte, so at stream completion the total minus that duration is token generation, not LiteLLM overhead.""" logging_obj = _logging_obj_stub() + logging_obj.stream = True logging_obj.model_call_details = {"litellm_params": {}, "llm_api_duration_ms": 200.0} logging_obj.caching_details = None + logging_obj.response_timing_metrics = {"_response_ms": 300.0, "litellm_overhead_time_ms": 100.0} class _CompletedEvent: def __init__(self) -> None: @@ -760,6 +762,9 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead(): assert iterator.completed_response._hidden_params["_response_ms"] == 10000.0 assert "litellm_overhead_time_ms" not in iterator.completed_response._hidden_params + logging_obj.set_response_timing_metrics.assert_called_once_with( + {"_response_ms": 10000.0, "litellm_overhead_time_ms": 100.0} + ) def _mock_config_with_completed_response(response: ResponsesAPIResponse) -> Mock: