From 5ed3629242ec572fcb8b5ee18fc766aa337b5a89 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 15 Oct 2026 19:05:22 -0700 Subject: [PATCH] fix(responses): write one spend-log row per native WebSocket turn (#46768) * fix(responses): write one spend-log row per native WebSocket turn Each response.create turn on a native Responses WebSocket connection logs and bills as its own call when its terminal event arrives, keyed by the resp_ id the client saw, instead of one row at socket close that sums every turn * test(responses): mock the websocket slot release and mark the per-turn start time naive * style(responses): keep the per-turn start time line within ruff format * refactor(responses): type websocket turn events as read-only mappings * fix(responses): settle the websocket budget reservation on the first turn only Every per-turn logging copy carried the connection's admission reservation, so two turns whose cost callbacks overlapped could both reconcile it. The first turn keeps the reservation and later turns charge their cost directly * fix(proxy): keep the legacy limiter's websocket slot until disconnect A websocket turn callback no longer decrements current_requests in the legacy limiter. The socket's slot is released once from arelease_max_parallel_requests_on_disconnect, which now dispatches to the legacy handler as well --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 18 ++ .../internal_call_metadata.py | 9 +- litellm/litellm_core_utils/litellm_logging.py | 51 ++-- .../proxy/hooks/parallel_request_limiter.py | 57 ++++- .../proxy/response_api_endpoints/endpoints.py | 2 + litellm/proxy/utils.py | 4 +- litellm/responses/streaming_iterator.py | 232 ++++++++++-------- litellm/types/utils.py | 1 + .../integration/_support/responses_ws_peer.py | 186 ++++++++++++++ .../test_openai_responses_websocket_wire.py | 226 +++-------------- .../test_responses_websocket_turn_rows.py | 114 +++++++++ .../test_litellm_logging.py | 52 ++++ .../hooks/test_parallel_request_limiter.py | 32 ++- .../response_api_endpoints/test_endpoints.py | 2 + .../test_responses_websocket_all_providers.py | 98 +++++++- tests/unit/test_cost_calculator.py | 20 ++ 16 files changed, 769 insertions(+), 335 deletions(-) create mode 100644 tests/integration/_support/responses_ws_peer.py create mode 100644 tests/integration/spend/test_responses_websocket_turn_rows.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 43ce71af7b8..0e315e97432 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -3008,6 +3008,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 @@ -3059,6 +3060,23 @@ 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( + response.id for _, response in _billable_responses_ws_events(results) if 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) + return LiteLLMRealtimeStreamLoggingObject( + id=ResponsesWebSocketTokenUsageProcessor.response_id_from_responses_ws_results(results), + usage=ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(results), + results=cast(OpenAIRealtimeStreamList, list(results)), # cast-ok: raw WS frames, unvalidated field + 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/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index c5e57bad050..1200d799018 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -157,7 +157,10 @@ def sanitize_user_api_key_auth(auth: object) -> object: return auth -def _sanitized(parent_metadata: Mapping[str, object]) -> dict[str, object]: # mutable-ok: SDK metadata kwarg +def without_budget_reservation( + parent_metadata: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: SDK metadata kwarg + """Parent metadata minus the admission reservation, also inside ``user_api_key_auth``.""" return { k: sanitize_user_api_key_auth(v) if k == _USER_API_KEY_AUTH_KEY else v for k, v in parent_metadata.items() @@ -176,7 +179,7 @@ def forwarded_internal_call_metadata( """ if not parent_metadata: return {} - return _sanitized(parent_metadata) | {INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin} + return without_budget_reservation(parent_metadata) | {INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin} def parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> Mapping[str, str]: @@ -202,4 +205,4 @@ def sanitized_forwardable_call_metadata( must not inherit per-request state such as its routing decision or logging payload. """ identity: Final = {k: v for k, v in parent_metadata.items() if k in FORWARDABLE_IDENTITY_METADATA_KEYS} - return _sanitized(identity) | {INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin} + return without_budget_reservation(identity) | {INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin} diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 7038cc531f5..e766a2be9c3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -19,7 +19,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._internal_context import post_response_phase @@ -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 @@ -771,6 +775,36 @@ 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": + """One turn of a long-lived connection: the connection's deployment, key and metadata under its own call id, + start time and chunk buffers, so the turn logs and bills as a request of its own.""" + connection_params: Final = _LITELLM_PARAMS.validate_python( + self.litellm_params # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped + ) + litellm_params: Final[dict[str, object]] = { + **connection_params, + **{ + key: _LITELLM_PARAMS.validate_python(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) @@ -2384,20 +2418,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.relays_llm_passthrough and isinstance(result, Response): from litellm.utils import ProviderConfigManager @@ -6726,7 +6749,7 @@ def get_standard_logging_object_payload( else raw_usage_dict ) - id = response_obj.get("id", kwargs.get("litellm_call_id")) + id = response_obj.get("id") 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/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 19251b53f13..47ea50debc9 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -1,6 +1,6 @@ import asyncio import sys -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn @@ -23,7 +23,7 @@ from litellm.proxy.auth.budget_throttle import throttled_limit from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup -from litellm.types.utils import Usage +from litellm.types.utils import CallTypes, Usage if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -48,6 +48,14 @@ def _response_total_tokens(response_obj: object) -> int: return response_usage.total_tokens if isinstance(response_usage, Usage) else 0 +_CONNECTION_SCOPED_CALL_TYPES: Final = frozenset({CallTypes.aresponses_websocket.value}) + + +def _slots_released(kwargs: Mapping[str, object]) -> int: + """A websocket turn callback releases nothing; the socket's slot is released once at disconnect.""" + return 0 if kwargs.get("call_type") in _CONNECTION_SCOPED_CALL_TYPES else 1 + + class CacheObject(TypedDict): current_global_requests: dict | None request_count_api_key: dict | None @@ -509,6 +517,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): litellm_parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(kwargs=kwargs) try: self.print_verbose("INSIDE parallel request limiter ASYNC SUCCESS LOGGING") + released: Final = _slots_released(kwargs) # pyright: ignore[reportUnknownArgumentType] # legacy callback kwargs are untyped global_max_parallel_requests: Final = kwargs["litellm_params"]["metadata"].get( "global_max_parallel_requests", None @@ -540,7 +549,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): # decrement await self.internal_usage_cache.async_increment_cache( key=_key, - value=-1, + value=-released, local_only=True, litellm_parent_otel_span=litellm_parent_otel_span, ) @@ -568,7 +577,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -609,7 +618,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -635,7 +644,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -661,7 +670,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -687,7 +696,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -709,6 +718,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): self.print_verbose("Inside Max Parallel Request Failure Hook") litellm_parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(kwargs=kwargs) _metadata: Final = kwargs["litellm_params"].get("metadata", {}) or {} + released: Final = _slots_released(kwargs) # pyright: ignore[reportUnknownArgumentType] # legacy callback kwargs are untyped global_max_parallel_requests: Final = _metadata.get("global_max_parallel_requests", None) user_api_key: Final = _metadata.get("user_api_key", None) self.print_verbose(f"user_api_key: [set={user_api_key is not None}]") @@ -736,7 +746,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): # decrement await self.internal_usage_cache.async_increment_cache( key=_key, - value=-1, + value=-released, local_only=True, litellm_parent_otel_span=litellm_parent_otel_span, ) @@ -758,7 +768,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): } new_val: Final = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - released, 0), "current_tpm": current["current_tpm"], "current_rpm": current["current_rpm"], } @@ -774,6 +784,33 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): verbose_proxy_logger.exception("Inside Parallel Request Limiter: An exception occurred - %s", e) @with_service_target("rate_limits") + async def async_release_max_parallel_requests_on_disconnect(self, user_api_key_dict: UserAPIKeyAuth) -> None: + """Release the slot ``async_pre_call_hook`` took for a connection whose per-turn callbacks released none.""" + from litellm.proxy.proxy_server import ( + general_settings, # pyright: ignore[reportUnknownVariableType] # module-level settings dict is untyped + ) + + now: Final = self._clock() + await self.async_log_success_event( # pyright: ignore[reportUnknownMemberType] # legacy callback signature is untyped + kwargs={ + "litellm_params": { + "metadata": { + "user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict), + "user_api_key_user_id": user_api_key_dict.user_id, + "user_api_key_team_id": user_api_key_dict.team_id, + "user_api_key_model_max_budget": user_api_key_dict.model_max_budget, # pyright: ignore[reportUnknownMemberType] # UserAPIKeyAuth.model_max_budget is a bare dict + "global_max_parallel_requests": general_settings.get( # pyright: ignore[reportUnknownMemberType] # module-level settings dict is untyped + "global_max_parallel_requests" + ), + } + }, + "user": user_api_key_dict.end_user_id, + }, + response_obj=None, + start_time=now, + end_time=now, + ) + async def get_internal_user_object( self, user_id: str, diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 4bfc4433815..b7f57d93a9b 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1905,6 +1905,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/proxy/utils.py b/litellm/proxy/utils.py index 9788c4635c9..5d2ebe73cac 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -231,7 +231,7 @@ from litellm.proxy.hooks.cache_control_check import ( # noqa: F401, RUF100 # l ) from litellm.proxy.hooks.parallel_request_limiter import ( # noqa: F401, RUF100 # legacy module exports PROXY_MaxParallelRequestsHandler, - _PROXY_MaxParallelRequestsHandler, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export + _PROXY_MaxParallelRequestsHandler, # pyright: ignore[reportPrivateUsage] # backwards-compatible package export ) from litellm.proxy.hooks.parallel_request_limiter_v3 import ( # noqa: F401, RUF100 # legacy module exports PROXY_MaxParallelRequestsHandler_v3, @@ -4405,7 +4405,7 @@ class ProxyLogging: double-decrement under the limiter's in-memory fallback. """ limiter: Final = self.get_proxy_hook("parallel_request_limiter") - if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): + if not isinstance(limiter, (PROXY_MaxParallelRequestsHandler_v3, _PROXY_MaxParallelRequestsHandler)): return await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 7532a528248..7cf68332c07 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 @@ -26,6 +28,7 @@ from litellm.constants import ( from litellm.exceptions import MidStreamFallbackError, RateLimitError from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( @@ -1902,6 +1905,77 @@ _HISTORY_TERMINAL_EVENT_TYPES: Final = frozenset({"response.completed", "respons RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output_text", "text"}) _RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"}) +_PER_TURN_METADATA_KEYS: Final = frozenset({"metadata", "litellm_metadata"}) + +_RESPONSES_WS_TERMINAL_EVENT_TYPES: Final = _HISTORY_TERMINAL_EVENT_TYPES | _RESPONSES_WS_FAILURE_EVENT_TYPES + + +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[Mapping[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) + + +def _record_usage_for_failure(turn_logging_obj: LiteLLMLoggingObj, events: Sequence[Mapping[str, object]]) -> None: + from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor + + 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) + + +@dataclass(frozen=True, slots=True) +class _ResponsesWsTurn: + """One ``response.create`` on a native Responses socket, logged and billed as its own call once a terminal + event (``response.completed``, ``response.incomplete``, ``response.failed`` or ``error``) closes it.""" + + logging_obj: LiteLLMLoggingObj + events: tuple[Mapping[str, object], ...] = () + + def with_event(self, event: Mapping[str, object]) -> _ResponsesWsTurn: + return _ResponsesWsTurn(logging_obj=self.logging_obj, events=(*self.events, event)) + + def failure(self, model: str, custom_llm_provider: str) -> Exception | None: + failed_event: Final = next( + (event for event in self.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(_ws_event_error(failed_event), model, custom_llm_provider) + _RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) @@ -1997,8 +2071,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._open_turns: tuple[_ResponsesWsTurn, ...] = () + self._failure: Exception | None = None self.first_message = first_message self.guardrail_callbacks: Sequence[PresidioGuardrailCallback] = guardrail_callbacks or [] self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or [] @@ -2008,105 +2082,69 @@ 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 _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 - - if self._should_store_event(event_obj): - self.messages.append(event_obj) - - 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 - - if msg_obj.get("type") != "response.create": - return - - 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: - failed_event: Final = next( - (event for event in self.messages if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None - ) - if failed_event is None: - return None - return _map_stream_error_to_exception( - _ws_event_error(failed_event), self.authorized_model or "", self.custom_llm_provider or "" + def _start_turn(self, client_frame: str | None) -> None: + start_time: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time + turn_logging_obj: Final = self.logging_obj.copy_for_new_call( + litellm_call_id=str(uuid.uuid4()), start_time=start_time ) + self._detach_budget_reservation() + 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)) - async def _log_messages(self) -> None: - if not self.logging_obj: + def _detach_budget_reservation(self) -> None: + """Only the first turn settles the connection's admission reservation; later turns charge their cost directly.""" + litellm_params: Final = _json_object_or_none( + self.logging_obj.litellm_params # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped + ) + if litellm_params is None: return - if self.input_messages: - self.logging_obj.model_call_details["messages"] = self.input_messages - if not self.messages: + self.logging_obj.litellm_params = { + key: without_budget_reservation(value) + if key in _PER_TURN_METADATA_KEYS and _is_json_object(value) + else value + for key, value in litellm_params.items() + } + + 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) + + 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 - exception: Final = self._failure_exception() + 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 _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) + + def _log_turn(self, turn: _ResponsesWsTurn) -> None: + events: Final = list(turn.events) + exception: Final = turn.failure(self.authorized_model or "", self.custom_llm_provider or "") if exception is None: - asyncio.create_task(self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True)) + asyncio.create_task(turn.logging_obj.dispatch_success_handlers(events, prefer_async_handlers=True)) return - self._record_usage_for_failure() + if self._failure is None: + self._failure = exception + _record_usage_for_failure(turn.logging_obj, events) traceback_exception: Final = "".join(traceback.format_exception(exception)) asyncio.create_task( - self.logging_obj.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True) + turn.logging_obj.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True) ) - def _record_usage_for_failure(self) -> 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) - def _wrap_response_event(self, response_str: str) -> str: try: event_obj: Final = _load_json_object(response_str) @@ -2164,7 +2202,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) @@ -2173,7 +2211,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: """ @@ -2518,8 +2556,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: @@ -2527,8 +2564,7 @@ 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: @@ -2551,7 +2587,7 @@ class ResponsesWebSocketStreaming: await self.backend_ws.close() except Exception: pass - return self._failure_exception() + return self._failure # --------------------------------------------------------------------------- diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a957e6957ee..0bcb732a0f5 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4587,6 +4587,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/integration/_support/responses_ws_peer.py b/tests/integration/_support/responses_ws_peer.py new file mode 100644 index 00000000000..1bf3e1fcbae --- /dev/null +++ b/tests/integration/_support/responses_ws_peer.py @@ -0,0 +1,186 @@ +"""A scripted provider-side Responses WebSocket peer. + +Answers every ``response.create`` frame with the events a script returns for that turn over +``wss://`` (the proxy only dials TLS provider sockets). Clients drive it through ``session``. +""" + +import asyncio +import itertools +import json +import ssl +import threading +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from typing import Final, TypeAlias + +import websockets +from integration._support.tls import server_context +from pydantic import JsonValue +from websockets.asyncio.server import ServerConnection, serve + +PROVIDER_MODEL: Final = "ws-peer-model" +USAGE: Final = {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7} +TERMINAL: Final = frozenset({"response.completed", "response.failed", "error"}) + +Script: TypeAlias = Callable[[str], tuple[dict[str, JsonValue], ...]] + + +@dataclass(frozen=True, slots=True) +class Credentials: + authorization: str | None + api_key: str | None + + +@dataclass(frozen=True, slots=True) +class Peer: + url: str + paths: SimpleQueue[str] + credentials: SimpleQueue[Credentials] + frames: SimpleQueue[dict[str, JsonValue]] + + +def events(response_id: str, text: str) -> tuple[dict[str, JsonValue], ...]: + message: Final = { + "type": "message", + "id": f"msg_{response_id}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + response: Final = {"id": response_id, "object": "response", "created_at": 1700000000, "model": PROVIDER_MODEL} + return ( + {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}, + { + "type": "response.output_text.delta", + "item_id": f"msg_{response_id}", + "output_index": 0, + "content_index": 0, + "delta": text, + }, + { + "type": "response.completed", + "response": {**response, "status": "completed", "output": [message], "usage": USAGE}, + }, + ) + + +def seven(response_id: str) -> tuple[dict[str, JsonValue], ...]: + return events(response_id, "seven") + + +async def _answer( + connection: ServerConnection, + paths: SimpleQueue[str], + credentials: SimpleQueue[Credentials], + frames: SimpleQueue[dict[str, JsonValue]], + script: Script, +) -> None: + request: Final = connection.request + paths.put(request.path if request is not None else "") + credentials.put( + Credentials(request.headers.get("Authorization"), request.headers.get("api-key")) + if request is not None + else Credentials(None, None) + ) + turns: Final = itertools.count(1) + async for raw in connection: + frame: Final = json.loads(raw) + frames.put(frame) + if frame.get("type") != "response.create": + continue + for event in script(f"resp_peer_{next(turns)}"): + await connection.send(json.dumps(event)) + + +async def _serve( + tls: ssl.SSLContext, + paths: SimpleQueue[str], + credentials: SimpleQueue[Credentials], + frames: SimpleQueue[dict[str, JsonValue]], + ports: SimpleQueue[int], + stop: asyncio.Event, + script: Script, +) -> None: + async with serve( + lambda connection: _answer(connection, paths, credentials, frames, script), "127.0.0.1", 0, ssl=tls + ) as server: + ports.put(next(iter(server.sockets)).getsockname()[1]) + await stop.wait() + + +@contextmanager +def responses_peer(cert: tuple[Path, Path], script: Script = seven) -> Iterator[Peer]: + loop: Final = asyncio.new_event_loop() + stop: Final = asyncio.Event() + paths: Final = SimpleQueue[str]() + credentials: Final = SimpleQueue[Credentials]() + frames: Final = SimpleQueue[dict[str, JsonValue]]() + ports: Final = SimpleQueue[int]() + thread: Final = threading.Thread( + target=loop.run_until_complete, + args=(_serve(server_context(*cert), paths, credentials, frames, ports, stop, script),), + daemon=True, + ) + thread.start() + try: + yield Peer(f"https://127.0.0.1:{ports.get(timeout=10)}/v1", paths, credentials, frames) + finally: + loop.call_soon_threadsafe(stop.set) + thread.join(timeout=10) + loop.close() + + +def create(model: str, text: str, previous_response_id: str | None = None) -> str: + return json.dumps( + { + "type": "response.create", + "model": model, + "store": True, + "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}], + **({} if previous_response_id is None else {"previous_response_id": previous_response_id}), + } + ) + + +async def turn(connection: websockets.ClientConnection, frame: str) -> tuple[dict[str, JsonValue], ...]: + await connection.send(frame) + return await until_terminal(connection, ()) + + +async def until_terminal( + connection: websockets.ClientConnection, received: tuple[dict[str, JsonValue], ...] +) -> tuple[dict[str, JsonValue], ...]: + event: Final = json.loads(await asyncio.wait_for(connection.recv(), timeout=20)) + collected: Final = (*received, event) + if event.get("type") in TERMINAL or len(collected) >= 50: + return collected + return await until_terminal(connection, collected) + + +async def session( + proxy_url: str, key: str, model: str, texts: tuple[str, ...] +) -> tuple[tuple[dict[str, JsonValue], ...], ...]: + proxy: Final = proxy_url.rstrip("/").replace("http://", "ws://") + async with websockets.connect( + f"{proxy}/v1/responses?model={model}", + additional_headers={"Authorization": f"Bearer {key}"}, + open_timeout=10, + ) as connection: + first: Final = await turn(connection, create(model, texts[0])) + if len(texts) == 1: + return (first,) + previous: Final = str(first[-1]["response"]["id"]) + second: Final = await turn(connection, create(model, texts[1], previous)) + return (first, second) + + +def completed(received: tuple[dict[str, JsonValue], ...]) -> dict[str, JsonValue]: + assert received[-1]["type"] == "response.completed", [event.get("type") for event in received] + return received[-1]["response"] + + +def drain(queue: SimpleQueue[dict[str, JsonValue]]) -> tuple[dict[str, JsonValue], ...]: + return tuple(queue.get_nowait() for _ in range(queue.qsize())) diff --git a/tests/integration/providers/test_openai_responses_websocket_wire.py b/tests/integration/providers/test_openai_responses_websocket_wire.py index 27d328369d3..4f9a119db2f 100644 --- a/tests/integration/providers/test_openai_responses_websocket_wire.py +++ b/tests/integration/providers/test_openai_responses_websocket_wire.py @@ -1,197 +1,37 @@ import asyncio import hashlib -import itertools -import json -import ssl -import threading import uuid -from collections.abc import Callable, Iterator -from contextlib import contextmanager -from dataclasses import dataclass +from collections.abc import Iterator from pathlib import Path -from queue import SimpleQueue -from typing import Final, TypeAlias +from typing import Final import pytest -import websockets from integration._support.client import Gateway, eventually, gateway_from_environment, object_value, string_value from integration._support.database import read_rows from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.responses_stream import RATE_LIMIT_MESSAGE, openai_error_event -from integration._support.tls import server_context, write_self_signed_cert +from integration._support.responses_ws_peer import ( + PROVIDER_MODEL, + USAGE, + completed, + drain, + events, + responses_peer, + session, +) +from integration._support.tls import write_self_signed_cert from pydantic import JsonValue -from websockets.asyncio.server import ServerConnection, serve pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) - -PROVIDER_MODEL: Final = "ws-peer-model" -USAGE: Final = {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7} -TERMINAL: Final = frozenset({"response.completed", "response.failed", "error"}) PROXY_OPENAI_KEY: Final = "sk-proxy-environment-openai-key" PROXY_AZURE_KEY: Final = "proxy-environment-azure-key" -_Script: TypeAlias = Callable[[str], tuple[dict[str, JsonValue], ...]] - - -@dataclass(frozen=True, slots=True) -class Credentials: - authorization: str | None - api_key: str | None - - -@dataclass(frozen=True, slots=True) -class Peer: - url: str - paths: SimpleQueue[str] - credentials: SimpleQueue[Credentials] - frames: SimpleQueue[dict[str, JsonValue]] - - -def _events(response_id: str, text: str) -> tuple[dict[str, JsonValue], ...]: - message: Final = { - "type": "message", - "id": f"msg_{response_id}", - "status": "completed", - "role": "assistant", - "content": [{"type": "output_text", "text": text, "annotations": []}], - } - response: Final = {"id": response_id, "object": "response", "created_at": 1700000000, "model": PROVIDER_MODEL} - return ( - {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}, - { - "type": "response.output_text.delta", - "item_id": f"msg_{response_id}", - "output_index": 0, - "content_index": 0, - "delta": text, - }, - { - "type": "response.completed", - "response": {**response, "status": "completed", "output": [message], "usage": USAGE}, - }, - ) - - -def _seven(response_id: str) -> tuple[dict[str, JsonValue], ...]: - return _events(response_id, "seven") - def _top_level_rate_limit(response_id: str) -> tuple[dict[str, JsonValue], ...]: - created: Final = _events(response_id, "")[0] + created: Final = events(response_id, "")[0] return (created, dict(openai_error_event("rate_limit_exceeded", RATE_LIMIT_MESSAGE))) -async def _answer( - connection: ServerConnection, - paths: SimpleQueue[str], - credentials: SimpleQueue[Credentials], - frames: SimpleQueue[dict[str, JsonValue]], - script: _Script, -) -> None: - request: Final = connection.request - paths.put(request.path if request is not None else "") - credentials.put( - Credentials(request.headers.get("Authorization"), request.headers.get("api-key")) - if request is not None - else Credentials(None, None) - ) - turns: Final = itertools.count(1) - async for raw in connection: - frame: Final = json.loads(raw) - frames.put(frame) - if frame.get("type") != "response.create": - continue - for event in script(f"resp_peer_{next(turns)}"): - await connection.send(json.dumps(event)) - - -async def _serve( - tls: ssl.SSLContext, - paths: SimpleQueue[str], - credentials: SimpleQueue[Credentials], - frames: SimpleQueue[dict[str, JsonValue]], - ports: SimpleQueue[int], - stop: asyncio.Event, - script: _Script, -) -> None: - async with serve( - lambda connection: _answer(connection, paths, credentials, frames, script), "127.0.0.1", 0, ssl=tls - ) as server: - ports.put(next(iter(server.sockets)).getsockname()[1]) - await stop.wait() - - -@contextmanager -def responses_peer(cert: tuple[Path, Path], script: _Script = _seven) -> Iterator[Peer]: - loop: Final = asyncio.new_event_loop() - stop: Final = asyncio.Event() - paths: Final = SimpleQueue[str]() - credentials: Final = SimpleQueue[Credentials]() - frames: Final = SimpleQueue[dict[str, JsonValue]]() - ports: Final = SimpleQueue[int]() - thread: Final = threading.Thread( - target=loop.run_until_complete, - args=(_serve(server_context(*cert), paths, credentials, frames, ports, stop, script),), - daemon=True, - ) - thread.start() - try: - yield Peer(f"https://127.0.0.1:{ports.get(timeout=10)}/v1", paths, credentials, frames) - finally: - loop.call_soon_threadsafe(stop.set) - thread.join(timeout=10) - loop.close() - - -def _create(model: str, text: str, previous_response_id: str | None = None) -> str: - return json.dumps( - { - "type": "response.create", - "model": model, - "store": True, - "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}], - **({} if previous_response_id is None else {"previous_response_id": previous_response_id}), - } - ) - - -async def _turn(connection: websockets.ClientConnection, frame: str) -> tuple[dict[str, JsonValue], ...]: - await connection.send(frame) - return await _until_terminal(connection, ()) - - -async def _until_terminal( - connection: websockets.ClientConnection, received: tuple[dict[str, JsonValue], ...] -) -> tuple[dict[str, JsonValue], ...]: - event: Final = json.loads(await asyncio.wait_for(connection.recv(), timeout=20)) - collected: Final = (*received, event) - if event.get("type") in TERMINAL or len(collected) >= 50: - return collected - return await _until_terminal(connection, collected) - - -async def _session( - proxy_url: str, key: str, model: str, texts: tuple[str, ...] -) -> tuple[tuple[dict[str, JsonValue], ...], ...]: - proxy: Final = proxy_url.rstrip("/").replace("http://", "ws://") - async with websockets.connect( - f"{proxy}/v1/responses?model={model}", - additional_headers={"Authorization": f"Bearer {key}"}, - open_timeout=10, - ) as connection: - first: Final = await _turn(connection, _create(model, texts[0])) - if len(texts) == 1: - return (first,) - previous: Final = str(first[-1]["response"]["id"]) - second: Final = await _turn(connection, _create(model, texts[1], previous)) - return (first, second) - - -def _completed(events: tuple[dict[str, JsonValue], ...]) -> dict[str, JsonValue]: - assert events[-1]["type"] == "response.completed", [event.get("type") for event in events] - return events[-1]["response"] - - @pytest.fixture(scope="module") def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: return write_self_signed_cert(tmp_path_factory.mktemp("responses-ws-cert")) @@ -208,10 +48,6 @@ def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) yield proxy -def _drain(queue: SimpleQueue[dict[str, JsonValue]]) -> tuple[dict[str, JsonValue], ...]: - return tuple(queue.get_nowait() for _ in range(queue.qsize())) - - def test_a_response_create_frame_streams_from_the_provider_socket_back_to_the_client( candidate: Gateway, cert: tuple[Path, Path] ) -> None: @@ -219,17 +55,17 @@ def test_a_response_create_frame_streams_from_the_provider_socket_back_to_the_cl model: Final = scenario.model(model=f"openai/{PROVIDER_MODEL}", api_base=peer.url) key: Final = scenario.key(models=[model]) text: Final = f"say seven {uuid.uuid4().hex}" - (events,) = asyncio.run(_session(str(candidate.client.base_url), key, model, (text,))) + (events,) = asyncio.run(session(str(candidate.client.base_url), key, model, (text,))) assert [event["type"] for event in events] == [ "response.created", "response.output_text.delta", "response.completed", ] - completed: Final = _completed(events) - assert completed["status"] == "completed" - assert completed["usage"] == USAGE + response: Final = completed(events) + assert response["status"] == "completed" + assert response["usage"] == USAGE assert peer.paths.get_nowait() == f"/v1/responses?model={PROVIDER_MODEL}" - (forwarded,) = _drain(peer.frames) + (forwarded,) = drain(peer.frames) assert forwarded["type"] == "response.create" assert forwarded["model"] == PROVIDER_MODEL assert forwarded["input"][0]["content"][0]["text"] == text @@ -242,12 +78,12 @@ def test_previous_response_id_from_the_first_turn_reaches_the_provider_as_its_ow model: Final = scenario.model(model=f"openai/{PROVIDER_MODEL}", api_base=peer.url) key: Final = scenario.key(models=[model]) texts: Final = (f"remember seven {uuid.uuid4().hex}", f"which number {uuid.uuid4().hex}") - first, second = asyncio.run(_session(str(candidate.client.base_url), key, model, texts)) - assert _completed(first)["status"] == "completed" - assert str(_completed(first)["id"]).startswith("resp_") and _completed(first)["id"] != "resp_peer_1" - assert _completed(second)["status"] == "completed" + first, second = asyncio.run(session(str(candidate.client.base_url), key, model, texts)) + assert completed(first)["status"] == "completed" + assert str(completed(first)["id"]).startswith("resp_") and completed(first)["id"] != "resp_peer_1" + assert completed(second)["status"] == "completed" assert peer.paths.qsize() == 1 - forwarded: Final = _drain(peer.frames) + forwarded: Final = drain(peer.frames) assert [frame["input"][0]["content"][0]["text"] for frame in forwarded] == list(texts) assert "previous_response_id" not in forwarded[0] assert forwarded[1]["previous_response_id"] == "resp_peer_1" @@ -260,10 +96,10 @@ def test_a_top_level_error_frame_reaches_the_client_and_logs_the_provider_code_a model: Final = scenario.model(model=f"openai/{PROVIDER_MODEL}", api_base=peer.url) key: Final = scenario.key(models=[model]) text: Final = f"say seven {uuid.uuid4().hex}" - (events,) = asyncio.run(_session(str(candidate.client.base_url), key, model, (text,))) + (events,) = asyncio.run(session(str(candidate.client.base_url), key, model, (text,))) assert [event["type"] for event in events] == ["response.created", "error"], events assert (events[-1]["code"], events[-1]["message"]) == ("rate_limit_exceeded", RATE_LIMIT_MESSAGE), events - (forwarded,) = _drain(peer.frames) + (forwarded,) = drain(peer.frames) assert forwarded["input"][0]["content"][0]["text"] == text (row,) = eventually( lambda: read_rows( @@ -280,7 +116,7 @@ def test_a_top_level_error_frame_reaches_the_client_and_logs_the_provider_code_a def _one_turn(candidate: Gateway, key: str, model: str) -> tuple[dict[str, JsonValue], ...]: - (events,) = asyncio.run(_session(str(candidate.client.base_url), key, model, (f"say seven {uuid.uuid4().hex}",))) + (events,) = asyncio.run(session(str(candidate.client.base_url), key, model, (f"say seven {uuid.uuid4().hex}",))) return events @@ -293,7 +129,7 @@ def test_a_bare_model_declared_as_openai_sends_its_deployment_key_to_the_provide model=PROVIDER_MODEL, custom_llm_provider="openai", api_base=peer.url, api_key=deployment_key ) events: Final = _one_turn(candidate, scenario.key(models=[model]), model) - assert _completed(events)["status"] == "completed" + assert completed(events)["status"] == "completed" assert peer.credentials.get_nowait().authorization == f"Bearer {deployment_key}" @@ -303,7 +139,7 @@ def test_an_openai_deployment_without_its_own_key_uses_the_proxy_openai_key( with responses_peer(cert) as peer, candidate.scenario() as scenario: model: Final = scenario.model(model=f"openai/{PROVIDER_MODEL}", api_base=peer.url, api_key=None) events: Final = _one_turn(candidate, scenario.key(models=[model]), model) - assert _completed(events)["status"] == "completed" + assert completed(events)["status"] == "completed" assert peer.credentials.get_nowait().authorization == f"Bearer {PROXY_OPENAI_KEY}" @@ -315,7 +151,7 @@ def test_an_azure_deployment_without_its_own_key_gets_the_azure_key_and_never_th model=f"azure/{PROVIDER_MODEL}", api_base=peer.url.removesuffix("/v1"), api_key=None ) events: Final = _one_turn(candidate, scenario.key(models=[model]), model) - assert _completed(events)["status"] == "completed" + assert completed(events)["status"] == "completed" credentials: Final = peer.credentials.get_nowait() assert credentials.api_key == PROXY_AZURE_KEY assert PROXY_OPENAI_KEY not in (credentials.authorization or "") @@ -329,8 +165,8 @@ def test_a_declared_openai_provider_wins_over_another_providers_model_prefix( model: Final = scenario.model(model=gateway_model, custom_llm_provider="openai", api_base=peer.url) key: Final = scenario.key(models=[model]) events: Final = _one_turn(candidate, key, model) - assert _completed(events)["status"] == "completed" - (forwarded,) = _drain(peer.frames) + assert completed(events)["status"] == "completed" + (forwarded,) = drain(peer.frames) assert forwarded["model"] == gateway_model (row,) = eventually( lambda: read_rows( diff --git a/tests/integration/spend/test_responses_websocket_turn_rows.py b/tests/integration/spend/test_responses_websocket_turn_rows.py new file mode 100644 index 00000000000..59b5925c6eb --- /dev/null +++ b/tests/integration/spend/test_responses_websocket_turn_rows.py @@ -0,0 +1,114 @@ +"""Every response.create turn on a native Responses WebSocket is its own spend row. + +OpenAI bills each turn as a request, so one socket carrying two turns has to leave two +rows, each keyed by the response id the client saw and priced from that turn's own usage, +instead of one row at socket close that sums both (LIT-8949). A turn the provider fails +leaves a failure row of its own and does not swallow the completed turn before it. +""" + +import asyncio +import hashlib +import uuid +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, gateway_from_environment, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.responses_stream import RATE_LIMIT_MESSAGE, openai_error_event +from integration._support.responses_ws_peer import PROVIDER_MODEL, USAGE, completed, events, responses_peer, session +from integration._support.tls import write_self_signed_cert +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 180) + +INPUT_RATE: Final = 0.001 +OUTPUT_RATE: Final = 0.002 +TURN_SPEND: Final = USAGE["input_tokens"] * INPUT_RATE + USAGE["output_tokens"] * OUTPUT_RATE +SPEND_ROWS_FOR_KEY: Final = ( + 'SELECT request_id, spend, total_tokens, call_type, status, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s' +) + + +def _second_turn_rate_limited(response_id: str) -> tuple[dict[str, JsonValue], ...]: + if response_id == "resp_peer_1": + return events(response_id, "seven") + return (events(response_id, "")[0], dict(openai_error_event("rate_limit_exceeded", RATE_LIMIT_MESSAGE))) + + +def _key_hash(key: str) -> str: + return hashlib.sha256(key.encode()).hexdigest() + + +@pytest.fixture(scope="module") +def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: + return write_self_signed_cert(tmp_path_factory.mktemp("responses-ws-rows-cert")) + + +@pytest.fixture(scope="module") +def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: + with gateway_from_environment() as base: + with owned_proxy(base, tmp_path_factory.mktemp("responses-ws-rows"), {"SSL_CERT_FILE": str(cert[0])}) as proxy: + yield proxy + + +def test_two_turns_on_one_socket_leave_two_rows_keyed_by_the_response_ids_the_client_saw( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert) as peer, candidate.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{PROVIDER_MODEL}", + api_base=peer.url, + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + ) + key: Final = scenario.key(models=[model]) + texts: Final = (f"first {uuid.uuid4().hex}", f"second {uuid.uuid4().hex}") + + first, second = asyncio.run(session(str(candidate.client.base_url), key, model, texts)) + + response_ids: Final = {str(completed(first)["id"]), str(completed(second)["id"])} + assert len(response_ids) == 2, response_ids + rows: Final = eventually( + lambda: read_rows(SPEND_ROWS_FOR_KEY, (_key_hash(key),)), + lambda found: len(found) >= 2, + seconds=70, + ) + assert {str(row["request_id"]) for row in rows} == response_ids, rows + assert [float(row["spend"]) for row in rows] == [pytest.approx(TURN_SPEND)] * 2, rows + assert {(row["call_type"], row["status"], row["total_tokens"]) for row in rows} == { + ("_aresponses_websocket", "success", USAGE["total_tokens"]) + }, rows + + +def test_a_turn_the_provider_rejects_leaves_its_own_failure_row_next_to_the_completed_turn( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert, _second_turn_rate_limited) as peer, candidate.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{PROVIDER_MODEL}", + api_base=peer.url, + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + ) + key: Final = scenario.key(models=[model]) + texts: Final = (f"first {uuid.uuid4().hex}", f"second {uuid.uuid4().hex}") + + first, second = asyncio.run(session(str(candidate.client.base_url), key, model, texts)) + + assert second[-1]["type"] == "error", second + rows: Final = eventually( + lambda: read_rows(SPEND_ROWS_FOR_KEY, (_key_hash(key),)), + lambda found: len(found) >= 2, + seconds=70, + ) + by_status: Final = {str(row["status"]): row for row in rows} + assert set(by_status) == {"success", "failure"}, rows + assert str(by_status["success"]["request_id"]) == str(completed(first)["id"]), rows + assert float(by_status["success"]["spend"]) == pytest.approx(TURN_SPEND), rows + assert float(by_status["failure"]["spend"]) == 0.0, rows + error: Final = object_value(object_value(by_status["failure"]["metadata"])["error_information"]) + assert str(error["error_code"]) == "429", error + assert RATE_LIMIT_MESSAGE in string_value(error["error_message"]), error diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 14334690624..df6593d29b1 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -11880,3 +11880,55 @@ def test_masking_function_not_in_metadata_when_not_provided(): assert "_langfuse_masking_function" not in result assert result["metadata"]["some_key"] == "some_value" + + +def test_copy_for_new_call_logs_a_turn_under_its_own_call_id_against_the_connections_reservation() -> None: + reservation: Final[dict[str, object]] = {"reserved_cost": 0.5, "entries": []} + connection: Final = LitellmLogging( + model="gpt-5.4", + messages=[], + stream=False, + call_type=CallTypes.aresponses_websocket.value, + start_time=time.time(), + litellm_call_id="connection-call", + function_id="connection-call", + ) + connection.litellm_params = { + "litellm_call_id": "connection-call", + "metadata": {"user_api_key": "hashed-key", "user_api_key_budget_reservation": reservation}, + } + connection.model_call_details["litellm_params"] = connection.litellm_params + turn_start: Final = datetime.datetime(2026, 1, 1, 12, 0, 0) + + turn: Final = connection.copy_for_new_call(litellm_call_id="turn-1", start_time=turn_start) + + assert (turn.litellm_call_id, turn.start_time) == ("turn-1", turn_start) + assert (turn.model_call_details["litellm_call_id"], turn.model_call_details["start_time"]) == ("turn-1", turn_start) + assert turn.model_call_details["litellm_params"]["litellm_call_id"] == "turn-1" + assert turn.litellm_params["metadata"]["user_api_key"] == "hashed-key" + assert turn.litellm_params["metadata"]["user_api_key_budget_reservation"] is reservation + turn.litellm_params["metadata"]["turn_only"] = True + assert "turn_only" not in connection.litellm_params["metadata"] + assert (connection.litellm_call_id, connection.model_call_details["litellm_call_id"]) == ( + "connection-call", + "connection-call", + ) + + +@pytest.mark.parametrize("response_id, expected_row_id", [(None, "turn-7"), ("resp_9", "resp_9")]) +def test_a_websocket_turn_row_is_keyed_by_the_response_id_or_else_its_call_id( + logging_obj: LitellmLogging, response_id: str | None, expected_row_id: str +) -> None: + from litellm.litellm_core_utils.litellm_logging import get_standard_logging_object_payload + + now: Final = datetime_unit_test(2026, 1, 1, 12, 0, 0) + payload: Final = get_standard_logging_object_payload( + kwargs={"litellm_call_id": "turn-7", "litellm_params": {"metadata": {}}}, + init_response_obj=LiteLLMRealtimeStreamLoggingObject(id=response_id, usage=Usage(), results=[]), + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + assert payload is not None + assert payload["id"] == expected_row_id diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter.py b/tests/unit/proxy/hooks/test_parallel_request_limiter.py index 89cb6bd1ef2..f17cad56751 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter.py @@ -16,7 +16,7 @@ from litellm.proxy.hooks.parallel_request_limiter import ( PROXY_MaxParallelRequestsHandler, ) from litellm.proxy.utils import InternalUsageCache, hash_token -from litellm.types.utils import EmbeddingResponse, ModelResponse, TextCompletionResponse, Usage +from litellm.types.utils import CallTypes, EmbeddingResponse, ModelResponse, TextCompletionResponse, Usage FROZEN_INSTANT: Final = datetime(2026, 1, 31, 23, 59, 30) LAST_MICROSECOND_OF_JANUARY: Final = datetime(2026, 1, 31, 23, 59, 59, 999999) @@ -323,3 +323,33 @@ async def test_retry_after_is_the_seconds_until_the_next_minute_of_the_injected_ ) assert exc_info.value.headers["retry-after"] == "14.5" + + +@pytest.mark.asyncio +async def test_websocket_turn_callbacks_keep_the_sockets_slot_until_the_disconnect_release(): + api_key = hash_token("sk-websocket") + handler = PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock) + bucket = f"{api_key}::{FROZEN_INSTANT.strftime('%Y-%m-%d-%H-%M')}::request_count" + await handler.internal_usage_cache.async_set_cache( + key=bucket, + value={"current_requests": 1, "current_tpm": 0, "current_rpm": 1}, + litellm_parent_otel_span=None, + ) + turn_kwargs = { + "call_type": CallTypes.aresponses_websocket.value, + "litellm_params": {"metadata": {"user_api_key": api_key}}, + "exception": RuntimeError("provider rejected the turn"), + } + + await handler.async_log_success_event( + kwargs=turn_kwargs, response_obj=None, start_time=FROZEN_INSTANT, end_time=FROZEN_INSTANT + ) + await handler.async_log_failure_event( + kwargs=turn_kwargs, response_obj=None, start_time=FROZEN_INSTANT, end_time=FROZEN_INSTANT + ) + after_two_turns = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) + assert after_two_turns["current_requests"] == 1, after_two_turns + + await handler.async_release_max_parallel_requests_on_disconnect(UserAPIKeyAuth(api_key=api_key)) + after_disconnect = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) + assert after_disconnect["current_requests"] == 0, after_disconnect diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 1af42947da7..0eaf619ba70 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1952,6 +1952,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 ( @@ -2016,6 +2017,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 ( diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 3625bbb7c04..2b78aa98b04 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1545,6 +1545,7 @@ class TestNativeWebSocketDeploymentDefaults: mock_logging = MagicMock() mock_logging.pre_call = MagicMock() mock_logging.dispatch_success_handlers = AsyncMock() + mock_logging.copy_for_new_call.return_value = mock_logging client_ws = MagicMock() client_ws.receive_text = AsyncMock(side_effect=RuntimeError("client closed")) @@ -1721,6 +1722,7 @@ class TestNativeWebSocketGuardrails: client_ws.send_text = AsyncMock() logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj delta_event = json.dumps( {"type": "response.output_text.delta", "delta": "alice@example.com"} @@ -1795,6 +1797,7 @@ class TestNativeWebSocketGuardrails: client_ws.send_text = AsyncMock() logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj done_events = [ json.dumps( @@ -2494,6 +2497,7 @@ class TestNativeWebSocketGuardrailMasking: ) logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj handler = _make_streaming( websocket=websocket, @@ -2550,6 +2554,7 @@ class TestNativeWebSocketGuardrailMasking: ) logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj handler = _make_streaming( websocket=websocket, @@ -2613,6 +2618,7 @@ class TestNativeWebSocketGuardrailMasking: ) logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj handler = _make_streaming( websocket=websocket, @@ -2676,6 +2682,7 @@ class TestNativeWebSocketGuardrailMasking: ) logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj handler = _make_streaming( websocket=websocket, @@ -3226,6 +3233,7 @@ class TestNativeWebSocketEncryptedContentAffinity: ) logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj handler = _make_streaming( websocket=websocket, backend_ws=backend_ws, @@ -3272,6 +3280,7 @@ class TestNativeWebSocketEncryptedContentAffinity: ) logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj handler = _make_streaming( websocket=websocket, backend_ws=backend_ws, @@ -3345,6 +3354,7 @@ class TestNativeWebSocketEncryptedContentAffinity: ) logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj logging_obj.dispatch_failure_handlers = AsyncMock() logging_obj.response_cost_calculator = MagicMock(return_value=0.0) handler = _make_streaming( @@ -3366,17 +3376,73 @@ class TestNativeWebSocketEncryptedContentAffinity: assert expected_message in str(exception) @pytest.mark.asyncio - async def test_backend_to_client_bills_completed_turns_before_a_failure(self): + async def test_backend_to_client_logs_each_turn_as_its_own_call(self, monkeypatch): import asyncio + import time from unittest.mock import AsyncMock import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import CallTypes + + pre_calls: list[tuple[str, object]] = [] + successes: list[tuple[str, list[object]]] = [] + failures: list[tuple[str, Exception]] = [] + partial_usage: list[tuple[str, int, float]] = [] + + def record_pre_call(self, input, api_key, model=None, additional_args=None): + pre_calls.append((self.litellm_call_id, self.model_call_details["messages"])) + + reservations: dict[str, tuple[object, object]] = {} + + def record_reservation(self): + metadata = self.litellm_params["metadata"] + reservations[self.litellm_call_id] = ( + metadata.get("user_api_key_budget_reservation"), + metadata["user_api_key_auth"].budget_reservation, + ) + + async def record_success(self, result, **kwargs): + record_reservation(self) + successes.append((self.litellm_call_id, [event["type"] for event in result])) + + async def record_failure(self, exception, traceback_exception, **kwargs): + record_reservation(self) + failures.append((self.litellm_call_id, exception)) + + def record_partial(self, usage, response_cost): + partial_usage.append((self.litellm_call_id, usage.total_tokens, response_cost)) + + monkeypatch.setattr(Logging, "pre_call", record_pre_call) + monkeypatch.setattr(Logging, "dispatch_success_handlers", record_success) + monkeypatch.setattr(Logging, "dispatch_failure_handlers", record_failure) + monkeypatch.setattr(Logging, "record_partial_usage_for_failure", record_partial) + monkeypatch.setattr(Logging, "response_cost_calculator", lambda self, result: 0.01) + connection_logging = Logging( + model="gpt-5.4", + messages=[], + stream=False, + call_type=CallTypes.aresponses_websocket.value, + start_time=time.time(), + litellm_call_id="connection-call", + function_id="connection-call", + ) + reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False} + connection_logging.litellm_params = { + "metadata": { + "user_api_key_budget_reservation": reservation, + "user_api_key_auth": UserAPIKeyAuth(api_key="sk-hashed", budget_reservation=reservation), + } + } + websocket = MagicMock() websocket.send_text = AsyncMock() backend_ws = MagicMock() backend_ws.recv = AsyncMock( side_effect=[ + json.dumps({"type": "response.created", "response": {"id": "resp_1", "status": "in_progress"}}), json.dumps( { "type": "response.completed", @@ -3392,21 +3458,27 @@ 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.01) - handler = _make_streaming(websocket=websocket, backend_ws=backend_ws, logging_obj=logging_obj, request_data={}) + handler = _make_streaming( + websocket=websocket, backend_ws=backend_ws, logging_obj=connection_logging, request_data={} + ) + handler._record_client_frame(json.dumps({"type": "response.create", "input": "first question"})) + handler._record_client_frame(json.dumps({"type": "response.create", "input": "second question"})) await handler.backend_to_client() await asyncio.sleep(0) - 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) - assert response_cost == 0.01 - logging_obj.dispatch_success_handlers.assert_not_awaited() - logging_obj.dispatch_failure_handlers.assert_awaited_once() + (first_turn, first_messages), (second_turn, second_messages) = pre_calls + assert len({first_turn, second_turn, "connection-call"}) == 3 + assert first_messages == [{"role": "user", "content": "first question"}] + assert second_messages == [{"role": "user", "content": "second question"}] + assert successes == [(first_turn, ["response.created", "response.completed"])] + ((failed_turn, exception),) = failures + assert failed_turn == second_turn + assert "bad turn" in str(exception) + assert partial_usage == [(second_turn, 0, 0.01)] + assert connection_logging.litellm_call_id == "connection-call" + assert reservations[first_turn] == (reservation, reservation) + assert reservations[second_turn] == (None, None) @pytest.mark.asyncio async def test_bidirectional_forward_returns_the_provider_failure(self): @@ -3450,6 +3522,7 @@ class TestNativeWebSocketEncryptedContentAffinity: backend_ws.close = AsyncMock() logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj logging_obj.dispatch_failure_handlers = AsyncMock() logging_obj.response_cost_calculator = MagicMock(return_value=0.0) handler = _make_streaming( @@ -3508,6 +3581,7 @@ class TestNativeWebSocketEncryptedContentAffinity: backend_ws.close = AsyncMock() logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj logging_obj.dispatch_failure_handlers = AsyncMock() handler = _make_streaming( websocket=websocket, diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 3db4432bd6c..cc2a1d62d52 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -9267,3 +9267,23 @@ def test_completion_cost_search_uses_search_tool_input_cost_per_query( assert map_cost == pytest.approx(expected_queries * litellm.model_cost["agentcore/search"]["input_cost_per_query"]) assert configured_cost == pytest.approx(expected_queries * 0.05) + + +def test_responses_ws_logging_object_carries_the_completed_responses_id_tier_and_usage() -> None: + events = [ + {"type": "response.created", "response": {"id": "resp_a", "status": "in_progress"}}, + { + "type": "response.completed", + "response": { + "id": "resp_a", + "service_tier": "priority", + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + }, + ] + + logged = ResponsesWebSocketTokenUsageProcessor.create_logging_object(events) + + assert (logged.id, logged.service_tier) == ("resp_a", "priority") + assert (logged.usage.prompt_tokens, logged.usage.completion_tokens, logged.usage.total_tokens) == (10, 5, 15) + assert ResponsesWebSocketTokenUsageProcessor.create_logging_object([{"type": "error", "error": {}}]).id is None