autogen_ext.auth.azure 源代码

from typing import List

from autogen_core import Component, ComponentBase
from pydantic import BaseModel
from typing_extensions import Self

from azure.core.credentials import TokenProvider
from azure.identity import DefaultAzureCredential, get_bearer_token_provider




class TokenProviderConfig(BaseModel):
    provider_kind: str
    scopes: List[str]





class AzureTokenProvider(ComponentBase[TokenProviderConfig], Component[TokenProviderConfig]):
    component_type = "token_provider"
    component_config_schema = TokenProviderConfig
    component_provider_override = "autogen_ext.auth.azure.AzureTokenProvider"

    def __init__(self, credential: TokenProvider, *scopes: str):
        self.credential = credential
        self.scopes = list(scopes)
        self.provider = get_bearer_token_provider(self.credential, *self.scopes)

    def __call__(self) -> str:
        return self.provider()

    def _to_config(self) -> TokenProviderConfig:
        """转储用于创建与当前实例配置匹配的组件新实例所需的配置。

        Returns:
            T: 该组件的配置信息。
        """

        if isinstance(self.credential, DefaultAzureCredential):
            # NOTE: we are not currently inspecting the chained credentials, so this could result in a loss of information
            return TokenProviderConfig(provider_kind="DefaultAzureCredential", scopes=self.scopes)
        else:
            raise ValueError("Only DefaultAzureCredential is supported")

    @classmethod
    def _from_config(cls, config: TokenProviderConfig) -> Self:
        """根据配置对象创建组件的新实例。

        Args:
            config (T): 配置对象。

        Returns:
            Self: 组件的新实例。
        """

        if config.provider_kind == "DefaultAzureCredential":
            return cls(DefaultAzureCredential(), *config.scopes)
        else:
            raise ValueError("Only DefaultAzureCredential is supported")