mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
fix(passthrough): record real TTFT and start_time for streaming requests
Pass-through streaming requests (/v1/messages, /vertex_ai/*, /gemini/*, /cohere/*, /assemblyai/*, /openai/*, /cursor/*) all share PassThroughStreamingHandler.chunk_processor, which had two timing bugs that interacted to collapse spend_logs.completionStartTime onto spend_logs.endTime — making the streaming phase (endTime - completionStartTime) round to roughly zero and TTFT effectively soak up the entire request duration for every pass-through streaming row. Root causes 1. `start_time` arg captured too late. The caller's start_time originates in BaseAnthropicMessagesStreamingIterator.__init__, which runs AFTER the upstream HTTP response has already been received. SpendLogs.startTime therefore reflects "moment we started reading the stream", not "moment the client request entered the proxy" — the real TTFT window is silently subtracted from Duration. 2. First-chunk arrival never recorded. The chunk loop yielded bytes without updating litellm_logging_obj.completion_start_time. The fallback at litellm_logging.py:1834-1837 sets completion_start_time = end_time, collapsing TTFT onto Duration. Fix - Use litellm_logging_obj.start_time when it's earlier than the caller-supplied start_time (true request-entry timestamp). - Record completion_start_time on the first emitted chunk. Both fields are write-once: existing values are preserved. Blast radius: all pass-through endpoints — they all flow through this single chunk_processor.
This commit is contained in:
parent
69b0dd2da0
commit
034658d7fb
2 changed files with 225 additions and 0 deletions
|
|
@ -35,6 +35,20 @@ class PassThroughStreamingHandler:
|
|||
passthrough_success_handler_obj: PassThroughEndpointLogging,
|
||||
url_route: str,
|
||||
):
|
||||
# Use the true request-entry timestamp held on the logging object
|
||||
# when it's earlier than the start_time passed in. The caller's
|
||||
# start_time is captured at the streaming-iterator constructor,
|
||||
# which runs AFTER the upstream HTTP response has already been
|
||||
# received — too late to represent when the client's request
|
||||
# entered the proxy. Without this override, SpendLogs.startTime
|
||||
# is artificially deflated by the full TTFT, making
|
||||
# `endTime - startTime` shorter than reality.
|
||||
true_start = getattr(litellm_logging_obj, "start_time", None)
|
||||
if isinstance(true_start, datetime) and (
|
||||
not isinstance(start_time, datetime) or true_start < start_time
|
||||
):
|
||||
start_time = true_start
|
||||
|
||||
raw_bytes: List[bytes] = []
|
||||
logging_scheduled = False
|
||||
model_name = PassThroughStreamingHandler._extract_model_for_cost_injection(
|
||||
|
|
@ -57,6 +71,15 @@ class PassThroughStreamingHandler:
|
|||
if not cost_injection_active:
|
||||
# Hot path: just buffer for end-of-stream logging and forward.
|
||||
async for chunk in response.aiter_bytes():
|
||||
# Record TTFT on the first chunk so spend_logs
|
||||
# `completionStartTime` reflects real time-to-first-token.
|
||||
# Without this, the fallback in litellm_logging.py sets
|
||||
# completion_start_time = end_time, making TTFT equal to
|
||||
# total Duration for every passthrough streaming request.
|
||||
if litellm_logging_obj.completion_start_time is None:
|
||||
litellm_logging_obj._update_completion_start_time(
|
||||
completion_start_time=datetime.now()
|
||||
)
|
||||
raw_bytes.append(chunk)
|
||||
yield chunk
|
||||
else:
|
||||
|
|
@ -66,6 +89,10 @@ class PassThroughStreamingHandler:
|
|||
assert model_name is not None
|
||||
resolved_model_name: str = model_name
|
||||
async for chunk in response.aiter_bytes():
|
||||
if litellm_logging_obj.completion_start_time is None:
|
||||
litellm_logging_obj._update_completion_start_time(
|
||||
completion_start_time=datetime.now()
|
||||
)
|
||||
raw_bytes.append(chunk)
|
||||
if endpoint_type == EndpointType.VERTEX_AI:
|
||||
if "streamRawPredict" in url_route or "rawPredict" in url_route:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,198 @@
|
|||
"""
|
||||
Tests for PassThroughStreamingHandler.chunk_processor timing corrections:
|
||||
|
||||
1. `start_time` is overridden by `litellm_logging_obj.start_time` when the
|
||||
latter is earlier (true request-entry timestamp).
|
||||
2. `completion_start_time` is recorded on the first emitted chunk so TTFT
|
||||
reflects real time-to-first-token instead of `endTime`.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime, timedelta
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.pass_through_endpoints.streaming_handler import (
|
||||
PassThroughStreamingHandler,
|
||||
)
|
||||
|
||||
|
||||
def _build_logging_obj(true_start: datetime, completion_start_time=None):
|
||||
"""A minimal LiteLLMLoggingObj-shaped stub the handler can read from."""
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.start_time = true_start
|
||||
logging_obj.completion_start_time = completion_start_time
|
||||
logging_obj._update_completion_start_time = MagicMock(
|
||||
side_effect=lambda completion_start_time: setattr(
|
||||
logging_obj, "completion_start_time", completion_start_time
|
||||
)
|
||||
)
|
||||
logging_obj.model_call_details = {}
|
||||
logging_obj.standard_logging_object = None
|
||||
return logging_obj
|
||||
|
||||
|
||||
async def _build_response_with_chunks(chunks):
|
||||
"""Wrap a list of bytes chunks into an httpx.Response-shaped mock that
|
||||
yields them via aiter_bytes()."""
|
||||
|
||||
async def _aiter_bytes():
|
||||
for c in chunks:
|
||||
yield c
|
||||
|
||||
response = MagicMock(spec=httpx.Response)
|
||||
response.aiter_bytes = _aiter_bytes
|
||||
response.headers = {}
|
||||
return response
|
||||
|
||||
|
||||
def _drain(gen):
|
||||
"""Synchronously drain an async generator to a list."""
|
||||
|
||||
async def _collect():
|
||||
return [x async for x in gen]
|
||||
|
||||
return asyncio.run(_collect())
|
||||
|
||||
|
||||
def test_chunk_processor_uses_logging_obj_start_time_when_earlier(monkeypatch):
|
||||
"""Caller's start_time is captured at the streaming iterator constructor,
|
||||
which runs after the upstream HTTP response is already in. The logging
|
||||
object's start_time reflects when the client request entered the proxy
|
||||
— if it's earlier, it should win.
|
||||
"""
|
||||
earlier = datetime(2026, 1, 1, 12, 0, 0)
|
||||
later = earlier + timedelta(seconds=1)
|
||||
|
||||
logging_obj = _build_logging_obj(true_start=earlier)
|
||||
|
||||
response = asyncio.run(_build_response_with_chunks([b"chunk-1", b"chunk-2"]))
|
||||
captured_start_time = {}
|
||||
|
||||
async def fake_log_streaming_request(*args, **kwargs):
|
||||
captured_start_time["start_time"] = kwargs.get("start_time")
|
||||
|
||||
monkeypatch.setattr(
|
||||
PassThroughStreamingHandler,
|
||||
"_route_streaming_logging_to_handler",
|
||||
AsyncMock(side_effect=fake_log_streaming_request),
|
||||
)
|
||||
|
||||
gen = PassThroughStreamingHandler.chunk_processor(
|
||||
response=response,
|
||||
request_body={},
|
||||
litellm_logging_obj=logging_obj,
|
||||
endpoint_type=MagicMock(),
|
||||
start_time=later,
|
||||
passthrough_success_handler_obj=MagicMock(),
|
||||
url_route="/v1/messages",
|
||||
)
|
||||
_drain(gen)
|
||||
|
||||
assert captured_start_time["start_time"] == earlier, (
|
||||
"Handler must override the later caller-supplied start_time with the "
|
||||
"earlier logging-obj start_time."
|
||||
)
|
||||
|
||||
|
||||
def test_chunk_processor_keeps_caller_start_time_when_earlier(monkeypatch):
|
||||
"""When the caller-supplied start_time is already earlier than the
|
||||
logging-obj's start_time (atypical but possible if logging was
|
||||
re-stamped), don't pull it forward."""
|
||||
earlier = datetime(2026, 1, 1, 12, 0, 0)
|
||||
later = earlier + timedelta(seconds=1)
|
||||
|
||||
logging_obj = _build_logging_obj(true_start=later)
|
||||
|
||||
response = asyncio.run(_build_response_with_chunks([b"chunk-1"]))
|
||||
captured_start_time = {}
|
||||
|
||||
async def fake_log_streaming_request(*args, **kwargs):
|
||||
captured_start_time["start_time"] = kwargs.get("start_time")
|
||||
|
||||
monkeypatch.setattr(
|
||||
PassThroughStreamingHandler,
|
||||
"_route_streaming_logging_to_handler",
|
||||
AsyncMock(side_effect=fake_log_streaming_request),
|
||||
)
|
||||
|
||||
gen = PassThroughStreamingHandler.chunk_processor(
|
||||
response=response,
|
||||
request_body={},
|
||||
litellm_logging_obj=logging_obj,
|
||||
endpoint_type=MagicMock(),
|
||||
start_time=earlier,
|
||||
passthrough_success_handler_obj=MagicMock(),
|
||||
url_route="/v1/messages",
|
||||
)
|
||||
_drain(gen)
|
||||
|
||||
assert (
|
||||
captured_start_time["start_time"] == earlier
|
||||
), "Caller-supplied start_time should win when it's already the earlier value."
|
||||
|
||||
|
||||
def test_chunk_processor_records_completion_start_on_first_chunk(monkeypatch):
|
||||
"""First chunk arrival populates litellm_logging_obj.completion_start_time
|
||||
so SpendLogs.completionStartTime reflects real TTFT instead of falling
|
||||
back to end_time."""
|
||||
logging_obj = _build_logging_obj(
|
||||
true_start=datetime(2026, 1, 1, 12, 0, 0),
|
||||
completion_start_time=None,
|
||||
)
|
||||
|
||||
response = asyncio.run(_build_response_with_chunks([b"chunk-1", b"chunk-2"]))
|
||||
|
||||
monkeypatch.setattr(
|
||||
PassThroughStreamingHandler,
|
||||
"_route_streaming_logging_to_handler",
|
||||
AsyncMock(),
|
||||
)
|
||||
|
||||
gen = PassThroughStreamingHandler.chunk_processor(
|
||||
response=response,
|
||||
request_body={},
|
||||
litellm_logging_obj=logging_obj,
|
||||
endpoint_type=MagicMock(),
|
||||
start_time=datetime(2026, 1, 1, 12, 0, 0),
|
||||
passthrough_success_handler_obj=MagicMock(),
|
||||
url_route="/v1/messages",
|
||||
)
|
||||
_drain(gen)
|
||||
|
||||
logging_obj._update_completion_start_time.assert_called_once()
|
||||
assert isinstance(logging_obj.completion_start_time, datetime)
|
||||
|
||||
|
||||
def test_chunk_processor_does_not_overwrite_existing_completion_start(monkeypatch):
|
||||
"""If a downstream layer has already set completion_start_time
|
||||
(e.g. wrapper iterator), the handler must not overwrite it."""
|
||||
pre_set = datetime(2026, 1, 1, 12, 0, 5)
|
||||
logging_obj = _build_logging_obj(
|
||||
true_start=datetime(2026, 1, 1, 12, 0, 0),
|
||||
completion_start_time=pre_set,
|
||||
)
|
||||
|
||||
response = asyncio.run(_build_response_with_chunks([b"chunk-1"]))
|
||||
|
||||
monkeypatch.setattr(
|
||||
PassThroughStreamingHandler,
|
||||
"_route_streaming_logging_to_handler",
|
||||
AsyncMock(),
|
||||
)
|
||||
|
||||
gen = PassThroughStreamingHandler.chunk_processor(
|
||||
response=response,
|
||||
request_body={},
|
||||
litellm_logging_obj=logging_obj,
|
||||
endpoint_type=MagicMock(),
|
||||
start_time=datetime(2026, 1, 1, 12, 0, 0),
|
||||
passthrough_success_handler_obj=MagicMock(),
|
||||
url_route="/v1/messages",
|
||||
)
|
||||
_drain(gen)
|
||||
|
||||
logging_obj._update_completion_start_time.assert_not_called()
|
||||
assert logging_obj.completion_start_time == pre_set
|
||||
Loading…
Add table
Reference in a new issue