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