alma-memory 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

alma/storage/azure_cosmos.py

@@ -17,20 +17,18 @@ Configuration (config.yaml):
         embedding_dim: 384
 """
 
-import json
 import logging
 from datetime import datetime, timezone
-from typing import Optional, List, Dict, Any, Tuple
-from dataclasses import asdict
+from typing import Any, Dict, List, Optional
 
+from alma.storage.base import StorageBackend
 from alma.types import (
+    AntiPattern,
+    DomainKnowledge,
     Heuristic,
     Outcome,
     UserPreference,
-    DomainKnowledge,
-    AntiPattern,
 )
-from alma.storage.base import StorageBackend
 
 logger = logging.getLogger(__name__)
 
@@ -39,6 +37,7 @@ try:
     from azure.cosmos import CosmosClient, PartitionKey, exceptions
     from azure.cosmos.container import ContainerProxy
     from azure.cosmos.database import DatabaseProxy
+
     AZURE_COSMOS_AVAILABLE = True
 except ImportError:
     AZURE_COSMOS_AVAILABLE = False
@@ -49,8 +48,7 @@ except ImportError:
     ContainerProxy = Any  # type: ignore
     DatabaseProxy = Any  # type: ignore
     logger.warning(
-        "azure-cosmos package not installed. "
-        "Install with: pip install azure-cosmos"
+        "azure-cosmos package not installed. Install with: pip install azure-cosmos"
     )
 
 
@@ -111,9 +109,7 @@ class AzureCosmosStorage(StorageBackend):
 
         # Get or create database
         if create_if_not_exists:
-            self.database = self.client.create_database_if_not_exists(
-                id=database_name
-            )
+            self.database = self.client.create_database_if_not_exists(id=database_name)
             self._init_containers()
         else:
             self.database = self.client.get_database_client(database_name)
@@ -195,7 +191,7 @@ class AzureCosmosStorage(StorageBackend):
             if cfg["vector_indexes"] and cfg["vector_path"]:
                 # Exclude vector path from regular indexing
                 indexing_policy["excludedPaths"].append(
-                    {"path": f'{cfg["vector_path"]}/*'}
+                    {"path": f"{cfg['vector_path']}/*"}
                 )
 
                 # Vector embedding policy for DiskANN
@@ -250,12 +246,14 @@ class AzureCosmosStorage(StorageBackend):
             "confidence": heuristic.confidence,
             "occurrence_count": heuristic.occurrence_count,
             "success_count": heuristic.success_count,
-            "last_validated": heuristic.last_validated.isoformat()
-            if heuristic.last_validated
-            else None,
-            "created_at": heuristic.created_at.isoformat()
-            if heuristic.created_at
-            else None,
+            "last_validated": (
+                heuristic.last_validated.isoformat()
+                if heuristic.last_validated
+                else None
+            ),
+            "created_at": (
+                heuristic.created_at.isoformat() if heuristic.created_at else None
+            ),
             "metadata": heuristic.metadata or {},
             "embedding": heuristic.embedding,
             "type": "heuristic",
@@ -301,9 +299,9 @@ class AzureCosmosStorage(StorageBackend):
             "preference": preference.preference,
             "source": preference.source,
             "confidence": preference.confidence,
-            "timestamp": preference.timestamp.isoformat()
-            if preference.timestamp
-            else None,
+            "timestamp": (
+                preference.timestamp.isoformat() if preference.timestamp else None
+            ),
             "metadata": preference.metadata or {},
             "type": "preference",
         }
@@ -324,9 +322,9 @@ class AzureCosmosStorage(StorageBackend):
             "fact": knowledge.fact,
             "source": knowledge.source,
             "confidence": knowledge.confidence,
-            "last_verified": knowledge.last_verified.isoformat()
-            if knowledge.last_verified
-            else None,
+            "last_verified": (
+                knowledge.last_verified.isoformat() if knowledge.last_verified else None
+            ),
             "metadata": knowledge.metadata or {},
             "embedding": knowledge.embedding,
             "type": "domain_knowledge",
@@ -348,12 +346,12 @@ class AzureCosmosStorage(StorageBackend):
             "why_bad": anti_pattern.why_bad,
             "better_alternative": anti_pattern.better_alternative,
             "occurrence_count": anti_pattern.occurrence_count,
-            "last_seen": anti_pattern.last_seen.isoformat()
-            if anti_pattern.last_seen
-            else None,
-            "created_at": anti_pattern.created_at.isoformat()
-            if anti_pattern.created_at
-            else None,
+            "last_seen": (
+                anti_pattern.last_seen.isoformat() if anti_pattern.last_seen else None
+            ),
+            "created_at": (
+                anti_pattern.created_at.isoformat() if anti_pattern.created_at else None
+            ),
             "metadata": anti_pattern.metadata or {},
             "embedding": anti_pattern.embedding,
             "type": "anti_pattern",
@@ -668,7 +666,7 @@ class AzureCosmosStorage(StorageBackend):
             return False
 
         doc = items[0]
-        project_id = doc["project_id"]
+        doc["project_id"]
 
         # Apply updates
         for key, value in updates.items():
@@ -710,6 +708,78 @@ class AzureCosmosStorage(StorageBackend):
         container.replace_item(item=heuristic_id, body=doc)
         return True
 
+    def update_heuristic_confidence(
+        self,
+        heuristic_id: str,
+        new_confidence: float,
+    ) -> bool:
+        """
+        Update confidence score for a heuristic.
+
+        Note: This requires a cross-partition query since we only have the ID.
+        For better performance, consider using update_heuristic() with the
+        project_id if available, which enables point reads.
+        """
+        container = self._get_container("heuristics")
+
+        # Find the heuristic (cross-partition query required without project_id)
+        query = "SELECT * FROM c WHERE c.id = @id"
+        items = list(
+            container.query_items(
+                query=query,
+                parameters=[{"name": "@id", "value": heuristic_id}],
+                enable_cross_partition_query=True,
+            )
+        )
+
+        if not items:
+            return False
+
+        doc = items[0]
+        doc["confidence"] = new_confidence
+
+        container.replace_item(item=heuristic_id, body=doc)
+        logger.debug(
+            f"Updated heuristic confidence: {heuristic_id} -> {new_confidence}"
+        )
+        return True
+
+    def update_knowledge_confidence(
+        self,
+        knowledge_id: str,
+        new_confidence: float,
+    ) -> bool:
+        """
+        Update confidence score for domain knowledge.
+
+        Note: This requires a cross-partition query since we only have the ID.
+        For better performance when project_id is known, fetch the document
+        directly using point read and update via save_domain_knowledge().
+        """
+        container = self._get_container("knowledge")
+
+        # Find the knowledge item (cross-partition query required without project_id)
+        query = "SELECT * FROM c WHERE c.id = @id"
+        items = list(
+            container.query_items(
+                query=query,
+                parameters=[{"name": "@id", "value": knowledge_id}],
+                enable_cross_partition_query=True,
+            )
+        )
+
+        if not items:
+            return False
+
+        doc = items[0]
+        doc["confidence"] = new_confidence
+
+        container.replace_item(item=knowledge_id, body=doc)
+        logger.debug(
+            f"Updated knowledge confidence: {knowledge_id} -> {new_confidence}"
+        )
+        return True
+
     # ==================== DELETE OPERATIONS ====================
 
     def delete_outcomes_older_than(

alma/storage/base.py

@@ -5,16 +5,15 @@ Abstract base class that all storage backends must implement.
 """
 
 from abc import ABC, abstractmethod
-from typing import Optional, List, Dict, Any
 from datetime import datetime
+from typing import Any, Dict, List, Optional
 
 from alma.types import (
+    AntiPattern,
+    DomainKnowledge,
     Heuristic,
     Outcome,
     UserPreference,
-    DomainKnowledge,
-    AntiPattern,
-    MemoryType,
 )
 
 
@@ -55,6 +54,22 @@ class StorageBackend(ABC):
         """Save an anti-pattern, return its ID."""
         pass
 
+    # ==================== BATCH WRITE OPERATIONS ====================
+
+    def save_heuristics(self, heuristics: List[Heuristic]) -> List[str]:
+        """Save multiple heuristics in a batch. Default implementation calls save_heuristic in a loop."""
+        return [self.save_heuristic(h) for h in heuristics]
+
+    def save_outcomes(self, outcomes: List[Outcome]) -> List[str]:
+        """Save multiple outcomes in a batch. Default implementation calls save_outcome in a loop."""
+        return [self.save_outcome(o) for o in outcomes]
+
+    def save_domain_knowledge_batch(
+        self, knowledge_items: List[DomainKnowledge]
+    ) -> List[str]:
+        """Save multiple domain knowledge items in a batch. Default implementation calls save_domain_knowledge in a loop."""
+        return [self.save_domain_knowledge(k) for k in knowledge_items]
+
     # ==================== READ OPERATIONS ====================
 
     @abstractmethod
@@ -339,6 +354,144 @@ class StorageBackend(ABC):
         """
         pass
 
+    # ==================== MULTI-AGENT MEMORY SHARING ====================
+
+    def get_heuristics_for_agents(
+        self,
+        project_id: str,
+        agents: List[str],
+        embedding: Optional[List[float]] = None,
+        top_k: int = 5,
+        min_confidence: float = 0.0,
+    ) -> List[Heuristic]:
+        """
+        Get heuristics from multiple agents in one call.
+
+        This enables multi-agent memory sharing where an agent can
+        read memories from agents it inherits from.
+
+        Args:
+            project_id: Project to query
+            agents: List of agent names to query
+            embedding: Query embedding for semantic search
+            top_k: Max results to return per agent
+            min_confidence: Minimum confidence threshold
+
+        Returns:
+            List of matching heuristics from all specified agents
+        """
+        # Default implementation: query each agent individually
+        results = []
+        for agent in agents:
+            agent_heuristics = self.get_heuristics(
+                project_id=project_id,
+                agent=agent,
+                embedding=embedding,
+                top_k=top_k,
+                min_confidence=min_confidence,
+            )
+            results.extend(agent_heuristics)
+        return results
+
+    def get_outcomes_for_agents(
+        self,
+        project_id: str,
+        agents: List[str],
+        task_type: Optional[str] = None,
+        embedding: Optional[List[float]] = None,
+        top_k: int = 5,
+        success_only: bool = False,
+    ) -> List[Outcome]:
+        """
+        Get outcomes from multiple agents in one call.
+
+        Args:
+            project_id: Project to query
+            agents: List of agent names to query
+            task_type: Filter by task type
+            embedding: Query embedding for semantic search
+            top_k: Max results to return per agent
+            success_only: Only return successful outcomes
+
+        Returns:
+            List of matching outcomes from all specified agents
+        """
+        results = []
+        for agent in agents:
+            agent_outcomes = self.get_outcomes(
+                project_id=project_id,
+                agent=agent,
+                task_type=task_type,
+                embedding=embedding,
+                top_k=top_k,
+                success_only=success_only,
+            )
+            results.extend(agent_outcomes)
+        return results
+
+    def get_domain_knowledge_for_agents(
+        self,
+        project_id: str,
+        agents: List[str],
+        domain: Optional[str] = None,
+        embedding: Optional[List[float]] = None,
+        top_k: int = 5,
+    ) -> List[DomainKnowledge]:
+        """
+        Get domain knowledge from multiple agents in one call.
+
+        Args:
+            project_id: Project to query
+            agents: List of agent names to query
+            domain: Filter by domain
+            embedding: Query embedding for semantic search
+            top_k: Max results to return per agent
+
+        Returns:
+            List of matching domain knowledge from all specified agents
+        """
+        results = []
+        for agent in agents:
+            agent_knowledge = self.get_domain_knowledge(
+                project_id=project_id,
+                agent=agent,
+                domain=domain,
+                embedding=embedding,
+                top_k=top_k,
+            )
+            results.extend(agent_knowledge)
+        return results
+
+    def get_anti_patterns_for_agents(
+        self,
+        project_id: str,
+        agents: List[str],
+        embedding: Optional[List[float]] = None,
+        top_k: int = 5,
+    ) -> List[AntiPattern]:
+        """
+        Get anti-patterns from multiple agents in one call.
+
+        Args:
+            project_id: Project to query
+            agents: List of agent names to query
+            embedding: Query embedding for semantic search
+            top_k: Max results to return per agent
+
+        Returns:
+            List of matching anti-patterns from all specified agents
+        """
+        results = []
+        for agent in agents:
+            agent_patterns = self.get_anti_patterns(
+                project_id=project_id,
+                agent=agent,
+                embedding=embedding,
+                top_k=top_k,
+            )
+            results.extend(agent_patterns)
+        return results
+
     # ==================== STATS ====================
 
     @abstractmethod