From 8edddf5506d825caeeeb61e46e392056bf284915 Mon Sep 17 00:00:00 2001 From: narmstrong-berri Date: Fri, 16 Oct 2026 04:55:59 +0000 Subject: [PATCH] refactor(chat): log bridged chat streams on the Responses API as chat rows 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. --- .../handler.py | 14 ++ litellm/responses/streaming_iterator.py | 9 +- ...itellm_responses_transformation_handler.py | 124 ++++++++++++++++++ 3 files changed, 146 insertions(+), 1 deletion(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index 082d9efbd74..12e2527934b 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -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, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index c5dc45566a9..5ca42c38622 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -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, diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py index 1c4636b6e12..4d0422051a4 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py @@ -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"], + )