diff --git a/gateway/routes/allowlist.py b/gateway/routes/allowlist.py index 3cfe9569b34..cced49f788a 100644 --- a/gateway/routes/allowlist.py +++ b/gateway/routes/allowlist.py @@ -121,6 +121,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = ( # Realtime / streaming "/v1/realtime", "/realtime", + "/v1/live/sessions", + "/live/sessions", # Health & ops "/health", "/metrics", diff --git a/litellm/constants.py b/litellm/constants.py index d9441f36a1f..f9cabd58bd4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -391,6 +391,11 @@ MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY: Final = "litellm.model_group_alias_resolve MCP_PEEKED_BODY_SCOPE_KEY: Final = "litellm_mcp_peeked_body" MCP_ADMISSION_BODY_PEEK_TIMEOUT_SECONDS: Final = 5.0 REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" +OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS: Final = 30 +OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS: Final = 5 +OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES: Final = frozenset( + {"response.completed", "response.incomplete", "response.failed"} +) REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged" # SSL/TLS cipher configuration for faster handshakes diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 533c6a9e00e..09603a881a3 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1,14 +1,16 @@ # What is this? ## File for 'response_cost' calculation in Logging import logging +import math import time from collections.abc import Mapping, Sequence +from datetime import datetime from functools import lru_cache from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast from httpx import Response -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, ValidationError from typing_extensions import ReadOnly, TypedDict, assert_never import litellm @@ -17,7 +19,9 @@ from litellm import verbose_logger from litellm.constants import ( DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND, + OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.cost_discount import resolve_cost_discount from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, is_decisions_call_type from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR @@ -110,7 +114,6 @@ from litellm.types.llms.openai import ( OpenAIModerationResponse, OpenAIRealtimeStreamList, OpenAIRealtimeStreamResponseBaseObject, - OpenAIRealtimeStreamSessionEvents, ResponseAPIUsage, ResponsesAPIResponse, ) @@ -2979,6 +2982,141 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor): _TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed" +class _LiveResponsePayload(BaseModel): + model_config = ConfigDict(frozen=True) + + id: str | None = None + model: str | None = None + usage: Mapping[str, object] | None = None + + +class _LiveResponseEvent(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str = "" + response: _LiveResponsePayload | None = None + + +class _LiveResponseEnvelope(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str = "" + event: _LiveResponseEvent | None = None + + +def _live_response_payload(result: Mapping[str, object]) -> _LiveResponsePayload | None: + try: + envelope: Final = _LiveResponseEnvelope.model_validate(result) + except ValidationError: + return None + if ( + envelope.type != "response.event" + or envelope.event is None + or envelope.event.type not in OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES + or envelope.event.response is None + or envelope.event.response.usage is None + ): + return None + return envelope.event.response + + +def _unique_live_response_payloads(results: Sequence[Mapping[str, object]]) -> tuple[_LiveResponsePayload, ...]: + response_payloads: Final = tuple( + response for result in results if (response := _live_response_payload(result)) is not None + ) + first_payload_by_id: Final = { + response.id: response for response in reversed(response_payloads) if response.id is not None + } + return tuple(reversed(first_payload_by_id.values())) + + +def _live_session_model(result: Mapping[str, object]) -> str | None: + session: Final = as_str_mapping(result.get("session")) + if session is None: + return None + model: Final = session.get("model") + return model if isinstance(model, str) else None + + +def _live_usage_seconds(result: Mapping[str, object]) -> float | None: + usage: Final = as_str_mapping(result.get("usage")) + if usage is None: + return None + seconds: Final = usage.get("seconds") + if isinstance(seconds, (int, float)) and not isinstance(seconds, bool) and math.isfinite(seconds): + return float(seconds) + return None + + +def _last_live_usage_seconds(results: Sequence[Mapping[str, object]], event_type: str) -> float | None: + return next( + ( + seconds + for result in reversed(results) + if result.get("type") == event_type and (seconds := _live_usage_seconds(result)) is not None + ), + None, + ) + + +def _live_session_wall_clock_seconds(litellm_logging_obj: LitellmLoggingObject | None) -> float: + if litellm_logging_obj is None: + return 0.0 + model_call_details: Final = litellm_logging_obj.model_call_details + start_time: Final = model_call_details.get("start_time") + end_time: Final = model_call_details.get("end_time") + if isinstance(start_time, datetime) and isinstance(end_time, datetime): + return max((end_time - start_time).total_seconds(), 0.0) + return 0.0 + + +def _live_input_cost_per_second(model_name: str, custom_llm_provider: str) -> float | None: + model_info: Final = _lookup_model_info_or_none(model_name, custom_llm_provider) + if model_info is None: + return None + rate: Final = model_info.get("input_cost_per_second") + if isinstance(rate, (int, float)) and not isinstance(rate, bool) and math.isfinite(rate): + return float(rate) + return None + + +def _first_live_input_cost_per_second( + potential_model_names: Sequence[str | None], + custom_llm_provider: str, +) -> float: + return next( + ( + rate + for model_name in potential_model_names + if model_name is not None + and (rate := _live_input_cost_per_second(model_name, custom_llm_provider)) is not None + ), + 0.0, + ) + + +def _live_response_token_costs( + response: _LiveResponsePayload, + custom_llm_provider: str, + data_residency: str | None, +) -> tuple[float, float]: + if response.id is None or response.model is None or response.usage is None: + return 0.0, 0.0 + try: + usage: Final = get_usage_object(response.model_dump()) + except (ValidationError, TypeError, ValueError, KeyError, AttributeError): + return 0.0, 0.0 + if usage is None: + return 0.0, 0.0 + costs: Final = _candidate_realtime_token_costs( + model_name=response.model, + combined_usage_object=usage, + custom_llm_provider=custom_llm_provider, + data_residency=data_residency, + ) + return (0.0, 0.0) if costs is None else costs + + def _candidate_realtime_token_costs( model_name: str, combined_usage_object: Usage, @@ -3065,14 +3203,21 @@ def handle_realtime_stream_cost_calculation( base_pricing_model: the deployment's resolved base_model, tried ahead of the session-reported model but after custom rates """ - received_model = None - potential_model_names: Final = [custom_pricing_model, base_pricing_model] - for result in results: - if result["type"] == "session.created": - received_model = cast(OpenAIRealtimeStreamSessionEvents, result)["session"].get("model", None) - potential_model_names.append(received_model) - - potential_model_names.append(litellm_model_name) + received_model: Final = next( + ( + model_name + for result in results + if result["type"] in ("session.created", "session.started") + and (model_name := _live_session_model(result)) is not None + ), + None, + ) + potential_model_names: Final = ( + custom_pricing_model, + base_pricing_model, + received_model, + litellm_model_name, + ) input_cost_per_token, output_cost_per_token = _first_priced_realtime_token_costs( potential_model_names=potential_model_names, combined_usage_object=combined_usage_object, @@ -3089,15 +3234,56 @@ def handle_realtime_stream_cost_calculation( if any(r.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE for r in results) else 0.0 ) - total_cost: Final = input_cost_per_token + output_cost_per_token + transcription_cost + live_lifecycle_event_types: Final = frozenset({"session.started", "session.usage.updated", "session.closed"}) + live_session_started: Final = any(result["type"] == "session.started" for result in results) + live_session_closed: Final = any(result["type"] == "session.closed" for result in results) + live_event_seen: Final = any(result["type"] in live_lifecycle_event_types for result in results) + closed_seconds: Final = _last_live_usage_seconds(results, "session.closed") + reported_seconds: Final = _last_live_usage_seconds(results, "session.usage.updated") + latest_reported_seconds: Final = 0.0 if reported_seconds is None else reported_seconds + voice_seconds: Final = ( + closed_seconds + if closed_seconds is not None + else max( + latest_reported_seconds, + _live_session_wall_clock_seconds(litellm_logging_obj), + ) + if live_session_started and not live_session_closed + else latest_reported_seconds + ) + live_model_names: Final = ( + (custom_pricing_model, base_pricing_model, received_model, litellm_model_name) if live_event_seen else () + ) + live_voice_cost: Final = voice_seconds * _first_live_input_cost_per_second( + live_model_names, + custom_llm_provider, + ) + live_response_costs: Final = tuple( + _live_response_token_costs( + response=response, + custom_llm_provider=custom_llm_provider, + data_residency=data_residency, + ) + for response in _unique_live_response_payloads(results) + ) + live_response_input_cost: Final = sum(costs[0] for costs in live_response_costs) + live_response_output_cost: Final = sum(costs[1] for costs in live_response_costs) + input_cost: Final = input_cost_per_token + live_response_input_cost + output_cost: Final = output_cost_per_token + live_response_output_cost + total_cost: Final = input_cost + output_cost + transcription_cost + live_voice_cost + additional_costs: Final = tuple( + (cost_name, cost) + for cost_name, cost in (("transcription_cost", transcription_cost), ("live_voice_cost", live_voice_cost)) + if cost > 0 + ) _store_cost_breakdown_in_logging_obj( litellm_logging_obj=litellm_logging_obj, - prompt_tokens_cost_usd_dollar=input_cost_per_token, - completion_tokens_cost_usd_dollar=output_cost_per_token, + prompt_tokens_cost_usd_dollar=input_cost, + completion_tokens_cost_usd_dollar=output_cost, cost_for_built_in_tools_cost_usd_dollar=0.0, total_cost_usd_dollar=total_cost, - additional_costs={"transcription_cost": transcription_cost} if transcription_cost > 0 else None, + additional_costs=dict(additional_costs) if additional_costs else None, data_residency=data_residency, ) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index e857461785c..c12eb4d458f 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -52,6 +52,10 @@ def safe_divide_seconds(seconds: float, denominator: float, default: float | Non _DROP_PARAMS_BOOL: Final = TypeAdapter(bool) +def as_str_mapping(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, Mapping) else None + + def normalize_drop_params(value: object) -> bool | None: if value is None or isinstance(value, bool): return value diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 57da3f8dabe..ff649953e69 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -466,7 +466,7 @@ class LoggingWorker: # Process all tasks concurrently for maximum speed await asyncio.gather(*[self._process_single_task(task) for task in tasks]) - def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine): + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, None, None]) -> None: """ Ensure the logging worker is initialized and enqueue the coroutine. """ diff --git a/litellm/llms/openai/live/__init__.py b/litellm/llms/openai/live/__init__.py new file mode 100644 index 00000000000..03d40156213 --- /dev/null +++ b/litellm/llms/openai/live/__init__.py @@ -0,0 +1,3 @@ +from .handler import LiveSessionStart, OpenAILiveSessions, parse_live_session_start + +__all__ = ("LiveSessionStart", "OpenAILiveSessions", "parse_live_session_start") diff --git a/litellm/llms/openai/live/handler.py b/litellm/llms/openai/live/handler.py new file mode 100644 index 00000000000..297a147a288 --- /dev/null +++ b/litellm/llms/openai/live/handler.py @@ -0,0 +1,555 @@ +import asyncio +import json +import math +import ssl as ssl_module +from collections.abc import AsyncIterator, Callable, Mapping +from contextlib import AbstractAsyncContextManager, suppress +from dataclasses import dataclass +from datetime import datetime +from itertools import count +from types import MappingProxyType +from typing import Final, Protocol, TypeAlias +from urllib.parse import urlsplit, urlunsplit + +from pydantic import TypeAdapter, ValidationError +from typing_extensions import NotRequired, ReadOnly, TypedDict + +from litellm._logging import redact_secrets +from litellm.constants import ( + OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS, + OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES, + REALTIME_SESSION_SUCCESS_LOGGED_KEY, + REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, +) +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker +from litellm.litellm_core_utils.realtime_errors import ( + close_after_upstream_handshake_refusal, + realtime_error_event, + websocket_close_reason, +) +from litellm.litellm_core_utils.realtime_streaming import backend_close_from +from litellm.llms.openai.realtime.handler import OpenAIRealtime +from litellm.types.utils import LiteLLMRealtimeStreamLoggingObject, Usage + +_LIVE_EVENT_FRAME_ADAPTER: Final = TypeAdapter(dict[str, object]) +_LIVE_USAGE_EVENT_TYPES: Final = frozenset({"session.usage.updated", "session.closed"}) +_LIVE_WEBSOCKET_SCHEMES: Final = MappingProxyType({"https": "wss", "http": "ws"}) +_LIVE_SESSION_CLOSE_FRAME: Final = '{"type": "session.close"}' + + +class _LiveSessionModel(TypedDict): + model: ReadOnly[str] + + +class _LiveSessionStartedEvent(TypedDict): + type: ReadOnly[str] + session: NotRequired[ReadOnly[_LiveSessionModel]] + + +class _LiveUsageSeconds(TypedDict): + seconds: ReadOnly[object] + + +class _LiveUsageEvent(TypedDict): + type: ReadOnly[str] + usage: NotRequired[ReadOnly[_LiveUsageSeconds]] + + +class _LiveResponseSummary(TypedDict): + id: ReadOnly[str] + model: ReadOnly[object] + usage: ReadOnly[object] + + +class _LiveTerminalResponseEvent(TypedDict): + type: ReadOnly[str] + response: ReadOnly[_LiveResponseSummary] + + +class _LiveResponseEvent(TypedDict): + type: ReadOnly[str] + event: ReadOnly[_LiveTerminalResponseEvent] + + +class _LivePreCallInput(TypedDict): + session_start: ReadOnly[Mapping[str, object]] + + +class _LivePreCallArgs(TypedDict): + api_base: ReadOnly[str] + headers: ReadOnly[Mapping[str, str]] + complete_input_dict: ReadOnly[_LivePreCallInput] + + +class LiveClientWebSocket(Protocol): + async def receive_text(self) -> str: ... + + async def send_text(self, data: str) -> None: ... + + async def close(self, code: int = 1000, reason: str | None = None) -> None: ... + + +class LiveBackendWebSocket(Protocol): + async def send(self, message: str, /) -> None: ... + + async def recv(self) -> str | bytes: ... + + async def close(self) -> None: ... + + +class LiveWebSocketConnector(Protocol): + def __call__( + self, + url: str, + *, + additional_headers: Mapping[str, str], + max_size: int | None, + ssl: bool | str | ssl_module.SSLContext | None, + ) -> AbstractAsyncContextManager[LiveBackendWebSocket]: ... + + +def _openai_live_websocket_connect( + url: str, + *, + additional_headers: Mapping[str, str], + max_size: int | None, + ssl: bool | str | ssl_module.SSLContext | None, +) -> AbstractAsyncContextManager[LiveBackendWebSocket]: + import websockets + + return websockets.connect( + url, + additional_headers=additional_headers, + max_size=max_size, + ssl=ssl, + ) + + +def _is_finite_usage_seconds(value: object) -> bool: + return ( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and (not isinstance(value, float) or math.isfinite(value)) + ) + + +def _session_started_event(event: Mapping[str, object]) -> _LiveSessionStartedEvent: + session: Final = as_str_mapping(event.get("session")) + session_model: Final = session.get("model") if session is not None else None + if isinstance(session_model, str): + started_with_model: Final[_LiveSessionStartedEvent] = { + "type": "session.started", + "session": {"model": session_model}, + } + return started_with_model + started: Final[_LiveSessionStartedEvent] = {"type": "session.started"} + return started + + +def _usage_event(event: Mapping[str, object], event_type: str) -> _LiveUsageEvent: + usage: Final = as_str_mapping(event.get("usage")) + seconds: Final = usage.get("seconds") if usage is not None else None + if _is_finite_usage_seconds(seconds): + usage_with_seconds: Final[_LiveUsageEvent] = {"type": event_type, "usage": {"seconds": seconds}} + return usage_with_seconds + usage_without_seconds: Final[_LiveUsageEvent] = {"type": event_type} + return usage_without_seconds + + +def _terminal_response_event(event: Mapping[str, object]) -> _LiveResponseEvent | None: + nested_event: Final = as_str_mapping(event.get("event")) + if nested_event is None: + return None + nested_event_type: Final = nested_event.get("type") + if not isinstance(nested_event_type, str) or nested_event_type not in OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES: + return None + response: Final = as_str_mapping(nested_event.get("response")) + response_id: Final = response.get("id") if response is not None else None + if not isinstance(response_id, str) or response is None: + return None + response_event: Final[_LiveResponseEvent] = { + "type": "response.event", + "event": { + "type": nested_event_type, + "response": { + "id": response_id, + "model": response.get("model"), + "usage": response.get("usage"), + }, + }, + } + return response_event + + +def _metering_event(message: str) -> Mapping[str, object] | None: + try: + event: Final = _LIVE_EVENT_FRAME_ADAPTER.validate_json(message) + except ValidationError: + return None + event_type: Final = event.get("type") + if event_type == "session.started": + return _session_started_event(event) + if isinstance(event_type, str) and event_type in _LIVE_USAGE_EVENT_TYPES: + return _usage_event(event, event_type) + if event_type == "response.event": + return _terminal_response_event(event) + return None + + +@dataclass(frozen=True, slots=True) +class LiveSessionStart: + frame: Mapping[str, object] + model: str + + +_LIVE_SESSION_START_ADAPTER: Final = TypeAdapter(dict[str, object]) + + +def parse_live_session_start(first_message: str) -> LiveSessionStart | None: + try: + frame: Final = _LIVE_SESSION_START_ADAPTER.validate_json(first_message.replace("\r", "").replace("\n", "")) + session: Final = _LIVE_SESSION_START_ADAPTER.validate_python(frame.get("session")) + except ValidationError: + return None + model: Final = session.get("model") + if frame.get("type") != "session.start" or not isinstance(model, str) or not model.strip(): + return None + return LiveSessionStart(frame=MappingProxyType(frame), model=model) + + +def _is_session_close_message(message: str) -> bool: + return _live_event_type(message) == "session.close" + + +def _is_session_closed_message(message: str) -> bool: + return _live_event_type(message) == "session.closed" + + +def _live_event_type(message: str) -> object: + try: + event: Final = _LIVE_EVENT_FRAME_ADAPTER.validate_json(message) + except ValidationError: + return None + return event.get("type") + + +_RetainedEvent: TypeAlias = tuple[int, Mapping[str, object]] + + +def _response_event_has_usage(event: Mapping[str, object]) -> bool: + nested_event: Final = as_str_mapping(event.get("event")) + response: Final = as_str_mapping(nested_event.get("response")) if nested_event is not None else None + return response is not None and as_str_mapping(response.get("usage")) is not None + + +class _LiveEventRetention: + def __init__(self) -> None: + self._event_sequence = count() + self._session_started: _RetainedEvent | None = None + self._latest_usage: _RetainedEvent | None = None + self._session_closed: _RetainedEvent | None = None + self._response_events: dict[str, _RetainedEvent] = {} # mutable-ok: one entry per response id, avoids copies + + def retain(self, event: Mapping[str, object]) -> None: + retained_item: Final = (next(self._event_sequence), event) + event_type: Final = event.get("type") + if event_type == "session.started": + if self._session_started is None: + self._session_started = retained_item + return + if event_type == "session.usage.updated": + usage: Final = as_str_mapping(event.get("usage")) + seconds: Final = usage.get("seconds") if usage is not None else None + if _is_finite_usage_seconds(seconds): + self._latest_usage = retained_item + return + if event_type == "session.closed": + self._session_closed = retained_item + return + nested_event: Final = as_str_mapping(event.get("event")) + response: Final = as_str_mapping(nested_event.get("response")) if nested_event is not None else None + response_id: Final = response.get("id") if response is not None else None + if not isinstance(response_id, str): + return + retained_response: Final = self._response_events.get(response_id) + if retained_response is None or not _response_event_has_usage(retained_response[1]): + self._response_events[response_id] = retained_item + + def results(self) -> tuple[Mapping[str, object], ...]: + retained_items: Final = tuple( + item for item in (self._session_started, self._latest_usage, self._session_closed) if item is not None + ) + tuple(self._response_events.values()) + return tuple(event for _, event in sorted(retained_items, key=lambda item: item[0])) + + +async def _forward_live_client_frames( + *, + websocket: LiveClientWebSocket, + backend_websocket: LiveBackendWebSocket, + client_gone: asyncio.Event, + client_close_sent: asyncio.Event, +) -> None: + async for message in _live_client_messages(websocket, client_gone): + await backend_websocket.send(message) + if _is_session_close_message(message): + client_close_sent.set() + + +async def _live_client_messages( + websocket: LiveClientWebSocket, + client_gone: asyncio.Event, +) -> AsyncIterator[str]: + from starlette.websockets import WebSocketDisconnect + + try: + while True: + yield await websocket.receive_text() + except (WebSocketDisconnect, RuntimeError): + client_gone.set() + + +async def _forward_live_backend_frames( + *, + websocket: LiveClientWebSocket, + backend_websocket: LiveBackendWebSocket, + client_gone: asyncio.Event, + event_retention: _LiveEventRetention, +) -> None: + from starlette.websockets import WebSocketDisconnect + from websockets.exceptions import ConnectionClosed + + event_loop: Final = asyncio.get_running_loop() + async for message, close_deadline in _live_backend_messages(websocket, backend_websocket): + if (event := _metering_event(message)) is not None: + event_retention.retain(event) + if client_gone.is_set(): + continue + try: + if close_deadline is None: + await websocket.send_text(message) + else: + await asyncio.wait_for( + websocket.send_text(message), + timeout=max(close_deadline - event_loop.time(), 0), + ) + except (TimeoutError, asyncio.TimeoutError, WebSocketDisconnect, RuntimeError, ConnectionClosed, OSError): + client_gone.set() + + +async def _live_backend_messages( + websocket: LiveClientWebSocket, + backend_websocket: LiveBackendWebSocket, +) -> AsyncIterator[tuple[str, float | None]]: + from starlette.websockets import WebSocketDisconnect + from websockets.exceptions import ConnectionClosed + + try: + while not _is_session_closed_message(message := _decode_live_websocket_frame(await backend_websocket.recv())): + yield message, None + except ConnectionClosed as error: + upstream_close: Final = backend_close_from(error) + with suppress(WebSocketDisconnect, RuntimeError, ConnectionClosed, OSError): + await websocket.close(code=upstream_close.code, reason=upstream_close.reason) + return + event_loop: Final = asyncio.get_running_loop() + close_deadline: Final = event_loop.time() + OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS + try: + yield message, close_deadline + while True: + yield ( + _decode_live_websocket_frame( + await asyncio.wait_for( + backend_websocket.recv(), + timeout=max(close_deadline - event_loop.time(), 0), + ) + ), + close_deadline, + ) + except (TimeoutError, asyncio.TimeoutError): + with suppress(WebSocketDisconnect, RuntimeError, ConnectionClosed, OSError): + await websocket.close(code=1000) + except ConnectionClosed as error: + with suppress(WebSocketDisconnect, RuntimeError, ConnectionClosed, OSError): + if error.rcvd is None: + await websocket.close(code=1000) + else: + await websocket.close(code=error.rcvd.code, reason=error.rcvd.reason) + + +def _decode_live_websocket_frame(frame: str | bytes) -> str: + return frame.decode("utf-8") if isinstance(frame, bytes) else frame + + +class OpenAILiveSessions(OpenAIRealtime): + def __init__( + self, + logging_worker: LoggingWorker = GLOBAL_LOGGING_WORKER, + websocket_connector: LiveWebSocketConnector | None = None, + clock: Callable[[], datetime] = datetime.now, + ) -> None: + super().__init__() + self._logging_worker = logging_worker + self._clock = clock + self._websocket_connector = ( + websocket_connector if websocket_connector is not None else _openai_live_websocket_connect + ) + + @staticmethod + def _construct_live_url(api_base: str) -> str: + parsed_url: Final = urlsplit(api_base) + websocket_scheme: Final = _LIVE_WEBSOCKET_SCHEMES.get(parsed_url.scheme, parsed_url.scheme) + return urlunsplit((websocket_scheme, parsed_url.netloc, "/v1/live/sessions", "", "")) + + async def async_live_session( + self, + *, + model: str, + websocket: LiveClientWebSocket, + logging_obj: LiteLLMLogging, + session_start: Mapping[str, object], + api_base: str | None, + api_key: str | None, + ) -> None: + if api_key is None: + raise ValueError("api_key is required for OpenAI Live session calls") + from websockets.exceptions import InvalidStatus + + resolved_api_base: Final = api_base or self._get_default_api_base() + url: Final = self._construct_live_url(resolved_api_base) + headers: Final = self._get_additional_headers(api_key) + ssl_config: Final = self._get_ssl_config(url) + session: Final = as_str_mapping(session_start.get("session")) + if session is None: + raise TypeError("session_start.session must be a mapping") + rewritten_start: Final = { + **session_start, + "session": {**session, "model": model}, + } + upstream_start: Final = json.dumps(rewritten_start) + pre_call_args: Final[_LivePreCallArgs] = { + "api_base": url, + "headers": headers, + "complete_input_dict": {"session_start": rewritten_start}, + } + event_retention: Final = _LiveEventRetention() + try: + logging_obj.pre_call(input=None, api_key=api_key, additional_args=pre_call_args) + async with self._websocket_connector( + url, + additional_headers=headers, + max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, + ssl=ssl_config, + ) as backend_websocket: + await backend_websocket.send(upstream_start) + try: + await self._bidirectional_forward( + websocket=websocket, + backend_websocket=backend_websocket, + event_retention=event_retention, + ) + finally: + session_end_time: Final = self._clock() + retained_events: Final = event_retention.results() + logging_result: Final = LiteLLMRealtimeStreamLoggingObject.model_validate( + {"usage": Usage(), "results": list(retained_events)} + ) + self._logging_worker.ensure_initialized_and_enqueue( + logging_obj.dispatch_success_handlers( + logging_result, end_time=session_end_time, prefer_async_handlers=True + ) + ) + logging_obj.model_call_details[ # rebind-ok: flags the shared logging object as already logged + REALTIME_SESSION_SUCCESS_LOGGED_KEY + ] = True + except InvalidStatus as error: + await close_after_upstream_handshake_refusal(websocket, error.response.status_code) + except Exception as error: # noqa: BLE001 # any relay failure must still close the client socket with a redacted reason + redacted_error: Final = redact_secrets(str(error)) + with suppress(Exception): + await websocket.send_text(realtime_error_event(redacted_error, error_type="server_error")) + try: + await websocket.close( + code=1011, + reason=websocket_close_reason(redacted_error, fallback="Internal server error"), + ) + except RuntimeError as close_error: + if "already completed" not in str(close_error) and "websocket.close" not in str(close_error): + raise RuntimeError(f"Unexpected error while closing WebSocket: {close_error}") from close_error + + async def _bidirectional_forward( + self, + *, + websocket: LiveClientWebSocket, + backend_websocket: LiveBackendWebSocket, + event_retention: _LiveEventRetention, + ) -> None: + from websockets.exceptions import ConnectionClosed + + client_gone: Final = asyncio.Event() + client_close_sent: Final = asyncio.Event() + client_forward_task: Final = asyncio.create_task( + _forward_live_client_frames( + websocket=websocket, + backend_websocket=backend_websocket, + client_gone=client_gone, + client_close_sent=client_close_sent, + ) + ) + backend_forward_task: Final = asyncio.create_task( + _forward_live_backend_frames( + websocket=websocket, + backend_websocket=backend_websocket, + client_gone=client_gone, + event_retention=event_retention, + ) + ) + client_gone_task: Final = asyncio.create_task(client_gone.wait()) + forwarding_tasks: Final = (client_forward_task, backend_forward_task, client_gone_task) + try: + completed_tasks: Final = await asyncio.wait( + forwarding_tasks, + return_when=asyncio.FIRST_COMPLETED, + ) + done_tasks: Final = completed_tasks[0] + if backend_forward_task in done_tasks: + backend_forward_task.result() + return + client_forward_error: Final = client_forward_task.exception() if client_forward_task in done_tasks else None + if client_forward_task not in done_tasks: + client_forward_task.cancel() + await asyncio.gather(client_forward_task, return_exceptions=True) + event_loop: Final = asyncio.get_running_loop() + close_deadline: Final = event_loop.time() + OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS + if not client_close_sent.is_set(): + remaining_send_timeout: Final = max(close_deadline - event_loop.time(), 0) + with suppress(ConnectionClosed, TimeoutError, asyncio.TimeoutError): + await asyncio.wait_for( + backend_websocket.send(_LIVE_SESSION_CLOSE_FRAME), + timeout=remaining_send_timeout, + ) + remaining_wait_timeout: Final = max(close_deadline - event_loop.time(), 0) + close_wait_result: Final = await asyncio.wait( + (backend_forward_task,), + timeout=remaining_wait_timeout, + ) + remaining_close_timeout: Final = max(close_deadline - event_loop.time(), 0) + with suppress(TimeoutError, asyncio.TimeoutError): + await asyncio.wait_for( + backend_websocket.close(), + timeout=remaining_close_timeout, + ) + if backend_forward_task in close_wait_result[0]: + backend_forward_task.result() + if client_forward_error is not None and not isinstance(client_forward_error, ConnectionClosed): + raise client_forward_error + return + backend_forward_task.cancel() + if client_forward_error is not None: + raise client_forward_error + finally: + for task in forwarding_tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*forwarding_tasks, return_exceptions=True) diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index 325169cfe15..853f707bc91 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -45,7 +45,7 @@ class OpenAIRealtime(OpenAIChatCompletion): api_key: str, *, openai_beta_realtime: bool = False, - ) -> dict: + ) -> dict[str, str]: """ Headers for the upstream OpenAI Realtime WebSocket. @@ -54,7 +54,7 @@ class OpenAIRealtime(OpenAIChatCompletion): so the legacy beta API is used. GA clients omit that header on the client connection and must send GA-shaped ``session.update`` payloads. """ - headers: Final[dict] = {"Authorization": f"Bearer {api_key}"} + headers: Final[dict[str, str]] = {"Authorization": f"Bearer {api_key}"} if openai_beta_realtime: headers["OpenAI-Beta"] = "realtime=v1" return headers diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b6c81f17761..74f853ee2ce 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -316,6 +316,9 @@ class LiteLLMRoutes(enum.Enum): "/realtime", "/v1/realtime", "/openai/v1/realtime", + "/live/sessions", + "/v1/live/sessions", + "/openai/v1/live/sessions", "/realtime?{model}", "/v1/realtime?{model}", "/openai/v1/realtime?{model}", diff --git a/litellm/proxy/agent_endpoints/auth/managed_authorization.py b/litellm/proxy/agent_endpoints/auth/managed_authorization.py index 64e50cd74e7..cbbd1548301 100644 --- a/litellm/proxy/agent_endpoints/auth/managed_authorization.py +++ b/litellm/proxy/agent_endpoints/auth/managed_authorization.py @@ -12,6 +12,7 @@ from litellm.types.agents import AgentResponse from litellm.types.proxy.agent_identity import AgentIdentityFailure, ManagedAgentContext _MANAGED_REALTIME_ROUTES: Final = frozenset(("/realtime", "/v1/realtime", "/openai/v1/realtime")) +_MANAGED_LIVE_ROUTES: Final = frozenset(("/live/sessions", "/v1/live/sessions", "/openai/v1/live/sessions")) _MANAGED_MODEL_ROUTES: Final = frozenset( f"{prefix}/{operation}" for prefix, operation in product( @@ -87,7 +88,7 @@ def managed_agent_route_allowed(route: str, method: str | None) -> bool: if route in ("/agents", "/v1/agents"): return method in (None, "GET", "HEAD") - if route in _MANAGED_REALTIME_ROUTES: + if route in _MANAGED_REALTIME_ROUTES or route in _MANAGED_LIVE_ROUTES: return method in (None, "GET") if route in _MANAGED_MODEL_ROUTES or RouteChecks.check_route_access(route, _MANAGED_MODEL_PATHS): return method in (None, "POST") @@ -106,6 +107,8 @@ def managed_inference_request( ) -> dict[str, object]: from litellm.proxy.auth.route_checks import RouteChecks + if route in _MANAGED_LIVE_ROUTES: + return dict(body) if route in _MANAGED_REALTIME_ROUTES: model: Final = query_model or body.get("model") if not isinstance(model, str) or not model: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 56773789059..38b1ffac9c8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1255,24 +1255,7 @@ async def common_checks( code=status.HTTP_400_BAD_REQUEST, ) - managed_policy: Final = managed_agent_policy(valid_token) - if _model and valid_token is not None and managed_policy is not None: - managed_models: Final = (managed_policy.object_permission or MappingProxyType({})).get("models", ()) - if not isinstance(managed_models, (list, tuple)) or not managed_models: - raise HTTPException(403, "This agent has no model grants") - can_object_call_model( - model=_resolve_team_alias( - _model, team_model_aliases_for_auth_check(valid_token), valid_token.team_id, llm_router - ), - llm_router=llm_router, - models=list(managed_models), - team_id=valid_token.team_id, - object_type="agent", - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - - await _check_agent_access_group_model_access(model=_model, valid_token=valid_token, llm_router=llm_router) - await _check_agent_caller_model_access( + await can_agent_call_model( model=_model, valid_token=valid_token, llm_router=llm_router, @@ -4763,7 +4746,7 @@ def _can_object_call_model( raise Exception(f"Unable to parse model, max fallback depth exceeded - received model: {model}") if isinstance(model, list): for m in model: - _can_object_call_model( + can_object_call_model( model=m, llm_router=llm_router, models=models, @@ -4863,6 +4846,40 @@ def _live_team_alias_target( return model if deleted_team_deployment else target +async def can_agent_call_model( + model: str | list[str] | None, # mutable-ok: the model checks it delegates to take list[str] + valid_token: UserAPIKeyAuth | None, + llm_router: Router | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> None: + managed_policy: Final = managed_agent_policy(valid_token) + if model and valid_token is not None and managed_policy is not None: + managed_models: Final = (managed_policy.object_permission or {}).get("models", ()) + if not isinstance(managed_models, (list, tuple)) or not managed_models: + raise HTTPException(403, "This agent has no model grants") + can_object_call_model( + model=_resolve_team_alias( + model, team_model_aliases_for_auth_check(valid_token), valid_token.team_id, llm_router + ), + llm_router=llm_router, + models=list(managed_models), + team_id=valid_token.team_id, + object_type="agent", + key_model_aliases=key_model_aliases_for_auth_check(valid_token), + ) + await _check_agent_access_group_model_access(model=model, valid_token=valid_token, llm_router=llm_router) + await _check_agent_caller_model_access( + model=model, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + async def _check_agent_access_group_model_access( model: str | list[str] | None, # mutable-ok: _can_object_call_model and the client message helper take list[str] valid_token: UserAPIKeyAuth | None, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2231706b282..09bb0443ac3 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1181,7 +1181,13 @@ async def user_api_key_auth_websocket(websocket: WebSocket) -> UserAPIKeyAuth: return await user_api_key_auth_websocket_for_model(websocket, model=websocket.query_params.get("model")) -async def user_api_key_auth_websocket_for_model(websocket: WebSocket, model: str | None) -> UserAPIKeyAuth: +async def user_api_key_auth_live_websocket(websocket: WebSocket) -> UserAPIKeyAuth: + return await user_api_key_auth_websocket_for_model(websocket, model=None) + + +async def user_api_key_auth_websocket_for_model( + websocket: WebSocket, model: str | None, *, close_on_auth_error: bool = True +) -> UserAPIKeyAuth: ws_scope: Final = websocket.scope or {} scope_headers: Final = list(ws_scope.get("headers") or []) # ``get_request_route`` falls back to ``request.url.path`` when @@ -1236,6 +1242,8 @@ async def user_api_key_auth_websocket_for_model(websocket: WebSocket, model: str try: return await user_api_key_auth(request=request, api_key=f"Bearer {api_key}") except Exception as e: + if not close_on_auth_error: + raise if is_invalid_virtual_key_error(e): raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION) log_model_access_denial(e) @@ -1243,6 +1251,16 @@ async def user_api_key_auth_websocket_for_model(websocket: WebSocket, model: str raise HTTPException(status_code=403, detail=str(e)) +async def user_api_key_auth_accepted_websocket_for_model(websocket: WebSocket, model: str) -> UserAPIKeyAuth: + """Re-authorize an already accepted WebSocket for a model it named after the upgrade. + + Runs the same key, team, member, user, project and agent checks, plus budget reservation, as + the upgrade does for ``?model=``, but raises auth errors without closing the socket so the + caller can reply with an error frame. + """ + return await user_api_key_auth_websocket_for_model(websocket, model, close_on_auth_error=False) + + def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_user_params: dict) -> UserAPIKeyAuth: valid_token.end_user_id = end_user_params.get("end_user_id") # Only overwrite token fields when the DB-derived value is not None. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1c7cf54ab4d..d0f11b9350d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -32,6 +32,7 @@ from itertools import chain from types import MappingProxyType, UnionType from typing import ( TYPE_CHECKING, + Annotated, Any, Final, Literal, @@ -73,6 +74,7 @@ from litellm.constants import ( LITELLM_SETTINGS_SAFE_DB_OVERRIDES, LITELLM_UI_ALLOW_HEADERS, LITELLM_UI_SESSION_DURATION, + OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS, RUNTIME_UPDATABLE_ROUTER_SETTINGS, ) from litellm.litellm_core_utils.asyncify import asyncify @@ -334,6 +336,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( mask_sensitive_keys, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.openai.live import parse_live_session_start from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._experimental.mcp_server.byok_credential_cache import byok_credential_cache @@ -347,6 +350,7 @@ from litellm.proxy.analytics_endpoints.analytics_endpoints import ( from litellm.proxy.auth.auth_checks import ( ROLE_BASED_PERMISSIONS_ADAPTER, ExperimentalUIJWTToken, + can_agent_call_model, can_key_call_resolved_model, get_team_object, log_db_metrics, @@ -391,6 +395,8 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401, RUF100 # lega _fetch_global_spend_with_event_coordination, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export fetch_global_spend_with_event_coordination, user_api_key_auth, + user_api_key_auth_accepted_websocket_for_model, + user_api_key_auth_live_websocket, user_api_key_auth_websocket, ) from litellm.proxy.batches_endpoints.endpoints import router as batches_router @@ -3917,14 +3923,8 @@ async def invalidate_spend_counter(counter_key: str) -> None: forget_spend_counter(counter_key) spend_counter_cache.in_memory_cache.delete_cache(key=counter_key) if spend_counter_cache.redis_cache is not None: - try: + with contextlib.suppress(Exception): await spend_counter_cache.redis_cache.async_delete_cache(key=counter_key) - except Exception: - verbose_proxy_logger.debug( - "Unable to delete stale spend counter %s after increment failure", - counter_key, - exc_info=True, - ) _invalidate_spend_counter: Final = invalidate_spend_counter @@ -13194,6 +13194,17 @@ async def _release_realtime_max_parallel_slot(user_api_key_dict: UserAPIKeyAuth) await release_like_http_disconnect(user_api_key_dict) +class _RealtimeErrorBody(TypedDict): + type: ReadOnly[str] + message: ReadOnly[str] + code: NotRequired[ReadOnly[str]] + + +class _RealtimeErrorEvent(TypedDict): + type: ReadOnly[str] + error: ReadOnly[_RealtimeErrorBody] + + async def _reject_realtime_session( websocket: WebSocket, user_api_key_dict: UserAPIKeyAuth, @@ -13201,13 +13212,19 @@ async def _reject_realtime_session( code: int, reason: str, error_message: str | None = None, + error_type: str = "guardrail_error", + error_code: str | None = None, ) -> None: try: if error_message is not None: try: - await websocket.send_text( - json.dumps({"type": "error", "error": {"type": "guardrail_error", "message": error_message}}) - ) + error_event: Final[_RealtimeErrorEvent] = { + "type": "error", + "error": {"type": error_type, "message": error_message} + if error_code is None + else {"type": error_type, "message": error_message, "code": error_code}, + } + await websocket.send_text(json.dumps(error_event)) except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below verbose_proxy_logger.debug("Could not send realtime pre-call error event to client; closing anyway") await websocket.close(code=code, reason=reason) @@ -13229,12 +13246,11 @@ async def realtime_websocket_endpoint( ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket), ): - requested_protocols: Final = [ - p.strip() for p in (websocket.headers.get("sec-websocket-protocol") or "").split(",") if p.strip() - ] - accept_kwargs: Final[dict] = {} - if requested_protocols: - accept_kwargs["subprotocol"] = requested_protocols[0] + requested_protocols: Final = tuple( + protocol.strip() + for protocol in (websocket.headers.get("sec-websocket-protocol") or "").split(",") + if protocol.strip() + ) route_model = model if route_model is None: @@ -13257,26 +13273,47 @@ async def realtime_websocket_endpoint( log_model_access_denial(e) await _reject_realtime_session(websocket, user_api_key_dict, code=1008, reason=e.message[:120]) return - await websocket.accept(**accept_kwargs) + await websocket.accept(subprotocol=requested_protocols[0] if requested_protocols else None) + await _route_realtime_websocket_session( + websocket, + user_api_key_dict, + route_model=route_model, + model=model, + intent=intent, + guardrails=guardrails, + ) + +async def _route_realtime_websocket_session( + websocket: WebSocket, + user_api_key_dict: UserAPIKeyAuth, + *, + route_model: str, + model: str | None, + intent: str | None, + guardrails: str | None, + request_path: str | None = None, + live_session_start: Mapping[str, object] | None = None, +) -> None: # Only use explicit parameters, not all query params query_params: Final = cast(RealtimeQueryParams, dict(_realtime_query_params_template(model, intent))) - data: dict[str, object] = { "model": route_model, "websocket": websocket, "query_params": query_params, # Only explicit params } - # Pass guardrails into data so pre-call guardrail processing picks them up if guardrails: data["guardrails"] = [g.strip() for g in guardrails.split(",") if g.strip()] + if live_session_start is not None: + data["live_session_start"] = live_session_start # Use raw ASGI headers (already lowercase bytes) to avoid extra work headers_list: Final = list(websocket.scope.get("headers") or []) - scope: Final = REALTIME_REQUEST_SCOPE_TEMPLATE.copy() scope["headers"] = headers_list + if request_path is not None: + scope["path"] = request_path request: Final = Request(scope=scope) @@ -13319,6 +13356,7 @@ async def realtime_websocket_endpoint( ) return except BaseException: + await _release_realtime_budget_reservation(user_api_key_dict) await _release_realtime_max_parallel_slot(user_api_key_dict) raise @@ -13356,6 +13394,141 @@ async def realtime_websocket_endpoint( await _release_realtime_max_parallel_slot(user_api_key_dict) +async def _reject_invalid_live_session_start( + websocket: WebSocket, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + await _reject_realtime_session( + websocket, + user_api_key_dict, + code=1008, + reason="Invalid session.start frame", + error_message="First message must be a session.start JSON object with a non-empty session.model.", + error_type="invalid_request_error", + error_code="invalid_session_start", + ) + + +LiveSessionModelAuthorizer: TypeAlias = Callable[[WebSocket, str], Awaitable[UserAPIKeyAuth]] + + +def live_session_model_authorizer() -> LiveSessionModelAuthorizer: + return user_api_key_auth_accepted_websocket_for_model + + +async def _authorize_live_session_model( + authorize_session_model: LiveSessionModelAuthorizer, + websocket: WebSocket, + model: str, +) -> UserAPIKeyAuth: + session_auth: Final = await authorize_session_model(websocket, model) + try: + await can_key_call_resolved_model( + model=model, + llm_model_list=llm_model_list, + valid_token=session_auth, + llm_router=llm_router, + ) + await can_agent_call_model( + model=model, + valid_token=session_auth, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except BaseException: + await _release_realtime_budget_reservation(session_auth) + raise + return session_auth + + +def _live_session_denial_code(error: ProxyException | HTTPException) -> str: + model_denied: Final = isinstance(error, ProxyException) and str(error.type).endswith("_model_access_denied") + return "model_not_allowed" if model_denied else "session_not_allowed" + + +@app.websocket("/openai/v1/live/sessions") +@app.websocket("/v1/live/sessions") +@app.websocket("/live/sessions") +async def live_sessions_websocket_endpoint( + websocket: WebSocket, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_live_websocket)], + authorize_session_model: Annotated[LiveSessionModelAuthorizer, Depends(live_session_model_authorizer)], +) -> None: + requested_protocols: Final = tuple( + protocol.strip() + for protocol in (websocket.headers.get("sec-websocket-protocol") or "").split(",") + if protocol.strip() + ) + try: + await websocket.accept(subprotocol=requested_protocols[0] if requested_protocols else None) + except BaseException: + await _release_realtime_budget_reservation(user_api_key_dict) + await _release_realtime_max_parallel_slot(user_api_key_dict) + raise + + try: + first_message: Final = await asyncio.wait_for( + websocket.receive_text(), + timeout=OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS, + ) + except asyncio.TimeoutError: + await _reject_invalid_live_session_start(websocket, user_api_key_dict) + return + except WebSocketDisconnect: + await _release_realtime_budget_reservation(user_api_key_dict) + await _release_realtime_max_parallel_slot(user_api_key_dict) + return + except RuntimeError: + await _reject_invalid_live_session_start(websocket, user_api_key_dict) + return + except BaseException: + await _release_realtime_budget_reservation(user_api_key_dict) + await _release_realtime_max_parallel_slot(user_api_key_dict) + raise + + session_start: Final = parse_live_session_start(first_message) + if session_start is None: + await _reject_invalid_live_session_start(websocket, user_api_key_dict) + return + + upgrade_auth: Final = user_api_key_dict.model_copy() + await _release_realtime_budget_reservation(upgrade_auth) + try: + session_auth: Final = await _authorize_live_session_model( + authorize_session_model, websocket, session_start.model + ) + except (ProxyException, HTTPException) as error: + denial_message: Final = error.message if isinstance(error, ProxyException) else str(error.detail) + if isinstance(error, ProxyException): + log_model_access_denial(error) + await _reject_realtime_session( + websocket, + upgrade_auth, + code=1008, + reason=websocket_close_reason(denial_message, fallback="Session not allowed"), + error_message=denial_message, + error_type="invalid_request_error", + error_code=_live_session_denial_code(error), + ) + return + except BaseException: + await _release_realtime_max_parallel_slot(upgrade_auth) + raise + + await _route_realtime_websocket_session( + websocket, + session_auth, + route_model=session_start.model, + model=None, + intent=None, + guardrails=None, + request_path=websocket.url.path, + live_session_start=session_start.frame, + ) + + ###################################################################### # /v1/assistant Endpoints diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 0e79cbfda13..d6bb44b1d6b 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -3,8 +3,11 @@ import asyncio import os from collections.abc import Mapping +from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, cast + +from typing_extensions import ReadOnly, TypedDict, assert_never import litellm from litellm.constants import ( @@ -37,6 +40,7 @@ from ..llms.azure.common_utils import get_azure_ad_token from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_protocol_for_client from ..llms.bedrock.realtime.handler import BedrockRealtime from ..llms.custom_httpx.http_handler import realtime_ssl_for_url +from ..llms.openai.live.handler import OpenAILiveSessions from ..llms.openai.realtime.handler import OpenAIRealtime from ..llms.vertex_ai.audio_transcription.realtime_transformation import is_vertex_speech_to_text_model from ..llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig, vertex_realtime_config @@ -51,6 +55,7 @@ if TYPE_CHECKING: azure_realtime: Final = AzureOpenAIRealtime() openai_realtime: Final = OpenAIRealtime() +openai_live_sessions: Final = OpenAILiveSessions() bedrock_realtime: Final = BedrockRealtime() xai_realtime: Final = XAIRealtime() vertex_llm_base: Final = VertexBase() @@ -75,9 +80,9 @@ def _with_resolved_session_model(session: dict[str, object], model_name: str) -> return {**session, "model": model_name} -def _build_litellm_metadata(kwargs: dict) -> dict: +def _build_litellm_metadata(kwargs: dict) -> dict[str, object]: """Build the litellm_metadata dict for guardrail checking (internal only, not forwarded to provider).""" - metadata: Final[dict] = {**(kwargs.get("litellm_metadata") or {})} + metadata: Final[dict[str, object]] = {**(kwargs.get("litellm_metadata") or {})} guardrails: Final = (kwargs.get("metadata") or {}).get("guardrails") or kwargs.get("guardrails") or [] if guardrails: metadata["guardrails"] = guardrails @@ -345,6 +350,84 @@ async def _resolve_vertex_access_token_bounded( ) from e +class _RealtimeGuardrailRequestData(TypedDict): + litellm_metadata: ReadOnly[Mapping[str, object]] + stream: ReadOnly[Literal[True]] + + +def _realtime_guardrail_configured(litellm_metadata: Mapping[str, object]) -> bool: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import GuardrailEventHooks + + event_hooks: Final = ( + GuardrailEventHooks.realtime_input_transcription, + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ) + request_data: Final[_RealtimeGuardrailRequestData] = {"litellm_metadata": litellm_metadata, "stream": True} + return any( + isinstance(callback, CustomGuardrail) + and any(callback.should_run_guardrail(data=request_data, event_type=hook) for hook in event_hooks) + for callback in litellm.callbacks + ) + + +@dataclass(frozen=True, slots=True) +class _LiveProviderUnsupported: + provider: str + + +@dataclass(frozen=True, slots=True) +class _LiveGuardrailsUnsupported: + pass + + +_LiveSessionRejection = _LiveProviderUnsupported | _LiveGuardrailsUnsupported + + +def _live_session_rejection( + custom_llm_provider: str, litellm_metadata: Mapping[str, object] +) -> _LiveSessionRejection | None: + if custom_llm_provider != "openai": + return _LiveProviderUnsupported(provider=custom_llm_provider) + if _realtime_guardrail_configured(litellm_metadata): + return _LiveGuardrailsUnsupported() + return None + + +def _raise_live_session_rejection(rejection: _LiveSessionRejection) -> NoReturn: + match rejection: + case _LiveProviderUnsupported(provider=provider): + raise ValueError(f"OpenAI Live sessions require the openai provider, got {provider}") + case _LiveGuardrailsUnsupported(): + raise ValueError("Guardrails are not supported on OpenAI Live sessions") + case _: + assert_never(rejection) + + +async def _alive_session( + model: str, + websocket: "WebSocket", + custom_llm_provider: str, + litellm_logging_obj: LiteLLMLogging, + session_start: Mapping[str, object], + api_base: str | None, + api_key: str | None, + litellm_metadata: Mapping[str, object], +) -> None: + rejection: Final = _live_session_rejection(custom_llm_provider, litellm_metadata) + if rejection is not None: + _raise_live_session_rejection(rejection) + await openai_live_sessions.async_live_session( + model=model, + websocket=websocket, + logging_obj=litellm_logging_obj, + session_start=session_start, + api_base=api_base or litellm.api_base or "https://api.openai.com/", + api_key=api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY"), + ) + + @wrapper_client async def _arealtime( model: str, @@ -356,6 +439,7 @@ async def _arealtime( client: object | None = None, timeout: float | None = None, query_params: RealtimeQueryParams | None = None, + live_session_start: Mapping[str, object] | None = None, **kwargs, ) -> None: """ @@ -398,6 +482,19 @@ async def _arealtime( custom_llm_provider=_custom_llm_provider, ) + if live_session_start is not None: + await _alive_session( + model=model, + websocket=websocket, + custom_llm_provider=_custom_llm_provider, + litellm_logging_obj=litellm_logging_obj, + session_start=live_session_start, + api_base=dynamic_api_base or litellm_params.api_base, + api_key=dynamic_api_key, + litellm_metadata=_build_litellm_metadata(kwargs), + ) + return + provider_config: BaseRealtimeConfig | None = None if _custom_llm_provider in LlmProviders._member_map_.values(): provider_config = ProviderConfigManager.get_provider_realtime_config( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e6103270ee3..eab7d52ac57 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1058,6 +1058,9 @@ API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = { "/realtime": [CallTypes.arealtime], "/v1/realtime": [CallTypes.arealtime], "/openai/v1/realtime": [CallTypes.arealtime], + "/live/sessions": (CallTypes.arealtime,), + "/v1/live/sessions": (CallTypes.arealtime,), + "/openai/v1/live/sessions": (CallTypes.arealtime,), # Provider-specific routes "/anthropic/v1/messages": [CallTypes.anthropic_messages], # Google GenAI routes diff --git a/tests/unit/llms/openai/live/__init__.py b/tests/unit/llms/openai/live/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/live/test_handler.py b/tests/unit/llms/openai/live/test_handler.py new file mode 100644 index 00000000000..ec1130376d0 --- /dev/null +++ b/tests/unit/llms/openai/live/test_handler.py @@ -0,0 +1,1208 @@ +import asyncio +import json +from collections.abc import Coroutine +from datetime import datetime +from typing import Final, NoReturn, cast + +import pytest +from pydantic import TypeAdapter +from starlette.websockets import WebSocketDisconnect +from websockets.datastructures import Headers +from websockets.exceptions import ConnectionClosed, ConnectionClosedError, ConnectionClosedOK, InvalidStatus +from websockets.frames import Close +from websockets.http11 import Response + +import litellm +from litellm.constants import REALTIME_SESSION_SUCCESS_LOGGED_KEY, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES +from litellm.cost_calculator import handle_realtime_stream_cost_calculation +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.openai.live import handler as live_handler +from litellm.llms.openai.live.handler import LiveSessionStart, OpenAILiveSessions, parse_live_session_start +from litellm.types.utils import LiteLLMRealtimeStreamLoggingObject, Usage + +_LIVE_EVENT_FRAME_ADAPTER: Final = TypeAdapter(dict[str, object]) + + +def _event_type(message: str) -> str | None: + event_type: Final = _LIVE_EVENT_FRAME_ADAPTER.validate_json(message).get("type") + return event_type if isinstance(event_type, str) else None + + +class _FakeClientWebSocket: + def __init__( + self, + messages: tuple[str, ...], + *, + disconnect_after_send_type: str | None = None, + disconnect_after_receive_type: str | None = None, + fail_on_send_type: str | None = None, + block_on_send_type: str | None = None, + receive_error_after_close: Exception | None = None, + ) -> None: + self._messages = iter(messages) + self._block_on_send_type = block_on_send_type + self._receive_error_after_close = receive_error_after_close + self._disconnect_after_send_type = disconnect_after_send_type + self._disconnect_after_receive_type = disconnect_after_receive_type + self._fail_on_send_type = fail_on_send_type + self.sent_text: asyncio.Queue[str] = asyncio.Queue() + self.closed = asyncio.Event() + self.close_info: tuple[int, str | None] | None = None + + async def receive_text(self) -> str: + try: + message: Final = next(self._messages) + if _event_type(message) == self._disconnect_after_receive_type: + self.closed.set() + return message + except StopIteration: + await self.closed.wait() + raise self._receive_error_after_close or WebSocketDisconnect(code=1000) + + async def send_text(self, data: str) -> None: + event_type: Final = _event_type(data) + if event_type == self._fail_on_send_type: + raise RuntimeError("client is disconnected") + if event_type == self._block_on_send_type: + await asyncio.Event().wait() + await self.sent_text.put(data) + if event_type == self._disconnect_after_send_type: + self.closed.set() + + async def close(self, code: int = 1000, reason: str | None = None) -> None: + self.close_info = (code, reason) + self.closed.set() + + +class _FakeUpstream: + def __init__( + self, + client_frames: tuple[str, ...], + upstream_frames: tuple[str, ...], + *, + close_response: str | None = None, + never_respond_after_close: bool = False, + block_session_close_send: bool = False, + block_close: bool = False, + fail_on_send_type: str | None = None, + send_error: Exception | None = None, + repeat_last_frame_forever: bool = False, + close_error: ConnectionClosed | None = None, + ) -> None: + self.client_frames = client_frames + self._close_error = close_error + self._repeat_last_frame_forever = repeat_last_frame_forever + self._last_frame: str | None = None + self._fail_on_send_type = fail_on_send_type + self._send_error = send_error + self._upstream_frames = iter(upstream_frames) + self._close_response = close_response + self._close_response_sent = False + self._never_respond_after_close = never_respond_after_close + self._block_session_close_send = block_session_close_send + self._block_close = block_close + self.sent_text: asyncio.Queue[str] = asyncio.Queue() + self.client_frames_sent = asyncio.Event() + self.session_close_sent = asyncio.Event() + self.closed_event = asyncio.Event() + self.closed = False + + async def send(self, data: str) -> None: + event_type: Final = _event_type(data) + if event_type == self._fail_on_send_type: + raise self._send_error or ConnectionClosedError(None, Close(code=1011, reason="upstream send failed"), None) + if event_type == "session.close" and self._block_session_close_send: + await asyncio.Event().wait() + await self.sent_text.put(data) + if event_type == "session.close": + self.session_close_sent.set() + if self.sent_text.qsize() == len(self.client_frames) + 1: + self.client_frames_sent.set() + + async def recv(self) -> str: + await self.client_frames_sent.wait() + try: + self._last_frame = next(self._upstream_frames) + return self._last_frame + except StopIteration as error: + if self._repeat_last_frame_forever and self._last_frame is not None: + await asyncio.sleep(0) + return self._last_frame + if self._close_response is not None and not self._close_response_sent: + await self.session_close_sent.wait() + self._close_response_sent = True + return self._close_response + if self._never_respond_after_close: + await self.session_close_sent.wait() + await self.closed_event.wait() + raise self._close_error or ConnectionClosedOK( + Close(code=1000, reason="upstream closed"), + Close(code=1000, reason="upstream closed"), + True, + ) from error + + async def close(self) -> None: + if self._block_close: + await asyncio.Event().wait() + self.closed = True + self.closed_event.set() + + async def __aenter__(self) -> "_FakeUpstream": + return self + + async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> None: + self.closed = True + self.closed_event.set() + + +class _FailingConnection: + def __init__(self, error: Exception) -> None: + self.error = error + + async def __aenter__(self) -> NoReturn: + raise self.error + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: object | None, + ) -> None: + return None + + +class _FakeLogging: + def __init__(self, pre_call_error: Exception | None = None) -> None: + self.model_call_details: dict[str, object] = {} + self.pre_call_kwargs: dict[str, object] = {} + self.dispatch_count = 0 + self.dispatched_result: object | None = None + self.dispatched_end_time: datetime | None = None + self.prefer_async_handlers = False + self._pre_call_error = pre_call_error + + def pre_call(self, **kwargs: object) -> None: + if self._pre_call_error is not None: + raise self._pre_call_error + self.pre_call_kwargs = kwargs + + async def dispatch_success_handlers( + self, + result: object, + *, + end_time: datetime | None = None, + prefer_async_handlers: bool = False, + ) -> None: + self.dispatch_count += 1 + self.dispatched_result = result + self.dispatched_end_time = end_time + self.prefer_async_handlers = prefer_async_handlers + + +class _FakeLoggingWorker: + def __init__(self) -> None: + self.coroutine: Coroutine[object, object, None] | None = None + + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: + self.coroutine = async_coroutine + + +class _FakeConnector: + def __init__(self, upstream: _FakeUpstream) -> None: + self.upstream = upstream + self.url: str | None = None + self.kwargs: dict[str, object] = {} + + def __call__(self, url: str, **kwargs: object) -> _FakeUpstream: + self.url = url + self.kwargs = kwargs + return self.upstream + + +@pytest.mark.parametrize( + ("api_base", "expected_url"), + [ + ("https://api.openai.com/", "wss://api.openai.com/v1/live/sessions"), + ("https://api.openai.com/v1", "wss://api.openai.com/v1/live/sessions"), + ("http://localhost:8080", "ws://localhost:8080/v1/live/sessions"), + ("https://api.openai.com/v1?api-version=1", "wss://api.openai.com/v1/live/sessions"), + ], +) +def test_construct_live_url_uses_exact_path_and_drops_query(api_base: str, expected_url: str) -> None: + assert OpenAILiveSessions._construct_live_url(api_base) == expected_url + + +@pytest.mark.asyncio +async def test_async_live_session_rewrites_model_and_relays_frames_without_mutating_input() -> None: + client_frames: Final = ('{"type":"input_audio.append","audio":"client-a"}', '{"type":"session.close"}') + upstream_frames: Final = ( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + '{"type":"response.audio.delta","delta":"server-a"}', + '{"type":"session.closed","usage":{"seconds":16}}', + ) + upstream: Final = _FakeUpstream(client_frames=client_frames, upstream_frames=upstream_frames) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=client_frames) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + session_start: Final[dict[str, object]] = { + "type": "session.start", + "session": { + "model": "proxy-live-alias", + "instructions": "Be concise", + "audio": {"input": {"format": "pcm16"}}, + }, + } + original_session_start: Final = json.loads(json.dumps(session_start)) + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start=session_start, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + sent_frames: Final = tuple( + await asyncio.gather(*(upstream.sent_text.get() for _ in range(upstream.sent_text.qsize()))) + ) + forwarded_frames: Final = tuple( + await asyncio.gather(*(client.sent_text.get() for _ in range(client.sent_text.qsize()))) + ) + rewritten_start: Final = json.loads(sent_frames[0]) + assert connector.url == "wss://api.openai.com/v1/live/sessions" + assert connector.kwargs["additional_headers"] == {"Authorization": "Bearer live-test-key"} + assert connector.kwargs["max_size"] == REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES + assert "ssl" in connector.kwargs + assert rewritten_start == { + "type": "session.start", + "session": { + "model": "gpt-live-1", + "instructions": "Be concise", + "audio": {"input": {"format": "pcm16"}}, + }, + } + assert sent_frames[1:] == client_frames + assert forwarded_frames == upstream_frames + assert session_start == original_session_start + assert client.close_info == (1000, "upstream closed") + assert logger.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] is True + assert worker.coroutine is not None + await worker.coroutine + assert logger.dispatch_count == 1 + assert logger.prefer_async_handlers is True + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert tuple(event["type"] for event in logged_result.results) == ( + "session.started", + "session.closed", + ) + + +@pytest.mark.asyncio +async def test_async_live_session_coalesces_cumulative_usage_updates() -> None: + client_frames: Final = ('{"type":"session.close"}',) + upstream_frames: Final = ( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + '{"type":"response.event","event":{"type":"response.completed","response":{"id":"resp-1"}}}', + '{"type":"response.event","event":{"type":"response.failed","response":{"id":"resp-2"}}}', + '{"type":"session.usage.updated","usage":{"seconds":1}}', + '{"type":"session.usage.updated","usage":{"seconds":2}}', + '{"type":"session.usage.updated","usage":{"seconds":3}}', + '{"type":"session.closed","usage":{"seconds":3}}', + ) + upstream: Final = _FakeUpstream(client_frames=client_frames, upstream_frames=upstream_frames) + connector: Final = _FakeConnector(upstream) + worker: Final = _FakeLoggingWorker() + logger: Final = _FakeLogging() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=_FakeClientWebSocket(messages=client_frames), + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert tuple((event["type"], event.get("usage")) for event in logged_result.results) == ( + ("session.started", None), + ("response.event", None), + ("response.event", None), + ("session.usage.updated", {"seconds": 3}), + ("session.closed", {"seconds": 3}), + ) + assert tuple(logged_result.results[1:3]) == ( + { + "type": "response.event", + "event": { + "type": "response.completed", + "response": {"id": "resp-1", "model": None, "usage": None}, + }, + }, + { + "type": "response.event", + "event": { + "type": "response.failed", + "response": {"id": "resp-2", "model": None, "usage": None}, + }, + }, + ) + + +@pytest.mark.asyncio +async def test_async_live_session_requires_api_key() -> None: + with pytest.raises(ValueError, match="api_key"): + await OpenAILiveSessions().async_live_session( + model="gpt-live-1", + websocket=_FakeClientWebSocket(messages=()), + logging_obj=cast(Logging, _FakeLogging()), + session_start={"type": "session.start", "session": {"model": "proxy-live-alias"}}, + api_base="https://api.openai.com/", + api_key=None, + ) + + +@pytest.mark.asyncio +async def test_async_live_session_does_not_log_success_on_handshake_refusal() -> None: + response: Final = Response(401, "Unauthorized", Headers()) + invalid_status: Final = InvalidStatus(response) + + def refuse_connection(url: str, **kwargs: object) -> _FailingConnection: + return _FailingConnection(invalid_status) + + client: Final = _FakeClientWebSocket(messages=()) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=refuse_connection).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + try: + assert worker.coroutine is None + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in logger.model_call_details + assert logger.dispatch_count == 0 + finally: + if worker.coroutine is not None: + await worker.coroutine + + +@pytest.mark.asyncio +async def test_async_live_session_does_not_log_success_when_connection_fails_before_connect() -> None: + def fail_connection(url: str, **kwargs: object) -> _FailingConnection: + return _FailingConnection(RuntimeError("connection failed")) + + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=fail_connection).async_live_session( + model="gpt-live-1", + websocket=_FakeClientWebSocket(messages=()), + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is None + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in logger.model_call_details + assert logger.dispatch_count == 0 + + +class _RefusingLoggingWorker: + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: + async_coroutine.close() + raise RuntimeError("logging queue unavailable") + + +@pytest.mark.asyncio +async def test_async_live_session_leaves_success_unmarked_when_enqueue_fails() -> None: + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + ) + client: Final = _FakeClientWebSocket(messages=()) + logger: Final = _FakeLogging() + + await OpenAILiveSessions( + logging_worker=_RefusingLoggingWorker(), websocket_connector=_FakeConnector(upstream) + ).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in logger.model_call_details + assert client.close_info is not None + assert client.close_info[0] == 1011 + + +@pytest.mark.parametrize( + "receive_error", + [None, RuntimeError('WebSocket is not connected. Need to call "accept" first.')], + ids=["clean-disconnect", "receive-runtime-error"], +) +@pytest.mark.asyncio +async def test_async_live_session_retains_final_usage_after_client_disconnect_and_bills_voice( + receive_error: Exception | None, +) -> None: + session_closed: Final = json.dumps({"type": "session.closed", "usage": {"seconds": 7}}) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + close_response=session_closed, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket( + messages=(), + disconnect_after_send_type="session.started", + receive_error_after_close=receive_error, + ) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert logged_result.results == [ + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.closed", "usage": {"seconds": 7}}, + ], client.sent_text.qsize() + sent_frames: Final = tuple( + await asyncio.gather(*(upstream.sent_text.get() for _ in range(upstream.sent_text.qsize()))) + ) + session_close: Final = json.dumps({"type": "session.close"}) + assert tuple(frame for frame in sent_frames if _event_type(frame) == "session.close") == (session_close,) + + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + cost: Final = handle_realtime_stream_cost_calculation( + results=logged_result.results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + ) + assert cost == pytest.approx(7 * rate) + + +@pytest.mark.asyncio +async def test_async_live_session_drains_session_closed_after_upstream_send_failure() -> None: + session_closed: Final = json.dumps({"type": "session.closed", "usage": {"seconds": 4}}) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + close_response=session_closed, + fail_on_send_type="conversation.item.create", + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=('{"type":"conversation.item.create"}',)) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert logged_result.results == [ + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.closed", "usage": {"seconds": 4}}, + ] + assert client.close_info is not None + assert client.close_info[0] == 1000 + assert upstream.session_close_sent.is_set() + + +@pytest.mark.asyncio +async def test_async_live_session_raises_after_draining_usage_when_upstream_send_fails_without_close() -> None: + session_closed: Final = json.dumps({"type": "session.closed", "usage": {"seconds": 4}}) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + close_response=session_closed, + fail_on_send_type="conversation.item.create", + send_error=OSError("upstream socket write failed"), + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=('{"type":"conversation.item.create"}',)) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert logged_result.results[-1] == {"type": "session.closed", "usage": {"seconds": 4}} + client_frames: Final = tuple( + await asyncio.gather(*(client.sent_text.get() for _ in range(client.sent_text.qsize()))) + ) + assert tuple(_event_type(frame) for frame in client_frames) == ("session.started", "session.closed", "error") + assert client.close_info is not None + assert client.close_info[0] == 1011 + + +@pytest.mark.asyncio +async def test_async_live_session_does_not_log_success_when_session_start_send_fails() -> None: + upstream: Final = _FakeUpstream(client_frames=(), upstream_frames=(), fail_on_send_type="session.start") + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=()) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is None + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in logger.model_call_details + assert client.close_info is not None + assert client.close_info[0] == 1011 + + +@pytest.mark.asyncio +async def test_async_live_session_ends_when_upstream_keeps_socket_open_after_session_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(live_handler, "OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS", 0.01) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + json.dumps({"type": "session.closed", "usage": {"seconds": 3}}), + ), + never_respond_after_close=True, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=()) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await asyncio.wait_for( + OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=2, + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert logged_result.results[-1] == {"type": "session.closed", "usage": {"seconds": 3}} + assert client.close_info is not None + assert client.close_info[0] == 1000 + assert upstream.closed + + +@pytest.mark.asyncio +async def test_async_live_session_closes_client_normally_when_upstream_drops_after_session_closed() -> None: + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + json.dumps({"type": "session.closed", "usage": {"seconds": 3}}), + ), + close_error=ConnectionClosedError(None, None), + ) + client: Final = _FakeClientWebSocket(messages=()) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await asyncio.wait_for( + OpenAILiveSessions(logging_worker=worker, websocket_connector=_FakeConnector(upstream)).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=2, + ) + + assert client.close_info == (1000, None) + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert logged_result.results[-1] == {"type": "session.closed", "usage": {"seconds": 3}} + + +@pytest.mark.asyncio +async def test_async_live_session_forwards_upstream_close_code_before_session_closed() -> None: + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + close_error=ConnectionClosedError(Close(code=4000, reason="invalid_request_error"), None), + ) + client: Final = _FakeClientWebSocket(messages=()) + + await asyncio.wait_for( + OpenAILiveSessions( + logging_worker=_FakeLoggingWorker(), websocket_connector=_FakeConnector(upstream) + ).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, _FakeLogging()), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=2, + ) + + assert client.close_info == (4000, "invalid_request_error") + + +@pytest.mark.parametrize( + ("first_message", "expected"), + [ + ( + '{"type":"session.start","session":{"model":"gpt-live-1","voice":"alloy"}}', + LiveSessionStart( + frame={"type": "session.start", "session": {"model": "gpt-live-1", "voice": "alloy"}}, + model="gpt-live-1", + ), + ), + ('{"type":"session.update","session":{"model":"gpt-live-1"}}', None), + ('{"type":"session.start","session":{"model":" "}}', None), + ('{"type":"session.start","session":"gpt-live-1"}', None), + ('{"type":"session.start"}', None), + ("not json", None), + ], +) +def test_parse_live_session_start(first_message: str, expected: LiveSessionStart | None) -> None: + parsed: Final = parse_live_session_start(first_message) + if expected is None: + assert parsed is None + return + assert parsed is not None + assert parsed.model == expected.model + assert dict(parsed.frame) == dict(expected.frame) + + +@pytest.mark.asyncio +async def test_async_live_session_bounds_total_time_when_upstream_keeps_streaming_after_session_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(live_handler, "OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS", 0.05) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + json.dumps({"type": "session.closed", "usage": {"seconds": 3}}), + json.dumps({"type": "session.usage.updated", "usage": {"seconds": 3}}), + ), + repeat_last_frame_forever=True, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=()) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await asyncio.wait_for( + OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=1, + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert {"type": "session.closed", "usage": {"seconds": 3}} in logged_result.results + assert client.close_info is not None + assert client.close_info[0] == 1000 + assert upstream.closed + + +@pytest.mark.asyncio +async def test_async_live_session_ends_when_client_stops_reading_after_session_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(live_handler, "OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS", 0.05) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + json.dumps({"type": "session.closed", "usage": {"seconds": 3}}), + ), + repeat_last_frame_forever=True, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=(), block_on_send_type="session.closed") + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await asyncio.wait_for( + OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=1, + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert {"type": "session.closed", "usage": {"seconds": 3}} in logged_result.results + assert client.close_info is not None + assert client.close_info[0] == 1000 + assert upstream.closed + + +@pytest.mark.asyncio +async def test_async_live_session_keeps_finite_usage_after_invalid_snapshot() -> None: + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + '{"type":"session.usage.updated","usage":{"seconds":7}}', + '{"type":"session.usage.updated","usage":{"seconds":"nan"}}', + ), + never_respond_after_close=True, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=(), disconnect_after_send_type="session.started") + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert logged_result.results == [ + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 7}}, + ] + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + cost: Final = handle_realtime_stream_cost_calculation( + results=logged_result.results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + ) + assert cost == pytest.approx(7 * rate) + + +@pytest.mark.asyncio +async def test_async_live_session_does_not_send_a_second_close_after_client_close( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(live_handler, "OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS", 0.01) + client_frames: Final = (json.dumps({"type": "session.close"}),) + upstream: Final = _FakeUpstream( + client_frames=client_frames, + upstream_frames=(), + never_respond_after_close=True, + ) + connector: Final = _FakeConnector(upstream) + worker: Final = _FakeLoggingWorker() + logger: Final = _FakeLogging() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=_FakeClientWebSocket( + messages=client_frames, + disconnect_after_receive_type="session.close", + ), + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + sent_frames: Final = tuple( + await asyncio.gather(*(upstream.sent_text.get() for _ in range(upstream.sent_text.qsize()))) + ) + assert tuple(frame for frame in sent_frames if _event_type(frame) == "session.close") == client_frames + assert worker.coroutine is not None + await worker.coroutine + + +@pytest.mark.asyncio +async def test_async_live_session_bounds_a_blocking_upstream_close_send( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(live_handler, "OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS", 0.05) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + never_respond_after_close=True, + block_session_close_send=True, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=(), disconnect_after_send_type="session.started") + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await asyncio.wait_for( + OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=1, + ) + + assert worker.coroutine is not None + await worker.coroutine + assert logger.dispatch_count == 1 + + +@pytest.mark.asyncio +async def test_async_live_session_bounds_a_blocking_upstream_close( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(live_handler, "OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS", 0.05) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + close_response='{"type":"session.closed","usage":{"seconds":7}}', + block_close=True, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=(), disconnect_after_send_type="session.started") + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await asyncio.wait_for( + OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=1, + ) + + assert worker.coroutine is not None + await worker.coroutine + assert logger.dispatch_count == 1 + + +@pytest.mark.asyncio +async def test_async_live_session_retains_closed_event_when_client_send_fails() -> None: + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + '{"type":"session.closed","usage":{"seconds":7}}', + ), + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=(), fail_on_send_type="session.closed") + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert {"type": "session.closed", "usage": {"seconds": 7}} in logged_result.results + + +@pytest.mark.asyncio +async def test_async_live_session_bounds_wait_for_usage_after_client_disconnect( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(live_handler, "OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS", 0.01) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + never_respond_after_close=True, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=(), disconnect_after_send_type="session.started") + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await asyncio.wait_for( + OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=1, + ) + + assert upstream.closed + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert all(event["type"] != "session.closed" for event in logged_result.results) + + +@pytest.mark.asyncio +async def test_async_live_session_compacts_and_deduplicates_terminal_response_events() -> None: + large_output: Final = "x" * 4096 + large_instructions: Final = "y" * 4096 + response_usage: Final[dict[str, object]] = {"input_tokens": 10, "output_tokens": 5} + duplicate_response_usage: Final[dict[str, object]] = {"input_tokens": 100, "output_tokens": 50} + first_response: Final = json.dumps( + { + "type": "response.event", + "event": { + "type": "response.completed", + "response": { + "id": "resp-1", + "model": "gpt-live-1", + "usage": response_usage, + "output": large_output, + "instructions": large_instructions, + }, + }, + } + ) + duplicate_response: Final = json.dumps( + { + "type": "response.event", + "event": { + "type": "response.completed", + "response": { + "id": "resp-1", + "model": "duplicate-model", + "usage": duplicate_response_usage, + "output": large_output, + "instructions": large_instructions, + }, + }, + } + ) + response_without_id: Final = json.dumps( + { + "type": "response.event", + "event": { + "type": "response.completed", + "response": {"model": "gpt-live-1", "usage": response_usage}, + }, + } + ) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + first_response, + duplicate_response, + response_without_id, + ), + ) + connector: Final = _FakeConnector(upstream) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=_FakeClientWebSocket(messages=()), + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + response_events: Final = tuple(event for event in logged_result.results if event["type"] == "response.event") + assert response_events == ( + { + "type": "response.event", + "event": { + "type": "response.completed", + "response": { + "id": "resp-1", + "model": "gpt-live-1", + "usage": response_usage, + }, + }, + }, + ) + + +@pytest.mark.asyncio +async def test_async_live_session_bills_a_later_response_event_with_usage_over_an_earlier_one_without() -> None: + response_usage: Final[dict[str, object]] = {"input_tokens": 10, "output_tokens": 5} + incomplete_without_usage: Final = json.dumps( + { + "type": "response.event", + "event": {"type": "response.incomplete", "response": {"id": "resp-1", "model": "gpt-live-1"}}, + } + ) + completed_with_usage: Final = json.dumps( + { + "type": "response.event", + "event": { + "type": "response.completed", + "response": {"id": "resp-1", "model": "gpt-live-1", "usage": response_usage}, + }, + } + ) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + incomplete_without_usage, + completed_with_usage, + ), + ) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=_FakeConnector(upstream)).async_live_session( + model="gpt-live-1", + websocket=_FakeClientWebSocket(messages=()), + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + response_events: Final = tuple(event for event in logged_result.results if event["type"] == "response.event") + assert response_events == ( + { + "type": "response.event", + "event": { + "type": "response.completed", + "response": {"id": "resp-1", "model": "gpt-live-1", "usage": response_usage}, + }, + }, + ), logged_result.results + + +@pytest.mark.asyncio +async def test_async_live_session_logs_the_relay_end_time_instead_of_the_logging_worker_time() -> None: + relay_end_time: Final = datetime(2026, 1, 1, 12, 0, 5) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + ) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions( + logging_worker=worker, + websocket_connector=_FakeConnector(upstream), + clock=lambda: relay_end_time, + ).async_live_session( + model="gpt-live-1", + websocket=_FakeClientWebSocket(messages=()), + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + assert logger.dispatched_end_time == relay_end_time + + +@pytest.mark.asyncio +async def test_async_live_session_reports_redacted_server_error_and_bounds_close_reason() -> None: + error_message: Final = "x" * 300 + client: Final = _FakeClientWebSocket(messages=()) + logger: Final = _FakeLogging(pre_call_error=RuntimeError(error_message)) + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "proxy-live-alias"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + error_event: Final = json.loads(client.sent_text.get_nowait()) + assert error_event == {"type": "error", "error": {"type": "server_error", "message": error_message}} + assert client.close_info is not None + assert client.close_info[0] == 1011 + assert client.close_info[1] is not None + assert len(client.close_info[1].encode("utf-8")) <= 123 + assert client.close_info[1] == "x" * 123 + try: + assert worker.coroutine is None + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in logger.model_call_details + assert logger.dispatch_count == 0 + finally: + if worker.coroutine is not None: + await worker.coroutine diff --git a/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py b/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py index eee985f0aca..fb676ae9a61 100644 --- a/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py +++ b/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py @@ -451,6 +451,20 @@ def test_managed_realtime_requires_a_model_and_ignores_completion_defaults(route ) +@pytest.mark.parametrize("route", ["/live/sessions", "/v1/live/sessions", "/openai/v1/live/sessions"]) +def test_managed_live_upgrade_defers_the_model_to_the_session_start_frame(route: str) -> None: + from litellm.proxy.agent_endpoints.auth.managed_authorization import ( + managed_agent_route_allowed, + managed_inference_request, + ) + + assert managed_agent_route_allowed(route, "GET") + assert managed_inference_request(route, {}, {"completion_model": "allowed-default"}, "cli") == {} + assert managed_inference_request( + route, {"model": "frame-model"}, {"completion_model": "allowed-default"}, "cli" + ) == {"model": "frame-model"} + + @pytest.mark.parametrize("mode,user", [("autonomous", None), ("delegated", "verified-human")]) def test_matching_identity_revision_and_execution_mode_pass_admission(mode: str, user: str | None) -> None: context: Final = ManagedAgentContext.model_validate( diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index da8b25c2cee..3660019cc2c 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -12,6 +12,7 @@ import pytest from starlette.datastructures import URL from starlette.types import Message from litellm._logging import verbose_proxy_logger +import json import logging import litellm from litellm.proxy.auth.user_api_key_auth import ( @@ -912,6 +913,54 @@ async def test_user_api_key_auth_websocket_carries_asgi_path(): assert request_arg.scope.get("root_path") == "" +@pytest.mark.asyncio +async def test_user_api_key_auth_live_websocket_ignores_query_model(): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_live_websocket + + mock_websocket = MagicMock(spec=WebSocket) + mock_websocket.query_params = {"model": "gpt-realtime"} + mock_websocket.headers = {"authorization": "Bearer some_api_key"} + mock_websocket.scope = {"headers": [(b"authorization", b"Bearer some_api_key")]} + mock_websocket.url = URL(url="/v1/live/sessions?model=gpt-realtime") + + with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True) as mock_user_api_key_auth: + await user_api_key_auth_live_websocket(mock_websocket) + + request_arg = mock_user_api_key_auth.call_args.kwargs["request"] + assert json.loads(await request_arg.body()) == {"model": ""} + mock_user_api_key_auth.assert_called_once() + + +@pytest.mark.asyncio +async def test_accepted_websocket_reauthorizes_the_frame_model_without_closing(): + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_accepted_websocket_for_model + + mock_websocket = MagicMock(spec=WebSocket) + mock_websocket.query_params = {} + mock_websocket.headers = {"authorization": "Bearer some_api_key"} + mock_websocket.scope = { + "type": "websocket", + "path": "/v1/live/sessions", + "headers": [(b"authorization", b"Bearer some_api_key")], + } + mock_websocket.url = URL(url="/v1/live/sessions") + denial = ProxyException("denied", "user_model_access_denied", "model", 403) + + with patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True, side_effect=denial + ) as mock_user_api_key_auth: + with pytest.raises(ProxyException) as raised: + await user_api_key_auth_accepted_websocket_for_model(mock_websocket, "gpt-live-1") + + assert raised.value is denial + request_arg = mock_user_api_key_auth.call_args.kwargs["request"] + assert json.loads(await request_arg.body()) == {"model": "gpt-live-1"} + assert request_arg.scope.get("path") == "/v1/live/sessions" + assert mock_user_api_key_auth.call_args.kwargs["api_key"] == "Bearer some_api_key" + mock_websocket.close.assert_not_called() + + @pytest.mark.parametrize("enforce_rbac", [True, False]) @pytest.mark.asyncio async def test_jwt_user_api_key_auth_builder_enforce_rbac(enforce_rbac, monkeypatch): diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 79f4edd4d62..80bc3ad1a34 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -1,9 +1,9 @@ import os import sys import traceback -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping from contextlib import asynccontextmanager -from typing import Final, Literal +from typing import Final, Literal, NoReturn, cast from unittest import mock from dotenv import load_dotenv @@ -35,7 +35,7 @@ logging.basicConfig( from unittest.mock import AsyncMock, MagicMock, patch -from fastapi import FastAPI, HTTPException +from fastapi import FastAPI, HTTPException, WebSocket # test /chat/completion request to the proxy from fastapi.testclient import TestClient @@ -3296,3 +3296,496 @@ def test_get_litellm_model_info(data): ): get_litellm_model_info(model=model) get_info_mock.assert_called_once_with(data["expected"]) + + +def test_live_session_accept_failure_releases_reservations( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + + released: Final[list[str]] = [] + + async def release_budget(user_api_key_dict: object) -> None: + released.append("budget") + + async def release_slot(user_api_key_dict: object) -> None: + released.append("slot") + + async def receive() -> dict[str, object]: + return {"type": "websocket.connect"} + + async def send_failing_accept(message: Mapping[str, object]) -> None: + if message["type"] == "websocket.accept": + raise RuntimeError("accept failed") + + scope: Final = { + "type": "websocket", + "path": "/v1/live/sessions", + "raw_path": b"/v1/live/sessions", + "root_path": "", + "scheme": "ws", + "query_string": b"", + "headers": [(b"host", b"testserver"), (b"authorization", b"Bearer test-api-key")], + "client": ("testclient", 50000), + "server": ("testserver", 80), + "subprotocols": [], + } + route_call: Final = AsyncMock() + monkeypatch.setattr(proxy_server, "route_request", route_call) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", release_slot) + + with pytest.raises(RuntimeError, match="accept failed"): + asyncio.run(client_no_auth.app(scope, receive, send_failing_accept)) + + assert released == ["budget", "slot"] + route_call.assert_not_awaited() + + +def _override_live_session_authorizer( + monkeypatch: pytest.MonkeyPatch, + authorizer: Callable[[WebSocket, str], Awaitable[UserAPIKeyAuth]], +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setitem( + proxy_server.app.dependency_overrides, + proxy_server.live_session_model_authorizer, + lambda: authorizer, + ) + + +def _start_live_session_and_wait_for_reply(client: TestClient) -> None: + with client.websocket_connect("/v1/live/sessions", headers={"Authorization": "Bearer test-api-key"}) as websocket: + websocket.send_json({"type": "session.start", "session": {"model": "gpt-3.5-turbo"}}) + websocket.receive_text() + + +def test_live_session_model_check_error_releases_reservations( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + + released: Final[list[str]] = [] + + async def release_budget(user_api_key_dict: object) -> None: + released.append("budget") + + async def release_slot(user_api_key_dict: object) -> None: + released.append("slot") + + async def failing_authorizer(websocket: object, model: str) -> UserAPIKeyAuth: + raise RuntimeError("agent lookup failed") + + route_call: Final = AsyncMock() + monkeypatch.setattr(proxy_server, "route_request", route_call) + _override_live_session_authorizer(monkeypatch, failing_authorizer) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", release_slot) + + with pytest.raises(RuntimeError, match="agent lookup failed"): + _start_live_session_and_wait_for_reply(client_no_auth) + + assert released == ["budget", "slot"] + route_call.assert_not_awaited() + + +class _PreCallAborted(BaseException): + pass + + +def test_live_session_pre_call_abort_releases_reservations( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + frame_auth: Final = UserAPIKeyAuth(token="frame-authorized") + released_for_frame: Final[list[str]] = [] + + async def release_budget(user_api_key_dict: UserAPIKeyAuth) -> None: + if user_api_key_dict is frame_auth: + released_for_frame.append("budget") + + async def release_slot(user_api_key_dict: UserAPIKeyAuth) -> None: + if user_api_key_dict is frame_auth: + released_for_frame.append("slot") + + async def authorize(websocket: object, model: str) -> UserAPIKeyAuth: + return frame_auth + + class _AbortingPreCallProcessing(ProxyBaseLLMRequestProcessing): + async def common_processing_pre_call_logic(self, *args: object, **kwargs: object) -> NoReturn: + raise _PreCallAborted + + route_call: Final = AsyncMock() + monkeypatch.setattr(proxy_server, "route_request", route_call) + _override_live_session_authorizer(monkeypatch, authorize) + monkeypatch.setattr(proxy_server, "ProxyBaseLLMRequestProcessing", _AbortingPreCallProcessing) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", release_slot) + + with pytest.raises(_PreCallAborted): + _start_live_session_and_wait_for_reply(client_no_auth) + + assert released_for_frame == ["budget", "slot"] + route_call.assert_not_awaited() + + +@pytest.mark.parametrize("path", ["/openai/v1/live/sessions", "/v1/live/sessions", "/live/sessions"]) +def test_live_session_invalid_start_is_rejected_without_routing( + path: str, + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + from starlette.websockets import WebSocketDisconnect + + async def no_op(user_api_key_dict: object) -> None: + return None + + route_call: Final = AsyncMock() + monkeypatch.setattr(proxy_server, "route_request", route_call) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op) + + with client_no_auth.websocket_connect(path, headers={"Authorization": "Bearer test-api-key"}) as websocket: + websocket.send_text('{"type":"session.start","session":{"model":""}}') + assert websocket.receive_json() == { + "type": "error", + "error": { + "type": "invalid_request_error", + "code": "invalid_session_start", + "message": "First message must be a session.start JSON object with a non-empty session.model.", + }, + } + with pytest.raises(WebSocketDisconnect) as error: + websocket.receive_json() + assert error.value.code == 1008 + + route_call.assert_not_awaited() + + +@pytest.mark.parametrize( + ("denial_type", "expected_code"), + (("user_model_access_denied", "model_not_allowed"), ("budget_exceeded", "session_not_allowed")), +) +def test_live_session_frame_authorization_denial_closes_with_policy_code( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, + denial_type: str, + expected_code: str, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._types import ProxyException + from starlette.websockets import WebSocketDisconnect + + async def no_op(user_api_key_dict: object) -> None: + return None + + authorized_models: Final[list[str]] = [] + + async def deny(websocket: object, model: str) -> UserAPIKeyAuth: + authorized_models.append(model) + raise ProxyException("model access denied for this user", denial_type, "model", 403) + + route_call: Final = AsyncMock() + monkeypatch.setattr(proxy_server, "route_request", route_call) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op) + _override_live_session_authorizer(monkeypatch, deny) + + with client_no_auth.websocket_connect( + "/v1/live/sessions", + headers={"Authorization": "Bearer test-api-key"}, + ) as websocket: + websocket.send_json({"type": "session.start", "session": {"model": "gpt-4o"}}) + assert websocket.receive_json() == { + "type": "error", + "error": { + "type": "invalid_request_error", + "code": expected_code, + "message": "model access denied for this user", + }, + } + with pytest.raises(WebSocketDisconnect) as error: + websocket.receive_json() + assert error.value.code == 1008 + + assert authorized_models == ["gpt-4o"] + route_call.assert_not_awaited() + + +def test_live_session_valid_start_routes_with_model_and_frame( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from types import SimpleNamespace + + from litellm.proxy import proxy_server + from starlette.websockets import WebSocket + + async def no_op(user_api_key_dict: object) -> None: + return None + + async def complete_call() -> None: + return None + + captured: Final = SimpleNamespace(data=None, route_type=None) + + async def route( + *, + data: dict[str, object], + route_type: str, + llm_router: object, + user_model: str | None, + ) -> Coroutine[object, object, None]: + captured.data = data + captured.route_type = route_type + await cast(WebSocket, data["websocket"]).send_text("routed") + return complete_call() + + frame_auth: Final = UserAPIKeyAuth(token="frame-authorized") + + async def authorize(websocket: object, model: str) -> UserAPIKeyAuth: + return frame_auth + + monkeypatch.setattr(proxy_server, "route_request", route) + _override_live_session_authorizer(monkeypatch, authorize) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op) + start_frame: Final[dict[str, object]] = { + "type": "session.start", + "session": {"model": "gpt-3.5-turbo", "instructions": "Stay brief"}, + } + + with client_no_auth.websocket_connect( + "/v1/live/sessions", + headers={"Authorization": "Bearer test-api-key"}, + ) as websocket: + websocket.send_json(start_frame) + assert websocket.receive_text() == "routed" + + routed_data: Final = cast(dict[str, object], captured.data) + assert captured.route_type == "_arealtime" + assert routed_data["model"] == "gpt-3.5-turbo" + assert routed_data["live_session_start"] == start_frame + assert routed_data["user_api_key_dict"] is frame_auth + + +def test_live_session_ignores_query_model_and_authorizes_frame_model( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastapi import Request + + from litellm.proxy import proxy_server + from starlette.websockets import WebSocket + + async def no_op(user_api_key_dict: object) -> None: + return None + + async def complete_call() -> None: + return None + + async def route( + *, + data: dict[str, object], + route_type: str, + llm_router: object, + user_model: str | None, + ) -> Coroutine[object, object, None]: + await cast(WebSocket, data["websocket"]).send_text("routed") + return complete_call() + + upgrade_auth: Final = UserAPIKeyAuth(token="upgrade", models=["gpt-live-1"]) + + async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth: + if await request.body() != b'{"model": ""}': + raise HTTPException(status_code=403, detail="Model is not allowed") + return upgrade_auth + + auth_call: Final = AsyncMock(side_effect=authenticate) + monkeypatch.setattr("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", auth_call) + frame_auth: Final = UserAPIKeyAuth(token="frame-authorized", models=["gpt-live-1"]) + authorize: Final = AsyncMock(return_value=frame_auth) + monkeypatch.setattr(proxy_server, "route_request", route) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op) + _override_live_session_authorizer(monkeypatch, authorize) + + with client_no_auth.websocket_connect( + "/v1/live/sessions?model=gpt-realtime", + headers={"Authorization": "Bearer test-api-key"}, + ) as websocket: + websocket.send_json({"type": "session.start", "session": {"model": "gpt-live-1"}}) + assert websocket.receive_text() == "routed" + + request_arg: Final = auth_call.await_args.kwargs["request"] + assert json.loads(asyncio.run(request_arg.body())) == {"model": ""} + assert authorize.await_args.args[1] == "gpt-live-1" + + +def test_live_session_reauthorized_token_still_needs_the_model_on_the_key( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from starlette.websockets import WebSocketDisconnect + + from litellm.proxy import proxy_server + + released_tokens: Final[list[str | None]] = [] + + async def release_budget(user_api_key_dict: UserAPIKeyAuth) -> None: + released_tokens.append(user_api_key_dict.token) + + async def no_op(user_api_key_dict: object) -> None: + return None + + async def custom_auth_without_common_checks(websocket: object, model: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth(token="custom-auth", models=["gpt-4o"]) + + route_call: Final = AsyncMock() + monkeypatch.setattr(proxy_server, "route_request", route_call) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op) + _override_live_session_authorizer(monkeypatch, custom_auth_without_common_checks) + + with client_no_auth.websocket_connect( + "/v1/live/sessions", + headers={"Authorization": "Bearer test-api-key"}, + ) as websocket: + websocket.send_json({"type": "session.start", "session": {"model": "gpt-3.5-turbo"}}) + error_event: Final = websocket.receive_json() + with pytest.raises(WebSocketDisconnect) as error: + websocket.receive_json() + + assert error_event["error"]["code"] == "model_not_allowed", error_event + assert error.value.code == 1008 + assert "custom-auth" in released_tokens + route_call.assert_not_awaited() + + +def test_live_session_releases_the_upgrade_reservation_before_the_frame_reserves( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from starlette.websockets import WebSocket + + from litellm.proxy import proxy_server + + released_reservations: Final[list[object]] = [] + + async def release_budget(user_api_key_dict: UserAPIKeyAuth) -> None: + released_reservations.append(user_api_key_dict.budget_reservation) + + async def no_op(user_api_key_dict: object) -> None: + return None + + async def complete_call() -> None: + return None + + async def route( + *, + data: dict[str, object], + route_type: str, + llm_router: object, + user_model: str | None, + ) -> Coroutine[object, object, None]: + await cast(WebSocket, data["websocket"]).send_text("routed") + return complete_call() + + frame_auth: Final = UserAPIKeyAuth(token="frame-authorized", budget_reservation={"reserved_for": "frame model"}) + + async def upgrade_auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth(token="upgrade", budget_reservation={"reserved_for": "query model"}) + + released_before_frame_auth: Final[list[tuple[object, ...]]] = [] + + async def authorize(websocket: object, model: str) -> UserAPIKeyAuth: + released_before_frame_auth.append(tuple(released_reservations)) + return frame_auth + + monkeypatch.setattr(proxy_server, "route_request", route) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op) + monkeypatch.setitem(proxy_server.app.dependency_overrides, proxy_server.user_api_key_auth_live_websocket, upgrade_auth) + _override_live_session_authorizer(monkeypatch, authorize) + + with client_no_auth.websocket_connect( + "/v1/live/sessions?model=gpt-4o", + headers={"Authorization": "Bearer test-api-key"}, + ) as websocket: + websocket.send_json({"type": "session.start", "session": {"model": "gpt-3.5-turbo"}}) + assert websocket.receive_text() == "routed" + + assert released_before_frame_auth == [({"reserved_for": "query model"},)], released_reservations + + +def test_live_session_preserves_pretty_printed_start_frame( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from types import SimpleNamespace + + from starlette.websockets import WebSocket + + from litellm.proxy import proxy_server + + async def no_op(user_api_key_dict: object) -> None: + return None + + async def complete_call() -> None: + return None + + captured: Final = SimpleNamespace(data=None, route_type=None) + + async def route( + *, + data: dict[str, object], + route_type: str, + llm_router: object, + user_model: str | None, + ) -> Coroutine[object, object, None]: + captured.data = data + captured.route_type = route_type + await cast(WebSocket, data["websocket"]).send_text("routed") + return complete_call() + + frame_auth: Final = UserAPIKeyAuth(token="frame-authorized") + + async def authorize(websocket: object, model: str) -> UserAPIKeyAuth: + return frame_auth + + monkeypatch.setattr(proxy_server, "route_request", route) + _override_live_session_authorizer(monkeypatch, authorize) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op) + + start_frame: Final = "\r\n".join( + [ + "{", + ' "type": "session.start",', + ' "session": {"model": "gpt-live-1", "instructions": "a\\nb"}', + "}", + ] + ) + expected_start: Final[dict[str, object]] = { + "type": "session.start", + "session": {"model": "gpt-live-1", "instructions": "a\nb"}, + } + + with client_no_auth.websocket_connect( + "/v1/live/sessions", + headers={"Authorization": "Bearer test-api-key"}, + ) as websocket: + websocket.send_text(start_frame) + assert websocket.receive_text() == "routed" + + routed_data: Final = cast(dict[str, object], captured.data) + assert captured.route_type == "_arealtime" + assert routed_data["model"] == "gpt-live-1" + assert routed_data["live_session_start"] == expected_start diff --git a/tests/unit/realtime_api/test_main.py b/tests/unit/realtime_api/test_main.py index 9d11d2191fa..52433665841 100644 --- a/tests/unit/realtime_api/test_main.py +++ b/tests/unit/realtime_api/test_main.py @@ -1,6 +1,7 @@ import asyncio import json import time +from collections.abc import Mapping from types import TracebackType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -33,6 +34,11 @@ class FakeLogging: pass +class _CallCapture: + def __init__(self) -> None: + self.values: Mapping[str, object] | None = None + + def test_resolves_top_level_session_model(): resolved = _with_resolved_session_model({"model": "alias/gpt-realtime"}, "gpt-realtime") assert resolved == {"model": "gpt-realtime"} @@ -238,6 +244,191 @@ def test_client_secret_forwards_nested_transcription_model_untouched(monkeypatch assert session["input_audio_transcription"]["model"] == "whisper-1" +@pytest.mark.asyncio +async def test_arealtime_live_session_dispatches_to_openai_live_handler(monkeypatch: pytest.MonkeyPatch) -> None: + captured: Final = _CallCapture() + + def mock_get_llm_provider( + model: str, + api_base: str | None, + api_key: str | None, + ) -> tuple[str, str, str | None, str | None]: + return "gpt-live-1", "openai", "resolved-key", "https://live-gateway.example/v1" + + class FakeLiveHandler: + async def async_live_session(self, **kwargs: object) -> None: + captured.values = kwargs + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + monkeypatch.setattr(realtime_main, "openai_live_sessions", FakeLiveHandler()) + start_frame: Final[dict[str, object]] = { + "type": "session.start", + "session": {"model": "live-alias", "instructions": "Stay brief"}, + } + websocket: Final = MagicMock() + logging_obj: Final = FakeLogging() + + await realtime_main._arealtime.__wrapped__( + model="live-alias", + websocket=websocket, + litellm_logging_obj=logging_obj, + live_session_start=start_frame, + ) + + assert captured.values is not None + assert captured.values["model"] == "gpt-live-1" + assert captured.values["websocket"] is websocket + assert captured.values["logging_obj"] is logging_obj + assert captured.values["session_start"] == start_frame + assert captured.values["api_base"] == "https://live-gateway.example/v1" + assert captured.values["api_key"] == "resolved-key" + + +@pytest.mark.asyncio +async def test_arealtime_live_session_rejects_matching_guardrail_before_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class RecordingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__( + guardrail_name="live-blocker", + event_hook=GuardrailEventHooks.realtime_input_transcription, + default_on=True, + ) + self.request_data: dict[str, object] | None = None + + def should_run_guardrail(self, data: dict[str, object], event_type: GuardrailEventHooks) -> bool: + self.request_data = data + return super().should_run_guardrail(data, event_type) + + class FakeLiveHandler: + def __init__(self) -> None: + self.called = False + + async def async_live_session(self, **kwargs: object) -> None: + self.called = True + + guardrail: Final = RecordingGuardrail() + live_handler: Final = FakeLiveHandler() + + def mock_get_llm_provider( + model: str, + api_base: str | None, + api_key: str | None, + ) -> tuple[str, str, str | None, str | None]: + return "gpt-live-1", "openai", "resolved-key", "https://live-gateway.example/v1" + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + monkeypatch.setattr(realtime_main, "openai_live_sessions", live_handler) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + with pytest.raises(ValueError, match="Guardrails are not supported on OpenAI Live sessions"): + await realtime_main._arealtime.__wrapped__( + model="live-alias", + websocket=MagicMock(), + litellm_logging_obj=FakeLogging(), + litellm_metadata={"session_id": "live-session"}, + live_session_start={"type": "session.start", "session": {"model": "live-alias"}}, + ) + + assert guardrail.request_data == {"litellm_metadata": {"session_id": "live-session"}, "stream": True} + assert live_handler.called is False + + +@pytest.mark.asyncio +async def test_arealtime_live_session_rejects_streaming_only_pre_call_guardrail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class FakeLiveHandler: + def __init__(self) -> None: + self.called = False + + async def async_live_session(self, **kwargs: object) -> None: + self.called = True + + guardrail: Final = CustomGuardrail( + guardrail_name="streaming-only-blocker", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + stream_scope="streaming", + ) + live_handler: Final = FakeLiveHandler() + + def mock_get_llm_provider( + model: str, + api_base: str | None, + api_key: str | None, + ) -> tuple[str, str, str | None, str | None]: + return "gpt-live-1", "openai", "resolved-key", "https://live-gateway.example/v1" + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + monkeypatch.setattr(realtime_main, "openai_live_sessions", live_handler) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + with pytest.raises(ValueError, match="Guardrails are not supported on OpenAI Live sessions"): + await realtime_main._arealtime.__wrapped__( + model="live-alias", + websocket=MagicMock(), + litellm_logging_obj=FakeLogging(), + litellm_metadata={"session_id": "live-session"}, + live_session_start={"type": "session.start", "session": {"model": "live-alias"}}, + ) + + assert live_handler.called is False + + +@pytest.mark.asyncio +async def test_arealtime_live_session_rejects_non_openai_provider(monkeypatch: pytest.MonkeyPatch) -> None: + def mock_get_llm_provider( + model: str, + api_base: str | None, + api_key: str | None, + ) -> tuple[str, str, str | None, str | None]: + return "claude-live", "anthropic", None, None + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + + with pytest.raises(ValueError, match="anthropic"): + await realtime_main._arealtime.__wrapped__( + model="anthropic/claude-live", + websocket=MagicMock(), + litellm_logging_obj=FakeLogging(), + live_session_start={"type": "session.start", "session": {"model": "claude-live"}}, + ) + + +@pytest.mark.asyncio +async def test_arealtime_without_live_start_preserves_openai_realtime_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: Final = _CallCapture() + + def mock_get_llm_provider( + model: str, + api_base: str | None, + api_key: str | None, + ) -> tuple[str, str, str | None, str | None]: + return "gpt-realtime", "openai", "resolved-key", "https://realtime-gateway.example/v1" + + class FakeRealtimeHandler: + async def async_realtime(self, **kwargs: object) -> None: + captured.values = kwargs + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + monkeypatch.setattr(realtime_main, "openai_realtime", FakeRealtimeHandler()) + + await realtime_main._arealtime.__wrapped__( + model="gpt-realtime", + websocket=MagicMock(), + litellm_logging_obj=FakeLogging(), + ) + + assert captured.values is not None + assert captured.values["model"] == "gpt-realtime" + assert captured.values["api_key"] == "resolved-key" + assert captured.values["api_base"] == "https://realtime-gateway.example/v1" + + class _CapturingConnect: def __init__(self) -> None: self.url: str | None = None diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 01cce6a3e0f..580658c3447 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -8649,3 +8649,264 @@ def test_batch_cost_calculator(): cost = completion_cost(**args) assert cost > 0 + + +def _live_session_results(*events: dict[str, object]) -> OpenAIRealtimeStreamList: + return cast(OpenAIRealtimeStreamList, list(events)) + + +def _live_logging_object( + start_time: datetime.datetime, + end_time: datetime.datetime, +) -> Logging: + return cast( + Logging, + SimpleNamespace(model_call_details={"start_time": start_time, "end_time": end_time}), + ) + + +def test_realtime_keeps_first_reported_session_model_for_pricing(_local_model_cost_map: None) -> None: + first_model: Final = "gpt-realtime-2" + later_model: Final = "gpt-realtime-mini" + combined_usage: Final = Usage(prompt_tokens=120, completion_tokens=60, total_tokens=180) + first_model_cost: Final = handle_realtime_stream_cost_calculation( + results=_live_session_results({"type": "session.created", "session": {"model": first_model}}), + combined_usage_object=combined_usage, + custom_llm_provider="openai", + litellm_model_name="unmapped-realtime-alias", + ) + first_model_with_later_update_cost: Final = handle_realtime_stream_cost_calculation( + results=_live_session_results( + {"type": "session.created", "session": {"model": first_model}}, + {"type": "session.started", "session": {"model": later_model}}, + ), + combined_usage_object=combined_usage, + custom_llm_provider="openai", + litellm_model_name="unmapped-realtime-alias", + ) + later_model_cost: Final = handle_realtime_stream_cost_calculation( + results=_live_session_results({"type": "session.started", "session": {"model": later_model}}), + combined_usage_object=combined_usage, + custom_llm_provider="openai", + litellm_model_name="unmapped-realtime-alias", + ) + + assert first_model_with_later_update_cost == pytest.approx(first_model_cost) + assert first_model_cost != pytest.approx(later_model_cost) + + +def test_live_session_charges_latest_closed_voice_usage(_local_model_cost_map: None) -> None: + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 12}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + {"type": "session.closed", "usage": {"seconds": 16}}, + ) + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + ) + + assert rate > 0 + assert cost == pytest.approx(16 * rate) + + +def test_live_session_uses_last_reported_duration_when_wall_clock_is_short(_local_model_cost_map: None) -> None: + start_time: Final = datetime.datetime(2025, 1, 1) + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + ) + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + litellm_logging_obj=_live_logging_object(start_time, start_time + datetime.timedelta(seconds=5)), + ) + + assert cost == pytest.approx(15 * rate) + + +def test_live_session_uses_latest_cumulative_usage_snapshot_not_sum(_local_model_cost_map: None) -> None: + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 12}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + ) + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + ) + + assert cost == pytest.approx(15 * rate) + + +def test_live_session_ignores_non_finite_usage_snapshot_after_finite_one(_local_model_cost_map: None) -> None: + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + {"type": "session.usage.updated", "usage": {"seconds": float("nan")}}, + ) + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + ) + + assert cost == pytest.approx(15 * rate) + + +def test_live_session_ignores_non_finite_closed_usage_snapshot(_local_model_cost_map: None) -> None: + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + {"type": "session.closed", "usage": {"seconds": float("inf")}}, + ) + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + ) + + assert cost == pytest.approx(15 * rate) + + +def test_live_session_skips_non_finite_rate_for_next_candidate( + _local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "live-non-finite-rate", + {"input_cost_per_second": float("nan"), "litellm_provider": "openai", "mode": "realtime"}, + ) + monkeypatch.setitem( + litellm.model_cost, + "live-fallback-rate", + {"input_cost_per_second": 0.25, "litellm_provider": "openai", "mode": "realtime"}, + ) + litellm.get_model_info.cache_clear() + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "live-non-finite-rate"}}, + {"type": "session.usage.updated", "usage": {"seconds": 4}}, + ) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="live-fallback-rate", + ) + + assert cost == pytest.approx(4 * 0.25) + + +def test_live_session_uses_wall_clock_duration_when_greater_than_reported(_local_model_cost_map: None) -> None: + start_time: Final = datetime.datetime(2025, 1, 1) + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + ) + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + litellm_logging_obj=_live_logging_object(start_time, start_time + datetime.timedelta(seconds=20)), + ) + + assert cost == pytest.approx(20 * rate) + + +def test_live_session_deduplicates_delegated_responses_and_charges_voice(_local_model_cost_map: None) -> None: + delegated_response: Final[dict[str, object]] = { + "type": "response.event", + "delegation_id": "delegation-1", + "event": { + "type": "response.completed", + "response": { + "id": "resp-1", + "model": "gpt-4o-mini", + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + }, + } + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.closed", "usage": {"seconds": 12}}, + delegated_response, + delegated_response, + ) + live_rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + backend_model: Final = litellm.model_cost["gpt-4o-mini"] + expected_backend_cost: Final = ( + 10 * cast(float, backend_model["input_cost_per_token"]) + + 5 * cast(float, backend_model["output_cost_per_token"]) + ) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + ) + + assert cost == pytest.approx(12 * live_rate + expected_backend_cost) + + +def test_live_session_zero_closed_seconds_do_not_use_usage_fallback(_local_model_cost_map: None) -> None: + start_time: Final = datetime.datetime(2025, 1, 1) + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + {"type": "session.closed", "usage": {"seconds": 0}}, + ) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + litellm_logging_obj=_live_logging_object(start_time, start_time + datetime.timedelta(seconds=20)), + ) + + assert cost == 0.0 + + +def test_live_session_uses_latest_snapshot_when_closed_usage_is_missing(_local_model_cost_map: None) -> None: + start_time: Final = datetime.datetime(2025, 1, 1) + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + {"type": "session.closed", "reason": "server_end"}, + ) + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + litellm_logging_obj=_live_logging_object(start_time, start_time + datetime.timedelta(seconds=20)), + ) + + assert cost == pytest.approx(15 * rate) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index fc2cc73ff10..bdad35822e5 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -9142,6 +9142,26 @@ export interface paths { patch?: never; trace?: never; }; + "/live/sessions": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * WebSocket: live_sessions_websocket_endpoint + * @description WebSocket connection endpoint + */ + get: operations["websocket_live_sessions_websocket_endpoint_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/login": { parameters: { query?: never; @@ -10822,6 +10842,26 @@ export interface paths { patch?: never; trace?: never; }; + "/openai/v1/live/sessions": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * WebSocket: live_sessions_websocket_endpoint + * @description WebSocket connection endpoint + */ + get: operations["websocket_live_sessions_websocket_endpoint_get_3"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/openai/v1/realtime": { parameters: { query?: never; @@ -20292,6 +20332,26 @@ export interface paths { patch?: never; trace?: never; }; + "/v1/live/sessions": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * WebSocket: live_sessions_websocket_endpoint + * @description WebSocket connection endpoint + */ + get: operations["websocket_live_sessions_websocket_endpoint_get_2"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v1/logs": { parameters: { query?: never; @@ -62824,6 +62884,24 @@ export interface operations { }; }; }; + websocket_live_sessions_websocket_endpoint_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description WebSocket Protocol Switched */ + 101: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; login_login_post: { parameters: { query?: never; @@ -65435,6 +65513,24 @@ export interface operations { }; }; }; + websocket_live_sessions_websocket_endpoint_get_3: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description WebSocket Protocol Switched */ + 101: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; websocket_realtime_websocket_endpoint_get_3: { parameters: { query?: never; @@ -77691,6 +77787,24 @@ export interface operations { }; }; }; + websocket_live_sessions_websocket_endpoint_get_2: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description WebSocket Protocol Switched */ + 101: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; ingest_otlp_traces_v1_logs_post: { parameters: { query?: never;