mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Unify Azure credential handling across all Python packages (#4088)
Replace ad_token, ad_token_provider, and get_entra_auth_token with a unified credential parameter across all Azure-related packages. Core changes: - Add AzureCredentialTypes (TokenCredential | AsyncTokenCredential) and AzureTokenProvider (Callable[[], str | Awaitable[str]]) type aliases - Add resolve_credential_to_token_provider() using azure.identity's get_bearer_token_provider for automatic token caching/refresh - Update AzureOpenAIChatClient, AzureOpenAIResponsesClient, and AzureOpenAIAssistantsClient to accept credential: AzureCredentialTypes | AzureTokenProvider - Remove ad_token, ad_token_provider params and get_entra_auth_token helpers Package updates: - azure-ai: Accept AzureCredentialTypes on AzureAIClient, AzureAIAgentClient, AzureAIProjectAgentProvider, AzureAIAgentsProvider - azure-ai-search: Accept AzureCredentialTypes on AzureAISearchContextProvider - purview: Accept AzureCredentialTypes | AzureTokenProvider on PurviewClient, PurviewPolicyMiddleware, PurviewChatPolicyMiddleware Fixes #3449 Fixes #3500 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
committed by
GitHub
Unverified
GitHub
parent
4c8f595019
commit
fd4e6e816c
@@ -30,6 +30,8 @@ _IMPORTS: dict[str, tuple[str, str]] = {
|
||||
"AzureAISearchSettings": ("agent_framework_azure_ai_search", "agent-framework-azure-ai-search"),
|
||||
"AzureAISettings": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
|
||||
"AzureAIAgentsProvider": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
|
||||
"AzureCredentialTypes": ("agent_framework.azure._entra_id_authentication", "agent-framework-core"),
|
||||
"AzureTokenProvider": ("agent_framework.azure._entra_id_authentication", "agent-framework-core"),
|
||||
"AzureOpenAIAssistantsClient": ("agent_framework.azure._assistants_client", "agent-framework-core"),
|
||||
"AzureOpenAIAssistantsOptions": ("agent_framework.azure._assistants_client", "agent-framework-core"),
|
||||
"AzureOpenAIChatClient": ("agent_framework.azure._chat_client", "agent-framework-core"),
|
||||
@@ -42,7 +44,6 @@ _IMPORTS: dict[str, tuple[str, str]] = {
|
||||
"DurableAIAgentClient": ("agent_framework_durabletask", "agent-framework-durabletask"),
|
||||
"DurableAIAgentOrchestrationContext": ("agent_framework_durabletask", "agent-framework-durabletask"),
|
||||
"DurableAIAgentWorker": ("agent_framework_durabletask", "agent-framework-durabletask"),
|
||||
"get_entra_auth_token": ("agent_framework.azure._entra_id_authentication", "agent-framework-core"),
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -21,7 +21,7 @@ from agent_framework_durabletask import (
|
||||
|
||||
from agent_framework.azure._assistants_client import AzureOpenAIAssistantsClient
|
||||
from agent_framework.azure._chat_client import AzureOpenAIChatClient
|
||||
from agent_framework.azure._entra_id_authentication import get_entra_auth_token
|
||||
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
|
||||
from agent_framework.azure._responses_client import AzureOpenAIResponsesClient
|
||||
from agent_framework.azure._shared import AzureOpenAISettings
|
||||
|
||||
@@ -37,13 +37,14 @@ __all__ = [
|
||||
"AzureAISearchContextProvider",
|
||||
"AzureAISearchSettings",
|
||||
"AzureAISettings",
|
||||
"AzureCredentialTypes",
|
||||
"AzureOpenAIAssistantsClient",
|
||||
"AzureOpenAIChatClient",
|
||||
"AzureOpenAIResponsesClient",
|
||||
"AzureOpenAISettings",
|
||||
"AzureTokenProvider",
|
||||
"DurableAIAgent",
|
||||
"DurableAIAgentClient",
|
||||
"DurableAIAgentOrchestrationContext",
|
||||
"DurableAIAgentWorker",
|
||||
"get_entra_auth_token",
|
||||
]
|
||||
|
||||
@@ -4,19 +4,16 @@ from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Generic
|
||||
from typing import Any, ClassVar, Generic
|
||||
|
||||
from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI
|
||||
from openai.lib.azure import AsyncAzureOpenAI
|
||||
|
||||
from .._settings import load_settings
|
||||
from ..exceptions import ServiceInitializationError
|
||||
from ..openai import OpenAIAssistantsClient
|
||||
from ..openai._assistants_client import OpenAIAssistantsOptions
|
||||
from ._entra_id_authentication import get_entra_auth_token
|
||||
from ._shared import DEFAULT_AZURE_TOKEN_ENDPOINT, AzureOpenAISettings, _apply_azure_defaults
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from azure.core.credentials import TokenCredential
|
||||
from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider, resolve_credential_to_token_provider
|
||||
from ._shared import AzureOpenAISettings, _apply_azure_defaults
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from typing import TypeVar # type: ignore # pragma: no cover
|
||||
@@ -61,10 +58,8 @@ class AzureOpenAIAssistantsClient(
|
||||
endpoint: str | None = None,
|
||||
base_url: str | None = None,
|
||||
api_version: str | None = None,
|
||||
ad_token: str | None = None,
|
||||
ad_token_provider: AsyncAzureADTokenProvider | None = None,
|
||||
token_endpoint: str | None = None,
|
||||
credential: TokenCredential | None = None,
|
||||
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
|
||||
default_headers: Mapping[str, str] | None = None,
|
||||
async_client: AsyncAzureOpenAI | None = None,
|
||||
env_file_path: str | None = None,
|
||||
@@ -93,11 +88,12 @@ class AzureOpenAIAssistantsClient(
|
||||
api_version: The deployment API version. If provided will override the value
|
||||
in the env vars or .env file.
|
||||
Can also be set via environment variable AZURE_OPENAI_API_VERSION.
|
||||
ad_token: The Azure Active Directory token.
|
||||
ad_token_provider: The Azure Active Directory token provider.
|
||||
token_endpoint: The token endpoint to request an Azure token.
|
||||
Can also be set via environment variable AZURE_OPENAI_TOKEN_ENDPOINT.
|
||||
credential: The Azure credential to use for authentication.
|
||||
credential: Azure credential or token provider for authentication. Accepts a
|
||||
``TokenCredential``, ``AsyncTokenCredential``, or a callable that returns a
|
||||
bearer token string (sync or async), for example from
|
||||
``azure.identity.get_bearer_token_provider()``.
|
||||
default_headers: The default headers mapping of string keys to
|
||||
string values for HTTP requests.
|
||||
async_client: An existing client to use.
|
||||
@@ -156,20 +152,15 @@ class AzureOpenAIAssistantsClient(
|
||||
"or 'AZURE_OPENAI_CHAT_DEPLOYMENT_NAME' environment variable."
|
||||
)
|
||||
|
||||
# Handle authentication: try API key first, then AD token, then Entra ID
|
||||
if (
|
||||
not async_client
|
||||
and not azure_openai_settings["api_key"]
|
||||
and not ad_token
|
||||
and not ad_token_provider
|
||||
and azure_openai_settings["token_endpoint"]
|
||||
and credential
|
||||
):
|
||||
token_ep = azure_openai_settings["token_endpoint"] or DEFAULT_AZURE_TOKEN_ENDPOINT
|
||||
ad_token = get_entra_auth_token(credential, token_ep)
|
||||
# Resolve credential to token provider
|
||||
ad_token_provider = None
|
||||
if not async_client and not azure_openai_settings["api_key"] and credential:
|
||||
ad_token_provider = resolve_credential_to_token_provider(
|
||||
credential, azure_openai_settings["token_endpoint"]
|
||||
)
|
||||
|
||||
if not async_client and not azure_openai_settings["api_key"] and not ad_token and not ad_token_provider:
|
||||
raise ServiceInitializationError("The Azure OpenAI API key, ad_token, or ad_token_provider is required.")
|
||||
if not async_client and not azure_openai_settings["api_key"] and not ad_token_provider:
|
||||
raise ServiceInitializationError("Please provide either api_key, credential, or a client.")
|
||||
|
||||
# Create Azure client if not provided
|
||||
if not async_client:
|
||||
@@ -180,8 +171,6 @@ class AzureOpenAIAssistantsClient(
|
||||
|
||||
if azure_openai_settings["api_key"]:
|
||||
client_params["api_key"] = azure_openai_settings["api_key"].get_secret_value()
|
||||
elif ad_token:
|
||||
client_params["azure_ad_token"] = ad_token
|
||||
elif ad_token_provider:
|
||||
client_params["azure_ad_token_provider"] = ad_token_provider
|
||||
|
||||
|
||||
@@ -8,8 +8,7 @@ import sys
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Generic
|
||||
|
||||
from azure.core.credentials import TokenCredential
|
||||
from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI
|
||||
from openai.lib.azure import AsyncAzureOpenAI
|
||||
from openai.types.chat.chat_completion import Choice
|
||||
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
|
||||
from pydantic import BaseModel
|
||||
@@ -29,6 +28,7 @@ from agent_framework.openai import OpenAIChatOptions
|
||||
from agent_framework.openai._chat_client import RawOpenAIChatClient
|
||||
|
||||
from .._settings import load_settings
|
||||
from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
|
||||
from ._shared import (
|
||||
AzureOpenAIConfigMixin,
|
||||
AzureOpenAISettings,
|
||||
@@ -169,10 +169,8 @@ class AzureOpenAIChatClient( # type: ignore[misc]
|
||||
endpoint: str | None = None,
|
||||
base_url: str | None = None,
|
||||
api_version: str | None = None,
|
||||
ad_token: str | None = None,
|
||||
ad_token_provider: AsyncAzureADTokenProvider | None = None,
|
||||
token_endpoint: str | None = None,
|
||||
credential: TokenCredential | None = None,
|
||||
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
|
||||
default_headers: Mapping[str, str] | None = None,
|
||||
async_client: AsyncAzureOpenAI | None = None,
|
||||
env_file_path: str | None = None,
|
||||
@@ -199,11 +197,12 @@ class AzureOpenAIChatClient( # type: ignore[misc]
|
||||
api_version: The deployment API version. If provided will override the value
|
||||
in the env vars or .env file.
|
||||
Can also be set via environment variable AZURE_OPENAI_API_VERSION.
|
||||
ad_token: The Azure Active Directory token.
|
||||
ad_token_provider: The Azure Active Directory token provider.
|
||||
token_endpoint: The token endpoint to request an Azure token.
|
||||
Can also be set via environment variable AZURE_OPENAI_TOKEN_ENDPOINT.
|
||||
credential: The Azure credential for authentication.
|
||||
credential: Azure credential or token provider for authentication. Accepts a
|
||||
``TokenCredential``, ``AsyncTokenCredential``, or a callable that returns a
|
||||
bearer token string (sync or async), for example from
|
||||
``azure.identity.get_bearer_token_provider()``.
|
||||
default_headers: The default headers mapping of string keys to
|
||||
string values for HTTP requests.
|
||||
async_client: An existing client to use.
|
||||
@@ -274,8 +273,6 @@ class AzureOpenAIChatClient( # type: ignore[misc]
|
||||
base_url=azure_openai_settings["base_url"],
|
||||
api_version=azure_openai_settings["api_version"], # type: ignore
|
||||
api_key=azure_openai_settings["api_key"].get_secret_value() if azure_openai_settings["api_key"] else None,
|
||||
ad_token=ad_token,
|
||||
ad_token_provider=ad_token_provider,
|
||||
token_endpoint=azure_openai_settings["token_endpoint"],
|
||||
credential=credential,
|
||||
default_headers=default_headers,
|
||||
|
||||
@@ -3,80 +3,66 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Union
|
||||
|
||||
from azure.core.exceptions import ClientAuthenticationError
|
||||
from azure.core.credentials import TokenCredential
|
||||
from azure.core.credentials_async import AsyncTokenCredential
|
||||
|
||||
from ..exceptions import ServiceInvalidAuthError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from azure.core.credentials import TokenCredential
|
||||
from azure.core.credentials_async import AsyncTokenCredential
|
||||
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
AzureTokenProvider = Callable[[], Union[str, Awaitable[str]]]
|
||||
"""A callable that returns a bearer token string, either synchronously or asynchronously."""
|
||||
|
||||
def get_entra_auth_token(
|
||||
credential: TokenCredential,
|
||||
token_endpoint: str,
|
||||
**kwargs: Any,
|
||||
) -> str | None:
|
||||
"""Retrieve a Microsoft Entra Auth Token for a given token endpoint.
|
||||
AzureCredentialTypes = Union[TokenCredential, AsyncTokenCredential]
|
||||
"""Union of Azure credential types.
|
||||
|
||||
The token endpoint may be specified as an environment variable, via the .env
|
||||
file or as an argument. If the token endpoint is not provided, the default is None.
|
||||
Accepts:
|
||||
- ``TokenCredential`` — synchronous Azure credential (e.g. ``DefaultAzureCredential()``)
|
||||
- ``AsyncTokenCredential`` — asynchronous Azure credential (e.g. ``azure.identity.aio.DefaultAzureCredential()``)
|
||||
"""
|
||||
|
||||
|
||||
def resolve_credential_to_token_provider(
|
||||
credential: AzureCredentialTypes | AzureTokenProvider,
|
||||
token_endpoint: str | None,
|
||||
) -> AzureTokenProvider:
|
||||
"""Convert an Azure credential or token provider into an ``ad_token_provider`` callable.
|
||||
|
||||
If the credential is already a callable token provider, it is returned as-is
|
||||
(``token_endpoint`` is not required in this case).
|
||||
If it is a ``TokenCredential`` or ``AsyncTokenCredential``, it is wrapped using
|
||||
``azure.identity.get_bearer_token_provider`` (sync or async variant) which
|
||||
handles token caching and automatic refresh.
|
||||
|
||||
Args:
|
||||
credential: The Azure credential to use for authentication.
|
||||
token_endpoint: The token endpoint to use to retrieve the authentication token.
|
||||
|
||||
Keyword Args:
|
||||
**kwargs: Additional keyword arguments to pass to the token retrieval method.
|
||||
credential: An Azure credential or token provider callable.
|
||||
token_endpoint: The token scope/endpoint
|
||||
(e.g. ``"https://cognitiveservices.azure.com/.default"``).
|
||||
Required when ``credential`` is a ``TokenCredential`` or ``AsyncTokenCredential``.
|
||||
|
||||
Returns:
|
||||
The Azure token or None if the token could not be retrieved.
|
||||
A callable that returns a bearer token string (sync or async).
|
||||
|
||||
Raises:
|
||||
ServiceInvalidAuthError: If the token endpoint is empty when needed for credential wrapping.
|
||||
"""
|
||||
# Already a token provider callable (not a credential object) — use directly
|
||||
if callable(credential) and not isinstance(credential, (TokenCredential, AsyncTokenCredential)):
|
||||
return credential
|
||||
|
||||
if not token_endpoint:
|
||||
raise ServiceInvalidAuthError(
|
||||
"A token endpoint must be provided either in settings, as an environment variable, or as an argument."
|
||||
)
|
||||
|
||||
try:
|
||||
auth_token = credential.get_token(token_endpoint, **kwargs)
|
||||
except ClientAuthenticationError as ex:
|
||||
logger.error(f"Failed to retrieve Azure token for the specified endpoint: `{token_endpoint}`, with error: {ex}")
|
||||
return None
|
||||
if isinstance(credential, AsyncTokenCredential):
|
||||
from azure.identity.aio import get_bearer_token_provider as get_async_bearer_token_provider
|
||||
|
||||
return auth_token.token if auth_token else None
|
||||
return get_async_bearer_token_provider(credential, token_endpoint)
|
||||
|
||||
from azure.identity import get_bearer_token_provider
|
||||
|
||||
async def get_entra_auth_token_async(
|
||||
credential: AsyncTokenCredential, token_endpoint: str, **kwargs: Any
|
||||
) -> str | None:
|
||||
"""Retrieve a async Microsoft Entra Auth Token for a given token endpoint.
|
||||
|
||||
The token endpoint may be specified as an environment variable, via the .env
|
||||
file or as an argument. If the token endpoint is not provided, the default is None.
|
||||
|
||||
Args:
|
||||
credential: The async Azure credential to use for authentication.
|
||||
token_endpoint: The token endpoint to use to retrieve the authentication token.
|
||||
|
||||
Keyword Args:
|
||||
**kwargs: Additional keyword arguments to pass to the token retrieval method.
|
||||
|
||||
Returns:
|
||||
The Azure token or None if the token could not be retrieved.
|
||||
"""
|
||||
if not token_endpoint:
|
||||
raise ServiceInvalidAuthError(
|
||||
"A token endpoint must be provided either in settings, as an environment variable, or as an argument."
|
||||
)
|
||||
|
||||
try:
|
||||
auth_token = await credential.get_token(token_endpoint, **kwargs)
|
||||
except ClientAuthenticationError as ex:
|
||||
logger.error(f"Failed to retrieve Azure token for the specified endpoint: `{token_endpoint}`, with error: {ex}")
|
||||
return None
|
||||
|
||||
return auth_token.token if auth_token else None
|
||||
return get_bearer_token_provider(credential, token_endpoint) # type: ignore[arg-type]
|
||||
|
||||
@@ -8,9 +8,7 @@ from typing import TYPE_CHECKING, Any, Generic
|
||||
from urllib.parse import urljoin, urlparse
|
||||
|
||||
from azure.ai.projects.aio import AIProjectClient
|
||||
from azure.core.credentials import TokenCredential
|
||||
from openai import AsyncOpenAI
|
||||
from openai.lib.azure import AsyncAzureADTokenProvider
|
||||
|
||||
from .._middleware import ChatMiddlewareLayer
|
||||
from .._settings import load_settings
|
||||
@@ -19,6 +17,7 @@ from .._tools import FunctionInvocationConfiguration, FunctionInvocationLayer
|
||||
from ..exceptions import ServiceInitializationError
|
||||
from ..observability import ChatTelemetryLayer
|
||||
from ..openai._responses_client import RawOpenAIResponsesClient
|
||||
from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
|
||||
from ._shared import (
|
||||
AzureOpenAIConfigMixin,
|
||||
AzureOpenAISettings,
|
||||
@@ -69,10 +68,8 @@ class AzureOpenAIResponsesClient( # type: ignore[misc]
|
||||
endpoint: str | None = None,
|
||||
base_url: str | None = None,
|
||||
api_version: str | None = None,
|
||||
ad_token: str | None = None,
|
||||
ad_token_provider: AsyncAzureADTokenProvider | None = None,
|
||||
token_endpoint: str | None = None,
|
||||
credential: TokenCredential | None = None,
|
||||
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
|
||||
default_headers: Mapping[str, str] | None = None,
|
||||
async_client: AsyncOpenAI | None = None,
|
||||
project_client: Any | None = None,
|
||||
@@ -109,11 +106,12 @@ class AzureOpenAIResponsesClient( # type: ignore[misc]
|
||||
api_version: The deployment API version. If provided will override the value
|
||||
in the env vars or .env file. Currently, the api_version must be "preview".
|
||||
Can also be set via environment variable AZURE_OPENAI_API_VERSION.
|
||||
ad_token: The Azure Active Directory token.
|
||||
ad_token_provider: The Azure Active Directory token provider.
|
||||
token_endpoint: The token endpoint to request an Azure token.
|
||||
Can also be set via environment variable AZURE_OPENAI_TOKEN_ENDPOINT.
|
||||
credential: The Azure credential for authentication.
|
||||
credential: Azure credential or token provider for authentication. Accepts a
|
||||
``TokenCredential``, ``AsyncTokenCredential``, or a callable that returns a
|
||||
bearer token string (sync or async), for example from
|
||||
``azure.identity.get_bearer_token_provider()``.
|
||||
default_headers: The default headers mapping of string keys to
|
||||
string values for HTTP requests.
|
||||
async_client: An existing client to use.
|
||||
@@ -230,8 +228,6 @@ class AzureOpenAIResponsesClient( # type: ignore[misc]
|
||||
base_url=azure_openai_settings["base_url"],
|
||||
api_version=azure_openai_settings["api_version"], # type: ignore
|
||||
api_key=azure_openai_settings["api_key"].get_secret_value() if azure_openai_settings["api_key"] else None,
|
||||
ad_token=ad_token,
|
||||
ad_token_provider=ad_token_provider,
|
||||
token_endpoint=azure_openai_settings["token_endpoint"],
|
||||
credential=credential,
|
||||
default_headers=default_headers,
|
||||
@@ -246,7 +242,7 @@ class AzureOpenAIResponsesClient( # type: ignore[misc]
|
||||
*,
|
||||
project_client: AIProjectClient | None,
|
||||
project_endpoint: str | None,
|
||||
credential: TokenCredential | None,
|
||||
credential: AzureCredentialTypes | AzureTokenProvider | None,
|
||||
) -> AsyncOpenAI:
|
||||
"""Create an AsyncOpenAI client from an Azure AI Foundry project.
|
||||
|
||||
|
||||
@@ -4,11 +4,10 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from collections.abc import Mapping
|
||||
from copy import copy
|
||||
from typing import Any, ClassVar, Final
|
||||
|
||||
from azure.core.credentials import TokenCredential
|
||||
from openai import AsyncOpenAI
|
||||
from openai.lib.azure import AsyncAzureOpenAI
|
||||
|
||||
@@ -16,7 +15,7 @@ from .._settings import SecretString
|
||||
from .._telemetry import APP_INFO, prepend_agent_framework_to_user_agent
|
||||
from ..exceptions import ServiceInitializationError
|
||||
from ..openai._shared import OpenAIBase
|
||||
from ._entra_id_authentication import get_entra_auth_token
|
||||
from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider, resolve_credential_to_token_provider
|
||||
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -135,10 +134,8 @@ class AzureOpenAIConfigMixin(OpenAIBase):
|
||||
base_url: str | None = None,
|
||||
api_version: str = DEFAULT_AZURE_API_VERSION,
|
||||
api_key: str | None = None,
|
||||
ad_token: str | None = None,
|
||||
ad_token_provider: Callable[[], str | Awaitable[str]] | None = None,
|
||||
token_endpoint: str | None = None,
|
||||
credential: TokenCredential | None = None,
|
||||
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
|
||||
default_headers: Mapping[str, str] | None = None,
|
||||
client: AsyncOpenAI | None = None,
|
||||
instruction_role: str | None = None,
|
||||
@@ -155,10 +152,10 @@ class AzureOpenAIConfigMixin(OpenAIBase):
|
||||
base_url: The base URL for Azure services.
|
||||
api_version: Azure API version. Defaults to the defined DEFAULT_AZURE_API_VERSION.
|
||||
api_key: API key for Azure services.
|
||||
ad_token: Azure AD token for authentication.
|
||||
ad_token_provider: A callable or coroutine function providing Azure AD tokens.
|
||||
token_endpoint: Azure AD token endpoint use to get the token.
|
||||
credential: Azure credential for authentication.
|
||||
token_endpoint: Azure AD token scope used to obtain a bearer token from a credential.
|
||||
credential: Azure credential or token provider for authentication. Accepts a
|
||||
``TokenCredential``, ``AsyncTokenCredential``, or a callable that returns a
|
||||
bearer token string (sync or async).
|
||||
default_headers: Default headers for HTTP requests.
|
||||
client: An existing client to use.
|
||||
instruction_role: The role to use for 'instruction' messages, for example, summarization
|
||||
@@ -172,16 +169,13 @@ class AzureOpenAIConfigMixin(OpenAIBase):
|
||||
merged_headers.update(APP_INFO)
|
||||
merged_headers = prepend_agent_framework_to_user_agent(merged_headers)
|
||||
if not client:
|
||||
# If the client is None, the api_key is none, the ad_token is none, and the ad_token_provider is none,
|
||||
# then we will attempt to get the ad_token using the default endpoint specified in the Azure OpenAI
|
||||
# settings.
|
||||
if not api_key and not ad_token_provider and not ad_token and token_endpoint and credential:
|
||||
ad_token = get_entra_auth_token(credential, token_endpoint)
|
||||
# Resolve credential to a token provider if needed
|
||||
ad_token_provider = None
|
||||
if not api_key and credential:
|
||||
ad_token_provider = resolve_credential_to_token_provider(credential, token_endpoint)
|
||||
|
||||
if not api_key and not ad_token and not ad_token_provider:
|
||||
raise ServiceInitializationError(
|
||||
"Please provide either api_key, ad_token or ad_token_provider or a client."
|
||||
)
|
||||
if not api_key and not ad_token_provider:
|
||||
raise ServiceInitializationError("Please provide either api_key, credential, or a client.")
|
||||
|
||||
if not endpoint and not base_url:
|
||||
raise ServiceInitializationError("Please provide an endpoint or a base_url")
|
||||
@@ -191,8 +185,6 @@ class AzureOpenAIConfigMixin(OpenAIBase):
|
||||
}
|
||||
if api_version:
|
||||
args["api_version"] = api_version
|
||||
if ad_token:
|
||||
args["azure_ad_token"] = ad_token
|
||||
if ad_token_provider:
|
||||
args["azure_ad_token_provider"] = ad_token_provider
|
||||
if api_key:
|
||||
|
||||
@@ -551,12 +551,16 @@ async def test_azure_assistants_client_agent_level_tool_persistence():
|
||||
|
||||
|
||||
def test_azure_assistants_client_entra_id_authentication() -> None:
|
||||
"""Test Entra ID authentication path with credential."""
|
||||
"""Test credential authentication path with sync credential."""
|
||||
mock_credential = MagicMock()
|
||||
mock_provider = MagicMock(return_value="token-string")
|
||||
|
||||
with (
|
||||
patch("agent_framework.azure._assistants_client.load_settings") as mock_load_settings,
|
||||
patch("agent_framework.azure._assistants_client.get_entra_auth_token") as mock_get_token,
|
||||
patch(
|
||||
"agent_framework.azure._assistants_client.resolve_credential_to_token_provider",
|
||||
return_value=mock_provider,
|
||||
) as mock_resolve,
|
||||
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
|
||||
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
|
||||
):
|
||||
@@ -564,28 +568,26 @@ def test_azure_assistants_client_entra_id_authentication() -> None:
|
||||
"chat_deployment_name": "test-deployment",
|
||||
"responses_deployment_name": None,
|
||||
"api_key": None,
|
||||
"token_endpoint": "https://login.microsoftonline.com/test",
|
||||
"token_endpoint": "https://cognitiveservices.azure.com/.default",
|
||||
"api_version": "2024-05-01-preview",
|
||||
"endpoint": "https://test-endpoint.openai.azure.com",
|
||||
"base_url": None,
|
||||
}
|
||||
mock_get_token.return_value = "entra-token-12345"
|
||||
|
||||
client = AzureOpenAIAssistantsClient(
|
||||
deployment_name="test-deployment",
|
||||
api_key="placeholder-key",
|
||||
endpoint="https://test-endpoint.openai.azure.com",
|
||||
credential=mock_credential,
|
||||
token_endpoint="https://login.microsoftonline.com/test",
|
||||
token_endpoint="https://cognitiveservices.azure.com/.default",
|
||||
)
|
||||
|
||||
# Verify Entra ID token was requested
|
||||
mock_get_token.assert_called_once_with(mock_credential, "https://login.microsoftonline.com/test")
|
||||
# Verify credential was resolved to a token provider
|
||||
mock_resolve.assert_called_once_with(mock_credential, "https://cognitiveservices.azure.com/.default")
|
||||
|
||||
# Verify client was created with the token
|
||||
# Verify client was created with the token provider
|
||||
mock_azure_client.assert_called_once()
|
||||
call_args = mock_azure_client.call_args[1]
|
||||
assert call_args["azure_ad_token"] == "entra-token-12345"
|
||||
assert call_args["azure_ad_token_provider"] is mock_provider
|
||||
|
||||
assert client is not None
|
||||
assert isinstance(client, AzureOpenAIAssistantsClient)
|
||||
@@ -605,7 +607,7 @@ def test_azure_assistants_client_no_authentication_error() -> None:
|
||||
}
|
||||
|
||||
# Test missing authentication raises error
|
||||
with pytest.raises(ServiceInitializationError, match="API key, ad_token, or ad_token_provider is required"):
|
||||
with pytest.raises(ServiceInitializationError, match="api_key, credential, or a client"):
|
||||
AzureOpenAIAssistantsClient(
|
||||
deployment_name="test-deployment",
|
||||
endpoint="https://test-endpoint.openai.azure.com",
|
||||
@@ -613,10 +615,16 @@ def test_azure_assistants_client_no_authentication_error() -> None:
|
||||
)
|
||||
|
||||
|
||||
def test_azure_assistants_client_ad_token_authentication() -> None:
|
||||
"""Test ad_token authentication client parameter path."""
|
||||
def test_azure_assistants_client_callable_credential() -> None:
|
||||
"""Test callable token provider as credential."""
|
||||
mock_provider = MagicMock(return_value="my-token")
|
||||
|
||||
with (
|
||||
patch("agent_framework.azure._assistants_client.load_settings") as mock_load_settings,
|
||||
patch(
|
||||
"agent_framework.azure._assistants_client.resolve_credential_to_token_provider",
|
||||
return_value=mock_provider,
|
||||
),
|
||||
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
|
||||
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
|
||||
):
|
||||
@@ -624,7 +632,7 @@ def test_azure_assistants_client_ad_token_authentication() -> None:
|
||||
"chat_deployment_name": "test-deployment",
|
||||
"responses_deployment_name": None,
|
||||
"api_key": None,
|
||||
"token_endpoint": None,
|
||||
"token_endpoint": "https://cognitiveservices.azure.com/.default",
|
||||
"api_version": "2024-05-01-preview",
|
||||
"endpoint": "https://test-endpoint.openai.azure.com",
|
||||
"base_url": None,
|
||||
@@ -633,49 +641,14 @@ def test_azure_assistants_client_ad_token_authentication() -> None:
|
||||
client = AzureOpenAIAssistantsClient(
|
||||
deployment_name="test-deployment",
|
||||
endpoint="https://test-endpoint.openai.azure.com",
|
||||
ad_token="test-ad-token-12345",
|
||||
credential=mock_provider,
|
||||
token_endpoint="https://cognitiveservices.azure.com/.default",
|
||||
)
|
||||
|
||||
# ad_token path
|
||||
# Verify client was created with the token provider
|
||||
mock_azure_client.assert_called_once()
|
||||
call_args = mock_azure_client.call_args[1]
|
||||
assert call_args["azure_ad_token"] == "test-ad-token-12345"
|
||||
|
||||
assert client is not None
|
||||
assert isinstance(client, AzureOpenAIAssistantsClient)
|
||||
|
||||
|
||||
def test_azure_assistants_client_ad_token_provider_authentication() -> None:
|
||||
"""Test ad_token_provider authentication client parameter path."""
|
||||
from openai.lib.azure import AsyncAzureADTokenProvider
|
||||
|
||||
mock_token_provider = MagicMock(spec=AsyncAzureADTokenProvider)
|
||||
|
||||
with (
|
||||
patch("agent_framework.azure._assistants_client.load_settings") as mock_load_settings,
|
||||
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
|
||||
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
|
||||
):
|
||||
mock_load_settings.return_value = {
|
||||
"chat_deployment_name": "test-deployment",
|
||||
"responses_deployment_name": None,
|
||||
"api_key": None,
|
||||
"token_endpoint": None,
|
||||
"api_version": "2024-05-01-preview",
|
||||
"endpoint": "https://test-endpoint.openai.azure.com",
|
||||
"base_url": None,
|
||||
}
|
||||
|
||||
client = AzureOpenAIAssistantsClient(
|
||||
deployment_name="test-deployment",
|
||||
endpoint="https://test-endpoint.openai.azure.com",
|
||||
ad_token_provider=mock_token_provider,
|
||||
)
|
||||
|
||||
# ad_token_provider path
|
||||
mock_azure_client.assert_called_once()
|
||||
call_args = mock_azure_client.call_args[1]
|
||||
assert call_args["azure_ad_token_provider"] is mock_token_provider
|
||||
assert call_args["azure_ad_token_provider"] is mock_provider
|
||||
|
||||
assert client is not None
|
||||
assert isinstance(client, AzureOpenAIAssistantsClient)
|
||||
|
||||
@@ -1,156 +1,61 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from azure.core.exceptions import ClientAuthenticationError
|
||||
from azure.core.credentials import TokenCredential
|
||||
from azure.core.credentials_async import AsyncTokenCredential
|
||||
|
||||
from agent_framework.azure._entra_id_authentication import (
|
||||
get_entra_auth_token,
|
||||
get_entra_auth_token_async,
|
||||
resolve_credential_to_token_provider,
|
||||
)
|
||||
from agent_framework.exceptions import ServiceInvalidAuthError
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_credential() -> MagicMock:
|
||||
"""Mock synchronous TokenCredential."""
|
||||
mock_cred = MagicMock()
|
||||
# Create a mock token object with a .token attribute
|
||||
mock_token = MagicMock()
|
||||
mock_token.token = "test-access-token-12345"
|
||||
mock_cred.get_token.return_value = mock_token
|
||||
return mock_cred
|
||||
TOKEN_ENDPOINT = "https://cognitiveservices.azure.com/.default"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_async_credential() -> MagicMock:
|
||||
"""Mock asynchronous AsyncTokenCredential."""
|
||||
mock_cred = MagicMock()
|
||||
# Create a mock token object with a .token attribute
|
||||
mock_token = MagicMock()
|
||||
mock_token.token = "test-async-access-token-12345"
|
||||
mock_cred.get_token = AsyncMock(return_value=mock_token)
|
||||
return mock_cred
|
||||
def test_resolve_sync_credential_returns_provider() -> None:
|
||||
"""Test that a sync TokenCredential is resolved via azure.identity.get_bearer_token_provider."""
|
||||
mock_credential = MagicMock(spec=TokenCredential)
|
||||
mock_provider = MagicMock(return_value="token-string")
|
||||
|
||||
with patch("azure.identity.get_bearer_token_provider", return_value=mock_provider) as mock_gbtp:
|
||||
result = resolve_credential_to_token_provider(mock_credential, TOKEN_ENDPOINT)
|
||||
|
||||
mock_gbtp.assert_called_once_with(mock_credential, TOKEN_ENDPOINT)
|
||||
assert result is mock_provider
|
||||
|
||||
|
||||
def test_get_entra_auth_token_success(mock_credential: MagicMock) -> None:
|
||||
"""Test successful token retrieval with sync function."""
|
||||
def test_resolve_async_credential_returns_provider() -> None:
|
||||
"""Test that an AsyncTokenCredential is resolved via azure.identity.aio.get_bearer_token_provider."""
|
||||
mock_credential = MagicMock(spec=AsyncTokenCredential)
|
||||
mock_provider = MagicMock(return_value="token-string")
|
||||
|
||||
token_endpoint = "https://test-endpoint.com/.default"
|
||||
with patch("azure.identity.aio.get_bearer_token_provider", return_value=mock_provider) as mock_gbtp:
|
||||
result = resolve_credential_to_token_provider(mock_credential, TOKEN_ENDPOINT)
|
||||
|
||||
result = get_entra_auth_token(mock_credential, token_endpoint)
|
||||
|
||||
# Assert - check the results
|
||||
assert result == "test-access-token-12345"
|
||||
mock_credential.get_token.assert_called_once_with(token_endpoint)
|
||||
mock_gbtp.assert_called_once_with(mock_credential, TOKEN_ENDPOINT)
|
||||
assert result is mock_provider
|
||||
|
||||
|
||||
async def test_get_entra_auth_token_async_success(mock_async_credential: MagicMock) -> None:
|
||||
"""Test successful token retrieval with async function."""
|
||||
def test_resolve_callable_provider_passthrough() -> None:
|
||||
"""Test that a callable token provider is returned as-is, without needing token_endpoint."""
|
||||
my_provider = lambda: "my-token" # noqa: E731
|
||||
|
||||
token_endpoint = "https://test-endpoint.com/.default"
|
||||
# Works with token_endpoint
|
||||
assert resolve_credential_to_token_provider(my_provider, TOKEN_ENDPOINT) is my_provider
|
||||
|
||||
result = await get_entra_auth_token_async(mock_async_credential, token_endpoint)
|
||||
|
||||
# Assert - check the results
|
||||
assert result == "test-async-access-token-12345"
|
||||
mock_async_credential.get_token.assert_called_once_with(token_endpoint)
|
||||
# Also works without token_endpoint
|
||||
assert resolve_credential_to_token_provider(my_provider, None) is my_provider
|
||||
assert resolve_credential_to_token_provider(my_provider, "") is my_provider
|
||||
|
||||
|
||||
def test_get_entra_auth_token_missing_endpoint(mock_credential: MagicMock) -> None:
|
||||
def test_resolve_missing_endpoint_raises() -> None:
|
||||
"""Test that missing token endpoint raises ServiceInvalidAuthError."""
|
||||
# Test with empty string
|
||||
mock_credential = MagicMock(spec=TokenCredential)
|
||||
|
||||
with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
|
||||
get_entra_auth_token(mock_credential, "")
|
||||
resolve_credential_to_token_provider(mock_credential, "")
|
||||
|
||||
# Test with None
|
||||
with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
|
||||
get_entra_auth_token(mock_credential, None) # type: ignore
|
||||
|
||||
|
||||
async def test_get_entra_auth_token_async_missing_endpoint(mock_async_credential: MagicMock) -> None:
|
||||
"""Test that missing token endpoint raises ServiceInvalidAuthError in async function."""
|
||||
# Test with empty string
|
||||
with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
|
||||
await get_entra_auth_token_async(mock_async_credential, "")
|
||||
|
||||
# Test with None
|
||||
with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
|
||||
await get_entra_auth_token_async(mock_async_credential, None) # type: ignore
|
||||
|
||||
|
||||
def test_get_entra_auth_token_auth_failure(mock_credential: MagicMock) -> None:
|
||||
"""Test that Azure authentication failure returns None."""
|
||||
|
||||
mock_credential.get_token.side_effect = ClientAuthenticationError("Auth failed")
|
||||
token_endpoint = "https://test-endpoint.com/.default"
|
||||
|
||||
result = get_entra_auth_token(mock_credential, token_endpoint)
|
||||
|
||||
# Assert - should return None on auth failure
|
||||
assert result is None
|
||||
mock_credential.get_token.assert_called_once_with(token_endpoint)
|
||||
|
||||
|
||||
async def test_get_entra_auth_token_async_auth_failure(mock_async_credential: MagicMock) -> None:
|
||||
"""Test that Azure authentication failure returns None in async function."""
|
||||
|
||||
mock_async_credential.get_token.side_effect = ClientAuthenticationError("Auth failed")
|
||||
token_endpoint = "https://test-endpoint.com/.default"
|
||||
|
||||
result = await get_entra_auth_token_async(mock_async_credential, token_endpoint)
|
||||
|
||||
# Assert - should return None on auth failure
|
||||
assert result is None
|
||||
mock_async_credential.get_token.assert_called_once_with(token_endpoint)
|
||||
|
||||
|
||||
def test_get_entra_auth_token_none_token_response(mock_credential: MagicMock) -> None:
|
||||
"""Test that None token response returns None."""
|
||||
mock_credential.get_token.return_value = None
|
||||
token_endpoint = "https://test-endpoint.com/.default"
|
||||
|
||||
result = get_entra_auth_token(mock_credential, token_endpoint)
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
mock_credential.get_token.assert_called_once_with(token_endpoint)
|
||||
|
||||
|
||||
async def test_get_entra_auth_token_async_none_token_response(mock_async_credential: MagicMock) -> None:
|
||||
"""Test that None token response returns None in async function."""
|
||||
mock_async_credential.get_token.return_value = None
|
||||
token_endpoint = "https://test-endpoint.com/.default"
|
||||
|
||||
result = await get_entra_auth_token_async(mock_async_credential, token_endpoint)
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
mock_async_credential.get_token.assert_called_once_with(token_endpoint)
|
||||
|
||||
|
||||
def test_get_entra_auth_token_with_kwargs(mock_credential: MagicMock) -> None:
|
||||
"""Test that kwargs are passed through to get_token."""
|
||||
|
||||
token_endpoint = "https://test-endpoint.com/.default"
|
||||
extra_kwargs = {"scopes": ["read", "write"], "tenant_id": "test-tenant"}
|
||||
|
||||
result = get_entra_auth_token(mock_credential, token_endpoint, **extra_kwargs)
|
||||
|
||||
# Assert
|
||||
assert result == "test-access-token-12345"
|
||||
mock_credential.get_token.assert_called_once_with(token_endpoint, **extra_kwargs)
|
||||
|
||||
|
||||
async def test_get_entra_auth_token_async_with_kwargs(mock_async_credential: MagicMock) -> None:
|
||||
"""Test that kwargs are passed through to async get_token."""
|
||||
|
||||
token_endpoint = "https://test-endpoint.com/.default"
|
||||
extra_kwargs = {"scopes": ["read", "write"], "tenant_id": "test-tenant"}
|
||||
|
||||
result = await get_entra_auth_token_async(mock_async_credential, token_endpoint, **extra_kwargs)
|
||||
|
||||
# Assert
|
||||
assert result == "test-async-access-token-12345"
|
||||
mock_async_credential.get_token.assert_called_once_with(token_endpoint, **extra_kwargs)
|
||||
resolve_credential_to_token_provider(mock_credential, None) # type: ignore[arg-type]
|
||||
|
||||
Reference in New Issue
Block a user