mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
feat(realtime): proxy OpenAI Live sessions on /v1/live/sessions
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
1b89c2af1a
commit
fabddfd71f
23 changed files with 3463 additions and 64 deletions
|
|
@ -121,6 +121,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
# Realtime / streaming
|
||||
"/v1/realtime",
|
||||
"/realtime",
|
||||
"/v1/live/sessions",
|
||||
"/live/sessions",
|
||||
# Health & ops
|
||||
"/health",
|
||||
"/metrics",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
"""
|
||||
|
|
|
|||
3
litellm/llms/openai/live/__init__.py
Normal file
3
litellm/llms/openai/live/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
from .handler import LiveSessionStart, OpenAILiveSessions, parse_live_session_start
|
||||
|
||||
__all__ = ("LiveSessionStart", "OpenAILiveSessions", "parse_live_session_start")
|
||||
555
litellm/llms/openai/live/handler.py
Normal file
555
litellm/llms/openai/live/handler.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
0
tests/unit/llms/openai/live/__init__.py
Normal file
0
tests/unit/llms/openai/live/__init__.py
Normal file
1208
tests/unit/llms/openai/live/test_handler.py
Normal file
1208
tests/unit/llms/openai/live/test_handler.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
114
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
114
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue