diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index a865de4118..aca32e1404 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -9,20 +9,26 @@ from typing import Any, Optional, cast from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from ....litellm_core_utils.realtime_streaming import RealTimeStreaming from ..openai import OpenAIChatCompletion +from litellm.types.realtime import RealtimeQueryParams class OpenAIRealtime(OpenAIChatCompletion): - def _construct_url(self, api_base: str, model: str) -> str: + def _construct_url(self, api_base: str, query_params: RealtimeQueryParams) -> str: """ - Example output: - "BACKEND_WS_URL = "wss://localhost:8080/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01""; + Construct the backend websocket URL with all query parameters (excluding 'model' if present). """ from httpx import URL api_base = api_base.replace("https://", "wss://") api_base = api_base.replace("http://", "ws://") - url = URL(api_base).join("/v1/realtime") - return str(url.copy_add_param("model", model)) + url = URL(api_base) + # Set the correct path + url = url.copy_with(path="/v1/realtime") + # Build query dict excluding 'model' + query_dict = {k: v for k, v in query_params.items() if k != "model"} + if query_dict: + url = url.copy_with(params=query_dict) + return str(url) async def async_realtime( self, @@ -33,6 +39,7 @@ class OpenAIRealtime(OpenAIChatCompletion): api_key: Optional[str] = None, client: Optional[Any] = None, timeout: Optional[float] = None, + query_params: Optional[RealtimeQueryParams] = None, ): import websockets from websockets.asyncio.client import ClientConnection @@ -42,7 +49,10 @@ class OpenAIRealtime(OpenAIChatCompletion): if api_key is None: raise ValueError("api_key is required for Azure OpenAI calls") - url = self._construct_url(api_base, model) + # Use all query params if provided, else fallback to just model + if query_params is None: + query_params = {"model": model} + url = self._construct_url(api_base, query_params) try: async with websockets.connect( # type: ignore diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9cc21f5538..7fca29f2bf 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -147,6 +147,7 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.types.realtime import RealtimeQueryParams from litellm.proxy._experimental.mcp_server.rest_endpoints import ( router as mcp_rest_endpoints_router, ) @@ -4680,15 +4681,20 @@ from litellm import _arealtime async def websocket_endpoint( websocket: WebSocket, model: str, + intent: str = fastapi.Query(None, description="The intent of the websocket connection."), 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, "intent": intent} + data = { "model": model, "websocket": websocket, + "query_params": query_params, # Only explicit params } headers = dict(websocket.headers.items()) # Convert headers to dict first diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index fcf6c21845..c69a058ea1 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -15,6 +15,7 @@ from ..litellm_core_utils.get_litellm_params import get_litellm_params 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 litellm.types.realtime import RealtimeQueryParams from ..utils import client as wrapper_client azure_realtime = AzureOpenAIRealtime() @@ -32,6 +33,7 @@ async def _arealtime( azure_ad_token: Optional[str] = None, client: Optional[Any] = None, timeout: Optional[float] = None, + query_params: Optional[RealtimeQueryParams] = None, **kwargs, ): """ @@ -132,6 +134,7 @@ async def _arealtime( api_key=api_key, client=None, timeout=timeout, + query_params=query_params, ) else: raise ValueError(f"Unsupported model: {model}") @@ -170,7 +173,7 @@ async def _realtime_health_check( ) elif custom_llm_provider == "openai": url = openai_realtime._construct_url( - api_base=api_base or "https://api.openai.com/", model=model + api_base=api_base or "https://api.openai.com/", query_params=RealtimeQueryParams(model=model) ) else: raise ValueError(f"Unsupported model: {model}") diff --git a/litellm/types/realtime.py b/litellm/types/realtime.py index c105983b1e..73b60a223a 100644 --- a/litellm/types/realtime.py +++ b/litellm/types/realtime.py @@ -41,3 +41,9 @@ class RealtimeModalityResponseTransformOutput(TypedDict): current_conversation_id: Optional[str] current_delta_chunks: Optional[List[OpenAIRealtimeResponseDelta]] current_delta_type: Optional[ALL_DELTA_TYPES] + + +class RealtimeQueryParams(TypedDict, total=False): + model: str + intent: Optional[str] + # Add more fields as needed 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 0681be6c22..e4378dbeae 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 @@ -19,9 +19,74 @@ def test_openai_realtime_handler_url_construction(api_base): handler = OpenAIRealtime() url = handler._construct_url( - api_base=api_base, model="gpt-4o-realtime-preview-2024-10-01" + api_base=api_base, query_params = { + "model": "gpt-4o-realtime-preview-2024-10-01", + } ) assert ( url - == f"wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01" + == f"wss://api.openai.com/v1/realtime" ) + + +def test_openai_realtime_handler_url_with_extra_params(): + from litellm.llms.openai.realtime.handler import OpenAIRealtime + from litellm.types.realtime import RealtimeQueryParams + + handler = OpenAIRealtime() + api_base = "https://api.openai.com/v1" + query_params: RealtimeQueryParams = { + "model": "gpt-4o-realtime-preview-2024-10-01", + "intent": "chat" + } + url = handler._construct_url(api_base=api_base, query_params=query_params) + # 'model' should be excluded from the query string + assert url.startswith("wss://api.openai.com/v1/realtime?") + assert "intent=chat" in url + + +import asyncio + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +@pytest.mark.asyncio +async def test_async_realtime_success(): + from litellm.llms.openai.realtime.handler import OpenAIRealtime + from litellm.types.realtime import RealtimeQueryParams + + handler = OpenAIRealtime() + api_base = "https://api.openai.com/v1" + api_key = "test-key" + model = "gpt-4o-realtime-preview-2024-10-01" + query_params: RealtimeQueryParams = {"model": model, "intent": "chat"} + + dummy_websocket = AsyncMock() + dummy_logging_obj = MagicMock() + mock_backend_ws = AsyncMock() + + class DummyAsyncContextManager: + def __init__(self, value): + self.value = value + async def __aenter__(self): + return self.value + async def __aexit__(self, exc_type, exc, tb): + return None + + 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() + mock_realtime_streaming.return_value = mock_streaming_instance + mock_streaming_instance.bidirectional_forward = AsyncMock() + + await handler.async_realtime( + model=model, + websocket=dummy_websocket, + logging_obj=dummy_logging_obj, + api_base=api_base, + api_key=api_key, + query_params=query_params, + ) + + mock_realtime_streaming.assert_called_once() + mock_streaming_instance.bidirectional_forward.assert_awaited_once()