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:
co-authored by
Copilot
parent
4c8f595019
commit
fd4e6e816c
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user