hexdag 0.5.0.dev1__py3-none-any.whl

hexdag_plugins/azure/azure_cosmos_adapter.py

@@ -0,0 +1,383 @@
+"""Azure Cosmos DB adapter for hexDAG framework.
+
+Provides persistent memory and state storage for agents and pipelines.
+"""
+
+import os
+import time
+from typing import Any
+
+from hexdag.core.ports.healthcheck import HealthStatus
+from hexdag.core.ports.memory import Memory
+
+
+class AzureCosmosAdapter(Memory):
+    """Azure Cosmos DB adapter for agent memory and pipeline state.
+
+    Provides persistent, scalable storage for agent memory, conversation
+    history, and pipeline execution state using Azure Cosmos DB.
+
+    Parameters
+    ----------
+    endpoint : str
+        Azure Cosmos DB endpoint URL
+    key : str
+        Azure Cosmos DB primary or secondary key (auto-resolved from AZURE_COSMOS_KEY)
+    database_name : str
+        Database name (default: "hexdag")
+    container_name : str
+        Container name (default: "memory")
+    partition_key : str
+        Partition key path (default: "/agent_id")
+    use_managed_identity : bool
+        Use Managed Identity instead of key auth (default: False)
+    throughput : int
+        Container throughput in RU/s (default: 400)
+
+    Examples
+    --------
+    YAML configuration::
+
+        spec:
+          ports:
+            memory:
+              adapter: hexdag_plugins.azure.AzureCosmosAdapter
+              config:
+                endpoint: ${AZURE_COSMOS_ENDPOINT}
+                database_name: "hexdag"
+                container_name: "agent_memory"
+
+    Python usage::
+
+        from hexdag_plugins.azure import AzureCosmosAdapter
+
+        adapter = AzureCosmosAdapter(
+            endpoint="https://my-cosmos.documents.azure.com:443/",
+            key="...",  # or auto-resolved from AZURE_COSMOS_KEY
+            database_name="hexdag",
+            container_name="agent_memory"
+        )
+
+        # Store agent memory
+        await adapter.astore("agent-123", {"context": "...", "history": [...]})
+
+        # Retrieve memory
+        memory = await adapter.aretrieve("agent-123")
+
+        # Search memories
+        results = await adapter.asearch("user query", top_k=5)
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+        key: str | None = None,
+        database_name: str = "hexdag",
+        container_name: str = "memory",
+        partition_key: str = "/agent_id",
+        use_managed_identity: bool = False,
+        throughput: int = 400,
+    ):
+        """Initialize Azure Cosmos DB adapter.
+
+        Args
+        ----
+            endpoint: Azure Cosmos DB endpoint URL
+            key: Cosmos DB key (or use managed identity)
+            database_name: Database name (default: "hexdag")
+            container_name: Container name (default: "memory")
+            partition_key: Partition key path (default: "/agent_id")
+            use_managed_identity: Use Managed Identity (default: False)
+            throughput: Container throughput RU/s (default: 400)
+        """
+        self.endpoint = endpoint
+        self.key = key or os.getenv("AZURE_COSMOS_KEY")
+        self.database_name = database_name
+        self.container_name = container_name
+        self.partition_key = partition_key
+        self.use_managed_identity = use_managed_identity
+        self.throughput = throughput
+
+        self._client = None
+        self._container = None
+
+    async def _get_container(self):
+        """Get or create Cosmos DB container."""
+        if self._container is None:
+            try:
+                from azure.cosmos import PartitionKey
+                from azure.cosmos.aio import CosmosClient
+            except ImportError as e:
+                raise ImportError(
+                    "Azure Cosmos SDK not installed. Install with: pip install azure-cosmos"
+                ) from e
+
+            if self.use_managed_identity:
+                try:
+                    from azure.identity.aio import DefaultAzureCredential
+
+                    credential = DefaultAzureCredential()
+                    self._client = CosmosClient(self.endpoint, credential=credential)
+                except ImportError as e:
+                    raise ImportError(
+                        "Azure Identity SDK not installed. Install with: pip install azure-identity"
+                    ) from e
+            else:
+                if not self.key:
+                    raise ValueError("key is required when use_managed_identity=False")
+                self._client = CosmosClient(self.endpoint, credential=self.key)
+
+            # Create database if not exists
+            database = await self._client.create_database_if_not_exists(id=self.database_name)
+
+            # Create container if not exists
+            self._container = await database.create_container_if_not_exists(
+                id=self.container_name,
+                partition_key=PartitionKey(path=self.partition_key),
+                offer_throughput=self.throughput,
+            )
+
+        return self._container
+
+    async def astore(
+        self,
+        key: str,
+        value: dict[str, Any],
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        """Store data in Cosmos DB.
+
+        Args
+        ----
+            key: Unique identifier for the data
+            value: Data to store
+            metadata: Optional metadata
+        """
+        container = await self._get_container()
+
+        # Extract agent_id from key for partitioning
+        agent_id = key.split(":")[0] if ":" in key else key
+
+        document = {
+            "id": key,
+            "agent_id": agent_id,
+            "data": value,
+            "metadata": metadata or {},
+            "created_at": time.time(),
+            "updated_at": time.time(),
+        }
+
+        await container.upsert_item(document)
+
+    async def aretrieve(self, key: str) -> dict[str, Any] | None:
+        """Retrieve data from Cosmos DB.
+
+        Args
+        ----
+            key: Unique identifier for the data
+
+        Returns
+        -------
+            Stored data or None if not found
+        """
+        container = await self._get_container()
+        agent_id = key.split(":")[0] if ":" in key else key
+
+        try:
+            item = await container.read_item(item=key, partition_key=agent_id)
+            return item.get("data")
+        except Exception:
+            return None
+
+    async def adelete(self, key: str) -> bool:
+        """Delete data from Cosmos DB.
+
+        Args
+        ----
+            key: Unique identifier for the data
+
+        Returns
+        -------
+            True if deleted, False if not found
+        """
+        container = await self._get_container()
+        agent_id = key.split(":")[0] if ":" in key else key
+
+        try:
+            await container.delete_item(item=key, partition_key=agent_id)
+            return True
+        except Exception:
+            return False
+
+    async def alist(self, prefix: str | None = None) -> list[str]:
+        """List all keys in memory.
+
+        Args
+        ----
+            prefix: Optional prefix to filter keys
+
+        Returns
+        -------
+            List of keys
+        """
+        container = await self._get_container()
+
+        if prefix:
+            query = f"SELECT c.id FROM c WHERE STARTSWITH(c.id, '{prefix}')"
+        else:
+            query = "SELECT c.id FROM c"
+
+        items = container.query_items(query=query, enable_cross_partition_query=True)
+        return [item["id"] async for item in items]
+
+    async def asearch(
+        self,
+        query: str,
+        top_k: int = 5,
+        filter_metadata: dict[str, Any] | None = None,
+    ) -> list[dict[str, Any]]:
+        """Search memories by content.
+
+        Note: This performs a simple text search. For vector similarity search,
+        consider using Azure Cognitive Search integration.
+
+        Args
+        ----
+            query: Search query string
+            top_k: Maximum number of results
+            filter_metadata: Optional metadata filters
+
+        Returns
+        -------
+            List of matching documents
+        """
+        container = await self._get_container()
+
+        # Build SQL query with CONTAINS for text search
+        sql_query = f"""
+            SELECT TOP {top_k} c.id, c.data, c.metadata, c.created_at
+            FROM c
+            WHERE CONTAINS(LOWER(c.data), LOWER('{query}'))
+        """
+
+        if filter_metadata:
+            for key, value in filter_metadata.items():
+                sql_query += f" AND c.metadata.{key} = '{value}'"
+
+        sql_query += " ORDER BY c.created_at DESC"
+
+        items = container.query_items(query=sql_query, enable_cross_partition_query=True)
+        return [item async for item in items]
+
+    async def astore_conversation(
+        self,
+        agent_id: str,
+        messages: list[dict[str, str]],
+        session_id: str | None = None,
+    ) -> None:
+        """Store conversation history.
+
+        Args
+        ----
+            agent_id: Agent identifier
+            messages: List of message dicts with 'role' and 'content'
+            session_id: Optional session identifier
+        """
+        key = f"{agent_id}:conversation:{session_id or 'default'}"
+        await self.astore(key, {"messages": messages}, {"type": "conversation"})
+
+    async def aretrieve_conversation(
+        self,
+        agent_id: str,
+        session_id: str | None = None,
+    ) -> list[dict[str, str]]:
+        """Retrieve conversation history.
+
+        Args
+        ----
+            agent_id: Agent identifier
+            session_id: Optional session identifier
+
+        Returns
+        -------
+            List of messages
+        """
+        key = f"{agent_id}:conversation:{session_id or 'default'}"
+        data = await self.aretrieve(key)
+        return data.get("messages", []) if data else []
+
+    async def aclear_agent(self, agent_id: str) -> int:
+        """Clear all data for an agent.
+
+        Args
+        ----
+            agent_id: Agent identifier
+
+        Returns
+        -------
+            Number of items deleted
+        """
+        keys = await self.alist(prefix=agent_id)
+        count = 0
+        for key in keys:
+            if await self.adelete(key):
+                count += 1
+        return count
+
+    async def ahealth_check(self) -> HealthStatus:
+        """Check Azure Cosmos DB connectivity.
+
+        Returns
+        -------
+            HealthStatus with connectivity details
+        """
+        try:
+            start_time = time.time()
+            container = await self._get_container()
+
+            # Simple query to verify connectivity
+            query = "SELECT VALUE COUNT(1) FROM c"
+            items = container.query_items(query=query, enable_cross_partition_query=True)
+            count = 0
+            async for item in items:
+                count = item
+
+            latency_ms = (time.time() - start_time) * 1000
+
+            return HealthStatus(
+                status="healthy",
+                adapter_name="AzureCosmos",
+                latency_ms=latency_ms,
+                details={
+                    "endpoint": self.endpoint,
+                    "database": self.database_name,
+                    "container": self.container_name,
+                    "document_count": count,
+                },
+            )
+
+        except Exception as e:
+            return HealthStatus(
+                status="unhealthy",
+                adapter_name="AzureCosmos",
+                latency_ms=0.0,
+                details={"error": str(e)},
+            )
+
+    async def aclose(self) -> None:
+        """Close the Cosmos DB client."""
+        if self._client:
+            await self._client.close()
+            self._client = None
+            self._container = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize adapter configuration (excluding secrets)."""
+        return {
+            "endpoint": self.endpoint,
+            "database_name": self.database_name,
+            "container_name": self.container_name,
+            "partition_key": self.partition_key,
+            "use_managed_identity": self.use_managed_identity,
+            "throughput": self.throughput,
+        }

hexdag_plugins/azure/azure_keyvault_adapter.py

@@ -0,0 +1,314 @@
+"""Azure Key Vault adapter for hexDAG framework.
+
+Provides secret resolution from Azure Key Vault for production deployments.
+"""
+
+from typing import Any
+
+from hexdag.core.ports.healthcheck import HealthStatus
+
+
+class AzureKeyVaultAdapter:
+    """Azure Key Vault adapter for secret resolution.
+
+    Supports both API key authentication and Azure Managed Identity for
+    secure, credential-free access in Azure environments.
+
+    Parameters
+    ----------
+    vault_url : str
+        Azure Key Vault URL (e.g., "https://my-vault.vault.azure.net")
+    use_managed_identity : bool, optional
+        Use Azure Managed Identity instead of explicit credentials (default: True)
+    tenant_id : str, optional
+        Azure AD tenant ID (only needed if use_managed_identity=False)
+    client_id : str, optional
+        Azure AD client ID (only needed if use_managed_identity=False)
+    client_secret : str, optional
+        Azure AD client secret (only needed if use_managed_identity=False)
+    cache_secrets : bool, optional
+        Cache retrieved secrets in memory (default: True)
+    cache_ttl : int, optional
+        Cache TTL in seconds (default: 300)
+
+    Examples
+    --------
+    YAML configuration with managed identity::
+
+        spec:
+          ports:
+            secret:
+              adapter: hexdag_plugins.azure.AzureKeyVaultAdapter
+              config:
+                vault_url: "https://my-vault.vault.azure.net"
+                use_managed_identity: true
+
+    YAML configuration with service principal::
+
+        spec:
+          ports:
+            secret:
+              adapter: hexdag_plugins.azure.AzureKeyVaultAdapter
+              config:
+                vault_url: "https://my-vault.vault.azure.net"
+                use_managed_identity: false
+                tenant_id: ${AZURE_TENANT_ID}
+                client_id: ${AZURE_CLIENT_ID}
+                client_secret: ${AZURE_CLIENT_SECRET}
+
+    Python usage::
+
+        from hexdag_plugins.azure import AzureKeyVaultAdapter
+
+        # With managed identity (recommended for Azure deployments)
+        adapter = AzureKeyVaultAdapter(
+            vault_url="https://my-vault.vault.azure.net",
+            use_managed_identity=True
+        )
+
+        # Retrieve secrets
+        api_key = await adapter.aget("OPENAI-API-KEY")
+        db_password = await adapter.aget("DB-PASSWORD")
+
+        # Batch retrieval
+        secrets = await adapter.aget_batch(["SECRET1", "SECRET2"])
+    """
+
+    def __init__(
+        self,
+        vault_url: str,
+        use_managed_identity: bool = True,
+        tenant_id: str | None = None,
+        client_id: str | None = None,
+        client_secret: str | None = None,
+        cache_secrets: bool = True,
+        cache_ttl: int = 300,
+    ):
+        """Initialize Azure Key Vault adapter.
+
+        Args
+        ----
+            vault_url: Azure Key Vault URL
+            use_managed_identity: Use Managed Identity for auth (default: True)
+            tenant_id: Azure AD tenant ID (for service principal auth)
+            client_id: Azure AD client ID (for service principal auth)
+            client_secret: Azure AD client secret (for service principal auth)
+            cache_secrets: Cache retrieved secrets (default: True)
+            cache_ttl: Cache TTL in seconds (default: 300)
+        """
+        self.vault_url = vault_url
+        self.use_managed_identity = use_managed_identity
+        self.tenant_id = tenant_id
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.cache_secrets = cache_secrets
+        self.cache_ttl = cache_ttl
+
+        self._client = None
+        self._cache: dict[str, tuple[str, float]] = {}
+
+    def _get_client(self):
+        """Get or create Key Vault client."""
+        if self._client is None:
+            try:
+                from azure.identity import (
+                    ClientSecretCredential,
+                    DefaultAzureCredential,
+                )
+                from azure.keyvault.secrets import SecretClient
+            except ImportError as e:
+                raise ImportError(
+                    "Azure SDK not installed. Install with: "
+                    "pip install azure-identity azure-keyvault-secrets"
+                ) from e
+
+            if self.use_managed_identity:
+                credential = DefaultAzureCredential()
+            else:
+                if not all([self.tenant_id, self.client_id, self.client_secret]):
+                    raise ValueError(
+                        "tenant_id, client_id, and client_secret are required "
+                        "when use_managed_identity=False"
+                    )
+                credential = ClientSecretCredential(
+                    tenant_id=self.tenant_id,
+                    client_id=self.client_id,
+                    client_secret=self.client_secret,
+                )
+
+            self._client = SecretClient(vault_url=self.vault_url, credential=credential)
+
+        return self._client
+
+    def _get_from_cache(self, secret_name: str) -> str | None:
+        """Get secret from cache if valid."""
+        if not self.cache_secrets:
+            return None
+
+        import time
+
+        if secret_name in self._cache:
+            value, timestamp = self._cache[secret_name]
+            if time.time() - timestamp < self.cache_ttl:
+                return value
+            del self._cache[secret_name]
+
+        return None
+
+    def _set_cache(self, secret_name: str, value: str) -> None:
+        """Set secret in cache."""
+        if self.cache_secrets:
+            import time
+
+            self._cache[secret_name] = (value, time.time())
+
+    async def aget(self, secret_name: str) -> str:
+        """Retrieve a secret from Azure Key Vault.
+
+        Args
+        ----
+            secret_name: Name of the secret in Key Vault
+
+        Returns
+        -------
+            Secret value as string
+
+        Raises
+        ------
+            ValueError: If secret not found
+            RuntimeError: If Key Vault access fails
+        """
+        # Check cache first
+        cached = self._get_from_cache(secret_name)
+        if cached is not None:
+            return cached
+
+        try:
+            client = self._get_client()
+            secret = client.get_secret(secret_name)
+
+            if secret.value is None:
+                raise ValueError(f"Secret '{secret_name}' has no value")
+
+            self._set_cache(secret_name, secret.value)
+            return secret.value
+
+        except Exception as e:
+            if "SecretNotFound" in str(e):
+                raise ValueError(f"Secret '{secret_name}' not found in Key Vault") from e
+            raise RuntimeError(f"Failed to retrieve secret '{secret_name}': {e}") from e
+
+    async def aget_batch(self, secret_names: list[str]) -> dict[str, str]:
+        """Retrieve multiple secrets from Azure Key Vault.
+
+        Args
+        ----
+            secret_names: List of secret names to retrieve
+
+        Returns
+        -------
+            Dictionary mapping secret names to values
+        """
+        results = {}
+        for name in secret_names:
+            try:
+                results[name] = await self.aget(name)
+            except ValueError:  # noqa: SIM105 - contextlib.suppress doesn't work with async
+                # Skip secrets that don't exist
+                pass
+        return results
+
+    async def aset(self, secret_name: str, value: str) -> None:
+        """Set a secret in Azure Key Vault.
+
+        Args
+        ----
+            secret_name: Name of the secret
+            value: Secret value to store
+        """
+        try:
+            client = self._get_client()
+            client.set_secret(secret_name, value)
+            self._set_cache(secret_name, value)
+        except Exception as e:
+            raise RuntimeError(f"Failed to set secret '{secret_name}': {e}") from e
+
+    async def adelete(self, secret_name: str) -> None:
+        """Delete a secret from Azure Key Vault.
+
+        Args
+        ----
+            secret_name: Name of the secret to delete
+        """
+        try:
+            client = self._get_client()
+            client.begin_delete_secret(secret_name)
+            if secret_name in self._cache:
+                del self._cache[secret_name]
+        except Exception as e:
+            raise RuntimeError(f"Failed to delete secret '{secret_name}': {e}") from e
+
+    async def alist(self) -> list[str]:
+        """List all secret names in the Key Vault.
+
+        Returns
+        -------
+            List of secret names
+        """
+        try:
+            client = self._get_client()
+            return [secret.name for secret in client.list_properties_of_secrets()]
+        except Exception as e:
+            raise RuntimeError(f"Failed to list secrets: {e}") from e
+
+    def clear_cache(self) -> None:
+        """Clear the secret cache."""
+        self._cache.clear()
+
+    async def ahealth_check(self) -> HealthStatus:
+        """Check Azure Key Vault connectivity.
+
+        Returns
+        -------
+            HealthStatus with connectivity details
+        """
+        import time
+
+        try:
+            start_time = time.time()
+            client = self._get_client()
+
+            # Try to list secrets (limited to 1) to verify connectivity
+            list(client.list_properties_of_secrets())
+
+            latency_ms = (time.time() - start_time) * 1000
+
+            return HealthStatus(
+                status="healthy",
+                adapter_name="AzureKeyVault",
+                latency_ms=latency_ms,
+                details={
+                    "vault_url": self.vault_url,
+                    "auth_method": "managed_identity"
+                    if self.use_managed_identity
+                    else "service_principal",
+                    "cache_enabled": self.cache_secrets,
+                },
+            )
+
+        except Exception as e:
+            return HealthStatus(
+                status="unhealthy",
+                adapter_name="AzureKeyVault",
+                latency_ms=0.0,
+                details={"error": str(e), "vault_url": self.vault_url},
+            )
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize adapter configuration (excluding secrets)."""
+        return {
+            "vault_url": self.vault_url,
+            "use_managed_identity": self.use_managed_identity,
+            "cache_secrets": self.cache_secrets,
+            "cache_ttl": self.cache_ttl,
+        }