diff --git a/litellm/constants.py b/litellm/constants.py index b312a15892..2110a3b37a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -206,6 +206,7 @@ REPEATED_STREAMING_CHUNK_LIMIT = int( os.getenv("REPEATED_STREAMING_CHUNK_LIMIT", 100) ) # catch if model starts looping the same chunk while streaming. Uses high default to prevent false positives. DEFAULT_MAX_LRU_CACHE_SIZE = int(os.getenv("DEFAULT_MAX_LRU_CACHE_SIZE", 16)) +_REALTIME_BODY_CACHE_SIZE = 1000 # Keep realtime helper caches bounded; workloads rarely exceed 1k models/intents INITIAL_RETRY_DELAY = float(os.getenv("INITIAL_RETRY_DELAY", 0.5)) MAX_RETRY_DELAY = float(os.getenv("MAX_RETRY_DELAY", 8.0)) JITTER = float(os.getenv("JITTER", 0.75)) diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index 23c04e640c..8e5581206d 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -10,6 +10,7 @@ from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from ....litellm_core_utils.realtime_streaming import RealTimeStreaming +from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..azure import AzureChatCompletion # BACKEND_WS_URL = "ws://localhost:8080/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01" @@ -61,12 +62,14 @@ class AzureOpenAIRealtime(AzureChatCompletion): url = self._construct_url(api_base, model, api_version) try: + ssl_context = get_shared_realtime_ssl_context() async with websockets.connect( # type: ignore url, extra_headers={ "api-key": api_key, # type: ignore }, max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, + ssl=ssl_context, ) as backend_ws: realtime_streaming = RealTimeStreaming( websocket, cast(ClientConnection, backend_ws), logging_obj diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index dbeceab5e0..c35e910ab0 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -215,6 +215,20 @@ def get_ssl_configuration( return ssl_verify +_shared_realtime_ssl_context: Optional[Union[bool, str, ssl.SSLContext]] = None + + +def get_shared_realtime_ssl_context() -> Union[bool, str, ssl.SSLContext]: + """ + Lazily create the SSL context reused by realtime websocket clients so we avoid + import-order cycles during startup while keeping a single shared configuration. + """ + global _shared_realtime_ssl_context + if _shared_realtime_ssl_context is None: + _shared_realtime_ssl_context = get_ssl_configuration() + return _shared_realtime_ssl_context + + def mask_sensitive_info(error_message): # Find the start of the key parameter if isinstance(error_message, str): diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 6b0aef31ff..a0e8190cc6 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -38,6 +38,7 @@ from litellm.llms.base_llm.google_genai.transformation import ( BaseGoogleGenAIGenerateContentConfig, ) from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from .http_handler import get_shared_realtime_ssl_context from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -3612,10 +3613,12 @@ class BaseLLMHTTPHandler: ) try: + ssl_context = get_shared_realtime_ssl_context() async with websockets.connect( # type: ignore url, extra_headers=headers, max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, + ssl=ssl_context, ) as backend_ws: realtime_streaming = RealTimeStreaming( websocket, diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index e1fb3f1260..882309bb2f 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -11,6 +11,7 @@ from litellm.types.realtime import RealtimeQueryParams from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from ....litellm_core_utils.realtime_streaming import RealTimeStreaming +from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..openai import OpenAIChatCompletion @@ -55,6 +56,7 @@ class OpenAIRealtime(OpenAIChatCompletion): url = self._construct_url(api_base, query_params) try: + ssl_context = get_shared_realtime_ssl_context() async with websockets.connect( # type: ignore url, extra_headers={ @@ -62,6 +64,7 @@ class OpenAIRealtime(OpenAIChatCompletion): "OpenAI-Beta": "realtime=v1", }, max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, + ssl=ssl_context, ) as backend_ws: realtime_streaming = RealTimeStreaming( websocket, cast(ClientConnection, backend_ws), logging_obj diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2e8513e1b9..a8d5c35ebb 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -54,6 +54,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, ) +from litellm.proxy.common_utils.realtime_utils import _realtime_request_body from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.secret_managers.main import get_secret_bool from litellm.types.services import ServiceTypes @@ -157,14 +158,8 @@ def _apply_budget_limits_to_end_user_params( async def user_api_key_auth_websocket(websocket: WebSocket): # Accept the WebSocket connection - request = Request( - scope={ - "type": "http", - "headers": [ - (k.lower().encode(), v.encode()) for k, v in websocket.headers.items() - ], - } - ) + scope_headers = list(websocket.scope.get("headers") or []) + request = Request(scope={"type": "http", "headers": scope_headers}) request._url = websocket.url @@ -172,13 +167,13 @@ async def user_api_key_auth_websocket(websocket: WebSocket): model = query_params.get("model") + async def return_body(): - return_string = f'{{"model": "{model}"}}' - # return string as bytes - return return_string.encode() - + return _realtime_request_body(model) + request.body = return_body # type: ignore + authorization = websocket.headers.get("authorization") # If no Authorization header, try the api-key header if not authorization: diff --git a/litellm/proxy/common_utils/realtime_utils.py b/litellm/proxy/common_utils/realtime_utils.py new file mode 100644 index 0000000000..4af7ad2514 --- /dev/null +++ b/litellm/proxy/common_utils/realtime_utils.py @@ -0,0 +1,15 @@ +from functools import lru_cache +from typing import Optional + +from litellm.constants import _REALTIME_BODY_CACHE_SIZE + + +@lru_cache(maxsize=_REALTIME_BODY_CACHE_SIZE) +def _realtime_request_body(model: Optional[str]) -> bytes: + """ + Generate the realtime websocket request body. Cached with LRU semantics to avoid repeated + string formatting work while keeping memory usage bounded. + """ + return f'{{"model": "{model or ""}"}}'.encode() + + diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 91353e0162..dbd802a9f2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14,6 +14,7 @@ from datetime import datetime, timedelta from typing import ( TYPE_CHECKING, Any, + Dict, List, Literal, Optional, @@ -48,6 +49,7 @@ from litellm.types.utils import ( TokenCountResponse, ) from litellm.utils import load_credentials_from_list +from litellm.proxy.common_utils.realtime_utils import _realtime_request_body if TYPE_CHECKING: from aiohttp import ClientSession @@ -60,6 +62,12 @@ else: Span = Any OpenTelemetry = Any +REALTIME_REQUEST_SCOPE_TEMPLATE: Dict[str, Any] = { + "type": "http", + "method": "POST", + "path": "/v1/realtime", +} + def showwarning(message, category, filename, lineno, file=None, line=None): traceback_info = f"{filename}:{lineno}: {category.__name__}: {message}\n" @@ -131,6 +139,7 @@ def generate_feedback_box(): from collections import defaultdict from contextlib import asynccontextmanager +from functools import lru_cache import litellm from litellm import Router @@ -151,6 +160,7 @@ from litellm.constants import ( PROXY_BATCH_WRITE_AT, PROXY_BUDGET_RESCHEDULER_MAX_TIME, PROXY_BUDGET_RESCHEDULER_MIN_TIME, + _REALTIME_BODY_CACHE_SIZE, ) from litellm.exceptions import RejectedRequestError from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting @@ -5574,12 +5584,24 @@ async def vertex_ai_live_passthrough_endpoint( # /v1/realtime Endpoints ###################################################################### -from litellm import _arealtime + +@lru_cache(maxsize=_REALTIME_BODY_CACHE_SIZE) +def _realtime_query_params_template( + model: str, intent: Optional[str] +) -> Tuple[Tuple[str, str], ...]: + """ + Build a hashable representation of the realtime query params so we can cache + the repetitive model/intent combinations. + """ + params: List[Tuple[str, str]] = [("model", model)] + if intent is not None: + params.append(("intent", intent)) + return tuple(params) @app.websocket("/v1/realtime") @app.websocket("/realtime") -async def websocket_endpoint( +async def realtime_websocket_endpoint( websocket: WebSocket, model: str, intent: str = fastapi.Query( @@ -5587,14 +5609,13 @@ async def websocket_endpoint( ), user_api_key_dict=Depends(user_api_key_auth_websocket), ): - import websockets await websocket.accept() # Only use explicit parameters, not all query params - query_params: RealtimeQueryParams = {"model": model} - if intent is not None: - query_params["intent"] = intent + query_params = cast( + RealtimeQueryParams, dict(_realtime_query_params_template(model, intent)) + ) data = { "model": model, @@ -5602,24 +5623,19 @@ async def websocket_endpoint( "query_params": query_params, # Only explicit params } - headers = dict(websocket.headers.items()) # Convert headers to dict first + # Use raw ASGI headers (already lowercase bytes) to avoid extra work + headers_list = list(websocket.scope.get("headers") or []) - request = Request( - scope={ - "type": "http", - "headers": [(k.lower().encode(), v.encode()) for k, v in headers.items()], - "method": "POST", - "path": "/v1/realtime", - } - ) + scope = REALTIME_REQUEST_SCOPE_TEMPLATE.copy() + scope["headers"] = headers_list + + request = Request(scope=scope) request._url = websocket.url - + async def return_body(): - return_string = f'{{"model": "{model}"}}' - # return string as bytes - return return_string.encode() - + return _realtime_request_body(model) + request.body = return_body # type: ignore ### ROUTE THE REQUEST ### diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 8978dbca17..93d6269b51 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -18,6 +18,7 @@ from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from ..llms.azure.realtime.handler import AzureOpenAIRealtime from ..llms.openai.realtime.handler import OpenAIRealtime from ..utils import client as wrapper_client +from ..llms.custom_httpx.http_handler import get_shared_realtime_ssl_context azure_realtime = AzureOpenAIRealtime() openai_realtime = OpenAIRealtime() @@ -178,11 +179,13 @@ async def _realtime_health_check( ) else: raise ValueError(f"Unsupported model: {model}") + ssl_context = get_shared_realtime_ssl_context() async with websockets.connect( # type: ignore url, extra_headers={ "api-key": api_key, # type: ignore }, max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, + ssl=ssl_context, ): return True diff --git a/tests/proxy_unit_tests/test_realtime_cache.py b/tests/proxy_unit_tests/test_realtime_cache.py new file mode 100644 index 0000000000..05688024c2 --- /dev/null +++ b/tests/proxy_unit_tests/test_realtime_cache.py @@ -0,0 +1,62 @@ +from typing import Any, cast + +import pytest + +from litellm.proxy.common_utils.realtime_utils import _realtime_request_body +from litellm.proxy.proxy_server import _realtime_query_params_template + + +@pytest.fixture(autouse=True) +def clear_realtime_caches(): + _realtime_request_body.cache_clear() + _realtime_query_params_template.cache_clear() + yield + _realtime_request_body.cache_clear() + _realtime_query_params_template.cache_clear() + + +def test_realtime_request_body_returns_immutable_bytes(): + cached_body = _realtime_request_body("gpt-4o") + + with pytest.raises(TypeError): + cast(Any, cached_body)[0] = ord("x") + + +def test_realtime_query_params_template_returns_immutable_tuples(): + cached_tuple = _realtime_query_params_template("gpt-4o", "intent-a") + + with pytest.raises(TypeError): + cast(Any, cached_tuple)[0] = ("model", "mutated") + + +def test_realtime_request_body_caches_each_model_separately(): + gpt4o_body_first = _realtime_request_body("gpt-4o") + gpt4o_body_second = _realtime_request_body("gpt-4o") + gpt4o_mini_body = _realtime_request_body("gpt-4o-mini") + + assert gpt4o_body_first is gpt4o_body_second + assert gpt4o_body_first == b'{"model": "gpt-4o"}' + assert gpt4o_mini_body == b'{"model": "gpt-4o-mini"}' + assert gpt4o_body_first is not gpt4o_mini_body + + +def test_realtime_query_params_template_caches_each_pair_separately(): + params_with_intent_first = _realtime_query_params_template("gpt-4o", "intent-a") + params_with_intent_second = _realtime_query_params_template("gpt-4o", "intent-a") + params_without_intent = _realtime_query_params_template("gpt-4o", None) + + assert params_with_intent_first is params_with_intent_second + assert params_with_intent_first == (("model", "gpt-4o"), ("intent", "intent-a")) + assert params_without_intent == (("model", "gpt-4o"),) + assert params_with_intent_first is not params_without_intent + + +def test_realtime_query_params_dict_copies_do_not_leak_state(): + params_dict_one = dict(_realtime_query_params_template("gpt-4o", "intent-a")) + params_dict_one["new"] = "value" + + params_dict_two = dict(_realtime_query_params_template("gpt-4o", "intent-a")) + + assert "new" not in params_dict_two + assert params_dict_two == {"model": "gpt-4o", "intent": "intent-a"} + diff --git a/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py b/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py index db165bb5fb..7bcbe37156 100644 --- a/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py +++ b/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py @@ -4,6 +4,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.llms.custom_httpx.http_handler import get_shared_realtime_ssl_context + sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path @@ -38,6 +40,7 @@ async def test_async_realtime_uses_max_size_parameter(): async def __aexit__(self, exc_type, exc, tb): return None + shared_context = get_shared_realtime_ssl_context() with patch("websockets.connect", return_value=DummyAsyncContextManager(mock_backend_ws)) as mock_ws_connect, \ patch("litellm.llms.azure.realtime.handler.RealTimeStreaming") as mock_realtime_streaming: @@ -61,6 +64,7 @@ async def test_async_realtime_uses_max_size_parameter(): # Verify max_size is set (default None for unlimited, matching OpenAI's SDK) assert "max_size" in called_kwargs assert called_kwargs["max_size"] is None + assert called_kwargs["ssl"] is shared_context # Default should be None (unlimited) to match OpenAI's official agents SDK # https://github.com/openai/openai-agents-python/blob/cf1b933660e44fd37b4350c41febab8221801409/src/agents/realtime/openai_realtime.py#L235 diff --git a/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py b/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py index ca55a3a2c1..bd973692c1 100644 --- a/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py +++ b/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py @@ -6,6 +6,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from litellm.llms.custom_httpx.http_handler import get_shared_realtime_ssl_context + sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path @@ -119,6 +121,7 @@ async def test_async_realtime_success(): async def __aexit__(self, exc_type, exc, tb): return None + shared_context = get_shared_realtime_ssl_context() with patch("websockets.connect", return_value=DummyAsyncContextManager(mock_backend_ws)) as mock_ws_connect, \ patch("litellm.llms.openai.realtime.handler.RealTimeStreaming") as mock_realtime_streaming: mock_streaming_instance = MagicMock() @@ -195,6 +198,7 @@ async def test_async_realtime_url_contains_model(): extra_headers = called_kwargs["extra_headers"] assert extra_headers["Authorization"] == f"Bearer {api_key}" assert extra_headers["OpenAI-Beta"] == "realtime=v1" + assert called_kwargs["ssl"] is shared_context mock_realtime_streaming.assert_called_once() mock_streaming_instance.bidirectional_forward.assert_awaited_once() @@ -230,6 +234,7 @@ async def test_async_realtime_uses_max_size_parameter(): async def __aexit__(self, exc_type, exc, tb): return None + shared_context = get_shared_realtime_ssl_context() with patch("websockets.connect", return_value=DummyAsyncContextManager(mock_backend_ws)) as mock_ws_connect, \ patch("litellm.llms.openai.realtime.handler.RealTimeStreaming") as mock_realtime_streaming: @@ -253,6 +258,7 @@ async def test_async_realtime_uses_max_size_parameter(): # Verify max_size is set (default None for unlimited, matching OpenAI's SDK) assert "max_size" in called_kwargs assert called_kwargs["max_size"] is None + assert called_kwargs["ssl"] is shared_context # Default should be None (unlimited) to match OpenAI's official agents SDK # https://github.com/openai/openai-agents-python/blob/cf1b933660e44fd37b4350c41febab8221801409/src/agents/realtime/openai_realtime.py#L235