diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index da420ab0578..6256bd5b8d9 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2892,6 +2892,7 @@ _RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "re class _ResponsesWsEventResponse(LiteLLMBaseModel): + id: str | None = None usage: Mapping[str, object] | None = None service_tier: str | None = None @@ -2943,6 +2944,26 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor): ) return ResponsesWebSocketTokenUsageProcessor.combine_usage_objects(list(collected_usage_objects)) + @staticmethod + def response_id_from_responses_ws_results(results: Sequence[Mapping[str, object]]) -> str | None: + response_ids: Final = tuple( + event.response.id + for event in (_ResponsesWsEvent.model_validate(result) for result in results) + if event.response is not None and event.response.id + ) + return response_ids[-1] if response_ids else None + + @staticmethod + def create_logging_object(results: Sequence[Mapping[str, object]]) -> LiteLLMRealtimeStreamLoggingObject: + tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(results) + raw_events: Final = cast(OpenAIRealtimeStreamList, list(results)) # cast-ok: raw WS frames, unvalidated field + return LiteLLMRealtimeStreamLoggingObject( + id=ResponsesWebSocketTokenUsageProcessor.response_id_from_responses_ws_results(results), + usage=ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(results), + results=raw_events, + service_tier=next(iter(tier_partition)) if len(tier_partition) == 1 else None, + ) + _TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 715e492b159..8c1b55b38c6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -18,7 +18,7 @@ from types import MappingProxyType, TracebackType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast from httpx import Response -from pydantic import BaseModel, JsonValue +from pydantic import BaseModel, JsonValue, TypeAdapter import litellm from litellm import _custom_logger_compatible_callbacks_literal @@ -536,6 +536,10 @@ def _resolve_mantle_region_for_cost( return resolve_mantle_region(litellm_params or MappingProxyType({})) +_PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"}) +_LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) + + def _provider_response_id(source: object) -> str | None: candidate: Final = source.get("id") if isinstance(source, dict) else getattr(source, "id", None) return candidate if isinstance(candidate, str) and candidate else None @@ -760,6 +764,32 @@ class Logging(LiteLLMLoggingBaseClass): self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None + def copy_for_new_call(self, litellm_call_id: str, start_time: datetime.datetime) -> "Logging": + connection_params: Final = _LITELLM_PARAMS.validate_python(self.litellm_params) + litellm_params: Final[dict[str, object]] = { + **connection_params, + **{ + key: {**value} + for key, value in connection_params.items() + if key in _PER_CALL_LITELLM_PARAM_DICTS and isinstance(value, dict) + }, + "litellm_call_id": litellm_call_id, + } + new_call: Final = copy.copy(self) + new_call.litellm_call_id = litellm_call_id + new_call.start_time = start_time + new_call.litellm_params = litellm_params + new_call.model_call_details = { + **self.model_call_details, + "litellm_call_id": litellm_call_id, + "start_time": start_time, + "first_api_call_start_time": None, + "litellm_params": litellm_params, + } + new_call.streaming_chunks = [] + new_call.sync_streaming_chunks = [] + return new_call + def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None: """Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``.""" self.response_timing_metrics = dict(timing_metrics) @@ -2362,20 +2392,9 @@ class Logging(LiteLLMLoggingBaseClass): ) elif self.call_type == CallTypes.aresponses_websocket.value and isinstance(result, list): # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped - combined_ws_usage: Final = ( - ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results( - results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream - ) - ) - ws_tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier( + logging_result = ResponsesWebSocketTokenUsageProcessor.create_logging_object( results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream ) - ws_service_tier: Final = next(iter(ws_tier_partition)) if len(ws_tier_partition) == 1 else None - logging_result = LiteLLMRealtimeStreamLoggingObject( - usage=combined_ws_usage, - results=result, # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream - service_tier=ws_service_tier, - ) elif ( self.call_type == CallTypes.llm_passthrough_route.value @@ -6628,7 +6647,7 @@ def get_standard_logging_object_payload( else raw_usage_dict ) - id = response_obj.get("id", kwargs.get("litellm_call_id")) + id = _provider_response_id(response_obj) or kwargs.get("litellm_call_id") _model_id: Final = metadata.get("model_info", {}).get("id", "") _model_group: Final = metadata.get("model_group", "") diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4a9b8fee2e8..7a766ebd904 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -155,7 +155,10 @@ from litellm.types.llms.openai import ( from litellm.types.realtime import RealtimeQueryParams from litellm.types.rerank import RerankResponse from litellm.types.responses.main import DeleteResponseResult -from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults +from litellm.types.responses.streaming_websocket import ( + ResponsesWebSocketRequestDefaults, + ResponsesWebSocketTurnFailureHook, +) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( CallTypes, @@ -6745,8 +6748,9 @@ class BaseLLMHTTPHandler: custom_llm_provider: str | None = None, first_message: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + on_turn_failure: ResponsesWebSocketTurnFailureHook | None = None, **kwargs: object, - ) -> Exception | None: + ) -> None: """ Handles Responses API WebSocket mode. @@ -6780,7 +6784,7 @@ class BaseLLMHTTPHandler: **kwargs, ) await handler.run() - return None + return import websockets from websockets.asyncio.client import ClientConnection @@ -6901,8 +6905,9 @@ class BaseLLMHTTPHandler: authorized_model=model, custom_llm_provider=custom_llm_provider, request_defaults=request_defaults, + on_turn_failure=on_turn_failure, ) - return await streaming.bidirectional_forward() + await streaming.bidirectional_forward() except websockets.exceptions.InvalidStatus as e: verbose_logger.exception("Error connecting to responses WS backend: %s", e) @@ -6916,7 +6921,6 @@ class BaseLLMHTTPHandler: pass else: raise Exception(f"Unexpected error while closing WebSocket: {close_error}") - return None def image_edit_handler( self, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 5b865f3b5c3..54c30073dcd 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -5377,7 +5377,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): that usage instead of refunding it. """ try: - stash: Final = get_request_stash() + failed_request: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(request_data) + stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(failed_request)) if stash is None: return await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0938d3f35c3..58ec0413d3a 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -469,11 +469,7 @@ class _ProxyDBLogger(CustomLogger): await _release_budget_reservation(budget_reservation=budget_reservation) # Non-model call types (health checks, afile_delete) have no model or standard_logging_object. # Use .get() for "stream" to avoid KeyError on health checks. - # WS session wrappers (_aresponses_websocket, _arealtime) also reach here with - # result=None; their per-turn costs are tracked on the inner aresponses/realtime calls. - if sl_object is None and ( - not kwargs.get("model") or kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime") - ): + if sl_object is None and (not kwargs.get("model") or kwargs.get("call_type") == "_arealtime"): verbose_proxy_logger.warning( "Cost tracking - skipping, no standard_logging_object for call_type=%s", kwargs.get("call_type", "unknown"), diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 7f9fcca369b..eaa6ec4c204 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -44,9 +44,12 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ) from litellm.types.responses.main import DeleteResponseResult +from litellm.types.responses.streaming_websocket import ResponsesWebSocketTurnFailureHook from litellm.types.utils import TokenCountResponse if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy.utils import ProxyLogging from litellm.router import Router router: Final = APIRouter() @@ -1436,6 +1439,40 @@ def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, objec return MappingProxyType({key: value for key, value in hints.items() if value is not None}) +_RESPONSES_WS_FIRST_FRAME_KEYS: Final = frozenset({"first_message", "input", "previous_response_id"}) +_RESPONSES_WS_PER_TURN_DICT_KEYS: Final = frozenset({"metadata", "litellm_metadata"}) + + +def _responses_ws_turn_request_data( + connection_request_data: Mapping[str, object], turn_logging_obj: "LiteLLMLoggingObj" +) -> dict[str, object]: + return { + **{key: value for key, value in connection_request_data.items() if key not in _RESPONSES_WS_FIRST_FRAME_KEYS}, + **{ + key: {**value} + for key, value in connection_request_data.items() + if key in _RESPONSES_WS_PER_TURN_DICT_KEYS and isinstance(value, dict) + }, + "litellm_call_id": turn_logging_obj.litellm_call_id, + "litellm_logging_obj": turn_logging_obj, + } + + +def _responses_ws_turn_failure_hook( + proxy_logging_obj: "ProxyLogging", + user_api_key_dict: UserAPIKeyAuth, + connection_request_data: Mapping[str, object], +) -> ResponsesWebSocketTurnFailureHook: + async def book_turn_failure(turn_logging_obj: "LiteLLMLoggingObj", exception: Exception) -> None: + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=exception, + request_data=_responses_ws_turn_request_data(connection_request_data, turn_logging_obj), + ) + + return book_turn_failure + + def _responses_ws_failure_frame(failure: Exception) -> str: raw_status: Final = getattr(failure, "status_code", None) status: Final = raw_status if isinstance(raw_status, int) and not isinstance(raw_status, bool) else 500 @@ -1584,7 +1621,10 @@ async def _responses_websocket_session( return routed_data: Final = dict( - data, user_api_key_dict=user_api_key_dict, **_routing_hints_from_first_ws_frame(first_message) + data, + user_api_key_dict=user_api_key_dict, + responses_ws_turn_failure_hook=_responses_ws_turn_failure_hook(proxy_logging_obj, user_api_key_dict, data), + **_routing_hints_from_first_ws_frame(first_message), ) # Phase 2: route to upstream provider try: @@ -1594,13 +1634,7 @@ async def _responses_websocket_session( llm_router=llm_router, user_model=user_model, ) - failure: Final = await llm_call - if isinstance(failure, Exception): - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=failure, - request_data=routed_data, - ) + await llm_call except Exception as e: verbose_proxy_logger.exception("Responses WebSocket error") with contextlib.suppress(Exception): @@ -1611,6 +1645,8 @@ async def _responses_websocket_session( request_data=routed_data, ) await websocket.close(code=1011, reason="Internal server error") + finally: + await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(user_api_key_dict) @router.websocket("/v1/responses") diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f6bd4bfba54..5df161fc366 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -59,7 +59,10 @@ from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.llms.openai.data_residency import infer_openai_data_residency from litellm.secret_managers.main import get_secret_str from litellm.types.responses.main import * -from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults +from litellm.types.responses.streaming_websocket import ( + ResponsesWebSocketRequestDefaults, + ResponsesWebSocketTurnFailureHook, +) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import all_litellm_params from litellm.utils import ( @@ -2422,12 +2425,13 @@ async def _aresponses_websocket( api_base: str | None = None, api_key: str | None = None, timeout: float | None = None, + responses_ws_turn_failure_hook: ResponsesWebSocketTurnFailureHook | None = None, **kwargs, -) -> Exception | None: +) -> None: """ Private function to handle the Responses API WebSocket mode. - For PROXY use only. Returns the provider failure that ended the connection, if any. + For PROXY use only. ``responses_ws_turn_failure_hook`` books each failed turn of a native socket. Resolves the LLM provider from ``model``, looks up the matching ``BaseResponsesAPIConfig``, and hands off to @@ -2500,7 +2504,7 @@ async def _aresponses_websocket( deployment_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _RESPONSES_WS_ROUTING_HINT_KEYS} first_message: Final = kwargs.get("first_message") - return await base_llm_http_handler.async_responses_websocket( + await base_llm_http_handler.async_responses_websocket( model=resolved_model, websocket=websocket, logging_obj=litellm_logging_obj, @@ -2517,5 +2521,6 @@ async def _aresponses_websocket( litellm_metadata=_build_litellm_metadata_for_ws(kwargs), custom_llm_provider=_custom_llm_provider, request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs), + on_turn_failure=responses_ws_turn_failure_hook, **remaining_kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 32a4f4284be..e4291bd366c 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -7,8 +7,10 @@ import time import traceback import uuid from collections.abc import AsyncIterable, Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime from functools import lru_cache +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Protocol, cast, overload, runtime_checkable @@ -56,6 +58,7 @@ if TYPE_CHECKING: ResponsesBackendWebSocket, ResponsesClientWebSocket, ResponsesWebSocketRequestDefaults, + ResponsesWebSocketTurnFailureHook, ) from litellm.types.router import LiteLLM_Params @@ -1831,9 +1834,59 @@ RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output _RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"}) +_RESPONSES_WS_TERMINAL_EVENT_TYPES: Final = _HISTORY_TERMINAL_EVENT_TYPES | _RESPONSES_WS_FAILURE_EVENT_TYPES + _RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) +def _json_object_or_none(message: object) -> Mapping[str, object] | None: + if _is_json_object(message): + return message + if not isinstance(message, (str, bytes)): + return None + try: + parsed: Final = _load_json_value(message) + except json.JSONDecodeError: + return None + return parsed if _is_json_object(parsed) else None + + +def _user_texts_in_input_item(item: object) -> tuple[str, ...]: + if not _is_json_object(item) or item.get("type") != "message" or item.get("role") != "user": + return () + content: Final = item.get("content", []) + if isinstance(content, str): + return (content,) + return tuple( + text + for block in _json_array_or_empty(content) + if _is_json_object(block) + and block.get("type") == "input_text" + and isinstance(text := block.get("text"), str) + and text + ) + + +def _logged_user_messages(client_frame: object) -> tuple[dict[str, object], ...]: + frame: Final = _json_object_or_none(client_frame) + if frame is None or frame.get("type") != "response.create": + return () + input_items: Final = frame.get("input", []) + if isinstance(input_items, str): + return ({"role": "user", "content": input_items},) + texts: Final = chain.from_iterable(_user_texts_in_input_item(item) for item in _json_array_or_empty(input_items)) + return tuple({"role": "user", "content": text} for text in texts) + + +@dataclass(frozen=True, slots=True) +class _ResponsesWsTurn: + logging_obj: LiteLLMLoggingObj + events: tuple[_MutableJsonObject, ...] = () + + def with_event(self, event: _MutableJsonObject) -> _ResponsesWsTurn: + return _ResponsesWsTurn(logging_obj=self.logging_obj, events=(*self.events, event)) + + def _ws_event_error(event: Mapping[str, object]) -> object: if event.get("type") == "error": return event.get("error") @@ -1916,6 +1969,7 @@ class ResponsesWebSocketStreaming: authorized_model: str | None = None, custom_llm_provider: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + on_turn_failure: ResponsesWebSocketTurnFailureHook | None = None, ): self.websocket = websocket self.backend_ws = backend_ws @@ -1925,8 +1979,8 @@ class ResponsesWebSocketStreaming: litellm_metadata: Final = self.request_data.get("litellm_metadata") self.litellm_metadata: dict[str, object] = litellm_metadata if _is_json_object(litellm_metadata) else {} self.custom_llm_provider: str | None = custom_llm_provider - self.messages: list[_MutableJsonObject] = [] - self.input_messages: list[dict[str, object]] = [] + self.on_turn_failure: ResponsesWebSocketTurnFailureHook | None = on_turn_failure + self._open_turns: tuple[_ResponsesWsTurn, ...] = () self.first_message = first_message self.guardrail_callbacks: Sequence[PresidioGuardrailCallback] = guardrail_callbacks or [] self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or [] @@ -1936,104 +1990,70 @@ class ResponsesWebSocketStreaming: self.authorized_model: str | None = authorized_model self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults - def _should_store_event(self, event_obj: _MutableJsonObject) -> bool: - return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES + def _start_turn(self, client_frame: str | None) -> None: + turn_start: Final = datetime.now() # noqa: DTZ005 # logging pipeline uses naive datetimes + turn_logging_obj: Final = self.logging_obj.copy_for_new_call( + litellm_call_id=str(uuid.uuid4()), start_time=turn_start + ) + turn_logging_obj.model_call_details["messages"] = list(_logged_user_messages(client_frame)) + turn_logging_obj.pre_call(input=client_frame, api_key="") + self._open_turns = (*self._open_turns, _ResponsesWsTurn(logging_obj=turn_logging_obj)) - def _store_event(self, event: str | bytes | dict[str, object]) -> None: - if isinstance(event, bytes): - event = event.decode("utf-8") - if isinstance(event, str): - try: - event_obj = _load_json_object(event) - except (json.JSONDecodeError, TypeError): - return - else: - event_obj = event + def _record_client_frame(self, client_frame: str) -> None: + frame: Final = _json_object_or_none(client_frame) + if frame is not None and frame.get("type") == "response.create": + self._start_turn(client_frame) - if self._should_store_event(event_obj): - self.messages.append(event_obj) + def _record_backend_event(self, event: str) -> None: + event_obj: Final = _json_object_or_none(event) + if event_obj is None or event_obj.get("type") not in RESPONSES_WS_LOGGED_EVENT_TYPES: + return + if not self._open_turns: + self._start_turn(None) + turn: Final = self._open_turns[0].with_event(event_obj) + if event_obj.get("type") in _RESPONSES_WS_TERMINAL_EVENT_TYPES: + self._open_turns = self._open_turns[1:] + self._log_turn(turn) + return + self._open_turns = (turn, *self._open_turns[1:]) - def _collect_input_from_client_event(self, message: object) -> None: - """Extract user input content from response.create for logging.""" - try: - if isinstance(message, str): - msg_obj = _load_json_object(message) - elif _is_json_object(message): - msg_obj = message - else: - return + def _log_open_turns(self) -> None: + open_turns: Final = self._open_turns + self._open_turns = () + for turn in open_turns: + if turn.events: + self._log_turn(turn) - if msg_obj.get("type") != "response.create": - return + def _log_turn(self, turn: _ResponsesWsTurn) -> None: + _spawn_logging_task( + asyncio.get_running_loop(), self._dispatch_turn(turn), task_name="Responses WS turn logging" + ) - input_items: Final = msg_obj.get("input", []) - if isinstance(input_items, str): - self.input_messages.append({"role": "user", "content": input_items}) - return - - if _is_json_array(input_items): - for item in input_items: - if not _is_json_object(item): - continue - if item.get("type") == "message" and item.get("role") == "user": - content = item.get("content", []) - if isinstance(content, str): - self.input_messages.append({"role": "user", "content": content}) - elif _is_json_array(content): - for c in content: - if _is_json_object(c) and c.get("type") == "input_text": - text = c.get("text", "") - if text: - self.input_messages.append({"role": "user", "content": text}) - except (json.JSONDecodeError, AttributeError, TypeError): - pass - - def _store_input(self, message: object) -> None: - self._collect_input_from_client_event(message) - if self.logging_obj: - self.logging_obj.pre_call(input=message, api_key="") - - def _failure_exception(self) -> Exception | None: + async def _dispatch_turn(self, turn: _ResponsesWsTurn) -> None: + events: Final = list(turn.events) failed_event: Final = next( - (event for event in self.messages if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None + (event for event in events if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None ) if failed_event is None: - return None - return _map_stream_error_to_exception( + await turn.logging_obj.dispatch_success_handlers(events, prefer_async_handlers=True) + return + exception: Final = _map_stream_error_to_exception( _ws_event_error(failed_event), self.authorized_model or "", self.custom_llm_provider or "" ) - - async def _log_messages(self) -> None: - if not self.logging_obj: - return - if self.input_messages: - self.logging_obj.model_call_details["messages"] = self.input_messages - if not self.messages: - return - exception: Final = self._failure_exception() - if exception is None: - asyncio.create_task(self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True)) - return - self._record_usage_for_failure() - traceback_exception: Final = "".join(traceback.format_exception(exception)) - asyncio.create_task( - self.logging_obj.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True) + self._record_usage_for_failure(turn.logging_obj, events) + await turn.logging_obj.dispatch_failure_handlers( + exception, "".join(traceback.format_exception(exception)), prefer_async_handlers=True ) + if self.on_turn_failure is not None: + await self.on_turn_failure(turn.logging_obj, exception) - def _record_usage_for_failure(self) -> None: + @staticmethod + def _record_usage_for_failure(turn_logging_obj: LiteLLMLoggingObj, events: Sequence[_MutableJsonObject]) -> None: from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor - from litellm.types.utils import LiteLLMRealtimeStreamLoggingObject - usage: Final = ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results( - self.messages - ) - tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(self.messages) - service_tier: Final = next(iter(tier_partition)) if len(tier_partition) == 1 else None - logging_result: Final = LiteLLMRealtimeStreamLoggingObject( - usage=usage, results=self.messages, service_tier=service_tier - ) - response_cost: Final = self.logging_obj.response_cost_calculator(result=logging_result) or 0.0 - self.logging_obj.record_partial_usage_for_failure(usage, response_cost) + logging_result: Final = ResponsesWebSocketTokenUsageProcessor.create_logging_object(events) + response_cost: Final = turn_logging_obj.response_cost_calculator(result=logging_result) or 0.0 + turn_logging_obj.record_partial_usage_for_failure(logging_result.usage, response_cost) def _wrap_response_event(self, response_str: str) -> str: try: @@ -2092,7 +2112,7 @@ class ResponsesWebSocketStreaming: # Log the output-masked form so PII redacted by apply_to_output # guardrails does not appear in success logs. - self._store_event(wrapped_str) + self._record_backend_event(wrapped_str) await self.websocket.send_text(wrapped_str) @@ -2101,7 +2121,7 @@ class ResponsesWebSocketStreaming: except Exception as e: verbose_logger.exception("Error in responses WS backend_to_client: %s", e) finally: - await self._log_messages() + self._log_open_turns() def _enforce_authorized_model(self, msg_obj: _MutableJsonObject) -> bool: """ @@ -2446,8 +2466,7 @@ class ResponsesWebSocketStreaming: try: if self.first_message is not None and await self._enforce_or_reject_frame(self.first_message): masked_first: Final = await self._mask_response_create(self.first_message) - self._store_input(masked_first) - self._store_event(masked_first) + self._record_client_frame(masked_first) await self.backend_ws.send(masked_first) while True: @@ -2455,14 +2474,13 @@ class ResponsesWebSocketStreaming: if not await self._enforce_or_reject_frame(message): continue masked = await self._mask_response_create(message) - self._store_input(masked) - self._store_event(masked) + self._record_client_frame(masked) await self.backend_ws.send(masked) except Exception as e: verbose_logger.debug("Responses WS client_to_backend ended: %s", e) - async def bidirectional_forward(self) -> Exception | None: + async def bidirectional_forward(self) -> None: forward_task: Final = asyncio.create_task(self.backend_to_client()) try: await self.client_to_backend() @@ -2479,7 +2497,6 @@ class ResponsesWebSocketStreaming: await self.backend_ws.close() except Exception: pass - return self._failure_exception() # --------------------------------------------------------------------------- diff --git a/litellm/router.py b/litellm/router.py index f173cdb37b4..309e8c2bd1e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8406,8 +8406,7 @@ class Router: Track remaining tpm/rpm quota for model in model_list """ try: - # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. - if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): + if kwargs.get("call_type") == "_arealtime": return if is_batch_retrieve_call_type(kwargs.get("call_type")): return diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 8ba1da824c6..8f81ec9414d 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -495,8 +495,7 @@ class RouterBudgetLimiting(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """Original method now uses helper functions""" verbose_router_logger.debug("in RouterBudgetLimiting.async_log_success_event") - # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. - if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): + if kwargs.get("call_type") == "_arealtime": return standard_logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_payload is None: diff --git a/litellm/types/responses/streaming_websocket.py b/litellm/types/responses/streaming_websocket.py index f369cbcebf8..de76398b5cc 100644 --- a/litellm/types/responses/streaming_websocket.py +++ b/litellm/types/responses/streaming_websocket.py @@ -2,10 +2,13 @@ from __future__ import annotations from collections.abc import Mapping from dataclasses import dataclass -from typing import Protocol +from typing import TYPE_CHECKING, Protocol from litellm.types.guardrails import PresidioPerRequestConfig +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging + class ResponsesClientWebSocket(Protocol): """Client-facing websocket surface used by the Responses API websocket handlers.""" @@ -43,6 +46,12 @@ class PresidioGuardrailCallback(Protocol): ) -> str: ... +class ResponsesWebSocketTurnFailureHook(Protocol): + """Books one failed ``response.create`` turn of a native Responses API socket from that turn's logging object.""" + + async def __call__(self, turn_logging_obj: Logging, exception: Exception) -> None: ... + + @dataclass(frozen=True, slots=True) class ResponsesWebSocketRequestDefaults: """Deployment-level request parameters merged into every ``response.create`` frame relayed over a native websocket.""" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4a962e5b3bd..90d6f86986e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4478,6 +4478,7 @@ class LiteLLMRealtimeStreamLoggingObject(LiteLLMPydanticObjectBase): results: SkipValidation[OpenAIRealtimeStreamList] usage: Usage service_tier: str | None = None + id: str | None = None _hidden_params: dict = {} @property diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 771ed5d2e15..8ea3f36068a 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -1998,34 +1998,34 @@ async def test_arealtime_marks_litellm_params_async(monkeypatch): @pytest.mark.asyncio -async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log( +async def test_aresponses_websocket_writes_no_connection_level_success_log( monkeypatch: pytest.MonkeyPatch, ): from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.responses.main import base_llm_http_handler - success_events = [] + logged_call_ids = [] class CaptureLogger(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - success_events.append(response_obj) + logged_call_ids.append(kwargs.get("litellm_call_id")) monkeypatch.setattr(litellm, "callbacks", [CaptureLogger()]) monkeypatch.setattr(litellm, "failure_callback", []) monkeypatch.setattr(litellm, "_async_failure_callback", []) monkeypatch.setattr(litellm, "success_callback", []) monkeypatch.setattr(litellm, "_async_success_callback", []) - failure = litellm.BadRequestError(message="invalid_encrypted_content", model="gpt-4o", llm_provider="openai") - with patch.object( # test-quality-ok: the provider socket is the seam; how the wrapper treats the relay's outcome is under test - base_llm_http_handler, "async_responses_websocket", AsyncMock(return_value=failure) + with patch.object( # test-quality-ok: the provider socket is the seam; each turn logs inside it, so the wrapper must not + base_llm_http_handler, "async_responses_websocket", AsyncMock(return_value=None) ): - outcome = await litellm._aresponses_websocket(model="openai/gpt-4o", websocket=MagicMock(), api_key="sk-test") + await litellm._aresponses_websocket( + model="openai/gpt-4o", websocket=MagicMock(), api_key="sk-test", litellm_call_id="socket-call" + ) await asyncio.sleep(0) with contextlib.suppress(asyncio.TimeoutError): await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) - assert outcome is failure - assert success_events == [] + assert "socket-call" not in logged_call_ids @pytest.mark.asyncio @@ -8078,6 +8078,74 @@ def test_normalize_logging_result_prices_responses_websocket_at_returned_service assert priority_http_cost > default_http_cost +@pytest.mark.parametrize( + ("completed_response_id", "expected_request_id"), + (("resp_wrapped_turn_2", "resp_wrapped_turn_2"), (None, "turn-2-call")), +) +def test_responses_websocket_turn_is_logged_under_the_response_id_the_client_saw( + completed_response_id: str | None, expected_request_id: str +): + events = [ + {"type": "response.created", "response": {"id": completed_response_id}}, + { + "type": "response.completed", + "response": { + "id": completed_response_id, + "usage": {"input_tokens": 100, "output_tokens": 40, "total_tokens": 140}, + }, + }, + ] + logging_obj = _responses_ws_logging_obj() + now = datetime.datetime.now() + + payload = litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload( + kwargs={"litellm_call_id": "turn-2-call", "model": "gpt-4o", "messages": []}, + init_response_obj=logging_obj.normalize_logging_result(result=events), + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["id"] == expected_request_id + + +def test_copy_for_new_call_gives_each_call_its_own_id_clock_and_metadata(): + reservation = {"reserved_cost": 0.5, "finalized": False} + connection = _responses_ws_logging_obj() + connection_start = connection.start_time + connection.litellm_params = { + "metadata": {"user_api_key_alias": "codex", "budget_reservation": reservation}, + "litellm_metadata": {"model_info": {"id": "dep-1"}}, + "litellm_call_id": connection.litellm_call_id, + } + connection.model_call_details = { + **connection.model_call_details, + "litellm_params": connection.litellm_params, + "first_api_call_start_time": connection_start, + } + turn_start = datetime.datetime(2026, 10, 9, 17, 24, 47) + + turn = connection.copy_for_new_call(litellm_call_id="turn-1-call", start_time=turn_start) + turn.litellm_params["metadata"]["status"] = "failure" + turn.litellm_params["litellm_metadata"]["model_group"] = "gpt-4o" + turn.model_call_details["standard_logging_object"] = {"id": "resp_1"} + + assert (turn.litellm_call_id, turn.start_time) == ("turn-1-call", turn_start) + assert turn.model_call_details["litellm_call_id"] == "turn-1-call" + assert turn.model_call_details["start_time"] == turn_start + assert turn.model_call_details["litellm_params"]["litellm_call_id"] == "turn-1-call" + assert turn.model_call_details["first_api_call_start_time"] is None + assert turn.litellm_params["metadata"]["budget_reservation"] is reservation + assert (connection.litellm_call_id, connection.start_time) == ("responses-ws-usage-test", connection_start) + assert connection.litellm_params["litellm_call_id"] == "responses-ws-usage-test" + assert "status" not in connection.litellm_params["metadata"] + assert "model_group" not in connection.litellm_params["litellm_metadata"] + assert "standard_logging_object" not in connection.model_call_details + assert connection.model_call_details["first_api_call_start_time"] == connection_start + + def test_get_standard_logging_object_payload_reads_overhead_from_logging_obj_for_dict_results(logging_obj): """LIT-5466: /v1/messages returns a plain dict with no _hidden_params, so the overhead recorded on the logging object must reach hidden_params.litellm_overhead_time_ms (SpendLogs).""" diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index d9cf8536912..28589a2f400 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4294,6 +4294,40 @@ async def test_post_call_failure_hook_releases_parallel_slot_v3(): ) +@pytest.mark.asyncio +async def test_post_call_failure_hook_for_another_call_keeps_the_owner_slot_v3(): + """ + A native Responses WebSocket holds one parallel slot for the whole socket and books each failed turn + under that turn's own call id. A failed turn must not free the slot while the socket is still open. + """ + _api_key = hash_token("sk-ws-turn-failure") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) + counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + connection_data: Dict[str, Any] = {"model": "gpt-4o-mini", "litellm_call_id": "socket-call"} + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=local_cache, data=connection_data, call_type="" + ) + + await handler.async_post_call_failure_hook( + request_data={"model": "gpt-4o-mini", "litellm_call_id": "turn-2-call"}, + original_exception=Exception("previous_response_not_found"), + user_api_key_dict=user_api_key_dict, + ) + in_flight_after_turn_failure = handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) + await handler.async_post_call_failure_hook( + request_data=connection_data, + original_exception=Exception("socket closed"), + user_api_key_dict=user_api_key_dict, + ) + + assert in_flight_after_turn_failure == 1 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 + + @pytest.mark.asyncio async def test_success_event_releases_parallel_slot_v3(monkeypatch): """ diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 5dac1d25572..ce967403265 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -542,29 +542,6 @@ class TestResponsesWSStreamingFirstMessage: class TestWSSessionCostTracking: - @pytest.mark.asyncio - async def test_router_budget_limiter_skips_aresponses_websocket_call_type(self): - """ - RouterBudgetLimiting.async_log_success_event must not raise when - call_type='_aresponses_websocket', even when standard_logging_object is None. - Per-turn costs are tracked by individual aresponses calls inside the session; - the outer session wrapper fires with result=None. - """ - from litellm.router_strategy.budget_limiter import RouterBudgetLimiting - - limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting) - kwargs = { - "call_type": "_aresponses_websocket", - "standard_logging_object": None, - "litellm_params": {"custom_llm_provider": "vertex_ai"}, - } - await limiter.async_log_success_event( - kwargs=kwargs, - response_obj=None, - start_time=None, - end_time=None, - ) - @pytest.mark.asyncio async def test_router_budget_limiter_skips_arealtime_call_type(self): """Same guard applies to _arealtime WS session wrappers.""" @@ -820,8 +797,7 @@ class TestResponsesWSFirstFrameModelAuth: ws.close.assert_not_awaited() @pytest.mark.asyncio - @pytest.mark.parametrize("provider_rejected", [True, False]) - async def test_endpoint_books_a_provider_rejected_connection_as_a_failed_request(self, provider_rejected: bool): + async def test_endpoint_books_each_failed_turn_as_its_own_failed_request(self): from litellm.proxy.response_api_endpoints.endpoints import ( responses_websocket_endpoint, ) @@ -833,23 +809,36 @@ class TestResponsesWSFirstFrameModelAuth: ws.url = "ws://testserver/v1/responses" ws.accept = AsyncMock() ws.receive_text = AsyncMock( - return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": []}) + return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": "first turn"}) ) ws.close = AsyncMock() + connection_metadata = {"user_api_key_alias": "codex", "model_group": "gpt-4o-mini"} processor = MagicMock() processor.common_processing_pre_call_logic = AsyncMock( - return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock()) + return_value=( + { + "model": "gpt-4o-mini", + "metadata": connection_metadata, + "litellm_metadata": {}, + "litellm_call_id": "connection-call", + }, + MagicMock(), + ) ) + turn_logging_obj = MagicMock() + turn_logging_obj.litellm_call_id = "turn-2-call" failure = litellm.BadRequestError( - message="invalid_encrypted_content", model="gpt-4o-mini", llm_provider="openai" + message="previous_response_not_found", model="gpt-4o-mini", llm_provider="openai" ) - async def fake_llm_call(): - return failure if provider_rejected else None + async def relay_with_one_failed_turn(): + routed = mock_route_request.await_args.kwargs["data"] + await routed["responses_ws_turn_failure_hook"](turn_logging_obj, failure) proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock() user_api_key_dict = MagicMock() with ( @@ -857,15 +846,15 @@ class TestResponsesWSFirstFrameModelAuth: "litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth", new_callable=AsyncMock, ), - patch( # test-quality-ok: the pre-call processor needs a live proxy; what the endpoint does with the relay's outcome is under test + patch( # test-quality-ok: the pre-call processor needs a live proxy; how the endpoint books a failed turn is under test "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", return_value=processor, ), - patch( # test-quality-ok: routing is the seam that hands back the relay's outcome + patch( # test-quality-ok: routing is the seam that hands the turn failure hook to the relay "litellm.proxy.route_llm_request.route_request", new_callable=AsyncMock, - return_value=fake_llm_call(), - ), + side_effect=lambda **_: relay_with_one_failed_turn(), + ) as mock_route_request, patch( # test-quality-ok: the failure hook is the proxy's only path to a failed spend log row "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj, @@ -878,14 +867,19 @@ class TestResponsesWSFirstFrameModelAuth: ) ws.close.assert_not_awaited() - if not provider_rejected: - proxy_logging_obj.post_call_failure_hook.assert_not_awaited() - return proxy_logging_obj.post_call_failure_hook.assert_awaited_once() booked = proxy_logging_obj.post_call_failure_hook.await_args.kwargs assert booked["original_exception"] is failure assert booked["user_api_key_dict"] is user_api_key_dict - assert booked["request_data"]["model"] == "gpt-4o-mini" + turn_request = booked["request_data"] + assert turn_request["litellm_call_id"] == "turn-2-call" + assert turn_request["litellm_logging_obj"] is turn_logging_obj + assert turn_request["model"] == "gpt-4o-mini" + assert turn_request["metadata"] == connection_metadata + assert turn_request["metadata"] is not connection_metadata + assert "input" not in turn_request + assert "first_message" not in turn_request + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once_with(user_api_key_dict) @pytest.mark.asyncio async def test_endpoint_sends_an_error_frame_when_routing_rejects_the_connection(self): @@ -914,6 +908,7 @@ class TestResponsesWSFirstFrameModelAuth: ) proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock(return_value=None) user_api_key_dict = MagicMock() with ( @@ -951,6 +946,7 @@ class TestResponsesWSFirstFrameModelAuth: assert booked["original_exception"] is rejection assert booked["user_api_key_dict"] is user_api_key_dict assert booked["request_data"]["model"] == "gpt-4o-mini" + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once_with(user_api_key_dict) @pytest.mark.asyncio async def test_reruns_model_auth_for_first_frame_model(self): diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index b29e9a6b7db..5b95c472480 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1898,6 +1898,29 @@ def _make_streaming(**kwargs): return ResponsesWebSocketStreaming(**kwargs) +class _TurnLoggingRecorder: + """Stands in for a socket's connection-level Logging object and keeps every per-turn copy it hands out.""" + + def __init__(self, response_cost: float = 0.0) -> None: + self.response_cost = response_cost + self.turns: list[MagicMock] = [] + self.connection = MagicMock() + self.connection.copy_for_new_call = MagicMock(side_effect=self._copy_for_new_call) + + def _copy_for_new_call(self, litellm_call_id: str, start_time: object) -> MagicMock: + from unittest.mock import AsyncMock + + turn = MagicMock() + turn.litellm_call_id = litellm_call_id + turn.start_time = start_time + turn.model_call_details = {} + turn.dispatch_success_handlers = AsyncMock() + turn.dispatch_failure_handlers = AsyncMock() + turn.response_cost_calculator = MagicMock(return_value=self.response_cost) + self.turns.append(turn) + return turn + + class TestNativeWebSocketGuardrailMasking: """Exercises the input/output PII masking hooks on ResponsesWebSocketStreaming.""" @@ -3051,12 +3074,11 @@ class TestNativeWebSocketEncryptedContentAffinity: Exception("stop"), ] ) - logging_obj = MagicMock() - logging_obj.dispatch_success_handlers = AsyncMock() + recorder = _TurnLoggingRecorder() handler = _make_streaming( websocket=websocket, backend_ws=backend_ws, - logging_obj=logging_obj, + logging_obj=recorder.connection, request_data={"litellm_metadata": dict(_AFFINITY_METADATA)}, custom_llm_provider="openai", ) @@ -3075,7 +3097,7 @@ class TestNativeWebSocketEncryptedContentAffinity: ) assert completed["response"]["output"][0]["encrypted_content"] == wrapped_content await asyncio.sleep(0) - logged = logging_obj.dispatch_success_handlers.await_args[0][0] + logged = recorder.turns[0].dispatch_success_handlers.await_args[0][0] assert logged[0]["response"]["id"] == completed["response"]["id"] @pytest.mark.asyncio @@ -3163,14 +3185,11 @@ class TestNativeWebSocketEncryptedContentAffinity: Exception("stop"), ] ) - logging_obj = MagicMock() - logging_obj.dispatch_success_handlers = AsyncMock() - logging_obj.dispatch_failure_handlers = AsyncMock() - logging_obj.response_cost_calculator = MagicMock(return_value=0.0) + recorder = _TurnLoggingRecorder() handler = _make_streaming( websocket=websocket, backend_ws=backend_ws, - logging_obj=logging_obj, + logging_obj=recorder.connection, request_data={}, authorized_model="gpt-5.6", custom_llm_provider="openai", @@ -3179,163 +3198,257 @@ class TestNativeWebSocketEncryptedContentAffinity: await handler.backend_to_client() await asyncio.sleep(0) - logging_obj.dispatch_success_handlers.assert_not_awaited() - logging_obj.dispatch_failure_handlers.assert_awaited_once() - exception = logging_obj.dispatch_failure_handlers.await_args[0][0] + (turn,) = recorder.turns + turn.dispatch_success_handlers.assert_not_awaited() + turn.dispatch_failure_handlers.assert_awaited_once() + exception = turn.dispatch_failure_handlers.await_args[0][0] assert exception.status_code == expected_status assert failure_frame.get("error", failure_frame.get("response", {}).get("error"))["message"] in str(exception) - @pytest.mark.asyncio - async def test_backend_to_client_bills_completed_turns_before_a_failure(self): + +def _completed_turn(response_id: str, input_tokens: int, output_tokens: int) -> list[dict[str, object]]: + return [ + {"type": "response.created", "response": {"id": response_id, "status": "in_progress"}}, + {"type": "response.output_text.delta", "delta": "hi"}, + { + "type": "response.completed", + "response": { + "id": response_id, + "status": "completed", + "output": [], + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + }, + }, + }, + ] + + +def _response_create(text: str) -> str: + return json.dumps( + { + "type": "response.create", + "model": "gpt-5.4-mini", + "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}], + } + ) + + +class _ScriptedBackend: + """Answers each forwarded response.create with the next scripted turn of upstream events.""" + + def __init__(self, turns: list[list[dict[str, object]]]) -> None: import asyncio + + self.turns = list(turns) + self.events: asyncio.Queue[str] = asyncio.Queue() + self.sent: list[str] = [] + + async def send(self, frame: str) -> None: + self.sent.append(frame) + if json.loads(frame).get("type") == "response.create" and self.turns: + for event in self.turns.pop(0): + self.events.put_nowait(json.dumps(event)) + + async def recv(self, decode: bool = False) -> str: + return await self.events.get() + + async def close(self) -> None: + return None + + +async def _until(condition) -> bool: + import asyncio + + for _ in range(200): + if condition(): + return True + await asyncio.sleep(0) + return False + + +class _ScriptedClient: + """Sends each frame only after the previous turn's log was dispatched, then hangs up.""" + + def __init__(self, frames: list[str], recorder: _TurnLoggingRecorder) -> None: + self.frames = list(frames) + self.recorder = recorder + self.sent_text: list[str] = [] + self.logged_turns_before_each_frame: list[int] = [] + self.logged_turns_before_hang_up: int | None = None + + def _logged_turns(self) -> int: + return sum( + 1 + for turn in self.recorder.turns + if turn.dispatch_success_handlers.await_count or turn.dispatch_failure_handlers.await_count + ) + + async def receive_text(self) -> str: + sent_so_far = len(self.logged_turns_before_each_frame) + await _until(lambda: self._logged_turns() >= sent_so_far) + if self.frames: + self.logged_turns_before_each_frame.append(self._logged_turns()) + return self.frames.pop(0) + self.logged_turns_before_hang_up = self._logged_turns() + raise Exception("client hung up") + + async def send_text(self, text: str) -> None: + self.sent_text.append(text) + + +class TestNativeWebSocketPerTurnLogging: + """One native Responses socket carries many response.create turns, and each turn is its own logged call.""" + + @pytest.mark.asyncio + async def test_each_turn_is_logged_as_its_own_call_while_the_socket_stays_open(self): + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + from litellm.responses.utils import ResponsesAPIRequestUtils + + recorder = _TurnLoggingRecorder() + client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder) + backend = _ScriptedBackend( + [_completed_turn("resp_1", 10, 1), _completed_turn("resp_2", 20, 2), _completed_turn("resp_3", 30, 3)] + ) + handler = _make_streaming( + websocket=client, backend_ws=backend, logging_obj=recorder.connection, request_data={} + ) + + await handler.bidirectional_forward() + + assert client.logged_turns_before_each_frame == [0, 1, 2] + assert client.logged_turns_before_hang_up == 3 + assert len({turn.litellm_call_id for turn in recorder.turns}) == 3 + assert recorder.connection.litellm_call_id not in {turn.litellm_call_id for turn in recorder.turns} + logged_events = [turn.dispatch_success_handlers.await_args[0][0] for turn in recorder.turns] + sent_frames = [json.loads(text) for text in client.sent_text] + ids_the_client_saw = [frame["response"]["id"] for frame in sent_frames if frame["type"] == "response.completed"] + assert [events[-1]["response"]["id"] for events in logged_events] == ids_the_client_saw + assert [ + ResponsesAPIRequestUtils.decode_responses_api_response_id(response_id)["response_id"] + for response_id in ids_the_client_saw + ] == ["resp_1", "resp_2", "resp_3"] + assert [len(events) for events in logged_events] == [2, 2, 2] + assert [events[-1]["response"]["usage"]["input_tokens"] for events in logged_events] == [10, 20, 30] + assert [turn.model_call_details["messages"] for turn in recorder.turns] == [ + [{"role": "user", "content": "one"}], + [{"role": "user", "content": "two"}], + [{"role": "user", "content": "three"}], + ] + for turn in recorder.turns: + turn.dispatch_failure_handlers.assert_not_awaited() + + @pytest.mark.asyncio + async def test_a_failed_turn_is_booked_alone_and_leaves_the_other_turns_successful(self): from unittest.mock import AsyncMock import websockets.exceptions # noqa: F401 (lazy submodule must be importable) - websocket = MagicMock() - websocket.send_text = AsyncMock() - backend_ws = MagicMock() - backend_ws.recv = AsyncMock( - side_effect=[ - json.dumps( + recorder = _TurnLoggingRecorder() + on_turn_failure = AsyncMock() + client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder) + backend = _ScriptedBackend( + [ + _completed_turn("resp_1", 10, 1), + [ { - "type": "response.completed", - "response": { - "id": "resp_1", - "status": "completed", - "output": [], - "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + "type": "error", + "status": 400, + "error": { + "type": "invalid_request_error", + "code": "previous_response_not_found", + "message": "Previous response with id 'resp_x' not found.", }, } - ), - json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": "bad turn"}}), - Exception("stop"), + ], + _completed_turn("resp_3", 30, 3), ] ) - logging_obj = MagicMock() - logging_obj.dispatch_success_handlers = AsyncMock() - logging_obj.dispatch_failure_handlers = AsyncMock() - logging_obj.response_cost_calculator = MagicMock(return_value=0.01) - handler = _make_streaming(websocket=websocket, backend_ws=backend_ws, logging_obj=logging_obj, request_data={}) + handler = _make_streaming( + websocket=client, + backend_ws=backend, + logging_obj=recorder.connection, + request_data={}, + authorized_model="gpt-5.4-mini", + custom_llm_provider="openai", + on_turn_failure=on_turn_failure, + ) - await handler.backend_to_client() - await asyncio.sleep(0) + await handler.bidirectional_forward() - logging_obj.record_partial_usage_for_failure.assert_called_once() - usage, response_cost = logging_obj.record_partial_usage_for_failure.call_args[0] - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15) + first, failed, third = recorder.turns + for successful in (first, third): + successful.dispatch_success_handlers.assert_awaited_once() + successful.dispatch_failure_handlers.assert_not_awaited() + failed.dispatch_success_handlers.assert_not_awaited() + failed.dispatch_failure_handlers.assert_awaited_once() + on_turn_failure.assert_awaited_once() + booked_turn, booked_exception = on_turn_failure.await_args[0] + assert booked_turn is failed + assert booked_exception is failed.dispatch_failure_handlers.await_args[0][0] + assert booked_exception.status_code == 400 + assert "not found" in str(booked_exception) + + @pytest.mark.asyncio + async def test_a_failed_turn_records_only_its_own_partial_usage(self): + import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + + recorder = _TurnLoggingRecorder(response_cost=0.01) + client = _ScriptedClient([_response_create(text) for text in ("one", "two")], recorder) + backend = _ScriptedBackend( + [ + _completed_turn("resp_1", 10, 5), + [ + {"type": "response.created", "response": {"id": "resp_2", "status": "in_progress"}}, + { + "type": "response.failed", + "response": { + "id": "resp_2", + "status": "failed", + "error": {"code": "server_error", "message": "upstream blew up"}, + "usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9}, + }, + }, + ], + ] + ) + handler = _make_streaming( + websocket=client, backend_ws=backend, logging_obj=recorder.connection, request_data={} + ) + + await handler.bidirectional_forward() + + first, failed = recorder.turns + first.record_partial_usage_for_failure.assert_not_called() + usage, response_cost = failed.record_partial_usage_for_failure.call_args[0] + assert (usage.prompt_tokens, usage.completion_tokens) == (0, 0) assert response_cost == 0.01 - logging_obj.dispatch_success_handlers.assert_not_awaited() - logging_obj.dispatch_failure_handlers.assert_awaited_once() @pytest.mark.asyncio - async def test_bidirectional_forward_returns_the_provider_failure(self): - import asyncio - from unittest.mock import AsyncMock - + async def test_a_turn_still_in_flight_when_the_socket_closes_is_logged_with_what_it_received(self): import websockets.exceptions # noqa: F401 (lazy submodule must be importable) - backend_drained = asyncio.Event() - backend_events = [ - json.dumps({"type": "response.created", "response": {"id": "resp_1", "status": "in_progress"}}), - json.dumps( - { - "type": "error", - "status": 400, - "error": { - "type": "invalid_request_error", - "code": "invalid_encrypted_content", - "message": "could not be verified", - }, - } - ), - ] - - async def recv(decode=False): - if backend_events: - return backend_events.pop(0) - backend_drained.set() - raise Exception("stop") - - async def receive_text(): - await backend_drained.wait() - raise Exception("client gone") - - websocket = MagicMock() - websocket.send_text = AsyncMock() - websocket.receive_text = receive_text - backend_ws = MagicMock() - backend_ws.recv = recv - backend_ws.send = AsyncMock() - backend_ws.close = AsyncMock() - logging_obj = MagicMock() - logging_obj.dispatch_success_handlers = AsyncMock() - logging_obj.dispatch_failure_handlers = AsyncMock() - logging_obj.response_cost_calculator = MagicMock(return_value=0.0) + recorder = _TurnLoggingRecorder() + client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder) + backend = _ScriptedBackend( + [ + _completed_turn("resp_1", 10, 1), + [{"type": "response.created", "response": {"id": "resp_2", "status": "in_progress"}}], + ] + ) handler = _make_streaming( - websocket=websocket, - backend_ws=backend_ws, - logging_obj=logging_obj, - request_data={}, - authorized_model="gpt-5.6", - custom_llm_provider="openai", + websocket=client, backend_ws=backend, logging_obj=recorder.connection, request_data={} ) - failure = await handler.bidirectional_forward() + await handler.bidirectional_forward() + await _until(lambda: recorder.turns[1].dispatch_success_handlers.await_count == 1) - assert isinstance(failure, Exception) - assert failure.status_code == 400 - assert "could not be verified" in str(failure) - - @pytest.mark.asyncio - async def test_bidirectional_forward_returns_none_after_a_completed_turn(self): - import asyncio - from unittest.mock import AsyncMock - - import websockets.exceptions # noqa: F401 (lazy submodule must be importable) - - backend_drained = asyncio.Event() - backend_events = [ - json.dumps( - { - "type": "response.completed", - "response": { - "id": "resp_1", - "status": "completed", - "output": [], - "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, - }, - } - ), - ] - - async def recv(decode=False): - if backend_events: - return backend_events.pop(0) - backend_drained.set() - raise Exception("stop") - - async def receive_text(): - await backend_drained.wait() - raise Exception("client gone") - - websocket = MagicMock() - websocket.send_text = AsyncMock() - websocket.receive_text = receive_text - backend_ws = MagicMock() - backend_ws.recv = recv - backend_ws.send = AsyncMock() - backend_ws.close = AsyncMock() - logging_obj = MagicMock() - logging_obj.dispatch_success_handlers = AsyncMock() - logging_obj.dispatch_failure_handlers = AsyncMock() - handler = _make_streaming( - websocket=websocket, - backend_ws=backend_ws, - logging_obj=logging_obj, - request_data={}, - authorized_model="gpt-5.6", - custom_llm_provider="openai", - ) - - assert await handler.bidirectional_forward() is None + first, in_flight, unanswered = recorder.turns + first.dispatch_success_handlers.assert_awaited_once() + in_flight.dispatch_success_handlers.assert_awaited_once() + assert [event["type"] for event in in_flight.dispatch_success_handlers.await_args[0][0]] == ["response.created"] + unanswered.dispatch_success_handlers.assert_not_awaited() + unanswered.dispatch_failure_handlers.assert_not_awaited() diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 90da579d910..522f4987e13 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -56,7 +56,9 @@ async def _log_success(limiter: RouterBudgetLimiting, kwargs: dict[str, object]) @pytest.mark.asyncio -@pytest.mark.parametrize("call_type", ["aresponses", "anthropic_messages", "aembedding", "arerank"]) +@pytest.mark.parametrize( + "call_type", ["aresponses", "_aresponses_websocket", "anthropic_messages", "aembedding", "arerank"] +) async def test_provider_spend_tracked_when_litellm_params_omits_provider(disable_budget_sync, call_type): """Non-chat surfaces carry the provider only on the standard logging payload.""" limiter = RouterBudgetLimiting(