mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
c19d49d919
* fix(responses): keep provider response headers in streaming logging callbacks The responses streaming iterator captures the provider's HTTP response headers into its own _hidden_params, but never puts them on the completed response, and the model_validate(model_dump()) copy made for logging drops pydantic private attributes. Success callbacks and StandardLoggingPayload.hidden_params.additional_headers therefore saw an empty dict for streaming /v1/responses, so Azure's apim-request-id was unreadable from the callback payload. Restore the headers on the nested response of the logging copy, preferring any the provider transform already set (the fake_stream path) and falling back to the ones the iterator captured from the stream. Skipped when the copy fell back to the original event, so a serialization failure never leaves logging-only state on the caller's object. * fix: satisfy LIT002 mutable-collection gate in header restore --------- Co-authored-by: Yucheng Zhu <yucheng@berri.ai>
539 lines
19 KiB
Python
539 lines
19 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 typing import Optional
|
|
from unittest.mock import Mock, patch
|
|
|
|
import httpx
|
|
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,
|
|
SyncResponsesAPIStreamingIterator,
|
|
)
|
|
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 _mock_config() -> Mock:
|
|
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 mock_config
|
|
|
|
|
|
def _make_iterator(
|
|
*,
|
|
sse_events: list[bytes],
|
|
logging_obj: LiteLLMLoggingObj,
|
|
trailing_error: Optional[Exception] = None,
|
|
) -> ResponsesAPIStreamingIterator:
|
|
async def aiter_bytes():
|
|
for evt in sse_events:
|
|
yield evt
|
|
if trailing_error is not None:
|
|
raise trailing_error
|
|
|
|
mock_response = Mock()
|
|
mock_response.headers = {}
|
|
mock_response.aiter_bytes = aiter_bytes
|
|
|
|
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",
|
|
)
|
|
|
|
|
|
def _make_sync_iterator(
|
|
*,
|
|
sse_events: list[bytes],
|
|
logging_obj: LiteLLMLoggingObj,
|
|
trailing_error: Optional[Exception] = None,
|
|
) -> SyncResponsesAPIStreamingIterator:
|
|
def iter_bytes():
|
|
for evt in sse_events:
|
|
yield evt
|
|
if trailing_error is not None:
|
|
raise trailing_error
|
|
|
|
mock_response = Mock()
|
|
mock_response.headers = {}
|
|
mock_response.iter_bytes = iter_bytes
|
|
|
|
return SyncResponsesAPIStreamingIterator(
|
|
response=mock_response,
|
|
model="gpt-4o-mini",
|
|
responses_api_provider_config=_mock_config(),
|
|
logging_obj=logging_obj,
|
|
litellm_metadata={},
|
|
custom_llm_provider="openai",
|
|
)
|
|
|
|
|
|
def _logging_obj_stub() -> Mock:
|
|
logging_obj = Mock(spec=LiteLLMLoggingObj)
|
|
logging_obj.completion_start_time = None
|
|
logging_obj.model_call_details = {"litellm_params": {}}
|
|
return logging_obj
|
|
|
|
|
|
@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
|
|
|
|
|
|
_COMPLETE_STREAM_EVENTS = [
|
|
_sse_event({"type": "response.created"}),
|
|
_sse_event({"type": "response.output_text.delta", "delta": "hi"}),
|
|
_sse_event({"type": "response.completed"}),
|
|
]
|
|
|
|
_TRAILING_ERRORS = [
|
|
httpx.ReadError("Response payload is not completed"),
|
|
httpx.RemoteProtocolError("peer closed connection without sending complete message body"),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type)
|
|
async def test_transport_error_after_completed_event_ends_stream_cleanly(trailing_error):
|
|
"""A sloppy connection close after `response.completed` must not turn a
|
|
complete stream into an error (regression guard for the transport no longer
|
|
swallowing ClientPayloadError/TransferEncodingError)."""
|
|
iterator = _make_iterator(
|
|
sse_events=_COMPLETE_STREAM_EVENTS,
|
|
logging_obj=_logging_obj_stub(),
|
|
trailing_error=trailing_error,
|
|
)
|
|
|
|
seen = [event.type async for event in iterator]
|
|
|
|
assert ResponsesAPIStreamEvents.RESPONSE_COMPLETED in seen
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_transport_error_before_completed_event_raises():
|
|
"""A connection lost before any terminal event is a real failure and must
|
|
surface, not end the stream as if it completed."""
|
|
iterator = _make_iterator(
|
|
sse_events=_COMPLETE_STREAM_EVENTS[:-1],
|
|
logging_obj=_logging_obj_stub(),
|
|
trailing_error=httpx.ReadError("Response payload is not completed"),
|
|
)
|
|
|
|
with pytest.raises(httpx.ReadError):
|
|
async for _ in iterator:
|
|
pass
|
|
|
|
|
|
@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type)
|
|
def test_sync_transport_error_after_completed_event_ends_stream_cleanly(trailing_error):
|
|
iterator = _make_sync_iterator(
|
|
sse_events=_COMPLETE_STREAM_EVENTS,
|
|
logging_obj=_logging_obj_stub(),
|
|
trailing_error=trailing_error,
|
|
)
|
|
|
|
seen = [event.type for event in iterator]
|
|
|
|
assert ResponsesAPIStreamEvents.RESPONSE_COMPLETED in seen
|
|
|
|
|
|
def test_sync_transport_error_before_completed_event_raises():
|
|
iterator = _make_sync_iterator(
|
|
sse_events=_COMPLETE_STREAM_EVENTS[:-1],
|
|
logging_obj=_logging_obj_stub(),
|
|
trailing_error=httpx.ReadError("Response payload is not completed"),
|
|
)
|
|
|
|
with pytest.raises(httpx.ReadError):
|
|
for _ in iterator:
|
|
pass
|
|
|
|
|
|
def test_stream_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypatch):
|
|
"""
|
|
Regression test for LIT-6184 on the /v1/responses streaming surface: the
|
|
completed-stream cache write was dispatched as a bare fire-and-forget task,
|
|
so asyncio.run cancelled it at loop close before the write landed. The
|
|
write must survive loop shutdown just like the chat-completions one.
|
|
"""
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
import litellm
|
|
from litellm.types.utils import CallTypes
|
|
|
|
writes = []
|
|
|
|
class _SlowWriteCache:
|
|
async def async_add_cache(self, result, dynamic_cache_object=None, **kwargs):
|
|
await asyncio.sleep(0.2)
|
|
writes.append(result)
|
|
|
|
def add_cache(self, *args, **kwargs):
|
|
raise AssertionError("sync write must not run on the async path")
|
|
|
|
caching_handler = SimpleNamespace(
|
|
request_kwargs={
|
|
"model": "test-model",
|
|
"input": "hello",
|
|
"stream": True,
|
|
"caching": True,
|
|
"metadata": None,
|
|
"custom_llm_provider": "openai",
|
|
},
|
|
preset_cache_key="responses-stream-cache-key",
|
|
original_function=litellm.aresponses,
|
|
dual_cache=None,
|
|
_should_store_result_in_cache=lambda original_function, kwargs: True,
|
|
)
|
|
logging_obj = SimpleNamespace(
|
|
model_call_details={"litellm_params": {}},
|
|
_llm_caching_handler=caching_handler,
|
|
)
|
|
iterator = ResponsesAPIStreamingIterator(
|
|
response=httpx.Response(200),
|
|
model="test-model",
|
|
responses_api_provider_config=Mock(spec=BaseResponsesAPIConfig),
|
|
logging_obj=logging_obj,
|
|
request_data=caching_handler.request_kwargs,
|
|
call_type=CallTypes.aresponses.value,
|
|
)
|
|
iterator.completed_response = ResponseCompletedEvent(
|
|
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
|
response=ResponsesAPIResponse(
|
|
id="resp_lit6184",
|
|
created_at=int(datetime.now().timestamp()),
|
|
status="completed",
|
|
model="test-model",
|
|
object="response",
|
|
output=[],
|
|
),
|
|
)
|
|
monkeypatch.setattr(litellm, "cache", _SlowWriteCache())
|
|
|
|
async def _short_lived_script():
|
|
iterator._persist_completed_response_to_cache(is_async=True)
|
|
|
|
asyncio.run(_short_lived_script())
|
|
|
|
assert len(writes) == 1
|
|
|
|
|
|
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.model_call_details = {"litellm_params": {}, "llm_api_duration_ms": 200.0}
|
|
logging_obj.caching_details = None
|
|
|
|
class _CompletedEvent:
|
|
def __init__(self) -> None:
|
|
self._hidden_params: dict = {}
|
|
|
|
iterator = _make_iterator(sse_events=[], logging_obj=logging_obj)
|
|
iterator.completed_response = _CompletedEvent()
|
|
iterator.start_time = datetime(2025, 1, 1, 0, 0, 0)
|
|
|
|
iterator._run_post_success_hooks(datetime(2025, 1, 1, 0, 0, 10))
|
|
|
|
assert iterator.completed_response._hidden_params["_response_ms"] == 10000.0
|
|
assert "litellm_overhead_time_ms" not in iterator.completed_response._hidden_params
|
|
|
|
|
|
def _responses_api_response_with_usage() -> ResponsesAPIResponse:
|
|
from litellm.types.llms.openai import ResponseAPIUsage
|
|
|
|
return ResponsesAPIResponse(
|
|
id="resp_lit6427",
|
|
created_at=int(datetime(2025, 1, 1).timestamp()),
|
|
status="completed",
|
|
model="mantle-claude",
|
|
object="response",
|
|
output=[],
|
|
usage=ResponseAPIUsage(input_tokens=20, output_tokens=60, total_tokens=80),
|
|
)
|
|
|
|
|
|
def test_stamp_responses_usage_cost_stamps_computed_cost():
|
|
from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
|
|
|
|
response = _responses_api_response_with_usage()
|
|
logging_obj = Mock(spec=LiteLLMLoggingObj)
|
|
logging_obj._response_cost_calculator.return_value = 0.000704
|
|
|
|
_stamp_responses_usage_cost(response, logging_obj)
|
|
|
|
assert getattr(response.usage, "cost", None) == pytest.approx(0.000704)
|
|
logging_obj._response_cost_calculator.assert_called_once_with(result=response)
|
|
|
|
|
|
def test_stamp_responses_usage_cost_keeps_provider_reported_cost():
|
|
from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
|
|
|
|
response = _responses_api_response_with_usage()
|
|
setattr(response.usage, "cost", 0.5)
|
|
logging_obj = Mock(spec=LiteLLMLoggingObj)
|
|
|
|
_stamp_responses_usage_cost(response, logging_obj)
|
|
|
|
assert getattr(response.usage, "cost", None) == pytest.approx(0.5)
|
|
logging_obj._response_cost_calculator.assert_not_called()
|
|
|
|
|
|
def test_stamp_responses_usage_cost_survives_calculator_failure():
|
|
from litellm.responses.streaming_iterator import _stamp_responses_usage_cost
|
|
|
|
response = _responses_api_response_with_usage()
|
|
logging_obj = Mock(spec=LiteLLMLoggingObj)
|
|
logging_obj._response_cost_calculator.side_effect = RuntimeError("cost map unavailable")
|
|
|
|
_stamp_responses_usage_cost(response, logging_obj)
|
|
|
|
assert getattr(response.usage, "cost", None) is None
|
|
|
|
|
|
def _capture_dispatch(logged: list):
|
|
"""Record the object handed to the success handlers.
|
|
|
|
``Mock(spec=LiteLLMLoggingObj).dispatch_success_handlers`` is an AsyncMock whose side effect
|
|
only runs when the coroutine is awaited, so capture with a plain function instead.
|
|
"""
|
|
|
|
async def _noop() -> None:
|
|
return None
|
|
|
|
def _dispatch(result, **kwargs):
|
|
logged.append(result)
|
|
return _noop()
|
|
|
|
return _dispatch
|
|
|
|
|
|
def _headers_config(*, transform_hidden_params: Optional[dict] = None) -> Mock:
|
|
"""Config whose completed event carries a real ResponsesAPIResponse, so the logging copy
|
|
performs a genuine model_dump/model_validate round trip."""
|
|
mock_config = Mock(spec=BaseResponsesAPIConfig)
|
|
|
|
def _transform(model, parsed_chunk, logging_obj):
|
|
evt_type = parsed_chunk.get("type")
|
|
if evt_type != "response.completed":
|
|
stub = Mock()
|
|
stub.type = evt_type
|
|
return stub
|
|
response = ResponsesAPIResponse(
|
|
id="resp_headers",
|
|
created_at=1,
|
|
output=[],
|
|
parallel_tool_calls=False,
|
|
tool_choice="auto",
|
|
tools=[],
|
|
)
|
|
if transform_hidden_params is not None:
|
|
response._hidden_params.update(transform_hidden_params)
|
|
return ResponseCompletedEvent(
|
|
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
|
response=response,
|
|
)
|
|
|
|
mock_config.transform_streaming_response.side_effect = _transform
|
|
return mock_config
|
|
|
|
|
|
def _make_header_iterator(
|
|
*,
|
|
headers: dict,
|
|
config: Mock,
|
|
logging_obj: LiteLLMLoggingObj,
|
|
) -> ResponsesAPIStreamingIterator:
|
|
async def aiter_bytes():
|
|
yield _sse_event({"type": "response.completed"})
|
|
|
|
mock_response = Mock()
|
|
mock_response.headers = headers
|
|
mock_response.aiter_bytes = aiter_bytes
|
|
|
|
return ResponsesAPIStreamingIterator(
|
|
response=mock_response,
|
|
model="gpt-4o-mini",
|
|
responses_api_provider_config=config,
|
|
logging_obj=logging_obj,
|
|
litellm_metadata={},
|
|
custom_llm_provider="azure",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_logging_response_carries_provider_response_headers():
|
|
"""LIT-6055: the provider headers the iterator captured must reach the logged response, so
|
|
custom loggers can read Azure's apim-request-id from the callback payload."""
|
|
logging_obj = _logging_obj_stub()
|
|
logged: list[object] = []
|
|
logging_obj.dispatch_success_handlers = _capture_dispatch(logged)
|
|
|
|
logging_obj._on_deferred_stream_complete = None
|
|
|
|
iterator = _make_header_iterator(
|
|
headers={"apim-request-id": "azure-correlation-1", "x-ms-region": "East US 2"},
|
|
config=_headers_config(),
|
|
logging_obj=logging_obj,
|
|
)
|
|
async for _ in iterator:
|
|
pass
|
|
|
|
assert len(logged) == 1
|
|
hidden_params = logged[0].response._hidden_params
|
|
assert hidden_params["additional_headers"]["llm_provider-apim-request-id"] == "azure-correlation-1"
|
|
assert hidden_params["additional_headers"]["llm_provider-x-ms-region"] == "East US 2"
|
|
assert hidden_params["headers"]["apim-request-id"] == "azure-correlation-1"
|
|
# the proxy builds the client's response headers from the iterator's own dict, so the logged
|
|
# response must hold copies rather than alias it
|
|
assert hidden_params["additional_headers"] is not iterator._hidden_params["additional_headers"]
|
|
assert hidden_params["headers"] is not iterator._raw_response_headers
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_logging_copy_preserves_transform_hidden_params():
|
|
"""LIT-6055: model_validate(model_dump()) drops pydantic private attributes, so headers a
|
|
provider transform already set on the response (fake_stream) must be re-applied."""
|
|
logging_obj = _logging_obj_stub()
|
|
logged: list[object] = []
|
|
logging_obj.dispatch_success_handlers = _capture_dispatch(logged)
|
|
|
|
logging_obj._on_deferred_stream_complete = None
|
|
|
|
iterator = _make_header_iterator(
|
|
headers={},
|
|
config=_headers_config(
|
|
transform_hidden_params={
|
|
"additional_headers": {"llm_provider-apim-request-id": "from-transform"},
|
|
"headers": {"apim-request-id": "from-transform"},
|
|
"response_cost": 0.5,
|
|
}
|
|
),
|
|
logging_obj=logging_obj,
|
|
)
|
|
async for _ in iterator:
|
|
pass
|
|
|
|
assert len(logged) == 1
|
|
hidden_params = logged[0].response._hidden_params
|
|
assert hidden_params["additional_headers"]["llm_provider-apim-request-id"] == "from-transform"
|
|
assert hidden_params["headers"]["apim-request-id"] == "from-transform"
|
|
assert iterator.completed_response is not logged[0]
|
|
# only the header keys travel: response_cost would short-circuit the cost calculator
|
|
assert "response_cost" not in hidden_params
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_logging_copy_fallback_leaves_caller_event_untouched():
|
|
"""LIT-6055: when the logging copy falls back to the original event, the header restore must
|
|
not stamp logging-only state onto the object the caller is iterating."""
|
|
logging_obj = _logging_obj_stub()
|
|
logged: list[object] = []
|
|
logging_obj.dispatch_success_handlers = _capture_dispatch(logged)
|
|
logging_obj._on_deferred_stream_complete = None
|
|
|
|
iterator = _make_header_iterator(
|
|
headers={"apim-request-id": "azure-correlation-1"},
|
|
config=_headers_config(),
|
|
logging_obj=logging_obj,
|
|
)
|
|
async for _ in iterator:
|
|
pass
|
|
|
|
assert len(logged) == 1
|
|
iterator._completed_response_logged = False
|
|
logged.clear()
|
|
with patch.object(type(iterator.completed_response), "model_dump", side_effect=ValueError("cannot serialize")):
|
|
iterator._log_completed_response(is_async=True)
|
|
|
|
assert logged == [iterator.completed_response]
|
|
assert iterator.completed_response.response._hidden_params == {}
|