Make api_key callable to enable token refresh for openai client

Author: johanste

src/openai/__init__.py

@@ -117,9 +117,7 @@
 
 from ._base_client import DEFAULT_TIMEOUT, DEFAULT_MAX_RETRIES
 
-api_key: str | None = None
-
-bearer_token_provider: _t.Callable[[], str] | None = None
+api_key: str | _t.Callable[[], str] | None = None
 
 organization: str | None = None
 
@@ -158,25 +156,14 @@ class _ModuleClient(OpenAI):
 
     @property  # type: ignore
     @override
-    def api_key(self) -> str | None:
+    def api_key(self) -> str | _t.Callable[[], str] | None:
         return api_key
 
     @api_key.setter  # type: ignore
-    def api_key(self, value: str | None) -> None:  # type: ignore
+    def api_key(self, value: str | _t.Callable[[], str] | None) -> None:  # type: ignore
         global api_key
-
         api_key = value
 
-    @property  # type: ignore
-    @override
-    def bearer_token_provider(self) -> _t.Callable[[], str] | None:
-        return bearer_token_provider
-
-    @bearer_token_provider.setter  # type: ignore
-    def bearer_token_provider(self, value: _t.Callable[[], str] | None) -> None:  # type: ignore
-        global bearer_token_provider
-
-        bearer_token_provider = value
 
     @property  # type: ignore
     @override
@@ -346,7 +333,7 @@ def _load_client() -> OpenAI:  # type: ignore[reportUnusedFunction]
             _client = _AzureModuleClient(  # type: ignore
                 api_version=api_version,
                 azure_endpoint=azure_endpoint,
-                api_key=api_key,
+                api_key=bearer_token_provider or api_key,
                 azure_ad_token=azure_ad_token,
                 azure_ad_token_provider=azure_ad_token_provider,
                 organization=organization,
@@ -360,8 +347,7 @@ def _load_client() -> OpenAI:  # type: ignore[reportUnusedFunction]
             return _client
 
         _client = _ModuleClient(
-            api_key=api_key,
-            bearer_token_provider=bearer_token_provider,
+            api_key=api_key or bearer_token_provider,
             organization=organization,
             project=project,
             webhook_secret=webhook_secret,

src/openai/_client.py

@@ -81,6 +81,7 @@
 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
@@ -96,8 +97,7 @@ class OpenAI(SyncAPIClient):
     def __init__(
         self,
         *,
-        api_key: str | None = None,
-        bearer_token_provider: Callable[[], str] | None = None,
+        api_key: str | None | Callable[[], str] = None,
         organization: str | None = None,
         project: str | None = None,
         webhook_secret: str | None = None,
@@ -129,16 +129,17 @@ def __init__(
         - `project` from `OPENAI_PROJECT_ID`
         - `webhook_secret` from `OPENAI_WEBHOOK_SECRET`
         """
-        if api_key and bearer_token_provider:
-            raise ValueError("The `api_key` and `bearer_token_provider` arguments are mutually exclusive")
         if api_key is None:
             api_key = os.environ.get("OPENAI_API_KEY")
-        if api_key is None and bearer_token_provider is None:
+        if api_key is None:
             raise OpenAIError(
-                "The api_key or bearer_token_provider client option must be set either by passing api_key or bearer_token_provider to the client or by setting the OPENAI_API_KEY environment variable"
+                "The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable"
             )
-        self.bearer_token_provider = bearer_token_provider
-        self.api_key = api_key or ""
+        if callable(api_key):
+            self.bearer_token_provider = api_key
+            self.api_key = ""
+        else:
+            self.api_key = api_key or ""
 
         if organization is None:
             organization = os.environ.get("OPENAI_ORG_ID")
@@ -328,8 +329,7 @@ def default_headers(self) -> dict[str, str | Omit]:
     def copy(
         self,
         *,
-        api_key: str | None = None,
-        bearer_token_provider: Callable[[], str] | None = None,
+        api_key: str | Callable[[], str] | None = None,
         organization: str | None = None,
         project: str | None = None,
         webhook_secret: str | None = None,
@@ -365,13 +365,9 @@ def copy(
         elif set_default_query is not None:
             params = set_default_query
 
-        bearer_token_provider = bearer_token_provider or self.bearer_token_provider
-        if bearer_token_provider is not None:
-            _extra_kwargs = {**_extra_kwargs, "bearer_token_provider": bearer_token_provider}
-
         http_client = http_client or self._client
         return self.__class__(
-            api_key=api_key or self.api_key,
+            api_key=api_key or self.api_key or self.bearer_token_provider,
             organization=organization or self.organization,
             project=project or self.project,
             webhook_secret=webhook_secret or self.webhook_secret,
@@ -427,6 +423,7 @@ 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
@@ -442,8 +439,7 @@ class AsyncOpenAI(AsyncAPIClient):
     def __init__(
         self,
         *,
-        api_key: str | None = None,
-        bearer_token_provider: Callable[[], Awaitable[str]] | None = None,
+        api_key: str | Callable[[], Awaitable[str]] | None = None,
         organization: str | None = None,
         project: str | None = None,
         webhook_secret: str | None = None,
@@ -475,16 +471,18 @@ def __init__(
         - `project` from `OPENAI_PROJECT_ID`
         - `webhook_secret` from `OPENAI_WEBHOOK_SECRET`
         """
-        if api_key and bearer_token_provider:
-            raise ValueError("The `api_key` and `bearer_token_provider` arguments are mutually exclusive")
         if api_key is None:
             api_key = os.environ.get("OPENAI_API_KEY")
-        if api_key is None and bearer_token_provider is None:
+        if api_key is None:
             raise OpenAIError(
-                "The api_key or bearer_token_provider client option must be set either by passing api_key or bearer_token_provider to the client or by setting the OPENAI_API_KEY environment variable"
+                "The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable"
             )
-        self.bearer_token_provider = bearer_token_provider
-        self.api_key = api_key or ""
+        if callable(api_key):
+            self.bearer_token_provider = api_key
+            self.api_key = ""
+        else:
+            self.bearer_token_provider = None
+            self.api_key = api_key or ""
 
         if organization is None:
             organization = os.environ.get("OPENAI_ORG_ID")
@@ -677,8 +675,7 @@ def default_headers(self) -> dict[str, str | Omit]:
     def copy(
         self,
         *,
-        api_key: str | None = None,
-        bearer_token_provider: Callable[[], Awaitable[str]] | None = None,
+        api_key: str | Callable[[], Awaitable[str]] | None = None,
         organization: str | None = None,
         project: str | None = None,
         webhook_secret: str | None = None,
@@ -714,13 +711,9 @@ def copy(
         elif set_default_query is not None:
             params = set_default_query
 
-        bearer_token_provider = bearer_token_provider or self.bearer_token_provider
-        if bearer_token_provider is not None:
-            _extra_kwargs = {**_extra_kwargs, "bearer_token_provider": bearer_token_provider}
-
         http_client = http_client or self._client
         return self.__class__(
-            api_key=api_key or self.api_key,
+            api_key=api_key or self.api_key or self.bearer_token_provider,
             organization=organization or self.organization,
             project=project or self.project,
             webhook_secret=webhook_secret or self.webhook_secret,

tests/test_client.py

@@ -946,7 +946,7 @@ def test_follow_redirects_disabled(self, respx_mock: MockRouter) -> None:
         assert exc_info.value.response.headers["Location"] == f"{base_url}/redirected"
 
     def test_refresh_auth_headers_token(self) -> None:
-        client = OpenAI(base_url=base_url, bearer_token_provider=lambda: "test_bearer_token")
+        client = OpenAI(base_url=base_url, api_key=lambda: "test_bearer_token")
         client.refresh_auth_headers()
         assert client.auth_headers.get("Authorization") == "Bearer test_bearer_token"
 
@@ -976,7 +976,7 @@ def token_provider() -> str:
 
             return "second"
 
-        client = OpenAI(base_url=base_url, bearer_token_provider=token_provider)
+        client = OpenAI(base_url=base_url, api_key=token_provider)
         client.chat.completions.create(messages=[], model="gpt-4")
 
         calls = cast("list[MockRequestCall]", respx_mock.calls)
@@ -985,23 +985,14 @@ def token_provider() -> str:
         assert calls[0].request.headers.get("Authorization") == "Bearer first"
         assert calls[1].request.headers.get("Authorization") == "Bearer second"
 
-    def test_auth_mutually_exclusive(self) -> None:
-        with pytest.raises(ValueError) as exc_info:
-            OpenAI(base_url=base_url, api_key=api_key, bearer_token_provider=lambda: "test_bearer_token")
-        assert str(exc_info.value) == "The `api_key` and `bearer_token_provider` arguments are mutually exclusive"
 
     def test_copy_auth(self) -> None:
-        client = OpenAI(base_url=base_url, bearer_token_provider=lambda: "test_bearer_token_1").copy(
-            bearer_token_provider=lambda: "test_bearer_token_2"
+        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()
         assert client.auth_headers == {"Authorization": "Bearer test_bearer_token_2"}
 
-    def test_copy_auth_mutually_exclusive(self) -> None:
-        with pytest.raises(ValueError) as exc_info:
-            OpenAI(base_url=base_url, api_key=api_key).copy(bearer_token_provider=lambda: "test_bearer_token")
-        assert str(exc_info.value) == "The `api_key` and `bearer_token_provider` arguments are mutually exclusive"
-
 
 class TestAsyncOpenAI:
     client = AsyncOpenAI(base_url=base_url, api_key=api_key, _strict_response_validation=True)
@@ -1957,7 +1948,7 @@ async def test_refresh_auth_headers_token_async(self) -> None:
         async def token_provider() -> str:
             return "test_bearer_token"
 
-        client = AsyncOpenAI(base_url=base_url, bearer_token_provider=token_provider)
+        client = AsyncOpenAI(base_url=base_url, api_key=token_provider)
         await client.refresh_auth_headers()
         assert client.auth_headers.get("Authorization") == "Bearer test_bearer_token"
 
@@ -1989,7 +1980,7 @@ async def token_provider() -> str:
 
             return "second"
 
-        client = AsyncOpenAI(base_url=base_url, bearer_token_provider=token_provider)
+        client = AsyncOpenAI(base_url=base_url, api_key=token_provider)
         await client.chat.completions.create(messages=[], model="gpt-4")
 
         calls = cast("list[MockRequestCall]", respx_mock.calls)
@@ -1998,14 +1989,6 @@ async def token_provider() -> str:
         assert calls[0].request.headers.get("Authorization") == "Bearer first"
         assert calls[1].request.headers.get("Authorization") == "Bearer second"
 
-    def test_auth_mutually_exclusive_async(self) -> None:
-        async def token_provider() -> str:
-            return "test_bearer_token"
-
-        with pytest.raises(ValueError) as exc_info:
-            AsyncOpenAI(base_url=base_url, api_key=api_key, bearer_token_provider=token_provider)
-        assert str(exc_info.value) == "The `api_key` and `bearer_token_provider` arguments are mutually exclusive"
-
     @pytest.mark.asyncio
     async def test_copy_auth(self) -> None:
         async def token_provider_1() -> str:
@@ -2014,16 +1997,8 @@ async def token_provider_1() -> str:
         async def token_provider_2() -> str:
             return "test_bearer_token_2"
 
-        client = AsyncOpenAI(base_url=base_url, bearer_token_provider=token_provider_1).copy(
-            bearer_token_provider=token_provider_2
+        client = AsyncOpenAI(base_url=base_url, api_key=token_provider_1).copy(
+            api_key=token_provider_2
         )
         await client.refresh_auth_headers()
         assert client.auth_headers == {"Authorization": "Bearer test_bearer_token_2"}
-
-    def test_copy_auth_mutually_exclusive_async(self) -> None:
-        async def token_provider() -> str:
-            return "test_bearer_token"
-
-        with pytest.raises(ValueError) as exc_info:
-            AsyncOpenAI(base_url=base_url, api_key=api_key).copy(bearer_token_provider=token_provider)
-        assert str(exc_info.value) == "The `api_key` and `bearer_token_provider` arguments are mutually exclusive"

tests/test_module_client.py

@@ -99,12 +99,8 @@ def test_http_client_option() -> None:
 
 
 def test_bearer_token_provider_option() -> None:
-    assert openai.bearer_token_provider is None
-    assert openai.completions._client.bearer_token_provider is None
+    openai.api_key = lambda: "foo"
 
-    openai.bearer_token_provider = lambda: "foo"
-
-    assert openai.bearer_token_provider() == "foo"
     assert openai.completions._client.bearer_token_provider
     assert openai.completions._client.bearer_token_provider() == "foo"
 
@@ -138,23 +134,19 @@ def test_only_api_key_results_in_openai_api() -> None:
 def test_only_bearer_token_provider_in_openai_api() -> None:
     with fresh_env():
         openai.api_type = None
-        openai.api_key = None
-        openai.bearer_token_provider = lambda: "example bearer token"
+        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:
     with fresh_env():
-        openai.api_key = "example API key"
-        openai.bearer_token_provider = lambda: "example bearer token"
+        openai.api_key = lambda: "example bearer token"
 
-        with pytest.raises(
-            ValueError,
-            match=r"The `api_key` and `bearer_token_provider` arguments are mutually exclusive",
-        ):
-            openai.completions._client  # noqa: B018
+        assert(openai.api_key() == "example bearer token")
 
+        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():