From eb5031da1e1aa59f5521b2b67397b3d8f0717ea5 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 22 Nov 2025 10:01:02 -0800 Subject: [PATCH] [Perf] Fix bottlenecks degrading realtime endpoint performance (#16670) * Cache realtime websocket request body Move the realtime request payload builder out of the websocket handler and wrap it with an LRU cache so repeated connections reuse the same bytes object. This keeps the JSON formatting cost down while bounding memory usage. * Optimize realtime websocket caching Refactored /v1/realtime to use cached helpers for both the JSON body and query params, introduced a reusable request-scope template, and optimized header handling to avoid redundant work. * Refine realtime websocket header handling * Reuse websocket scope headers in auth * Refactor realtime request body helper Move the realtime request body formatter into proxy common utils so it can be reused across modules. Reuse it in the websocket auth flow to share LRU caching and avoid ad hoc byte builders. * fix: revert to old pattern The old pattern was necessary, we can just return the optimized function instead. * Reuse SSL context for realtime Create a shared SSLContext for OpenAI realtime websocket dials and pass it into websockets.connect so we stop re-reading verify paths on every session. * feat: reuse shared TLS context for realtime websockets - add `SHARED_REALTIME_SSL_CONTEXT` helper so all realtime websocket clients share the same TLS settings - wire the shared context into OpenAI, Azure, custom HTTPX handlers, and realtime health checks - update realtime tests to assert that the expected SSL context is passed to `websockets.connect` This keeps TLS configuration consistent and avoids recreating SSL contexts per connection. * Reuse HTTP SSL context for realtime Remove the standalone realtime SSL helper, expose a shared context directly from the HTTP handler, and point all realtime websocket clients and tests to it. Add the websocket header comparison tool. * Lazy-load shared realtime SSL context Fix circular imports introduced by eagerly instantiating the shared TLS context. Make the HTTP handler lazily create the context and have realtime clients/tests fetch it on demand, keeping configuration consistent without breaking startup. * add: unit test for realtime LRU caches * fix: merge conflict with imports --- litellm/constants.py | 1 + litellm/llms/azure/realtime/handler.py | 3 + litellm/llms/custom_httpx/http_handler.py | 14 +++++ litellm/llms/custom_httpx/llm_http_handler.py | 3 + litellm/llms/openai/realtime/handler.py | 3 + litellm/proxy/auth/user_api_key_auth.py | 19 +++--- litellm/proxy/common_utils/realtime_utils.py | 15 +++++ litellm/proxy/proxy_server.py | 56 +++++++++++------ litellm/realtime_api/main.py | 3 + tests/proxy_unit_tests/test_realtime_cache.py | 62 +++++++++++++++++++ .../realtime/test_azure_realtime_handler.py | 4 ++ .../realtime/test_openai_realtime_handler.py | 6 ++ 12 files changed, 157 insertions(+), 32 deletions(-) create mode 100644 litellm/proxy/common_utils/realtime_utils.py create mode 100644 tests/proxy_unit_tests/test_realtime_cache.py diff --git a/litellm/constants.py b/litellm/constants.py index b312a15892b..2110a3b37a0 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 23c04e640c4..8e5581206de 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 dbeceab5e04..c35e910ab08 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 6b0aef31ff6..a0e8190cc65 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 e1fb3f12602..882309bb2fa 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 2e8513e1b9c..a8d5c35ebbd 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 00000000000..4af7ad2514f --- /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 91353e01624..dbd802a9f2f 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 8978dbca175..93d6269b517 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 00000000000..05688024c25 --- /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 db165bb5fbc..7bcbe371568 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 ca55a3a2c1d..bd973692c1e 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