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..6995998899a 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,216 @@ 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"], + ) + + +_BUILT_IN_TOOL_CALLS: Final[Mapping[str, tuple[Sequence[Mapping[str, object]], Sequence[Mapping[str, object]]]]] = { + "web_search": ( + tuple( + { + "type": "web_search_call", + "id": f"ws_{i}", + "status": "completed", + "action": {"type": "search", "query": "q"}, + } + for i in range(3) + ), + (), + ), + "image_generation": ( + ( + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "result": "aGk=", + "quality": "low", + "size": "1024x1024", + }, + ), + ({"type": "image_generation", "model": "gpt-image-1"},), + ), +} + + +def _responses_body_with(tool_call: str) -> dict[str, object]: + output_items, tools = _BUILT_IN_TOOL_CALLS[tool_call] + message: Final[dict[str, object]] = { + "type": "message", + "id": "msg_tool", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "done", "annotations": []}], + } + return { + "id": "resp_tool", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4.1", + "output": [*output_items, message], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": list(tools), + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + + +async def _rows_for(recorder: _SuccessRowRecorder, count: int) -> tuple[dict[str, object], ...]: + for _ in range(200): + if len(recorder.rows) >= count: + break + await asyncio.sleep(0.01) + return recorder.rows + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tool_call", sorted(_BUILT_IN_TOOL_CALLS)) +async def test_bridged_chat_call_is_billed_for_the_built_in_tool_calls_the_provider_made( + monkeypatch: pytest.MonkeyPatch, tool_call: str +): + recorder: Final = _SuccessRowRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=False) as upstream: + upstream.post("https://example.invalid/v1/responses").mock( + return_value=httpx.Response(200, json=_responses_body_with(tool_call)) + ) + bridged = await litellm.acompletion( + model="openai/responses/gpt-4.1", + api_base="https://example.invalid/v1", + api_key="x", + messages=[{"role": "user", "content": "hi"}], + ) + await _rows_for(recorder, 1) + await litellm.aresponses(model="openai/gpt-4.1", api_base="https://example.invalid/v1", api_key="x", input="hi") + rows: Final = await _rows_for(recorder, 2) + + bridged_row, native_row = rows + assert bridged_row["call_type"] == "acompletion" + assert native_row["call_type"] == "aresponses" + native_cost: Final = native_row["response_cost"] + assert isinstance(native_cost, float) + assert bridged_row["response_cost"] == pytest.approx(native_cost) + assert bridged._hidden_params["response_cost"] == pytest.approx(native_cost) + breakdown: Final = _LOGGED_ROW.validate_python(bridged_row["cost_breakdown"]) + tool_usage_cost: Final = breakdown["tool_usage_cost"] + assert isinstance(tool_usage_cost, float) and tool_usage_cost > 0