Review feedback + make sure you can swap between callable and string api key values for module level client.

Author: johanste

src/openai/__init__.py

@@ -157,15 +157,26 @@ class _ModuleClient(OpenAI):
     @property  # type: ignore
     @override
     def api_key(self) -> str | _t.Callable[[], str] | None:
-        return api_key
+        return api_key() if callable(api_key) else api_key
 
     @api_key.setter  # type: ignore
     def api_key(self, value: str | _t.Callable[[], str] | None) -> None:  # type: ignore
         global api_key
         api_key = value
 
+    @property
+    def _api_key_provider(self) -> _t.Callable[[], str] | None: # type: ignore
+        return None
 
-    @property  # type: ignore
+    @_api_key_provider.setter
+    def _api_key_provider(self, value: _t.Callable[[], str] | None) -> None: # type: ignore
+        global api_key
+        # Yes, setting the api_key is intentional. The module level client accepts callables
+        # for the module level api_key and will call it to retrieve the value
+        # if it is a callable.
+        api_key = value
+
+    @property
     @override
     def organization(self) -> str | None:
         return organization

src/openai/_client.py

@@ -81,7 +81,6 @@
 class OpenAI(SyncAPIClient):
     # client options
     api_key: str
-    bearer_token_provider: Callable[[], str] | None = None
     organization: str | None
     project: str | None
     webhook_secret: str | None
@@ -137,10 +136,10 @@ def __init__(
             )
         if callable(api_key):
             self.api_key = ""
-            self.bearer_token_provider = api_key
+            self._api_key_provider: Callable[[], str] | None = api_key
         else:
             self.api_key = api_key or ""
-            self.bearer_token_provider = None
+            self._api_key_provider = None
 
         if organization is None:
             organization = os.environ.get("OPENAI_ORG_ID")
@@ -173,7 +172,6 @@ def __init__(
         )
 
         self._default_stream_cls = Stream
-        self._auth_headers: dict[str, str] = {}
 
     @cached_property
     def completions(self) -> Completions:
@@ -296,28 +294,23 @@ def with_streaming_response(self) -> OpenAIWithStreamedResponse:
     def qs(self) -> Querystring:
         return Querystring(array_format="brackets")
 
-    def refresh_auth_headers(self) -> None:
-        if self.bearer_token_provider:
-            secret = self.bearer_token_provider()
-        else:
-            secret = self.api_key
-        if not secret:
-            # if the api key is an empty string, encoding the header will fail
-            # so we set it to an empty dict
-            # this is to avoid sending an invalid Authorization header
-            self._auth_headers = {}
-        else:
-            self._auth_headers = {"Authorization": f"Bearer {secret}"}
+    def _refresh_api_key(self) -> None:
+        if self._api_key_provider:
+            self.api_key = self._api_key_provider()
 
     @override
     def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOptions:
-        self.refresh_auth_headers()
+        self._refresh_api_key()
         return super()._prepare_options(options)
 
     @property
     @override
     def auth_headers(self) -> dict[str, str]:
-        return self._auth_headers
+        api_key = self.api_key
+        if not api_key:
+            # if the api key is an empty string, encoding the header will fail
+            return {}
+        return {"Authorization": f"Bearer {api_key}"}
 
     @property
     @override
@@ -371,7 +364,7 @@ def copy(
 
         http_client = http_client or self._client
         return self.__class__(
-            api_key=api_key or self.bearer_token_provider or self.api_key,
+            api_key=api_key or self._api_key_provider or self.api_key,
             organization=organization or self.organization,
             project=project or self.project,
             webhook_secret=webhook_secret or self.webhook_secret,
@@ -427,7 +420,6 @@ def _make_status_error(
 class AsyncOpenAI(AsyncAPIClient):
     # client options
     api_key: str
-    bearer_token_provider: Callable[[], Awaitable[str]] | None = None
     organization: str | None
     project: str | None
     webhook_secret: str | None
@@ -483,10 +475,10 @@ def __init__(
             )
         if callable(api_key):
             self.api_key = ""
-            self.bearer_token_provider = api_key
+            self._api_key_provider: Callable[[], Awaitable[str]] | None = api_key
         else:
             self.api_key = api_key or ""
-            self.bearer_token_provider = None
+            self._api_key_provider = None
 
         if organization is None:
             organization = os.environ.get("OPENAI_ORG_ID")
@@ -519,7 +511,6 @@ def __init__(
         )
 
         self._default_stream_cls = AsyncStream
-        self._auth_headers: dict[str, str] = {}
 
     @cached_property
     def completions(self) -> AsyncCompletions:
@@ -642,28 +633,23 @@ def with_streaming_response(self) -> AsyncOpenAIWithStreamedResponse:
     def qs(self) -> Querystring:
         return Querystring(array_format="brackets")
 
-    async def refresh_auth_headers(self) -> None:
-        if self.bearer_token_provider:
-            secret = await self.bearer_token_provider()
-        else:
-            secret = self.api_key
-        if not secret:
-            # if the api key is an empty string, encoding the header will fail
-            # so we set it to an empty dict
-            # this is to avoid sending an invalid Authorization header
-            self._auth_headers = {}
-        else:
-            self._auth_headers = {"Authorization": f"Bearer {secret}"}
+    async def _refresh_api_key(self) -> None:
+        if self._api_key_provider:
+            self.api_key = await self._api_key_provider()
 
     @override
     async def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOptions:
-        await self.refresh_auth_headers()
+        await self._refresh_api_key()
         return await super()._prepare_options(options)
 
     @property
     @override
     def auth_headers(self) -> dict[str, str]:
-        return self._auth_headers
+        api_key = self.api_key
+        if not api_key:
+            # if the api key is an empty string, encoding the header will fail
+            return {}
+        return {"Authorization": f"Bearer {api_key}"}
 
     @property
     @override
@@ -717,7 +703,7 @@ def copy(
 
         http_client = http_client or self._client
         return self.__class__(
-            api_key=api_key or self.bearer_token_provider or self.api_key,
+            api_key=api_key or self._api_key_provider or self.api_key,
             organization=organization or self.organization,
             project=project or self.project,
             webhook_secret=webhook_secret or self.webhook_secret,

src/openai/lib/azure.py

@@ -255,10 +255,10 @@ def __init__(
         self._azure_endpoint = httpx.URL(azure_endpoint) if azure_endpoint else None
 
     @override
-    def copy(  # type: ignore
+    def copy(
         self,
         *,
-        api_key: str | None = None,
+        api_key: str | Callable[[], str] | None = None,
         organization: str | None = None,
         project: str | None = None,
         webhook_secret: str | None = None,
@@ -301,7 +301,7 @@ def copy(  # type: ignore
             },
         )
 
-    with_options = copy  # type: ignore
+    with_options = copy
 
     def _get_azure_ad_token(self) -> str | None:
         if self._azure_ad_token is not None:
@@ -435,7 +435,7 @@ def __init__(
         azure_endpoint: str | None = None,
         azure_deployment: str | None = None,
         api_version: str | None = None,
-        api_key: str | None = None,
+        api_key: str | Callable[[], Awaitable[str]] | None = None,
         azure_ad_token: str | None = None,
         azure_ad_token_provider: AsyncAzureADTokenProvider | None = None,
         organization: str | None = None,
@@ -536,10 +536,10 @@ def __init__(
         self._azure_endpoint = httpx.URL(azure_endpoint) if azure_endpoint else None
 
     @override
-    def copy(  # type: ignore
+    def copy(
         self,
         *,
-        api_key: str | None = None,
+        api_key: str | Callable[[], Awaitable[str]] | None = None,
         organization: str | None = None,
         project: str | None = None,
         webhook_secret: str | None = None,
@@ -582,7 +582,7 @@ def copy(  # type: ignore
             },
         )
 
-    with_options = copy  # type: ignore
+    with_options = copy
 
     async def _get_azure_ad_token(self) -> str | None:
         if self._azure_ad_token is not None:

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

@@ -358,7 +358,7 @@ async def __aenter__(self) -> AsyncRealtimeConnection:
             raise OpenAIError("You need to install `openai[realtime]` to use this method") from exc
 
         extra_query = self.__extra_query
-        await self.__client.refresh_auth_headers()
+        await self.__client._refresh_api_key()
         auth_headers = self.__client.auth_headers
         if is_async_azure_client(self.__client):
             url, auth_headers = await self.__client._configure_realtime(self.__model, extra_query)
@@ -541,7 +541,7 @@ def __enter__(self) -> RealtimeConnection:
             raise OpenAIError("You need to install `openai[realtime]` to use this method") from exc
 
         extra_query = self.__extra_query
-        self.__client.refresh_auth_headers()
+        self.__client._refresh_api_key()
         auth_headers = self.__client.auth_headers
         if is_azure_client(self.__client):
             url, auth_headers = self.__client._configure_realtime(self.__model, extra_query)

tests/test_client.py

@@ -945,14 +945,24 @@ def test_follow_redirects_disabled(self, respx_mock: MockRouter) -> None:
         assert exc_info.value.response.status_code == 302
         assert exc_info.value.response.headers["Location"] == f"{base_url}/redirected"
 
-    def test_refresh_auth_headers_token(self) -> None:
+    def test_api_key_before_after_refresh_provider(self) -> None:
         client = OpenAI(base_url=base_url, api_key=lambda: "test_bearer_token")
-        client.refresh_auth_headers()
+
+        assert client.api_key == ""
+        assert 'Authorization' not in client.auth_headers
+
+        client._refresh_api_key()
+
+        assert client.api_key == "test_bearer_token"
         assert client.auth_headers.get("Authorization") == "Bearer test_bearer_token"
 
-    def test_refresh_auth_headers_key(self) -> None:
+
+    def test_api_key_before_after_refresh_str(self) -> None:
         client = OpenAI(base_url=base_url, api_key="test_api_key")
-        client.refresh_auth_headers()
+
+        assert client.auth_headers.get("Authorization") == "Bearer test_api_key"
+        client._refresh_api_key()
+
         assert client.auth_headers.get("Authorization") == "Bearer test_api_key"
 
     @pytest.mark.respx()
@@ -985,12 +995,11 @@ def token_provider() -> str:
         assert calls[0].request.headers.get("Authorization") == "Bearer first"
         assert calls[1].request.headers.get("Authorization") == "Bearer second"
 
-
     def test_copy_auth(self) -> None:
         client = OpenAI(base_url=base_url, api_key=lambda: "test_bearer_token_1").copy(
             api_key=lambda: "test_bearer_token_2"
         )
-        client.refresh_auth_headers()
+        client._refresh_api_key()
         assert client.auth_headers == {"Authorization": "Bearer test_bearer_token_2"}
 
 
@@ -1944,18 +1953,28 @@ async def test_follow_redirects_disabled(self, respx_mock: MockRouter) -> None:
         assert exc_info.value.response.headers["Location"] == f"{base_url}/redirected"
 
     @pytest.mark.asyncio
-    async def test_refresh_auth_headers_token_async(self) -> None:
-        async def token_provider() -> str:
+    async def test_api_key_before_after_refresh_provider(self) -> None:
+        async def mock_api_key_provider():
             return "test_bearer_token"
+        
+        client = AsyncOpenAI(base_url=base_url, api_key=mock_api_key_provider)
 
-        client = AsyncOpenAI(base_url=base_url, api_key=token_provider)
-        await client.refresh_auth_headers()
+        assert client.api_key == ""
+        assert 'Authorization' not in client.auth_headers
+
+        await client._refresh_api_key()
+
+        assert client.api_key == "test_bearer_token"
         assert client.auth_headers.get("Authorization") == "Bearer test_bearer_token"
 
+
     @pytest.mark.asyncio
-    async def test_refresh_auth_headers_key_async(self) -> None:
+    async def test_api_key_before_after_refresh_str(self) -> None:
         client = AsyncOpenAI(base_url=base_url, api_key="test_api_key")
-        await client.refresh_auth_headers()
+
+        assert client.auth_headers.get("Authorization") == "Bearer test_api_key"
+        await client._refresh_api_key()
+
         assert client.auth_headers.get("Authorization") == "Bearer test_api_key"
 
     @pytest.mark.asyncio
@@ -1997,8 +2016,6 @@ async def token_provider_1() -> str:
         async def token_provider_2() -> str:
             return "test_bearer_token_2"
 
-        client = AsyncOpenAI(base_url=base_url, api_key=token_provider_1).copy(
-            api_key=token_provider_2
-        )
-        await client.refresh_auth_headers()
+        client = AsyncOpenAI(base_url=base_url, api_key=token_provider_1).copy(api_key=token_provider_2)
+        await client._refresh_api_key()
         assert client.auth_headers == {"Authorization": "Bearer test_bearer_token_2"}

tests/test_module_client.py

@@ -97,12 +97,22 @@ def test_http_client_option() -> None:
     assert openai.completions._client._client is new_client
 
 
-def test_bearer_token_provider_option() -> None:
-    openai.api_key = lambda: "foo"
+def test_api_key_callable() -> None:
+    openai.api_key = lambda: "1"
+    assert openai.completions._client.api_key == "1"
 
-    assert openai.completions._client.bearer_token_provider
-    assert openai.completions._client.bearer_token_provider() == "foo"
+def test_api_key_overridable() -> None:
+    openai.api_key = lambda: "1"
+    assert openai.completions._client.api_key == "1"
+    assert openai.completions._client._api_key_provider is None
 
+    openai.api_key = "2"
+    assert openai.completions._client.api_key == "2"
+    assert openai.completions._client._api_key_provider is None
+
+    openai.api_key = lambda: "3"
+    assert openai.completions._client.api_key == "3"
+    assert openai.completions._client._api_key_provider is None
 
 import contextlib
 from typing import Iterator
@@ -130,22 +140,23 @@ def test_only_api_key_results_in_openai_api() -> None:
         assert type(openai.completions._client).__name__ == "_ModuleClient"
 
 
-def test_only_bearer_token_provider_in_openai_api() -> None:
+def test_only_api_key_in_openai_api() -> None:
     with fresh_env():
         openai.api_type = None
         openai.api_key = lambda: "example bearer token"
 
         assert type(openai.completions._client).__name__ == "_ModuleClient"
 
 
-def test_both_api_key_and_bearer_token_provider_in_openai_api() -> None:
+def test_both_api_key_and_api_key_provider_in_openai_api() -> None:
     with fresh_env():
         openai.api_key = lambda: "example bearer token"
 
-        assert(openai.api_key() == "example bearer token")
+        assert openai.api_key() == "example bearer token"
 
         openai.api_key = "example API key"
-        assert(openai.api_key == "example API key")
+        assert openai.api_key == "example API key"
+
 
 def test_azure_api_key_env_without_api_version() -> None:
     with fresh_env():