Skip to content

Python: Unify Azure credential handling across all packages #4087

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
eavanvalkenburg wants to merge 1 commit into from
Closed

Python: Unify Azure credential handling across all packages #4087

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
8 changes: 5 additions & 3 deletions python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, Message
from agent_framework._sessions import AgentSession, BaseContextProvider, SessionContext
from agent_framework._settings import SecretString, load_settings
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
from agent_framework.exceptions import ServiceInitializationError
from azure.core.credentials import AzureKeyCredential
from azure.core.credentials_async import AsyncTokenCredential
Expand Down Expand Up @@ -148,7 +149,7 @@ def __init__(
endpoint: str | None = None,
index_name: str | None = None,
api_key: str | AzureKeyCredential | None = None,
credential: AsyncTokenCredential | None = None,
credential: AzureCredentialTypes | None = None,
*,
mode: Literal["semantic", "agentic"] = "semantic",
top_k: int = 5,
Expand All @@ -175,7 +176,8 @@ def __init__(
endpoint: Azure AI Search endpoint URL.
index_name: Name of the search index to query.
api_key: API key for authentication.
credential: AsyncTokenCredential for managed identity authentication.
credential: Azure credential for managed identity authentication.
Accepts a TokenCredential, AsyncTokenCredential, or a callable token provider.
mode: Search mode - "semantic" or "agentic". Default: "semantic".
top_k: Maximum number of documents to retrieve. Default: 5.
semantic_configuration_name: Name of semantic configuration in the index.
Expand Down Expand Up @@ -223,7 +225,7 @@ def __init__(

resolved_credential: AzureKeyCredential | AsyncTokenCredential
if credential:
resolved_credential = credential
resolved_credential = credential # type: ignore[assignment]
Copy link

Copilot AI Feb 19, 2026

Choose a reason for hiding this comment

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

AzureAISearchContextProvider now types credential as AzureCredentialTypes (includes sync TokenCredential) and the docstring also mentions callable token providers, but the implementation constructs azure.search.documents.aio.SearchClient and annotates resolved_credential as AzureKeyCredential | AsyncTokenCredential. This is internally inconsistent (note the # type: ignore[assignment]) and may allow passing unsupported credential types to the async SearchClient. Recommend (a) removing the “callable token provider” claim, and (b) narrowing credential back to the credential type(s) that the async SearchClient actually supports, so the type: ignore and mismatched annotations can be removed.

Suggested change
resolved_credential = credential # type: ignore[assignment]
if isinstance(credential, (AzureKeyCredential, AsyncTokenCredential)):
resolved_credential = credential
else:
raise ServiceInitializationError(
"Unsupported credential type for Azure AI Search async client. "
"Use AzureKeyCredential or an AsyncTokenCredential."
)

Copilot uses AI. Check for mistakes.
elif isinstance(api_key, AzureKeyCredential):
resolved_credential = api_key
elif settings.get("api_key"):
Expand Down
9 changes: 5 additions & 4 deletions python/packages/azure-ai/agent_framework_azure_ai/_agent_provider.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,11 @@
from agent_framework._mcp import MCPTool
from agent_framework._settings import load_settings
from agent_framework._tools import ToolTypes
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
from agent_framework.exceptions import ServiceInitializationError
from azure.ai.agents.aio import AgentsClient
from azure.ai.agents.models import Agent as AzureAgent
from azure.ai.agents.models import ResponseFormatJsonSchema, ResponseFormatJsonSchemaType
from azure.core.credentials_async import AsyncTokenCredential
from pydantic import BaseModel

from ._chat_client import AzureAIAgentClient, AzureAIAgentOptions
Expand Down Expand Up @@ -93,7 +93,7 @@ def __init__(
agents_client: AgentsClient | None = None,
*,
project_endpoint: str | None = None,
credential: AsyncTokenCredential | None = None,
credential: AzureCredentialTypes | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
Expand All @@ -106,7 +106,8 @@ def __init__(
Keyword Args:
project_endpoint: The Azure AI Project endpoint URL.
Can also be set via AZURE_AI_PROJECT_ENDPOINT environment variable.
credential: Azure async credential for authentication.
credential: Azure credential for authentication. Accepts a TokenCredential,
AsyncTokenCredential, or a callable token provider.
Comment on lines +109 to +110
Copy link

Copilot AI Feb 19, 2026

Choose a reason for hiding this comment

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

The docstring states credential can be a “callable token provider”, but the parameter type is AzureCredentialTypes | None (credentials only). Please update the docstring to reflect the actual supported types, or expand the accepted types and add explicit support (right now the callable case isn’t handled).

Suggested change
credential: Azure credential for authentication. Accepts a TokenCredential,
AsyncTokenCredential, or a callable token provider.
credential: Azure credential for authentication. Accepts types compatible with
AzureCredentialTypes (for example, a TokenCredential or AsyncTokenCredential).

Copilot uses AI. Check for mistakes.
Required if agents_client is not provided.
env_file_path: Path to .env file for loading settings.
env_file_encoding: Encoding of the .env file.
Expand Down Expand Up @@ -137,7 +138,7 @@ def __init__(
raise ServiceInitializationError("Azure credential is required when agents_client is not provided.")
self._agents_client = AgentsClient(
endpoint=resolved_endpoint,
credential=credential,
credential=credential, # type: ignore[arg-type]
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
self._should_close_client = True
Expand Down
9 changes: 5 additions & 4 deletions python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
)
from agent_framework._settings import load_settings
from agent_framework._tools import ToolTypes
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
from agent_framework.exceptions import ServiceInitializationError, ServiceInvalidRequestError, ServiceResponseException
from agent_framework.observability import ChatTelemetryLayer
from azure.ai.agents.aio import AgentsClient
Expand Down Expand Up @@ -84,7 +85,6 @@
ToolDefinition,
ToolOutput,
)
from azure.core.credentials_async import AsyncTokenCredential
from pydantic import BaseModel

from ._shared import AzureAISettings, to_azure_ai_agent_tools
Expand Down Expand Up @@ -415,7 +415,7 @@ def __init__(
thread_id: str | None = None,
project_endpoint: str | None = None,
model_deployment_name: str | None = None,
credential: AsyncTokenCredential | None = None,
credential: AzureCredentialTypes | None = None,
should_cleanup_agent: bool = True,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
Expand All @@ -439,7 +439,8 @@ def __init__(
Ignored when a agents_client is passed.
model_deployment_name: The model deployment name to use for agent creation.
Can also be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME.
credential: Azure async credential to use for authentication.
credential: Azure credential for authentication. Accepts a TokenCredential,
AsyncTokenCredential, or a callable token provider.
Comment on lines +442 to +443
Copy link

Copilot AI Feb 19, 2026

Choose a reason for hiding this comment

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

The docstring claims credential accepts a “callable token provider”, but the parameter type is AzureCredentialTypes | None. This mismatch is user-facing and can lead to runtime errors if callers pass a callable. Either remove the callable claim or update the signature and implement a supported adapter type.

Suggested change
credential: Azure credential for authentication. Accepts a TokenCredential,
AsyncTokenCredential, or a callable token provider.
credential: Azure credential for authentication. Accepts a TokenCredential
or AsyncTokenCredential.

Copilot uses AI. Check for mistakes.
should_cleanup_agent: Whether to cleanup (delete) agents created by this client when
the client is closed or context is exited. Defaults to True. Only affects agents
created by this client instance; existing agents passed via agent_id are never deleted.
Expand Down Expand Up @@ -513,7 +514,7 @@ class MyOptions(AzureAIAgentOptions, total=False):
raise ServiceInitializationError("Azure credential is required when agents_client is not provided.")
agents_client = AgentsClient(
endpoint=resolved_endpoint,
credential=credential,
credential=credential, # type: ignore[arg-type]
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
should_close_client = True
Expand Down
14 changes: 8 additions & 6 deletions python/packages/azure-ai/agent_framework_azure_ai/_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
)
from agent_framework._settings import load_settings
from agent_framework._tools import ToolTypes
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.observability import ChatTelemetryLayer
from agent_framework.openai import OpenAIResponsesOptions
Expand All @@ -41,7 +42,6 @@
WebSearchPreviewTool,
)
from azure.ai.projects.models import FileSearchTool as ProjectsFileSearchTool
from azure.core.credentials_async import AsyncTokenCredential
from azure.core.exceptions import ResourceNotFoundError

from ._shared import AzureAISettings, create_text_format_config
Expand Down Expand Up @@ -109,7 +109,7 @@ def __init__(
conversation_id: str | None = None,
project_endpoint: str | None = None,
model_deployment_name: str | None = None,
credential: AsyncTokenCredential | None = None,
credential: AzureCredentialTypes | None = None,
use_latest_version: bool | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
Expand All @@ -132,7 +132,8 @@ def __init__(
Ignored when a project_client is passed.
model_deployment_name: The model deployment name to use for agent creation.
Can also be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME.
credential: Azure async credential to use for authentication.
credential: Azure credential for authentication. Accepts a TokenCredential,
AsyncTokenCredential, or a callable token provider.
Comment on lines +135 to +136
Copy link

Copilot AI Feb 19, 2026

Choose a reason for hiding this comment

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

The docstring says credential accepts “TokenCredential, AsyncTokenCredential, or a callable token provider”, but the type is AzureCredentialTypes | None (no token provider). Either update the docstring to match (credential objects only) or expand the accepted types and implement a safe adapter for callables.

Suggested change
credential: Azure credential for authentication. Accepts a TokenCredential,
AsyncTokenCredential, or a callable token provider.
credential: Azure credential for authentication. Accepts an AzureCredentialTypes
instance, such as a TokenCredential or AsyncTokenCredential.

Copilot uses AI. Check for mistakes.
use_latest_version: Boolean flag that indicates whether to use latest agent version
if it exists in the service.
env_file_path: Path to environment file for loading settings.
Expand Down Expand Up @@ -197,7 +198,7 @@ class MyOptions(ChatOptions, total=False):
raise ServiceInitializationError("Azure credential is required when project_client is not provided.")
project_client = AIProjectClient(
endpoint=resolved_endpoint,
credential=credential,
credential=credential, # type: ignore[arg-type]
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
should_close_client = True
Expand Down Expand Up @@ -950,7 +951,7 @@ def __init__(
conversation_id: str | None = None,
project_endpoint: str | None = None,
model_deployment_name: str | None = None,
credential: AsyncTokenCredential | None = None,
credential: AzureCredentialTypes | None = None,
use_latest_version: bool | None = None,
middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
Expand All @@ -972,7 +973,8 @@ def __init__(
Ignored when a project_client is passed.
model_deployment_name: The model deployment name to use for agent creation.
Can also be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME.
credential: Azure async credential to use for authentication.
credential: Azure credential for authentication. Accepts a TokenCredential
or AsyncTokenCredential.
use_latest_version: Boolean flag that indicates whether to use latest agent version
if it exists in the service.
middleware: Optional sequence of chat middlewares to include.
Expand Down
9 changes: 5 additions & 4 deletions python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from agent_framework._mcp import MCPTool
from agent_framework._settings import load_settings
from agent_framework._tools import ToolTypes
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
from agent_framework.exceptions import ServiceInitializationError
from azure.ai.projects.aio import AIProjectClient
from azure.ai.projects.models import (
Expand All @@ -29,7 +30,6 @@
from azure.ai.projects.models import (
FunctionTool as AzureFunctionTool,
)
from azure.core.credentials_async import AsyncTokenCredential

from ._client import AzureAIClient, AzureAIProjectAgentOptions
from ._shared import AzureAISettings, create_text_format_config, from_azure_ai_tools, to_azure_ai_tools
Expand Down Expand Up @@ -103,7 +103,7 @@ def __init__(
*,
project_endpoint: str | None = None,
model: str | None = None,
credential: AsyncTokenCredential | None = None,
credential: AzureCredentialTypes | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
Expand All @@ -116,7 +116,8 @@ def __init__(
Ignored when a project_client is passed.
model: The default model deployment name to use for agent creation.
Can also be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME.
credential: Azure async credential to use for authentication.
credential: Azure credential for authentication. Accepts a TokenCredential,
AsyncTokenCredential, or a callable token provider.
Comment on lines +119 to +120
Copy link

Copilot AI Feb 19, 2026

Choose a reason for hiding this comment

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

The docstring says credential accepts “TokenCredential, AsyncTokenCredential, or a callable token provider”, but the parameter is typed as AzureCredentialTypes | None (no token provider). Please align the docs with the actual supported types to avoid misleading users.

Suggested change
credential: Azure credential for authentication. Accepts a TokenCredential,
AsyncTokenCredential, or a callable token provider.
credential: Azure credential for authentication. Accepts credentials supported by
AzureCredentialTypes (for example, TokenCredential or AsyncTokenCredential).

Copilot uses AI. Check for mistakes.
Required when project_client is not provided.
env_file_path: Path to environment file for loading settings.
env_file_encoding: Encoding of the environment file.
Expand Down Expand Up @@ -149,7 +150,7 @@ def __init__(

project_client = AIProjectClient(
endpoint=resolved_endpoint,
credential=credential,
credential=credential, # type: ignore[arg-type]
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
self._should_close_client = True
Expand Down
3 changes: 2 additions & 1 deletion python/packages/core/agent_framework/azure/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@
"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"),
Expand All @@ -42,7 +44,6 @@
"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"),
}


Expand Down
5 changes: 3 additions & 2 deletions python/packages/core/agent_framework/azure/__init__.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -37,13 +37,14 @@ __all__ = [
"AzureAISearchContextProvider",
"AzureAISearchSettings",
"AzureAISettings",
"AzureCredentialTypes",
"AzureOpenAIAssistantsClient",
"AzureOpenAIChatClient",
"AzureOpenAIResponsesClient",
"AzureOpenAISettings",
"AzureTokenProvider",
"DurableAIAgent",
"DurableAIAgentClient",
"DurableAIAgentOrchestrationContext",
"DurableAIAgentWorker",
"get_entra_auth_token",
]
47 changes: 18 additions & 29 deletions python/packages/core/agent_framework/azure/_assistants_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,19 +4,16 @@

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
Expand Down Expand Up @@ -61,10 +58,8 @@ def __init__(
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,
Expand Down Expand Up @@ -93,11 +88,12 @@ def __init__(
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.
Expand Down Expand Up @@ -156,20 +152,15 @@ class MyOptions(AzureOpenAIAssistantsOptions, total=False):
"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)

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.")
# 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_provider:
raise ServiceInitializationError("Please provide either api_key, credential, or a client.")

# Create Azure client if not provided
if not async_client:
Expand All @@ -180,8 +171,6 @@ class MyOptions(AzureOpenAIAssistantsOptions, total=False):

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

Expand Down
Loading