diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 83d6fcc0bee..5ece02893fe 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1709,6 +1709,9 @@ class Logging(LiteLLMLoggingBaseClass): results=result, ) + elif self.call_type == CallTypes.aresponses_websocket.value and isinstance(result, list): + logging_result = ResponseAPILoggingUtils.build_response_from_websocket_events(events=result) or result + elif ( self.call_type == CallTypes.llm_passthrough_route.value or self.call_type == CallTypes.allm_passthrough_route.value diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 429ddeef36a..c14d6313987 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -14,7 +14,7 @@ from typing import ( overload, ) -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError import litellm from litellm._logging import verbose_logger @@ -1091,3 +1091,58 @@ class ResponseAPILoggingUtils: setattr(chat_usage, "cost", response_api_usage.cost) return chat_usage + + @staticmethod + def build_response_from_websocket_events( + events: Iterable[Any], + ) -> ResponsesAPIResponse | None: + """ + Build a single ``ResponsesAPIResponse`` from the events collected on a + Responses API WebSocket connection so token usage and cost are logged + the same way as the HTTP ``/v1/responses`` path. + + WebSocket logging dispatches the raw list of forwarded events. Without + this the logging pipeline sees a bare ``list``, extracts no usage, and + records zero tokens even though ``response.completed`` carries real + usage. Usage is summed across every terminal event (a connection may + serve multiple sequential requests); the remaining fields come from the + last terminal response. Returns ``None`` when no terminal event with a + response object is present. + """ + terminal_responses = tuple( + event["response"] + for event in events + if isinstance(event, dict) + and event.get("type") in ("response.completed", "response.incomplete") + and isinstance(event.get("response"), dict) + ) + if not terminal_responses: + return None + + usages = tuple(response["usage"] for response in terminal_responses if isinstance(response.get("usage"), dict)) + input_tokens = sum(usage.get("input_tokens") or 0 for usage in usages) + output_tokens = sum(usage.get("output_tokens") or 0 for usage in usages) + total_tokens = sum(usage.get("total_tokens") or 0 for usage in usages) + + response_dict = { + "id": "", + "created_at": 0, + "output": [], + **terminal_responses[-1], + **( + { + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": total_tokens or (input_tokens + output_tokens), + } + } + if usages + else {} + ), + } + try: + return ResponsesAPIResponse(**response_dict) + except ValidationError as e: + verbose_logger.debug("could not build ResponsesAPIResponse from websocket events: %s", e) + return None diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 4509abc7749..2d605d8c11e 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -2454,3 +2454,148 @@ class TestNativeWebSocketUrlConstruction: mock_config.get_websocket_url.assert_called_once() _, call_kwargs = mock_config.get_websocket_url.call_args assert call_kwargs["litellm_params"]["api_version"] == "2025-04-01-preview" + + +class TestResponsesWebSocketUsageLogging: + """Regression tests for LIT-4900. + + Responses API WebSocket logging dispatched the raw list of forwarded + events to the success pipeline, which extracted no usage and recorded + zero prompt/completion/total tokens even though ``response.completed`` + carried real usage (while ``response_cost`` was non-zero). + """ + + def _build_logging_obj(self): + import time + + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LitellmLogging, + ) + + logging_obj = LitellmLogging( + model="gpt-4o", + messages=[{"role": "user", "content": "Reply with one word: ok"}], + stream=False, + call_type="_aresponses_websocket", + start_time=time.time(), + litellm_call_id="lit-4900", + function_id="fn", + ) + logging_obj.update_environment_variables( + model="gpt-4o", + optional_params={}, + litellm_params={"metadata": {}}, + custom_llm_provider="openai", + ) + return logging_obj + + @pytest.mark.asyncio + async def test_websocket_completed_event_usage_is_logged(self): + from unittest.mock import AsyncMock + + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + from litellm.responses.streaming_iterator import ResponsesWebSocketStreaming + + class FakeBackendWS: + def __init__(self, events): + self._events = list(events) + + async def recv(self, decode=False): + if self._events: + return self._events.pop(0) + raise websockets.exceptions.ConnectionClosed(None, None) + + completed_event = json.dumps( + { + "type": "response.completed", + "response": { + "id": "resp_abc", + "object": "response", + "created_at": 1234567890, + "model": "gpt-4o", + "status": "completed", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok"}], + } + ], + "usage": { + "input_tokens": 12, + "output_tokens": 5, + "total_tokens": 17, + }, + }, + } + ) + + client_ws = MagicMock() + client_ws.send_text = AsyncMock() + logging_obj = self._build_logging_obj() + + handler = ResponsesWebSocketStreaming( + websocket=client_ws, + backend_ws=FakeBackendWS([completed_event]), + logging_obj=logging_obj, + ) + + await handler.backend_to_client() + + assert any(event.get("type") == "response.completed" for event in handler.messages) + + logging_obj._success_handler_helper_fn( + result=handler.messages, + cache_hit=False, + standard_logging_object=None, + ) + + slo = logging_obj.model_call_details["standard_logging_object"] + assert slo["prompt_tokens"] == 12 + assert slo["completion_tokens"] == 5 + assert slo["total_tokens"] == 17 + assert slo["response_cost"] > 0 + + def test_build_response_from_websocket_events_combines_usage(self): + from litellm.responses.utils import ResponseAPILoggingUtils + + events = [ + {"type": "response.created", "response": {"id": "r1"}}, + { + "type": "response.completed", + "response": { + "id": "r1", + "created_at": 1, + "output": [], + "usage": {"input_tokens": 12, "output_tokens": 5, "total_tokens": 17}, + }, + }, + { + "type": "response.completed", + "response": { + "id": "r2", + "created_at": 2, + "output": [], + "usage": {"input_tokens": 3, "output_tokens": 4, "total_tokens": 7}, + }, + }, + ] + + response = ResponseAPILoggingUtils.build_response_from_websocket_events(events=events) + + assert response is not None + assert response.id == "r2" + assert response.usage.input_tokens == 15 + assert response.usage.output_tokens == 9 + assert response.usage.total_tokens == 24 + + def test_build_response_from_websocket_events_without_completed_returns_none(self): + from litellm.responses.utils import ResponseAPILoggingUtils + + events = [ + {"type": "response.created", "response": {"id": "r1"}}, + {"type": "response.output_text.delta", "delta": "hi"}, + ] + + assert ResponseAPILoggingUtils.build_response_from_websocket_events(events=events) is None