refactor(chat): log bridged chat streams on the Responses API as chat rows
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled

The chat stream that wraps the native Responses stream now logs the request,
so the spend row carries the chat completion the client received instead of
the Responses object. The Responses iterator hands its completed response to
the bridge instead of dispatching success logging itself.
This commit is contained in:
narmstrong-berri 2026-10-16 04:55:59 +00:00
parent e3a4ec7387
commit 8edddf5506
3 changed files with 146 additions and 1 deletions

View file

@ -70,6 +70,18 @@ def _executor_logging_obj(
return native_logging_obj(logging_obj, request_data, request_data)
def _logged_by_client_stream(_completed: object) -> None:
return None
def _hand_completion_to_client_stream(result: object, executor_logging: "LiteLLMLoggingObj | None") -> None:
"""The chat stream wrapping the executor's Responses stream logs the request once, in the shape the client got."""
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
if executor_logging is not None and isinstance(result, BaseResponsesAPIStreamingIterator):
result.hand_completion_to(_logged_by_client_stream)
def _retagged_as_responses(
request_data: Mapping[str, object], logging_obj: "LiteLLMLoggingObj"
) -> Mapping[str, object]:
@ -303,6 +315,7 @@ class ResponsesToCompletionBridgeHandler:
return self._apply_post_stream_processing(result, model, custom_llm_provider)
from .transformation import is_structured_output_request
_hand_completion_to_client_stream(result, executor_logging)
completion_stream: Final = self.transformation_handler.get_model_response_iterator(
streaming_response=result,
sync_stream=True,
@ -411,6 +424,7 @@ class ResponsesToCompletionBridgeHandler:
return self._apply_post_stream_processing(result, model, custom_llm_provider)
from .transformation import is_structured_output_request
_hand_completion_to_client_stream(result, executor_logging)
completion_stream: Final = self.transformation_handler.get_model_response_iterator(
streaming_response=result,
sync_stream=False,

View file

@ -363,6 +363,7 @@ class BaseResponsesAPIStreamingIterator:
"""
_pending_logging_tasks: tuple[asyncio.Task[object], ...] = ()
_on_complete: Callable[[object], None] | None = None
def __init__(
self,
@ -416,6 +417,10 @@ class BaseResponsesAPIStreamingIterator:
) # GUARANTEE OPENAI HEADERS IN RESPONSE
self._raw_response_headers: Mapping[str, str] = MappingProxyType(dict(self.response.headers or {}))
def hand_completion_to(self, on_complete: Callable[[object], None]) -> None:
"""A bridge whose own client stream logs the request takes the completed response instead of the logger."""
self._on_complete = on_complete
def _check_max_streaming_duration(self) -> None:
"""Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS."""
if LITELLM_MAX_STREAMING_DURATION_SECONDS is None:
@ -615,7 +620,9 @@ class BaseResponsesAPIStreamingIterator:
self._restore_provider_response_headers(logging_response)
end_time: Final = datetime.now()
if is_async:
if self._on_complete is not None:
self._on_complete(logging_response)
elif is_async:
logging_coroutine: Final = self.logging_obj.dispatch_success_handlers(
logging_response,
start_time=self.start_time,

View file

@ -1,9 +1,16 @@
import asyncio
import json
import math
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import respx
from openai.types.chat import ChatCompletion
from pydantic import TypeAdapter
import litellm
from litellm.caching.caching import Cache
@ -428,3 +435,120 @@ def test_completion_keeps_responses_surface_for_what_only_it_serves(bridge_model
assert result is cached
assert responses_call.call_count == 1
assert executor_call.call_count == 0
_LOGGED_ROW: Final = TypeAdapter(dict[str, object])
_LOGGED_CHAT: Final = TypeAdapter(ChatCompletion)
class _SuccessRowRecorder(CustomLogger):
rows: tuple[dict[str, object], ...] = ()
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
self.rows = (*self.rows, _LOGGED_ROW.validate_python(kwargs["standard_logging_object"]))
def _responses_sse(events: Sequence[Mapping[str, object]]) -> bytes:
return "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events).encode()
@pytest.mark.asyncio
async def test_bridged_chat_stream_logs_one_chat_row_billed_from_provider_usage(monkeypatch: pytest.MonkeyPatch):
recorder: Final = _SuccessRowRecorder()
monkeypatch.setattr(litellm, "callbacks", [recorder])
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
prices: Final = {"input_cost_per_token": 1e-6, "output_cost_per_token": 4e-6, "cache_read_input_token_cost": 1e-7}
litellm.register_model(
{
"openai/bridge-stream-model": {
"litellm_provider": "openai",
"mode": "responses",
"max_tokens": 100,
**prices,
}
}
)
usage: Final[dict[str, object]] = {
"input_tokens": 100,
"input_tokens_details": {"cached_tokens": 40},
"output_tokens": 20,
"output_tokens_details": {"reasoning_tokens": 5},
"total_tokens": 120,
}
message: Final[dict[str, object]] = {
"type": "message",
"id": "msg_stream",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "streamed", "annotations": []}],
}
response: Final[dict[str, object]] = {
"id": "resp_stream",
"object": "response",
"created_at": 1,
"model": "bridge-stream-model",
"output": [],
"parallel_tool_calls": True,
"tool_choice": "auto",
"tools": [],
}
sse: Final = _responses_sse(
[
{"type": "response.created", "sequence_number": 0, "response": {**response, "status": "in_progress"}},
{"type": "response.output_item.added", "sequence_number": 1, "output_index": 0, "item": message},
{
"type": "response.output_text.delta",
"sequence_number": 2,
"item_id": "msg_stream",
"output_index": 0,
"content_index": 0,
"delta": "streamed",
},
{"type": "response.output_item.done", "sequence_number": 3, "output_index": 0, "item": message},
{
"type": "response.completed",
"sequence_number": 4,
"response": {**response, "status": "completed", "output": [message], "usage": usage},
},
]
)
with respx.mock(assert_all_called=False) as upstream:
responses_route = upstream.post("https://example.invalid/v1/responses").mock(
return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, content=sse)
)
stream = await litellm.acompletion(
model="openai/bridge-stream-model",
api_base="https://example.invalid/v1",
api_key="x",
messages=[{"role": "user", "content": "hi"}],
stream=True,
)
assert isinstance(stream, CustomStreamWrapper)
text = "".join([chunk.choices[0].delta.content or "" async for chunk in stream if chunk.choices])
for _ in range(100):
if recorder.rows:
break
await asyncio.sleep(0.01)
await asyncio.sleep(0.1)
assert text == "streamed"
assert responses_route.call_count == 1
assert len(recorder.rows) == 1
row: Final = recorder.rows[0]
logged: Final = _LOGGED_CHAT.validate_python(row["response"])
assert logged.object == "chat.completion"
assert logged.choices[0].message.content == "streamed"
assert row["prompt_tokens"] == 100
assert row["completion_tokens"] == 20
cost: Final = row["response_cost"]
assert isinstance(cost, float)
assert math.isclose(
cost,
60 * prices["input_cost_per_token"]
+ 40 * prices["cache_read_input_token_cost"]
+ 20 * prices["output_cost_per_token"],
)