claude-mpm 3.9.11__py3-none-any.whl → 4.0.3__py3-none-any.whl

claude_mpm/services/memory/indexed_memory.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 #!/usr/bin/env python3
 """Indexed memory service for high-performance memory queries.
 
@@ -16,25 +18,24 @@ WHY indexed memory:
 
 import bisect
 import hashlib
-import json
 import mmap
 import os
 import pickle
 import re
 import time
-from collections import defaultdict, Counter
+from collections import Counter, defaultdict
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta
-from pathlib import Path
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
-from ...core.logger import get_logger
 from ...core.cache import FileSystemCache, get_file_cache
+from ...core.logger import get_logger
 
 
 @dataclass
 class MemoryEntry:
     """Single memory entry with metadata."""
+
     id: str
     agent_id: str
     content: str
@@ -48,6 +49,7 @@ class MemoryEntry:
 @dataclass
 class QueryResult:
     """Result from a memory query."""
+
     entries: List[MemoryEntry]
     total_count: int
     query_time: float
@@ -57,14 +59,14 @@ class QueryResult:
 
 class InvertedIndex:
     """Inverted index for fast text searches.
-    
+
     WHY inverted index:
     - Maps words to document IDs for O(1) lookups
     - Supports boolean queries (AND, OR, NOT)
     - Enables relevance scoring with TF-IDF
     - Efficient for full-text search
     """
-    
+
     def __init__(self):
         # Word -> Set of memory IDs
         self.index: Dict[str, Set[str]] = defaultdict(set)
@@ -72,73 +74,69 @@ class InvertedIndex:
         self.doc_freqs: Dict[str, Counter] = {}
         # Total documents
         self.doc_count = 0
-        
+
         self.logger = get_logger("inverted_index")
-    
+
     def add_entry(self, entry_id: str, text: str):
         """Add entry to inverted index."""
         # Tokenize text
         words = self._tokenize(text)
-        
+
         # Update word frequencies
         self.doc_freqs[entry_id] = Counter(words)
-        
+
         # Update inverted index
         for word in set(words):
             self.index[word].add(entry_id)
-        
+
         self.doc_count += 1
-    
+
     def remove_entry(self, entry_id: str):
         """Remove entry from index."""
         if entry_id not in self.doc_freqs:
             return
-        
+
         # Remove from inverted index
         words = self.doc_freqs[entry_id].keys()
         for word in words:
             self.index[word].discard(entry_id)
             if not self.index[word]:
                 del self.index[word]
-        
+
         # Remove document frequency
         del self.doc_freqs[entry_id]
         self.doc_count -= 1
-    
-    def search(
-        self,
-        query: str,
-        operator: str = 'AND'
-    ) -> Set[str]:
+
+    def search(self, query: str, operator: str = "AND") -> Set[str]:
         """Search index for matching entries.
-        
+
         Args:
             query: Search query
             operator: Boolean operator (AND, OR, NOT)
-        
+
         Returns:
             Set of matching entry IDs
         """
         words = self._tokenize(query)
         if not words:
             return set()
-        
+
         # Get entry sets for each word
         entry_sets = [self.index.get(word, set()) for word in words]
-        
+
         if not entry_sets:
             return set()
-        
+
         # Apply boolean operator
-        if operator == 'AND':
+        if operator == "AND":
             result = entry_sets[0]
             for s in entry_sets[1:]:
                 result = result.intersection(s)
-        elif operator == 'OR':
+        elif operator == "OR":
             result = set()
             for s in entry_sets:
                 result = result.union(s)
-        elif operator == 'NOT':
+        elif operator == "NOT":
             # Return entries that don't contain any query words
             all_entries = set(self.doc_freqs.keys())
             excluded = set()
@@ -147,149 +145,154 @@ class InvertedIndex:
             result = all_entries - excluded
         else:
             result = entry_sets[0]
-        
+
         return result
-    
-    def calculate_relevance(
-        self,
-        entry_id: str,
-        query: str
-    ) -> float:
+
+    def calculate_relevance(self, entry_id: str, query: str) -> float:
         """Calculate TF-IDF relevance score.
-        
+
         Args:
             entry_id: Memory entry ID
             query: Search query
-        
+
         Returns:
             Relevance score (0.0 to 1.0)
         """
         if entry_id not in self.doc_freqs:
             return 0.0
-        
+
         query_words = self._tokenize(query)
         if not query_words:
             return 0.0
-        
+
         score = 0.0
         doc_freq = self.doc_freqs[entry_id]
-        
+
         for word in query_words:
             if word not in doc_freq:
                 continue
-            
+
             # Term frequency
             tf = doc_freq[word] / sum(doc_freq.values())
-            
+
             # Inverse document frequency
             if word in self.index:
                 idf = 1.0 + (self.doc_count / len(self.index[word]))
             else:
                 idf = 1.0
-            
+
             score += tf * idf
-        
+
         # Normalize score
         return min(1.0, score / len(query_words))
-    
+
     def _tokenize(self, text: str) -> List[str]:
         """Tokenize text into words."""
         # Convert to lowercase and split on non-alphanumeric
         text = text.lower()
-        words = re.findall(r'\b[a-z0-9]+\b', text)
-        
+        words = re.findall(r"\b[a-z0-9]+\b", text)
+
         # Remove stop words (simplified list)
-        stop_words = {'the', 'a', 'an', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for'}
+        stop_words = {
+            "the",
+            "a",
+            "an",
+            "and",
+            "or",
+            "but",
+            "in",
+            "on",
+            "at",
+            "to",
+            "for",
+        }
         return [w for w in words if w not in stop_words and len(w) > 2]
-    
+
     def save(self, path: Path):
         """Persist index to disk."""
         data = {
-            'index': dict(self.index),
-            'doc_freqs': dict(self.doc_freqs),
-            'doc_count': self.doc_count
+            "index": dict(self.index),
+            "doc_freqs": dict(self.doc_freqs),
+            "doc_count": self.doc_count,
         }
-        with open(path, 'wb') as f:
+        with open(path, "wb") as f:
             pickle.dump(data, f)
-    
+
     def load(self, path: Path):
         """Load index from disk."""
         if not path.exists():
             return
-        
-        with open(path, 'rb') as f:
+
+        with open(path, "rb") as f:
             data = pickle.load(f)
-        
-        self.index = defaultdict(set, {k: set(v) for k, v in data['index'].items()})
-        self.doc_freqs = data['doc_freqs']
-        self.doc_count = data['doc_count']
+
+        self.index = defaultdict(set, {k: set(v) for k, v in data["index"].items()})
+        self.doc_freqs = data["doc_freqs"]
+        self.doc_count = data["doc_count"]
 
 
 class BTreeIndex:
     """B-tree index for sorted queries.
-    
+
     WHY B-tree index:
     - Maintains sorted order for range queries
     - O(log n) search, insert, delete
     - Efficient for timestamp-based queries
     - Supports pagination
     """
-    
+
     def __init__(self, key_func=None):
         # Sorted list of (key, entry_id) tuples
         self.index: List[Tuple[Any, str]] = []
         self.key_func = key_func or (lambda x: x)
         self.logger = get_logger("btree_index")
-    
+
     def add_entry(self, entry_id: str, key: Any):
         """Add entry to B-tree index."""
         bisect.insort(self.index, (self.key_func(key), entry_id))
-    
+
     def remove_entry(self, entry_id: str):
         """Remove entry from index."""
         self.index = [(k, id) for k, id in self.index if id != entry_id]
-    
+
     def range_search(
-        self,
-        min_key: Any = None,
-        max_key: Any = None,
-        limit: int = None
+        self, min_key: Any = None, max_key: Any = None, limit: int = None
     ) -> List[str]:
         """Search for entries in key range.
-        
+
         Args:
             min_key: Minimum key value (inclusive)
             max_key: Maximum key value (inclusive)
             limit: Maximum results to return
-        
+
         Returns:
             List of matching entry IDs
         """
         # Find range boundaries
         if min_key is not None:
             min_key = self.key_func(min_key)
-            start = bisect.bisect_left(self.index, (min_key, ''))
+            start = bisect.bisect_left(self.index, (min_key, ""))
         else:
             start = 0
-        
+
         if max_key is not None:
             max_key = self.key_func(max_key)
-            end = bisect.bisect_right(self.index, (max_key, '\xff'))
+            end = bisect.bisect_right(self.index, (max_key, "\xff"))
         else:
             end = len(self.index)
-        
+
         # Extract entry IDs
         results = [entry_id for _, entry_id in self.index[start:end]]
-        
+
         if limit:
             results = results[:limit]
-        
+
         return results
-    
+
     def get_recent(self, n: int = 10) -> List[str]:
         """Get n most recent entries."""
         return [entry_id for _, entry_id in self.index[-n:]]
-    
+
     def get_oldest(self, n: int = 10) -> List[str]:
         """Get n oldest entries."""
         return [entry_id for _, entry_id in self.index[:n]]
@@ -297,38 +300,38 @@ class BTreeIndex:
 
 class IndexedMemoryService:
     """High-performance memory service with multiple indexes.
-    
+
     WHY this design:
     - Multiple specialized indexes for different query types
     - LRU cache for frequent queries
     - Incremental index updates for efficiency
     - Memory-mapped files for large datasets
-    
+
     Example:
         memory = IndexedMemoryService()
-        
+
         # Add memory entry
         memory.add_memory(
             agent_id='engineer',
             content='Use dependency injection for testability',
             category='pattern'
         )
-        
+
         # Fast text search
         results = memory.search('dependency injection')
-        
+
         # Range query by timestamp
         recent = memory.get_recent_memories(hours=24)
     """
-    
+
     def __init__(
         self,
         data_dir: Optional[Path] = None,
         cache_size_mb: int = 50,
-        enable_mmap: bool = False
+        enable_mmap: bool = False,
     ):
         """Initialize indexed memory service.
-        
+
         Args:
             data_dir: Directory for persisting indexes
             cache_size_mb: Cache size for query results
@@ -336,51 +339,51 @@ class IndexedMemoryService:
         """
         self.data_dir = data_dir or Path.home() / ".claude-mpm" / "memory"
         self.data_dir.mkdir(parents=True, exist_ok=True)
-        
+
         self.enable_mmap = enable_mmap
-        
+
         # Memory storage
         self.memories: Dict[str, MemoryEntry] = {}
-        
+
         # Indexes
         self.text_index = InvertedIndex()
         self.time_index = BTreeIndex(key_func=lambda dt: dt.timestamp())
         self.agent_index: Dict[str, Set[str]] = defaultdict(set)
         self.category_index: Dict[str, Set[str]] = defaultdict(set)
         self.tag_index: Dict[str, Set[str]] = defaultdict(set)
-        
+
         # Query cache
         self.cache = get_file_cache(max_size_mb=cache_size_mb, default_ttl=300)
-        
+
         # Logger
         self.logger = get_logger("indexed_memory")
-        
+
         # Load existing data
         self._load_indexes()
-    
+
     def add_memory(
         self,
         agent_id: str,
         content: str,
-        category: str = 'general',
+        category: str = "general",
         tags: List[str] = None,
-        metadata: Dict[str, Any] = None
+        metadata: Dict[str, Any] = None,
     ) -> str:
         """Add a new memory entry.
-        
+
         Args:
             agent_id: Agent that owns this memory
             content: Memory content
             category: Memory category
             tags: Optional tags
             metadata: Optional metadata
-        
+
         Returns:
             Memory entry ID
         """
         # Generate ID
         entry_id = self._generate_id(agent_id, content)
-        
+
         # Create entry
         entry = MemoryEntry(
             id=entry_id,
@@ -389,12 +392,12 @@ class IndexedMemoryService:
             category=category,
             timestamp=datetime.now(),
             tags=tags or [],
-            metadata=metadata or {}
+            metadata=metadata or {},
         )
-        
+
         # Store entry
         self.memories[entry_id] = entry
-        
+
         # Update indexes
         self.text_index.add_entry(entry_id, content)
         self.time_index.add_entry(entry_id, entry.timestamp)
@@ -402,13 +405,13 @@ class IndexedMemoryService:
         self.category_index[category].add(entry_id)
         for tag in entry.tags:
             self.tag_index[tag].add(entry_id)
-        
+
         # Invalidate cache
         self.cache.invalidate_pattern("query:*")
-        
+
         self.logger.debug(f"Added memory {entry_id} for agent {agent_id}")
         return entry_id
-    
+
     def search(
         self,
         query: str,
@@ -416,10 +419,10 @@ class IndexedMemoryService:
         category: Optional[str] = None,
         tags: Optional[List[str]] = None,
         limit: int = 50,
-        operator: str = 'AND'
+        operator: str = "AND",
     ) -> QueryResult:
         """Search memories with multiple filters.
-        
+
         Args:
             query: Text search query
             agent_id: Filter by agent
@@ -427,91 +430,87 @@ class IndexedMemoryService:
             tags: Filter by tags
             limit: Maximum results
             operator: Boolean operator for text search
-        
+
         Returns:
             Query results with metadata
         """
         start_time = time.time()
-        
+
         # Generate cache key
         cache_key = f"query:{hashlib.md5(f'{query}:{agent_id}:{category}:{tags}:{limit}:{operator}'.encode()).hexdigest()}"
-        
+
         # Check cache
         cached = self.cache.get(cache_key)
         if cached:
             return QueryResult(
-                entries=cached['entries'],
-                total_count=cached['total_count'],
+                entries=cached["entries"],
+                total_count=cached["total_count"],
                 query_time=0.001,
-                index_used='cache',
-                cache_hit=True
+                index_used="cache",
+                cache_hit=True,
             )
-        
+
         # Text search
         if query:
             matching_ids = self.text_index.search(query, operator)
         else:
             matching_ids = set(self.memories.keys())
-        
+
         # Apply filters
         if agent_id:
             matching_ids &= self.agent_index.get(agent_id, set())
-        
+
         if category:
             matching_ids &= self.category_index.get(category, set())
-        
+
         if tags:
             for tag in tags:
                 matching_ids &= self.tag_index.get(tag, set())
-        
+
         # Get entries and calculate relevance
         entries = []
         for entry_id in matching_ids:
             if entry_id in self.memories:
                 entry = self.memories[entry_id]
                 if query:
-                    entry.relevance_score = self.text_index.calculate_relevance(entry_id, query)
+                    entry.relevance_score = self.text_index.calculate_relevance(
+                        entry_id, query
+                    )
                 entries.append(entry)
-        
+
         # Sort by relevance and timestamp
         entries.sort(key=lambda e: (-e.relevance_score, -e.timestamp.timestamp()))
-        
+
         # Apply limit
         limited_entries = entries[:limit]
-        
+
         # Cache result
-        cache_data = {
-            'entries': limited_entries,
-            'total_count': len(entries)
-        }
+        cache_data = {"entries": limited_entries, "total_count": len(entries)}
         self.cache.put(cache_key, cache_data, ttl=300)
-        
+
         # Return result
         return QueryResult(
             entries=limited_entries,
             total_count=len(entries),
             query_time=time.time() - start_time,
-            index_used='text_index' if query else 'full_scan'
+            index_used="text_index" if query else "full_scan",
         )
-    
+
     def get_recent_memories(
-        self,
-        hours: Optional[int] = None,
-        days: Optional[int] = None,
-        limit: int = 50
+        self, hours: Optional[int] = None, days: Optional[int] = None, limit: int = 50
     ) -> QueryResult:
         """Get recent memories within time range.
-        
+
         Args:
             hours: Hours to look back
             days: Days to look back
             limit: Maximum results
-        
+
         Returns:
             Recent memories
         """
         start_time = time.time()
-        
+
         # Calculate time range
         now = datetime.now()
         if hours:
@@ -520,117 +519,120 @@ class IndexedMemoryService:
             min_time = now - timedelta(days=days)
         else:
             min_time = None
-        
+
         # Use time index for range search
         entry_ids = self.time_index.range_search(
-            min_key=min_time,
-            max_key=now,
-            limit=limit
+            min_key=min_time, max_key=now, limit=limit
         )
-        
+
         # Get entries
         entries = [self.memories[id] for id in entry_ids if id in self.memories]
-        
+
         return QueryResult(
             entries=entries,
             total_count=len(entries),
             query_time=time.time() - start_time,
-            index_used='time_index'
+            index_used="time_index",
         )
-    
-    def get_agent_memories(
-        self,
-        agent_id: str,
-        limit: int = 50
-    ) -> QueryResult:
+
+    def get_agent_memories(self, agent_id: str, limit: int = 50) -> QueryResult:
         """Get all memories for a specific agent.
-        
+
         Args:
             agent_id: Agent ID
             limit: Maximum results
-        
+
         Returns:
             Agent's memories
         """
         start_time = time.time()
-        
+
         # Use agent index
         entry_ids = list(self.agent_index.get(agent_id, set()))[:limit]
-        
+
         # Get entries
         entries = [self.memories[id] for id in entry_ids if id in self.memories]
-        
+
         # Sort by timestamp
         entries.sort(key=lambda e: -e.timestamp.timestamp())
-        
+
         return QueryResult(
             entries=entries,
             total_count=len(self.agent_index.get(agent_id, set())),
             query_time=time.time() - start_time,
-            index_used='agent_index'
+            index_used="agent_index",
         )
-    
+
     def _generate_id(self, agent_id: str, content: str) -> str:
         """Generate unique ID for memory entry."""
         timestamp = datetime.now().isoformat()
         hash_input = f"{agent_id}:{content[:100]}:{timestamp}"
         return hashlib.md5(hash_input.encode()).hexdigest()[:12]
-    
+
     def _save_indexes(self):
         """Persist all indexes to disk."""
         # Save text index
         self.text_index.save(self.data_dir / "text_index.pkl")
-        
+
         # Save other indexes
-        with open(self.data_dir / "indexes.pkl", 'wb') as f:
-            pickle.dump({
-                'memories': self.memories,
-                'time_index': self.time_index.index,
-                'agent_index': dict(self.agent_index),
-                'category_index': dict(self.category_index),
-                'tag_index': dict(self.tag_index)
-            }, f)
-        
+        with open(self.data_dir / "indexes.pkl", "wb") as f:
+            pickle.dump(
+                {
+                    "memories": self.memories,
+                    "time_index": self.time_index.index,
+                    "agent_index": dict(self.agent_index),
+                    "category_index": dict(self.category_index),
+                    "tag_index": dict(self.tag_index),
+                },
+                f,
+            )
+
         self.logger.info(f"Saved {len(self.memories)} memories to disk")
-    
+
     def _load_indexes(self):
         """Load indexes from disk."""
         # Load text index
         text_index_path = self.data_dir / "text_index.pkl"
         if text_index_path.exists():
             self.text_index.load(text_index_path)
-        
+
         # Load other indexes
         indexes_path = self.data_dir / "indexes.pkl"
         if indexes_path.exists():
-            with open(indexes_path, 'rb') as f:
+            with open(indexes_path, "rb") as f:
                 data = pickle.load(f)
-            
-            self.memories = data.get('memories', {})
-            self.time_index.index = data.get('time_index', [])
-            self.agent_index = defaultdict(set, {k: set(v) for k, v in data.get('agent_index', {}).items()})
-            self.category_index = defaultdict(set, {k: set(v) for k, v in data.get('category_index', {}).items()})
-            self.tag_index = defaultdict(set, {k: set(v) for k, v in data.get('tag_index', {}).items()})
-            
+
+            self.memories = data.get("memories", {})
+            self.time_index.index = data.get("time_index", [])
+            self.agent_index = defaultdict(
+                set, {k: set(v) for k, v in data.get("agent_index", {}).items()}
+            )
+            self.category_index = defaultdict(
+                set, {k: set(v) for k, v in data.get("category_index", {}).items()}
+            )
+            self.tag_index = defaultdict(
+                set, {k: set(v) for k, v in data.get("tag_index", {}).items()}
+            )
+
             self.logger.info(f"Loaded {len(self.memories)} memories from disk")
-    
+
     def get_stats(self) -> Dict[str, Any]:
         """Get memory service statistics."""
         return {
-            'total_memories': len(self.memories),
-            'agents': len(self.agent_index),
-            'categories': len(self.category_index),
-            'tags': len(self.tag_index),
-            'index_size': {
-                'text': len(self.text_index.index),
-                'time': len(self.time_index.index),
-                'agent': sum(len(v) for v in self.agent_index.values()),
-                'category': sum(len(v) for v in self.category_index.values()),
-                'tag': sum(len(v) for v in self.tag_index.values())
+            "total_memories": len(self.memories),
+            "agents": len(self.agent_index),
+            "categories": len(self.category_index),
+            "tags": len(self.tag_index),
+            "index_size": {
+                "text": len(self.text_index.index),
+                "time": len(self.time_index.index),
+                "agent": sum(len(v) for v in self.agent_index.values()),
+                "category": sum(len(v) for v in self.category_index.values()),
+                "tag": sum(len(v) for v in self.tag_index.values()),
             },
-            'cache_stats': self.cache.get_stats()
+            "cache_stats": self.cache.get_stats(),
         }
-    
+
     def cleanup(self):
         """Save indexes and cleanup resources."""
         self._save_indexes()
@@ -645,4 +647,4 @@ def get_indexed_memory() -> IndexedMemoryService:
     global _memory_service
     if _memory_service is None:
         _memory_service = IndexedMemoryService()
-    return _memory_service
+    return _memory_service