refactor(realtime): type the Azure protocol picker with the streaming module's websocket protocol

litellm/litellm_core_utils/realtime_streaming.py

@@ -90,12 +90,12 @@ class _ResponseDoneBody(TypedDict, total=False):
     output: ReadOnly[Sequence[Mapping[str, object]]]
 
 
-class _ScopedWebSocket(Protocol):
+class ScopedWebSocket(Protocol):
     @property
     def scope(self) -> _ASGIScope: ...
 
 
-class _ClientWebSocket(_ScopedWebSocket, Protocol):
+class _ClientWebSocket(ScopedWebSocket, Protocol):
     async def send_text(self, data: str) -> None: ...
     async def receive_text(self) -> str: ...
     async def close(self, code: int = 1000, reason: str | None = None) -> None: ...
@@ -1149,7 +1149,7 @@ class RealTimeStreaming:
         )
 
     @staticmethod
-    def _detect_beta_header(websocket: _ScopedWebSocket) -> bool:
+    def _detect_beta_header(websocket: ScopedWebSocket) -> bool:
         """Return True if the client sent 'OpenAI-Beta: realtime=v1'.
 
         Checks the raw ASGI scope headers so it works for both FastAPI WebSocket
@@ -1584,6 +1584,6 @@ class RealTimeStreaming:
             verbose_logger.debug("Could not relay the upstream close to the client: %s", e)
 
 
-def client_sent_openai_beta_realtime_header(websocket: _ScopedWebSocket) -> bool:
+def client_sent_openai_beta_realtime_header(websocket: ScopedWebSocket) -> bool:
     """True when the client WebSocket includes ``OpenAI-Beta: realtime=v1``."""
     return RealTimeStreaming._detect_beta_header(websocket)

litellm/llms/azure/realtime/handler.py

@@ -6,7 +6,7 @@ This requires websockets, and is currently only supported on LiteLLM Proxy.
 
 from collections.abc import Mapping
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, Protocol, cast
+from typing import Any, Final, Protocol, cast
 
 from litellm._logging import _redact_string, verbose_proxy_logger
 from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES
@@ -15,14 +15,12 @@ from litellm.types.realtime import RealtimeQueryParams
 from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
 from ....litellm_core_utils.realtime_streaming import (
     RealTimeStreaming,
+    ScopedWebSocket,
     client_sent_openai_beta_realtime_header,
 )
 from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context
 from ..azure import AzureChatCompletion
 
-if TYPE_CHECKING:
-    from fastapi import WebSocket
-
 # BACKEND_WS_URL = "ws://localhost:8080/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01"
 
 
@@ -41,7 +39,7 @@ def azure_realtime_protocol_for_client(
     configured_protocol: object,
     *,
     query_params: RealtimeQueryParams | None,
-    websocket: "WebSocket",
+    websocket: ScopedWebSocket,
 ) -> str:
     if isinstance(configured_protocol, str) and configured_protocol:
         return configured_protocol