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"], + )