Bearer token provider addition to openaiclient

src/openai/__init__.py

@@ -116,6 +116,8 @@
 
 api_key: str | None = None
 
+bearer_token_provider: _t.Callable[[], str] | None = None
+
 organization: str | None = None
 
 project: str | None = None
@@ -160,6 +162,17 @@ def api_key(self, value: str | None) -> None:  # type: ignore
 
         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
     def organization(self) -> str | None:
@@ -332,6 +345,7 @@ def _load_client() -> OpenAI:  # type: ignore[reportUnusedFunction]
 
         _client = _ModuleClient(
             api_key=api_key,
+            bearer_token_provider=bearer_token_provider,
             organization=organization,
             project=project,
             base_url=base_url,

src/openai/_client.py

@@ -3,11 +3,13 @@
 from __future__ import annotations
 
 import os
-from typing import TYPE_CHECKING, Any, Union, Mapping
+from typing import TYPE_CHECKING, Any, Union, Mapping, Callable, Awaitable
 from typing_extensions import Self, override
 
 import httpx
 
+from openai._models import FinalRequestOptions
+
 from . import _exceptions
 from ._qs import Querystring
 from ._types import (
@@ -91,6 +93,7 @@ def __init__(
         self,
         *,
         api_key: str | None = None,
+        bearer_token_provider: Callable[[], str] | None = None,
         organization: str | None = None,
         project: str | None = None,
         base_url: str | httpx.URL | None = None,
@@ -120,13 +123,16 @@ def __init__(
         - `organization` from `OPENAI_ORG_ID`
         - `project` from `OPENAI_PROJECT_ID`
         """
+        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:
+        if api_key is None and bearer_token_provider is None:
             raise OpenAIError(
-                "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"
+                "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"
             )
-        self.api_key = api_key
+        self.bearer_token_provider = bearer_token_provider
+        self.api_key = api_key or ""
 
         if organization is None:
             organization = os.environ.get("OPENAI_ORG_ID")
@@ -155,6 +161,7 @@ def __init__(
         )
 
         self._default_stream_cls = Stream
+        self._auth_headers: dict[str, str] = {}
 
     @cached_property
     def completions(self) -> Completions:
@@ -265,11 +272,25 @@ def with_streaming_response(self) -> OpenAIWithStreamedResponse:
     def qs(self) -> Querystring:
         return Querystring(array_format="brackets")
 
+    def refresh_auth_headers(self) -> None:
+        secret = self.bearer_token_provider() if self.bearer_token_provider else 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}"}
+
+    @override
+    def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOptions:
+        self.refresh_auth_headers()
+        return super()._prepare_options(options)
+
     @property
     @override
     def auth_headers(self) -> dict[str, str]:
-        api_key = self.api_key
-        return {"Authorization": f"Bearer {api_key}"}
+        return self._auth_headers
 
     @property
     @override
@@ -286,6 +307,7 @@ def copy(
         self,
         *,
         api_key: str | None = None,
+        bearer_token_provider: Callable[[], str] | None = None,
         organization: str | None = None,
         project: str | None = None,
         websocket_base_url: str | httpx.URL | None = None,
@@ -320,6 +342,10 @@ 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,
@@ -392,6 +418,7 @@ def __init__(
         self,
         *,
         api_key: str | None = None,
+        bearer_token_provider: Callable[[], Awaitable[str]] | None = None,
         organization: str | None = None,
         project: str | None = None,
         base_url: str | httpx.URL | None = None,
@@ -421,13 +448,16 @@ def __init__(
         - `organization` from `OPENAI_ORG_ID`
         - `project` from `OPENAI_PROJECT_ID`
         """
+        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:
+        if api_key is None and bearer_token_provider is None:
             raise OpenAIError(
-                "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"
+                "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"
             )
-        self.api_key = api_key
+        self.bearer_token_provider = bearer_token_provider
+        self.api_key = api_key or ""
 
         if organization is None:
             organization = os.environ.get("OPENAI_ORG_ID")
@@ -456,6 +486,7 @@ def __init__(
         )
 
         self._default_stream_cls = AsyncStream
+        self._auth_headers: dict[str, str] = {}
 
     @cached_property
     def completions(self) -> AsyncCompletions:
@@ -566,11 +597,28 @@ 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}"}
+
+    @override
+    async def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOptions:
+        await self.refresh_auth_headers()
+        return await super()._prepare_options(options)
+
     @property
     @override
     def auth_headers(self) -> dict[str, str]:
-        api_key = self.api_key
-        return {"Authorization": f"Bearer {api_key}"}
+        return self._auth_headers
 
     @property
     @override
@@ -587,6 +635,7 @@ def copy(
         self,
         *,
         api_key: str | None = None,
+        bearer_token_provider: Callable[[], Awaitable[str]] | None = None,
         organization: str | None = None,
         project: str | None = None,
         websocket_base_url: str | httpx.URL | None = None,
@@ -621,6 +670,10 @@ 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,

src/openai/lib/azure.py

@@ -250,7 +250,7 @@ def __init__(
         self._azure_endpoint = httpx.URL(azure_endpoint) if azure_endpoint else None
 
     @override
-    def copy(
+    def copy(  # type: ignore
         self,
         *,
         api_key: str | None = None,
@@ -294,7 +294,7 @@ def copy(
             },
         )
 
-    with_options = copy
+    with_options = copy  # type: ignore
 
     def _get_azure_ad_token(self) -> str | None:
         if self._azure_ad_token is not None:
@@ -524,7 +524,7 @@ def __init__(
         self._azure_endpoint = httpx.URL(azure_endpoint) if azure_endpoint else None
 
     @override
-    def copy(
+    def copy(  # type: ignore
         self,
         *,
         api_key: str | None = None,
@@ -568,7 +568,7 @@ def copy(
             },
         )
 
-    with_options = copy
+    with_options = copy  # type: ignore
 
     async def _get_azure_ad_token(self) -> str | None:
         if self._azure_ad_token is not None:
@@ -614,7 +614,7 @@ async def _configure_realtime(self, model: str, extra_query: Query) -> tuple[htt
             "api-version": self._api_version,
             "deployment": self._azure_deployment or model,
         }
-        if self.api_key != "<missing API key>":
+        if self.api_key and self.api_key != "<missing API key>":
             auth_headers = {"api-key": self.api_key}
         else:
             token = await self._get_azure_ad_token()

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

@@ -358,6 +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()
         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)
@@ -540,6 +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()
         auth_headers = self.__client.auth_headers
         if is_azure_client(self.__client):
             url, auth_headers = self.__client._configure_realtime(self.__model, extra_query)