mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
8449ecee6a
Streaming pass-through for native Anthropic /v1/messages and the /v1/responses streaming iterator never set logging_obj.completion_start_time, so _success_handler_helper_fn fell back to completion_start_time = end_time. Downstream TTFT consumers (Prometheus, OTEL, Langfuse, Admin UI, spend logs completionStartTime) then reported time-to-first-token equal to total request duration. Stamp completion_start_time on the first chunk in PassThroughStreamingHandler. chunk_processor and BaseResponsesAPIStreamingIterator._process_chunk, mirroring CustomStreamWrapper for /chat/completions. Resolves LIT-4185 Co-authored-by: yucheng <yucheng@yuchengs-MBP.localdomain>
124 lines
4.2 KiB
Python
124 lines
4.2 KiB
Python
"""Regression tests for LIT-4185 — /v1/responses streaming must stamp
|
|
completion_start_time on the first chunk so downstream TTFT consumers
|
|
(Prometheus, OTEL, SpendLogs completionStartTime) do not fall back to
|
|
completion_start_time = end_time."""
|
|
|
|
import json
|
|
from datetime import datetime
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
|
|
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
|
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
|
|
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
|
|
from litellm.types.llms.openai import (
|
|
ResponseCompletedEvent,
|
|
ResponsesAPIResponse,
|
|
ResponsesAPIStreamEvents,
|
|
)
|
|
|
|
|
|
def _sse_event(payload: dict) -> bytes:
|
|
return f"data: {json.dumps(payload)}\n\n".encode("utf-8")
|
|
|
|
|
|
def _make_iterator(
|
|
*,
|
|
sse_events: list[bytes],
|
|
logging_obj: LiteLLMLoggingObj,
|
|
) -> ResponsesAPIStreamingIterator:
|
|
async def aiter_bytes():
|
|
for evt in sse_events:
|
|
yield evt
|
|
|
|
mock_response = Mock()
|
|
mock_response.headers = {}
|
|
mock_response.aiter_bytes = aiter_bytes
|
|
|
|
mock_config = Mock(spec=BaseResponsesAPIConfig)
|
|
mock_responses_api_response = Mock(spec=ResponsesAPIResponse)
|
|
mock_responses_api_response.id = "resp_ttft"
|
|
|
|
def _transform(model, parsed_chunk, logging_obj):
|
|
evt_type = parsed_chunk.get("type")
|
|
if evt_type == "response.completed":
|
|
completed = Mock(spec=ResponseCompletedEvent)
|
|
completed.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED
|
|
completed.response = mock_responses_api_response
|
|
return completed
|
|
stub = Mock()
|
|
stub.type = evt_type
|
|
return stub
|
|
|
|
mock_config.transform_streaming_response.side_effect = _transform
|
|
|
|
return ResponsesAPIStreamingIterator(
|
|
response=mock_response,
|
|
model="gpt-4o-mini",
|
|
responses_api_provider_config=mock_config,
|
|
logging_obj=logging_obj,
|
|
litellm_metadata={},
|
|
custom_llm_provider="openai",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_responses_streaming_stamps_completion_start_time_on_first_chunk():
|
|
"""Without the fix, `logging_obj.completion_start_time` stays None across the
|
|
entire stream and _success_handler_helper_fn falls back to end_time — collapsing
|
|
the reported TTFT to full generation time."""
|
|
logging_obj = Mock(spec=LiteLLMLoggingObj)
|
|
logging_obj.completion_start_time = None
|
|
logging_obj.model_call_details = {"litellm_params": {}}
|
|
stamped: list[datetime] = []
|
|
|
|
def _update(*, completion_start_time):
|
|
stamped.append(completion_start_time)
|
|
logging_obj.completion_start_time = completion_start_time
|
|
logging_obj.model_call_details["completion_start_time"] = completion_start_time
|
|
|
|
logging_obj._update_completion_start_time.side_effect = _update
|
|
|
|
iterator = _make_iterator(
|
|
sse_events=[
|
|
_sse_event({"type": "response.created"}),
|
|
_sse_event({"type": "response.output_text.delta", "delta": "hi"}),
|
|
_sse_event({"type": "response.completed"}),
|
|
],
|
|
logging_obj=logging_obj,
|
|
)
|
|
|
|
async for _ in iterator:
|
|
pass
|
|
|
|
assert len(stamped) == 1, (
|
|
f"Expected exactly one first-chunk stamp; got {len(stamped)}. "
|
|
"Later chunks must not re-stamp completion_start_time."
|
|
)
|
|
assert isinstance(stamped[0], datetime)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_responses_streaming_does_not_reset_prior_completion_start_time():
|
|
"""If `completion_start_time` is already set (e.g. by an outer wrapper), the
|
|
iterator must not overwrite it — otherwise TTFT would collapse to
|
|
time-to-last-chunk under contention."""
|
|
prior = datetime(2020, 1, 1, 0, 0, 0)
|
|
logging_obj = Mock(spec=LiteLLMLoggingObj)
|
|
logging_obj.completion_start_time = prior
|
|
logging_obj.model_call_details = {"litellm_params": {}}
|
|
|
|
iterator = _make_iterator(
|
|
sse_events=[
|
|
_sse_event({"type": "response.created"}),
|
|
_sse_event({"type": "response.completed"}),
|
|
],
|
|
logging_obj=logging_obj,
|
|
)
|
|
|
|
async for _ in iterator:
|
|
pass
|
|
|
|
logging_obj._update_completion_start_time.assert_not_called()
|
|
assert logging_obj.completion_start_time == prior
|