extract azure logic out of enter

Author: kristapratico

src/openai/_utils/__init__.py

@@ -34,6 +34,8 @@
     maybe_coerce_boolean as maybe_coerce_boolean,
     maybe_coerce_integer as maybe_coerce_integer,
     is_async_azure_client as is_async_azure_client,
+    configure_azure_realtime as configure_azure_realtime,
+    configure_azure_realtime_async as configure_azure_realtime_async,
 )
 from ._typing import (
     is_list_type as is_list_type,

src/openai/_utils/_utils.py

@@ -22,7 +22,7 @@
 
 import sniffio
 
-from .._types import NotGiven, FileTypes, NotGivenOr, HeadersLike
+from .._types import Query, NotGiven, FileTypes, NotGivenOr, HeadersLike
 from .._compat import parse_date as parse_date, parse_datetime as parse_datetime
 
 _T = TypeVar("_T")
@@ -419,8 +419,46 @@ def json_safe(data: object) -> object:
 
 
 def is_azure_client(client: object) -> TypeGuard[AzureOpenAI]:
-    return hasattr(client, "_azure_ad_token_provider")
+    from ..lib.azure import AzureOpenAI
+
+    return isinstance(client, AzureOpenAI)
 
 
 def is_async_azure_client(client: object) -> TypeGuard[AsyncAzureOpenAI]:
-    return hasattr(client, "_azure_ad_token_provider")
+    from ..lib.azure import AsyncAzureOpenAI
+
+    return isinstance(client, AsyncAzureOpenAI)
+
+
+def configure_azure_realtime(client: AzureOpenAI, model: str, extra_query: Query) -> tuple[Query, dict[str, str]]:
+    auth_headers = {}
+    query = {
+        **extra_query,
+        "api-version": client._api_version,
+        "deployment": client._azure_deployment or model,
+    }
+    if client.api_key != "<missing API key>":
+        auth_headers = {"api-key": client.api_key}
+    else:
+        token = client._get_azure_ad_token()
+        if token:
+            auth_headers = {"Authorization": f"Bearer {token}"}
+    return query, auth_headers
+
+
+async def configure_azure_realtime_async(
+    client: AsyncAzureOpenAI, model: str, extra_query: Query
+) -> tuple[Query, dict[str, str]]:
+    auth_headers = {}
+    query = {
+        **extra_query,
+        "api-version": client._api_version,
+        "deployment": client._azure_deployment or model,
+    }
+    if client.api_key != "<missing API key>":
+        auth_headers = {"api-key": client.api_key}
+    else:
+        token = await client._get_azure_ad_token()
+        if token:
+            auth_headers = {"Authorization": f"Bearer {token}"}
+    return query, auth_headers

src/openai/resources/beta/realtime/realtime.py

@@ -26,6 +26,8 @@
     strip_not_given,
     async_maybe_transform,
     is_async_azure_client,
+    configure_azure_realtime,
+    configure_azure_realtime_async,
 )
 from ...._compat import cached_property
 from ...._models import construct_type_unchecked
@@ -321,20 +323,10 @@ async def __aenter__(self) -> AsyncRealtimeConnection:
         except ImportError as exc:
             raise OpenAIError("You need to install `openai[realtime]` to use this method") from exc
 
-        auth_headers = self.__client.auth_headers
         extra_query = self.__extra_query
+        auth_headers = self.__client.auth_headers
         if is_async_azure_client(self.__client):
-            extra_query = {
-                **self.__extra_query,
-                "api-version": self.__client._api_version,
-                "deployment": self.__client._azure_deployment or self.__model,
-            }
-            if self.__client.api_key != "<missing API key>":
-                auth_headers = {"api-key": self.__client.api_key}
-            else:
-                token = await self.__client._get_azure_ad_token()
-                if token:
-                    auth_headers = {"Authorization": f"Bearer {token}"}
+            extra_query, auth_headers = await configure_azure_realtime_async(self.__client, self.__model, extra_query)
 
         url = self._prepare_url().copy_with(
             params={
@@ -513,20 +505,10 @@ def __enter__(self) -> RealtimeConnection:
         except ImportError as exc:
             raise OpenAIError("You need to install `openai[realtime]` to use this method") from exc
 
-        auth_headers = self.__client.auth_headers
         extra_query = self.__extra_query
+        auth_headers = self.__client.auth_headers
         if is_azure_client(self.__client):
-            extra_query = {
-                **self.__extra_query,
-                "api-version": self.__client._api_version,
-                "deployment": self.__client._azure_deployment or self.__model,
-            }
-            if self.__client.api_key != "<missing API key>":
-                auth_headers = {"api-key": self.__client.api_key}
-            else:
-                token = self.__client._get_azure_ad_token()
-                if token:
-                    auth_headers = {"Authorization": f"Bearer {token}"}
+            extra_query, auth_headers = configure_azure_realtime(self.__client, self.__model, extra_query)
 
         url = self._prepare_url().copy_with(
             params={