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
2026-02-19 16:30:16 +00:00
committed by GitHub
co-authored by Copilot
parent 4c8f595019
commit fd4e6e816c
16 changed files with 200 additions and 348 deletions
@@ -11,6 +11,7 @@ from uuid import uuid4
import httpx
from agent_framework import AGENT_FRAMEWORK_USER_AGENT
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
from agent_framework.observability import get_tracer
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
@@ -39,18 +40,19 @@ logger = logging.getLogger("agent_framework.purview")
class PurviewClient:
"""Async client for calling Graph Purview endpoints.
Supports both synchronous TokenCredential and asynchronous AsyncTokenCredential implementations.
A sync credential will be invoked in a thread to avoid blocking the event loop.
Supports synchronous TokenCredential, asynchronous AsyncTokenCredential,
or callable token providers. A sync credential will be invoked in a thread
to avoid blocking the event loop.
"""
def __init__(
self,
credential: TokenCredential | AsyncTokenCredential,
credential: AzureCredentialTypes | AzureTokenProvider,
settings: PurviewSettings,
*,
timeout: float | None = 10.0,
):
self._credential: TokenCredential | AsyncTokenCredential = credential
self._credential: AzureCredentialTypes | AzureTokenProvider = credential
self._settings = settings
self._graph_uri = (settings.get("graph_base_uri") or "https://graph.microsoft.com/v1.0/").rstrip("/")
self._timeout = timeout
@@ -60,10 +62,14 @@ class PurviewClient:
await self._client.aclose()
async def _get_token(self, *, tenant_id: str | None = None) -> str:
"""Acquire an access token using either async or sync credential."""
scopes = get_purview_scopes(self._settings)
"""Acquire an access token using either async or sync credential, or callable token provider."""
cred = self._credential
token = cred.get_token(*scopes, tenant_id=tenant_id)
# Callable token provider — returns a token string directly
if callable(cred) and not isinstance(cred, (TokenCredential, AsyncTokenCredential)):
result = cred()
return await result if inspect.isawaitable(result) else result # type: ignore[return-value]
scopes = get_purview_scopes(self._settings)
token = cred.get_token(*scopes, tenant_id=tenant_id) # type: ignore[union-attr]
token = await token if inspect.isawaitable(token) else token
return token.token
@@ -4,8 +4,7 @@ import logging
from collections.abc import Awaitable, Callable
from agent_framework import AgentContext, AgentMiddleware, ChatContext, ChatMiddleware, MiddlewareTermination
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
from ._cache import CacheProvider
from ._client import PurviewClient
@@ -20,7 +19,7 @@ logger = logging.getLogger("agent_framework.purview")
class PurviewPolicyMiddleware(AgentMiddleware):
"""Agent middleware that enforces Purview policies on prompt and response.
Accepts either a synchronous TokenCredential or an AsyncTokenCredential.
Accepts a TokenCredential, AsyncTokenCredential, or callable token provider.
Usage:
@@ -28,14 +27,14 @@ class PurviewPolicyMiddleware(AgentMiddleware):
from agent_framework.microsoft import PurviewPolicyMiddleware, PurviewSettings
from agent_framework import Agent
credential = ... # TokenCredential or AsyncTokenCredential
credential = ... # TokenCredential, AsyncTokenCredential, or callable
settings = PurviewSettings(app_name="My App")
agent = Agent(client=client, instructions="...", middleware=[PurviewPolicyMiddleware(credential, settings)])
"""
def __init__(
self,
credential: TokenCredential | AsyncTokenCredential,
credential: AzureCredentialTypes | AzureTokenProvider,
settings: PurviewSettings,
cache_provider: CacheProvider | None = None,
) -> None:
@@ -153,14 +152,14 @@ class PurviewChatPolicyMiddleware(ChatMiddleware):
from agent_framework.microsoft import PurviewChatPolicyMiddleware, PurviewSettings
from agent_framework import ChatClient
credential = ... # TokenCredential or AsyncTokenCredential
credential = ... # TokenCredential, AsyncTokenCredential, or callable
settings = PurviewSettings(app_name="My App")
client = ChatClient(..., middleware=[PurviewChatPolicyMiddleware(credential, settings)])
"""
def __init__(
self,
credential: TokenCredential | AsyncTokenCredential,
credential: AzureCredentialTypes | AzureTokenProvider,
settings: PurviewSettings,
cache_provider: CacheProvider | None = None,
) -> None: