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

alma/storage/azure_cosmos.py

@@ -17,20 +17,19 @@ 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.storage.constants import AZURE_COSMOS_CONTAINER_NAMES, MemoryType
 from alma.types import (
+    AntiPattern,
+    DomainKnowledge,
     Heuristic,
     Outcome,
     UserPreference,
-    DomainKnowledge,
-    AntiPattern,
 )
-from alma.storage.base import StorageBackend
 
 logger = logging.getLogger(__name__)
 
@@ -39,6 +38,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 +49,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"
     )
 
 
@@ -63,21 +62,19 @@ class AzureCosmosStorage(StorageBackend):
     - DiskANN vector indexing for similarity search
     - Partition key: project_id for efficient queries
 
-    Container structure:
-    - alma-heuristics: Heuristics with vector embeddings
-    - alma-outcomes: Task outcomes with vector embeddings
-    - alma-preferences: User preferences (no vectors)
-    - alma-knowledge: Domain knowledge with vector embeddings
-    - alma-antipatterns: Anti-patterns with vector embeddings
+    Container structure (uses canonical memory type names with alma_ prefix):
+    - alma_heuristics: Heuristics with vector embeddings
+    - alma_outcomes: Task outcomes with vector embeddings
+    - alma_preferences: User preferences (no vectors)
+    - alma_domain_knowledge: Domain knowledge with vector embeddings
+    - alma_anti_patterns: Anti-patterns with vector embeddings
+
+    Container names are derived from alma.storage.constants.AZURE_COSMOS_CONTAINER_NAMES
+    for consistency across all storage backends.
     """
 
-    CONTAINER_NAMES = {
-        "heuristics": "alma-heuristics",
-        "outcomes": "alma-outcomes",
-        "preferences": "alma-preferences",
-        "knowledge": "alma-knowledge",
-        "antipatterns": "alma-antipatterns",
-    }
+    # Use canonical container names from constants
+    CONTAINER_NAMES = AZURE_COSMOS_CONTAINER_NAMES
 
     def __init__(
         self,
@@ -111,9 +108,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)
@@ -125,6 +120,14 @@ class AzureCosmosStorage(StorageBackend):
                 container_name
             )
 
+        # Cache for partition key mappings: {container_key: {doc_id: partition_key}}
+        # This reduces RU consumption by avoiding cross-partition queries
+        self._partition_key_cache: Dict[str, Dict[str, str]] = {
+            mt: {} for mt in MemoryType.ALL
+        }
+        # Maximum cache size per container to prevent memory issues
+        self._cache_max_size = 1000
+
         logger.info(f"Connected to Azure Cosmos DB: {database_name}")
 
     @classmethod
@@ -150,29 +153,29 @@ class AzureCosmosStorage(StorageBackend):
 
     def _init_containers(self):
         """Initialize containers with vector search indexing."""
-        # Container configs with indexing policies
+        # Container configs with indexing policies (using canonical memory types)
         container_configs = {
-            "heuristics": {
+            MemoryType.HEURISTICS: {
                 "partition_key": "/project_id",
                 "vector_path": "/embedding",
                 "vector_indexes": True,
             },
-            "outcomes": {
+            MemoryType.OUTCOMES: {
                 "partition_key": "/project_id",
                 "vector_path": "/embedding",
                 "vector_indexes": True,
             },
-            "preferences": {
+            MemoryType.PREFERENCES: {
                 "partition_key": "/user_id",
                 "vector_path": None,
                 "vector_indexes": False,
             },
-            "knowledge": {
+            MemoryType.DOMAIN_KNOWLEDGE: {
                 "partition_key": "/project_id",
                 "vector_path": "/embedding",
                 "vector_indexes": True,
             },
-            "antipatterns": {
+            MemoryType.ANTI_PATTERNS: {
                 "partition_key": "/project_id",
                 "vector_path": "/embedding",
                 "vector_indexes": True,
@@ -195,7 +198,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
@@ -235,11 +238,123 @@ class AzureCosmosStorage(StorageBackend):
         """Get container client by key."""
         return self._containers[container_key]
 
+    def _cache_partition_key(
+        self, container_key: str, doc_id: str, partition_key: str
+    ) -> None:
+        """
+        Cache the partition key mapping for a document.
+
+        This enables point reads for future operations, reducing RU consumption
+        by avoiding expensive cross-partition queries.
+        """
+        cache = self._partition_key_cache[container_key]
+
+        # Evict oldest entries if cache is full (simple FIFO eviction)
+        if len(cache) >= self._cache_max_size:
+            # Remove first 10% of entries
+            keys_to_remove = list(cache.keys())[: self._cache_max_size // 10]
+            for key in keys_to_remove:
+                del cache[key]
+
+        cache[doc_id] = partition_key
+
+    def _get_cached_partition_key(
+        self, container_key: str, doc_id: str
+    ) -> Optional[str]:
+        """
+        Get cached partition key for a document if available.
+
+        Returns None if the partition key is not cached.
+        """
+        return self._partition_key_cache.get(container_key, {}).get(doc_id)
+
+    def _invalidate_partition_key_cache(self, container_key: str, doc_id: str) -> None:
+        """Remove a document from the partition key cache."""
+        cache = self._partition_key_cache.get(container_key, {})
+        cache.pop(doc_id, None)
+
+    def _point_read_document(
+        self,
+        container_key: str,
+        doc_id: str,
+        partition_key: Optional[str] = None,
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Attempt to read a document using a point read (1 RU) instead of a query.
+
+        If partition_key is provided, performs a direct point read.
+        If partition_key is not provided but is cached, uses the cached value.
+        If neither is available, falls back to a cross-partition query.
+
+        Args:
+            container_key: The container key (e.g., 'heuristics', 'knowledge')
+            doc_id: The document ID
+            partition_key: Optional partition key for direct point read
+
+        Returns:
+            The document if found, None otherwise
+        """
+        container = self._get_container(container_key)
+
+        # Try to get partition key from cache if not provided
+        if partition_key is None:
+            partition_key = self._get_cached_partition_key(container_key, doc_id)
+
+        # If we have a partition key, use point read (1 RU)
+        if partition_key is not None:
+            try:
+                doc = container.read_item(item=doc_id, partition_key=partition_key)
+                # Refresh cache on successful read
+                self._cache_partition_key(container_key, doc_id, partition_key)
+                return doc
+            except exceptions.CosmosResourceNotFoundError:
+                # Document not found or partition key was wrong
+                self._invalidate_partition_key_cache(container_key, doc_id)
+                # Fall through to cross-partition query
+            except Exception as e:
+                logger.warning(f"Point read failed for {doc_id}: {e}")
+                # Fall through to cross-partition query
+
+        # Fallback: Cross-partition query (expensive but necessary without partition key)
+        logger.debug(
+            f"Using cross-partition query for {doc_id} in {container_key} "
+            "(consider providing project_id for better performance)"
+        )
+        query = "SELECT * FROM c WHERE c.id = @id"
+        items = list(
+            container.query_items(
+                query=query,
+                parameters=[{"name": "@id", "value": doc_id}],
+                enable_cross_partition_query=True,
+            )
+        )
+
+        if items:
+            doc = items[0]
+            # Cache the partition key for future operations
+            pk_field = self._get_partition_key_field(container_key)
+            if pk_field and pk_field in doc:
+                self._cache_partition_key(container_key, doc_id, doc[pk_field])
+            return doc
+
+        return None
+
+    def _get_partition_key_field(self, container_key: str) -> Optional[str]:
+        """Get the partition key field name for a container."""
+        partition_key_fields = {
+            MemoryType.HEURISTICS: "project_id",
+            MemoryType.OUTCOMES: "project_id",
+            MemoryType.PREFERENCES: "user_id",
+            MemoryType.DOMAIN_KNOWLEDGE: "project_id",
+            MemoryType.ANTI_PATTERNS: "project_id",
+        }
+        return partition_key_fields.get(container_key)
+
     # ==================== WRITE OPERATIONS ====================
 
     def save_heuristic(self, heuristic: Heuristic) -> str:
         """Save a heuristic."""
-        container = self._get_container("heuristics")
+        container = self._get_container(MemoryType.HEURISTICS)
 
         doc = {
             "id": heuristic.id,
@@ -250,24 +365,30 @@ 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",
         }
 
         container.upsert_item(doc)
+        # Cache partition key for efficient future updates
+        self._cache_partition_key(
+            MemoryType.HEURISTICS, heuristic.id, heuristic.project_id
+        )
         logger.debug(f"Saved heuristic: {heuristic.id}")
         return heuristic.id
 
     def save_outcome(self, outcome: Outcome) -> str:
         """Save an outcome."""
-        container = self._get_container("outcomes")
+        container = self._get_container(MemoryType.OUTCOMES)
 
         doc = {
             "id": outcome.id,
@@ -287,12 +408,14 @@ class AzureCosmosStorage(StorageBackend):
         }
 
         container.upsert_item(doc)
+        # Cache partition key for efficient future updates
+        self._cache_partition_key(MemoryType.OUTCOMES, outcome.id, outcome.project_id)
         logger.debug(f"Saved outcome: {outcome.id}")
         return outcome.id
 
     def save_user_preference(self, preference: UserPreference) -> str:
         """Save a user preference."""
-        container = self._get_container("preferences")
+        container = self._get_container(MemoryType.PREFERENCES)
 
         doc = {
             "id": preference.id,
@@ -301,20 +424,24 @@ 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",
         }
 
         container.upsert_item(doc)
+        # Cache partition key for efficient future updates
+        self._cache_partition_key(
+            MemoryType.PREFERENCES, preference.id, preference.user_id
+        )
         logger.debug(f"Saved preference: {preference.id}")
         return preference.id
 
     def save_domain_knowledge(self, knowledge: DomainKnowledge) -> str:
         """Save domain knowledge."""
-        container = self._get_container("knowledge")
+        container = self._get_container(MemoryType.DOMAIN_KNOWLEDGE)
 
         doc = {
             "id": knowledge.id,
@@ -324,21 +451,25 @@ 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",
         }
 
         container.upsert_item(doc)
+        # Cache partition key for efficient future updates
+        self._cache_partition_key(
+            MemoryType.DOMAIN_KNOWLEDGE, knowledge.id, knowledge.project_id
+        )
         logger.debug(f"Saved domain knowledge: {knowledge.id}")
         return knowledge.id
 
     def save_anti_pattern(self, anti_pattern: AntiPattern) -> str:
         """Save an anti-pattern."""
-        container = self._get_container("antipatterns")
+        container = self._get_container(MemoryType.ANTI_PATTERNS)
 
         doc = {
             "id": anti_pattern.id,
@@ -348,18 +479,22 @@ 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",
         }
 
         container.upsert_item(doc)
+        # Cache partition key for efficient future updates
+        self._cache_partition_key(
+            MemoryType.ANTI_PATTERNS, anti_pattern.id, anti_pattern.project_id
+        )
         logger.debug(f"Saved anti-pattern: {anti_pattern.id}")
         return anti_pattern.id
 
@@ -374,7 +509,7 @@ class AzureCosmosStorage(StorageBackend):
         min_confidence: float = 0.0,
     ) -> List[Heuristic]:
         """Get heuristics with optional vector search."""
-        container = self._get_container("heuristics")
+        container = self._get_container(MemoryType.HEURISTICS)
 
         if embedding:
             # Vector search query
@@ -426,6 +561,12 @@ class AzureCosmosStorage(StorageBackend):
             )
         )
 
+        # Cache partition keys for efficient future updates
+        for doc in items:
+            self._cache_partition_key(
+                MemoryType.HEURISTICS, doc["id"], doc["project_id"]
+            )
+
         return [self._doc_to_heuristic(doc) for doc in items]
 
     def get_outcomes(
@@ -438,7 +579,7 @@ class AzureCosmosStorage(StorageBackend):
         success_only: bool = False,
     ) -> List[Outcome]:
         """Get outcomes with optional vector search."""
-        container = self._get_container("outcomes")
+        container = self._get_container(MemoryType.OUTCOMES)
 
         if embedding:
             # Vector search query
@@ -496,6 +637,10 @@ class AzureCosmosStorage(StorageBackend):
             )
         )
 
+        # Cache partition keys for efficient future updates
+        for doc in items:
+            self._cache_partition_key(MemoryType.OUTCOMES, doc["id"], doc["project_id"])
+
         return [self._doc_to_outcome(doc) for doc in items]
 
     def get_user_preferences(
@@ -504,7 +649,7 @@ class AzureCosmosStorage(StorageBackend):
         category: Optional[str] = None,
     ) -> List[UserPreference]:
         """Get user preferences."""
-        container = self._get_container("preferences")
+        container = self._get_container(MemoryType.PREFERENCES)
 
         query = "SELECT * FROM c WHERE c.user_id = @user_id"
         parameters = [{"name": "@user_id", "value": user_id}]
@@ -522,6 +667,10 @@ class AzureCosmosStorage(StorageBackend):
             )
         )
 
+        # Cache partition keys for efficient future updates
+        for doc in items:
+            self._cache_partition_key(MemoryType.PREFERENCES, doc["id"], doc["user_id"])
+
         return [self._doc_to_preference(doc) for doc in items]
 
     def get_domain_knowledge(
@@ -533,7 +682,7 @@ class AzureCosmosStorage(StorageBackend):
         top_k: int = 5,
     ) -> List[DomainKnowledge]:
         """Get domain knowledge with optional vector search."""
-        container = self._get_container("knowledge")
+        container = self._get_container(MemoryType.DOMAIN_KNOWLEDGE)
 
         if embedding:
             query = """
@@ -585,6 +734,12 @@ class AzureCosmosStorage(StorageBackend):
             )
         )
 
+        # Cache partition keys for efficient future updates
+        for doc in items:
+            self._cache_partition_key(
+                MemoryType.DOMAIN_KNOWLEDGE, doc["id"], doc["project_id"]
+            )
+
         return [self._doc_to_domain_knowledge(doc) for doc in items]
 
     def get_anti_patterns(
@@ -595,7 +750,7 @@ class AzureCosmosStorage(StorageBackend):
         top_k: int = 5,
     ) -> List[AntiPattern]:
         """Get anti-patterns with optional vector search."""
-        container = self._get_container("antipatterns")
+        container = self._get_container(MemoryType.ANTI_PATTERNS)
 
         if embedding:
             query = """
@@ -641,6 +796,12 @@ class AzureCosmosStorage(StorageBackend):
             )
         )
 
+        # Cache partition keys for efficient future updates
+        for doc in items:
+            self._cache_partition_key(
+                MemoryType.ANTI_PATTERNS, doc["id"], doc["project_id"]
+            )
+
         return [self._doc_to_anti_pattern(doc) for doc in items]
 
     # ==================== UPDATE OPERATIONS ====================
@@ -649,26 +810,28 @@ class AzureCosmosStorage(StorageBackend):
         self,
         heuristic_id: str,
         updates: Dict[str, Any],
+        project_id: Optional[str] = None,
     ) -> bool:
-        """Update a heuristic's fields."""
-        container = self._get_container("heuristics")
+        """
+        Update a heuristic's fields.
 
-        # We need project_id to read the item (partition key)
-        # First try to find the heuristic
-        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,
-            )
-        )
+        Args:
+            heuristic_id: The heuristic document ID
+            updates: Dictionary of fields to update
+            project_id: Optional partition key for efficient point read (1 RU).
+                       If not provided, will attempt cache lookup, then
+                       fall back to cross-partition query (more expensive).
+
+        Returns:
+            True if update succeeded, False if document not found
+        """
+        container = self._get_container(MemoryType.HEURISTICS)
 
-        if not items:
-            return False
+        # Use optimized point read with cache fallback
+        doc = self._point_read_document(MemoryType.HEURISTICS, heuristic_id, project_id)
 
-        doc = items[0]
-        project_id = doc["project_id"]
+        if not doc:
+            return False
 
         # Apply updates
         for key, value in updates.items():
@@ -684,24 +847,29 @@ class AzureCosmosStorage(StorageBackend):
         self,
         heuristic_id: str,
         success: bool,
+        project_id: Optional[str] = None,
     ) -> bool:
-        """Increment heuristic occurrence count."""
-        container = self._get_container("heuristics")
+        """
+        Increment heuristic occurrence count.
 
-        # Find the heuristic
-        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,
-            )
-        )
+        Args:
+            heuristic_id: The heuristic document ID
+            success: Whether this occurrence was successful
+            project_id: Optional partition key for efficient point read (1 RU).
+                       If not provided, will attempt cache lookup, then
+                       fall back to cross-partition query (more expensive).
+
+        Returns:
+            True if update succeeded, False if document not found
+        """
+        container = self._get_container(MemoryType.HEURISTICS)
 
-        if not items:
+        # Use optimized point read with cache fallback
+        doc = self._point_read_document(MemoryType.HEURISTICS, heuristic_id, project_id)
+
+        if not doc:
             return False
 
-        doc = items[0]
         doc["occurrence_count"] = doc.get("occurrence_count", 0) + 1
         if success:
             doc["success_count"] = doc.get("success_count", 0) + 1
@@ -710,6 +878,88 @@ class AzureCosmosStorage(StorageBackend):
         container.replace_item(item=heuristic_id, body=doc)
         return True
 
+    def update_heuristic_confidence(
+        self,
+        heuristic_id: str,
+        new_confidence: float,
+        project_id: Optional[str] = None,
+    ) -> bool:
+        """
+        Update confidence score for a heuristic.
+
+        Args:
+            heuristic_id: The heuristic document ID
+            new_confidence: The new confidence value
+            project_id: Optional partition key for efficient point read (1 RU).
+                       If not provided, will attempt cache lookup, then
+                       fall back to cross-partition query (more expensive).
+
+        Returns:
+            True if update succeeded, False if document not found
+
+        Performance Note:
+            - With project_id: 1 RU for point read + write cost
+            - With cached partition key: 1 RU for point read + write cost
+            - Without either: Cross-partition query (variable, higher RUs)
+        """
+        container = self._get_container(MemoryType.HEURISTICS)
+
+        # Use optimized point read with cache fallback
+        doc = self._point_read_document(MemoryType.HEURISTICS, heuristic_id, project_id)
+
+        if not doc:
+            return False
+
+        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,
+        project_id: Optional[str] = None,
+    ) -> bool:
+        """
+        Update confidence score for domain knowledge.
+
+        Args:
+            knowledge_id: The knowledge document ID
+            new_confidence: The new confidence value
+            project_id: Optional partition key for efficient point read (1 RU).
+                       If not provided, will attempt cache lookup, then
+                       fall back to cross-partition query (more expensive).
+
+        Returns:
+            True if update succeeded, False if document not found
+
+        Performance Note:
+            - With project_id: 1 RU for point read + write cost
+            - With cached partition key: 1 RU for point read + write cost
+            - Without either: Cross-partition query (variable, higher RUs)
+        """
+        container = self._get_container(MemoryType.DOMAIN_KNOWLEDGE)
+
+        # Use optimized point read with cache fallback
+        doc = self._point_read_document(
+            MemoryType.DOMAIN_KNOWLEDGE, knowledge_id, project_id
+        )
+
+        if not doc:
+            return False
+
+        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(
@@ -719,7 +969,7 @@ class AzureCosmosStorage(StorageBackend):
         agent: Optional[str] = None,
     ) -> int:
         """Delete old outcomes."""
-        container = self._get_container("outcomes")
+        container = self._get_container(MemoryType.OUTCOMES)
 
         query = """
         SELECT c.id FROM c
@@ -762,7 +1012,7 @@ class AzureCosmosStorage(StorageBackend):
         agent: Optional[str] = None,
     ) -> int:
         """Delete low-confidence heuristics."""
-        container = self._get_container("heuristics")
+        container = self._get_container(MemoryType.HEURISTICS)
 
         query = """
         SELECT c.id FROM c
@@ -798,27 +1048,59 @@ class AzureCosmosStorage(StorageBackend):
         logger.info(f"Deleted {deleted} low-confidence heuristics")
         return deleted
 
-    def delete_heuristic(self, heuristic_id: str) -> bool:
-        """Delete a specific heuristic by ID."""
-        container = self._get_container("heuristics")
+    def delete_heuristic(
+        self, heuristic_id: str, project_id: Optional[str] = None
+    ) -> bool:
+        """
+        Delete a specific heuristic by ID.
 
-        # Find the heuristic to get project_id
-        query = "SELECT c.project_id FROM c WHERE c.id = @id"
-        items = list(
-            container.query_items(
-                query=query,
-                parameters=[{"name": "@id", "value": heuristic_id}],
-                enable_cross_partition_query=True,
+        Args:
+            heuristic_id: The heuristic document ID
+            project_id: Optional partition key for efficient point read (1 RU).
+                       If not provided, will attempt cache lookup, then
+                       fall back to cross-partition query (more expensive).
+
+        Returns:
+            True if deletion succeeded, False if document not found
+        """
+        container = self._get_container(MemoryType.HEURISTICS)
+
+        # Try to get partition key from cache if not provided
+        if project_id is None:
+            project_id = self._get_cached_partition_key(
+                MemoryType.HEURISTICS, heuristic_id
             )
+
+        # If we have a partition key, try direct delete
+        if project_id is not None:
+            try:
+                container.delete_item(item=heuristic_id, partition_key=project_id)
+                self._invalidate_partition_key_cache(
+                    MemoryType.HEURISTICS, heuristic_id
+                )
+                return True
+            except exceptions.CosmosResourceNotFoundError:
+                # Document not found or partition key was wrong
+                self._invalidate_partition_key_cache(
+                    MemoryType.HEURISTICS, heuristic_id
+                )
+                # Fall through to cross-partition lookup
+
+        # Fallback: Cross-partition query to find the document
+        logger.debug(
+            f"Using cross-partition query for delete {heuristic_id} "
+            "(consider providing project_id for better performance)"
         )
+        doc = self._point_read_document(MemoryType.HEURISTICS, heuristic_id, None)
 
-        if not items:
+        if not doc:
             return False
 
-        project_id = items[0]["project_id"]
+        project_id = doc["project_id"]
 
         try:
             container.delete_item(item=heuristic_id, partition_key=project_id)
+            self._invalidate_partition_key_cache(MemoryType.HEURISTICS, heuristic_id)
             return True
         except exceptions.CosmosResourceNotFoundError:
             return False
@@ -838,39 +1120,38 @@ class AzureCosmosStorage(StorageBackend):
             "database": self.database_name,
         }
 
-        # Count items in each container
-        container_keys = ["heuristics", "outcomes", "knowledge", "antipatterns"]
-        for key in container_keys:
-            container = self._get_container(key)
-            query = "SELECT VALUE COUNT(1) FROM c WHERE c.project_id = @project_id"
-            parameters = [{"name": "@project_id", "value": project_id}]
-
-            if agent and key != "preferences":
-                query = """
-                SELECT VALUE COUNT(1) FROM c
-                WHERE c.project_id = @project_id AND c.agent = @agent
-                """
-                parameters.append({"name": "@agent", "value": agent})
+        # Count items in each container using canonical memory types
+        for memory_type in MemoryType.ALL:
+            container = self._get_container(memory_type)
 
-            result = list(
-                container.query_items(
-                    query=query,
-                    parameters=parameters,
-                    enable_cross_partition_query=False,
-                    partition_key=project_id,
+            if memory_type == MemoryType.PREFERENCES:
+                # Preferences use user_id, not project_id
+                result = list(
+                    container.query_items(
+                        query="SELECT VALUE COUNT(1) FROM c",
+                        enable_cross_partition_query=True,
+                    )
                 )
-            )
-            stats[f"{key}_count"] = result[0] if result else 0
-
-        # Preferences count (no project_id filter)
-        container = self._get_container("preferences")
-        result = list(
-            container.query_items(
-                query="SELECT VALUE COUNT(1) FROM c",
-                enable_cross_partition_query=True,
-            )
-        )
-        stats["preferences_count"] = result[0] if result else 0
+            else:
+                query = "SELECT VALUE COUNT(1) FROM c WHERE c.project_id = @project_id"
+                parameters = [{"name": "@project_id", "value": project_id}]
+
+                if agent:
+                    query = """
+                    SELECT VALUE COUNT(1) FROM c
+                    WHERE c.project_id = @project_id AND c.agent = @agent
+                    """
+                    parameters.append({"name": "@agent", "value": agent})
+
+                result = list(
+                    container.query_items(
+                        query=query,
+                        parameters=parameters,
+                        enable_cross_partition_query=False,
+                        partition_key=project_id,
+                    )
+                )
+            stats[f"{memory_type}_count"] = result[0] if result else 0
 
         stats["total_count"] = sum(
             stats.get(k, 0) for k in stats if k.endswith("_count")