Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 3 additions & 10 deletions src/openai/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,8 @@

azure_endpoint: str | None = _os.environ.get("AZURE_OPENAI_ENDPOINT")

azure_ad_token: str | None = _os.environ.get("AZURE_OPENAI_AD_TOKEN")
# Keep explicit module configuration separate from Azure's environment fallback.
azure_ad_token: str | None = None

azure_ad_token_provider: _azure.AzureADTokenProvider | None = None

Expand Down Expand Up @@ -362,14 +363,11 @@ def _load_client() -> OpenAI: # type: ignore[reportUnusedFunction]
global _client

if _client is None:
global api_type, azure_endpoint, azure_ad_token, api_version
global api_type, azure_endpoint, api_version

if azure_endpoint is None:
azure_endpoint = _os.environ.get("AZURE_OPENAI_ENDPOINT")

if azure_ad_token is None:
azure_ad_token = _os.environ.get("AZURE_OPENAI_AD_TOKEN")

if api_version is None:
api_version = _os.environ.get("OPENAI_API_VERSION")

Expand All @@ -381,11 +379,6 @@ def _load_client() -> OpenAI: # type: ignore[reportUnusedFunction]
if has_openai and (has_azure or has_azure_ad):
raise _AmbiguousModuleClientUsageError()

if (azure_ad_token is not None or azure_ad_token_provider is not None) and _os.environ.get(
"AZURE_OPENAI_API_KEY"
) is not None:
raise _AmbiguousModuleClientUsageError()

if has_azure or has_azure_ad:
api_type = "azure"
else:
Expand Down
106 changes: 92 additions & 14 deletions src/openai/lib/azure.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,56 @@ def __init__(self) -> None:
)


_AzureApiKeyT = TypeVar("_AzureApiKeyT", bound=Union[str, Callable[..., Any]])
_AzureTokenProviderT = TypeVar("_AzureTokenProviderT", bound=Callable[..., Any])


def _resolve_azure_auth(
api_key: _AzureApiKeyT | None,
azure_ad_token: str | None,
azure_ad_token_provider: _AzureTokenProviderT | None,
) -> tuple[_AzureApiKeyT | str | None, str | None, _AzureTokenProviderT | None]:
# The sentinel is internal state carried by copies, not a second credential.
if (
sum(
(
api_key is not None and api_key != API_KEY_SENTINEL,
azure_ad_token is not None,
azure_ad_token_provider is not None,
Comment on lines +79 to +81

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Preserve explicit auth precedence in the module client

When the module-level client is forced to Azure, _load_client() materializes AZURE_OPENAI_AD_TOKEN into azure_ad_token before passing it to _AzureModuleClient; therefore an explicit openai.api_key or openai.azure_ad_token_provider is counted here alongside that ambient token and raises MutuallyExclusiveAuthError. For example, openai.api_type = "azure" with an explicit API key and an existing token environment variable now fails instead of honoring the explicit credential as the direct client does. Preserve whether the token came from the environment in this path and add a focused module-client authentication regression.

AGENTS.md reference: AGENTS.md:L31-L35

Useful? React with 👍 / 👎.

)
)
> 1
):
raise MutuallyExclusiveAuthError()

# Explicit credentials select the principal. Only an entirely unspecified
# configuration may consult the environment; retain AD-token precedence.
resolved_api_key: _AzureApiKeyT | str | None = api_key
if api_key is None and azure_ad_token is None and azure_ad_token_provider is None:
azure_ad_token = os.environ.get("AZURE_OPENAI_AD_TOKEN")
if azure_ad_token is None:
resolved_api_key = os.environ.get("AZURE_OPENAI_API_KEY")

return resolved_api_key, azure_ad_token, azure_ad_token_provider


def _copy_azure_auth(
api_key: _AzureApiKeyT | None,
azure_ad_token: str | None,
azure_ad_token_provider: _AzureTokenProviderT | None,
*,
current_api_key: _AzureApiKeyT | str,
current_token: str | None,
current_provider: _AzureTokenProviderT | None,
) -> tuple[_AzureApiKeyT | str, str | None, _AzureTokenProviderT | None]:
if api_key is None and azure_ad_token is None and azure_ad_token_provider is None:
return current_api_key or API_KEY_SENTINEL, current_token, current_provider

# Prevent OpenAI.copy() from inheriting the old API key when switching to AD.
key, token, provider = _resolve_azure_auth(api_key, azure_ad_token, azure_ad_token_provider)
return key or API_KEY_SENTINEL, token, provider


class BaseAzureClient(BaseClient[_HttpxClientT, _DefaultStreamT]):
_azure_endpoint: httpx2.URL | None
_azure_deployment: str | None
Expand Down Expand Up @@ -203,6 +253,10 @@ def __init__(
- `api_version` from `OPENAI_API_VERSION`
- `azure_endpoint` from `AZURE_OPENAI_ENDPOINT`

Pass at most one of `api_key`, `azure_ad_token`, or `azure_ad_token_provider`.
An explicit credential takes precedence over Azure credential environment variables.
With no explicit credential, `AZURE_OPENAI_AD_TOKEN` takes precedence over `AZURE_OPENAI_API_KEY`.

Args:
azure_endpoint: Your Azure endpoint, including the resource, e.g. `https://example-resource.azure.openai.com/`

Expand All @@ -216,11 +270,9 @@ def __init__(
if is_x509_workload_identity(workload_identity):
raise OpenAIError("X.509 workload identity is not supported by Azure clients")

if api_key is None:
api_key = os.environ.get("AZURE_OPENAI_API_KEY")

if azure_ad_token is None:
azure_ad_token = os.environ.get("AZURE_OPENAI_AD_TOKEN")
api_key, azure_ad_token, azure_ad_token_provider = _resolve_azure_auth(
api_key, azure_ad_token, azure_ad_token_provider
)

if _enforce_credentials and api_key is None and azure_ad_token is None and azure_ad_token_provider is None:
raise OpenAIError(
Expand Down Expand Up @@ -321,6 +373,15 @@ def copy(
if is_x509_workload_identity(workload_identity):
raise OpenAIError("X.509 workload identity is not supported by Azure clients")

api_key, azure_ad_token, azure_ad_token_provider = _copy_azure_auth(
api_key,
azure_ad_token,
azure_ad_token_provider,
current_api_key=self._api_key_provider or self.api_key,
current_token=self._azure_ad_token,
current_provider=self._azure_ad_token_provider,
)

return super().copy(
api_key=api_key,
admin_api_key=admin_api_key,
Expand All @@ -340,8 +401,8 @@ def copy(
_enforce_credentials=_enforce_credentials,
_extra_kwargs={
"api_version": api_version or self._api_version,
"azure_ad_token": azure_ad_token or self._azure_ad_token,
"azure_ad_token_provider": azure_ad_token_provider or self._azure_ad_token_provider,
"azure_ad_token": azure_ad_token,
"azure_ad_token_provider": azure_ad_token_provider,
**_extra_kwargs,
},
)
Expand Down Expand Up @@ -384,6 +445,9 @@ def _validate_headers(self, headers: Headers, custom_headers: Headers) -> None:

@override
def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOptions:
if self._api_key_provider is not None:
self._refresh_api_key()

headers: dict[str, str | Omit] = {**options.headers} if is_given(options.headers) else {}

options = model_copy(options)
Expand Down Expand Up @@ -536,6 +600,10 @@ def __init__(
- `api_version` from `OPENAI_API_VERSION`
- `azure_endpoint` from `AZURE_OPENAI_ENDPOINT`

Pass at most one of `api_key`, `azure_ad_token`, or `azure_ad_token_provider`.
An explicit credential takes precedence over Azure credential environment variables.
With no explicit credential, `AZURE_OPENAI_AD_TOKEN` takes precedence over `AZURE_OPENAI_API_KEY`.

Args:
azure_endpoint: Your Azure endpoint, including the resource, e.g. `https://example-resource.azure.openai.com/`

Expand All @@ -549,11 +617,9 @@ def __init__(
if is_x509_workload_identity(workload_identity):
raise OpenAIError("X.509 workload identity is not supported by Azure clients")

if api_key is None:
api_key = os.environ.get("AZURE_OPENAI_API_KEY")

if azure_ad_token is None:
azure_ad_token = os.environ.get("AZURE_OPENAI_AD_TOKEN")
api_key, azure_ad_token, azure_ad_token_provider = _resolve_azure_auth(
api_key, azure_ad_token, azure_ad_token_provider
)

if _enforce_credentials and api_key is None and azure_ad_token is None and azure_ad_token_provider is None:
raise OpenAIError(
Expand Down Expand Up @@ -654,6 +720,15 @@ def copy(
if is_x509_workload_identity(workload_identity):
raise OpenAIError("X.509 workload identity is not supported by Azure clients")

api_key, azure_ad_token, azure_ad_token_provider = _copy_azure_auth(
api_key,
azure_ad_token,
azure_ad_token_provider,
current_api_key=self._api_key_provider or self.api_key,
current_token=self._azure_ad_token,
current_provider=self._azure_ad_token_provider,
)

return super().copy(
api_key=api_key,
admin_api_key=admin_api_key,
Expand All @@ -673,8 +748,8 @@ def copy(
_enforce_credentials=_enforce_credentials,
_extra_kwargs={
"api_version": api_version or self._api_version,
"azure_ad_token": azure_ad_token or self._azure_ad_token,
"azure_ad_token_provider": azure_ad_token_provider or self._azure_ad_token_provider,
"azure_ad_token": azure_ad_token,
"azure_ad_token_provider": azure_ad_token_provider,
**_extra_kwargs,
},
)
Expand Down Expand Up @@ -719,6 +794,9 @@ def _validate_headers(self, headers: Headers, custom_headers: Headers) -> None:

@override
async def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOptions:
if self._api_key_provider is not None:
await self._refresh_api_key()

headers: dict[str, str | Omit] = {**options.headers} if is_given(options.headers) else {}

options = model_copy(options)
Expand Down
Loading
Loading