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

alma/storage/azure_cosmos.py

@@ -22,6 +22,7 @@ from datetime import datetime, timezone
 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,
@@ -61,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,
@@ -121,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
@@ -146,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,
@@ -231,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,
@@ -260,12 +379,16 @@ class AzureCosmosStorage(StorageBackend):
         }
 
         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,
@@ -285,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,
@@ -307,12 +432,16 @@ class AzureCosmosStorage(StorageBackend):
         }
 
         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,
@@ -331,12 +460,16 @@ class AzureCosmosStorage(StorageBackend):
         }
 
         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,
@@ -358,6 +491,10 @@ class AzureCosmosStorage(StorageBackend):
         }
 
         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
 
@@ -372,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
@@ -424,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(
@@ -436,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
@@ -494,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(
@@ -502,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}]
@@ -520,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(
@@ -531,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 = """
@@ -583,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(
@@ -593,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 = """
@@ -639,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 ====================
@@ -647,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]
-        doc["project_id"]
+        if not doc:
+            return False
 
         # Apply updates
         for key, value in updates.items():
@@ -682,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)
+
+        # Use optimized point read with cache fallback
+        doc = self._point_read_document(MemoryType.HEURISTICS, heuristic_id, project_id)
 
-        if not items:
+        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
@@ -712,30 +882,34 @@ class AzureCosmosStorage(StorageBackend):
         self,
         heuristic_id: str,
         new_confidence: float,
+        project_id: Optional[str] = None,
     ) -> 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.
+        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("heuristics")
+        container = self._get_container(MemoryType.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,
-            )
-        )
+        # Use optimized point read with cache fallback
+        doc = self._point_read_document(MemoryType.HEURISTICS, heuristic_id, project_id)
 
-        if not items:
+        if not doc:
             return False
 
-        doc = items[0]
         doc["confidence"] = new_confidence
 
         container.replace_item(item=heuristic_id, body=doc)
@@ -748,30 +922,36 @@ class AzureCosmosStorage(StorageBackend):
         self,
         knowledge_id: str,
         new_confidence: float,
+        project_id: Optional[str] = None,
     ) -> 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().
+        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("knowledge")
+        container = self._get_container(MemoryType.DOMAIN_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,
-            )
+        # Use optimized point read with cache fallback
+        doc = self._point_read_document(
+            MemoryType.DOMAIN_KNOWLEDGE, knowledge_id, project_id
         )
 
-        if not items:
+        if not doc:
             return False
 
-        doc = items[0]
         doc["confidence"] = new_confidence
 
         container.replace_item(item=knowledge_id, body=doc)
@@ -789,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
@@ -832,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
@@ -868,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
@@ -908,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")