first commit

agent_dhal/agentdhal_extensions/auth/azure/__init__.py

@@ -0,0 +1,56 @@
+from typing import List
+
+from agentdhal_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 = "agentdhal_extensions.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:
+        """Dump the configuration that would be requite to create a new instance of a component matching the configuration of this instance.
+
+        Returns:
+            T: The configuration of the component.
+        """
+
+        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:
+        """Create a new instance of the component from a configuration object.
+
+        Args:
+            config (T): The configuration object.
+
+        Returns:
+            Self: The new instance of the component.
+        """
+
+        if config.provider_kind == "DefaultAzureCredential":
+            return cls(DefaultAzureCredential(), *config.scopes)
+        else:
+            raise ValueError("Only DefaultAzureCredential is supported")