Update memory_manager.py

memory_manager.py

@@ -1,5 +1,13 @@
+# memory_manager.py - Production-ready memory management with FAISS
+
 import os
 import shutil
+import json
+import hashlib
+from datetime import datetime
+from typing import List, Dict, Optional, Any
+import logging
+
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 from langchain.docstore.document import Document
@@ -8,127 +16,341 @@ from langchain.docstore.document import Document
 MEMORY_DIR = "memory"
 INDEX_NAME = "faiss"
 MODEL_NAME = "all-MiniLM-L6-v2"
+METADATA_FILE = "memory_metadata.json"
+
+# Get logger
+logger = logging.getLogger(__name__)
 
 class MemoryManager:
-    def __init__(self):
-        self.embeddings = HuggingFaceEmbeddings(model_name=MODEL_NAME)
+    """
+    Manages long-term memory using FAISS vector store with semantic search capabilities.
+    """
+    
+    def __init__(self, memory_dir: str = MEMORY_DIR, model_name: str = MODEL_NAME):
+        """
+        Initialize the memory manager with embeddings and vector store.
+        
+        Args:
+            memory_dir: Directory to store memory files
+            model_name: Name of the HuggingFace embedding model
+        """
+        self.memory_dir = memory_dir
+        self.model_name = model_name
+        self.metadata_path = os.path.join(memory_dir, METADATA_FILE)
+        
+        logger.info(f"Initializing MemoryManager with model: {model_name}")
+        
+        # Initialize embeddings
+        try:
+            self.embeddings = HuggingFaceEmbeddings(
+                model_name=model_name,
+                model_kwargs={'device': 'cpu'},
+                encode_kwargs={'normalize_embeddings': True}
+            )
+            logger.info("Embeddings model loaded successfully")
+        except Exception as e:
+            logger.error(f"Failed to load embeddings model: {e}")
+            raise
+        
+        # Load or create vector store
         self.vector_store = self._load_or_create_vector_store()
-
-    def reset_memory(self):
-        """Removes the memory directory and re-initializes a new, empty index."""
-        if os.path.exists(MEMORY_DIR):
-            shutil.rmtree(MEMORY_DIR)
-            os.makedirs(MEMORY_DIR, exist_ok=True)
-            print("🧠 Memory reset successfully.")
-            self.vector_store = self._create_new_index()
-
-    def _load_or_create_vector_store(self):
-        """Loads FAISS index or creates a new one, handling potential corruption."""
-        index_path = os.path.join(MEMORY_DIR, f"{INDEX_NAME}.faiss")
-        if os.path.exists(index_path):
+        self.metadata = self._load_metadata()
+    
+    def _load_metadata(self) -> Dict[str, Any]:
+        """Load metadata about stored memories."""
+        if os.path.exists(self.metadata_path):
             try:
-                print("🧠 Loading existing memory from disk...")
-                return FAISS.load_local(
-                    folder_path=MEMORY_DIR, 
-                    embeddings=self.embeddings,
-                    index_name=INDEX_NAME,
-                    allow_dangerous_deserialization=True
-                )
+                with open(self.metadata_path, 'r', encoding='utf-8') as f:
+                    return json.load(f)
             except Exception as e:
-                print(f"⚠️ Error loading memory index: {e}. Rebuilding index.")
-                shutil.rmtree(MEMORY_DIR)
-                os.makedirs(MEMORY_DIR, exist_ok=True)
-                return self._create_new_index()
-        else:
-            print("🧠 No existing memory found. Creating a new one.")
-            return self._create_new_index()
-
-    def _create_new_index(self):
-        """Creates a fresh, empty FAISS index."""
-        dummy_doc = [Document(page_content="Initial memory entry.")]
-        # Note: If memory needs to be truly empty, use a small, persistent dummy doc
-        # or handle an empty index creation if FAISS allows it. Keeping dummy for robustness.
-        vs = FAISS.from_documents(dummy_doc, self.embeddings)
-        vs.save_local(folder_path=MEMORY_DIR, index_name=INDEX_NAME)
-        return vs
-
-    def add_to_memory(self, text_to_add: str, metadata: dict):
-        print(f"📝 Adding new memory: {text_to_add[:100]}...")
-        doc = Document(page_content=text_to_add, metadata=metadata)
-        self.vector_store.add_documents([doc])
-        self.vector_store.save_local(folder_path=MEMORY_DIR, index_name=INDEX_NAME)
-
-    def retrieve_relevant_memories(self, query: str, k: int = 5) -> list[Document]:
-        print(f"🔍 Searching memory for: {query[:50]}...")
-        return self.vector_store.similarity_search(query, k=k)
-
-# --- FIX: Instantiate the class globally to satisfy 'from memory_manager import memory_manager' ---
-# The name of the instance must match the name being imported by graph.py
-memory_manager = MemoryManager()
-import os
-import shutil
-from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_community.vectorstores import FAISS
-from langchain.docstore.document import Document
-
-# --- Configuration ---
-MEMORY_DIR = "memory"
-INDEX_NAME = "faiss"
-MODEL_NAME = "all-MiniLM-L6-v2"
-
-class MemoryManager:
-    def __init__(self):
-        self.embeddings = HuggingFaceEmbeddings(model_name=MODEL_NAME)
-        self.vector_store = self._load_or_create_vector_store()
-
-    def reset_memory(self):
-        """Removes the memory directory and re-initializes a new, empty index."""
-        if os.path.exists(MEMORY_DIR):
-            shutil.rmtree(MEMORY_DIR)
-            os.makedirs(MEMORY_DIR, exist_ok=True)
-            print("🧠 Memory reset successfully.")
+                logger.warning(f"Could not load metadata: {e}")
+        
+        return {
+            "created_at": datetime.now().isoformat(),
+            "total_memories": 0,
+            "last_updated": None,
+            "memory_hashes": set()
+        }
+    
+    def _save_metadata(self) -> None:
+        """Save metadata about stored memories."""
+        try:
+            # Convert set to list for JSON serialization
+            metadata_copy = self.metadata.copy()
+            if "memory_hashes" in metadata_copy and isinstance(metadata_copy["memory_hashes"], set):
+                metadata_copy["memory_hashes"] = list(metadata_copy["memory_hashes"])
+            
+            with open(self.metadata_path, 'w', encoding='utf-8') as f:
+                json.dump(metadata_copy, f, indent=2)
+        except Exception as e:
+            logger.warning(f"Could not save metadata: {e}")
+    
+    def reset_memory(self) -> bool:
+        """
+        Removes all memory data and re-initializes a new, empty index.
+        
+        Returns:
+            True if successful, False otherwise
+        """
+        try:
+            if os.path.exists(self.memory_dir):
+                shutil.rmtree(self.memory_dir)
+                logger.info("Removed existing memory directory")
+            
+            os.makedirs(self.memory_dir, exist_ok=True)
+            logger.info("Memory reset successfully")
+            
+            # Reinitialize
             self.vector_store = self._create_new_index()
-
-    def _load_or_create_vector_store(self):
-        """Loads FAISS index or creates a new one, handling potential corruption."""
-        index_path = os.path.join(MEMORY_DIR, f"{INDEX_NAME}.faiss")
+            self.metadata = self._load_metadata()
+            
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to reset memory: {e}")
+            return False
+    
+    def _load_or_create_vector_store(self) -> FAISS:
+        """
+        Loads FAISS index from disk or creates a new one if not found or corrupted.
+        
+        Returns:
+            FAISS vector store instance
+        """
+        index_path = os.path.join(self.memory_dir, f"{INDEX_NAME}.faiss")
+        
         if os.path.exists(index_path):
             try:
-                print("🧠 Loading existing memory from disk...")
-                return FAISS.load_local(
-                    folder_path=MEMORY_DIR, 
+                logger.info("Loading existing memory from disk...")
+                vector_store = FAISS.load_local(
+                    folder_path=self.memory_dir,
                     embeddings=self.embeddings,
                     index_name=INDEX_NAME,
-                    allow_dangerous_deserialization=True
+                    allow_dangerous_deserialization=True  # Required for loading
                 )
+                logger.info("Successfully loaded existing memory index")
+                return vector_store
+                
             except Exception as e:
-                print(f"⚠️ Error loading memory index: {e}. Rebuilding index.")
-                shutil.rmtree(MEMORY_DIR)
-                os.makedirs(MEMORY_DIR, exist_ok=True)
+                logger.warning(f"Error loading memory index: {e}. Creating new index...")
+                
+                # Backup corrupted index
+                backup_dir = os.path.join(self.memory_dir, f"backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}")
+                try:
+                    shutil.move(self.memory_dir, backup_dir)
+                    logger.info(f"Corrupted index backed up to: {backup_dir}")
+                except Exception as backup_error:
+                    logger.warning(f"Could not backup corrupted index: {backup_error}")
+                
+                os.makedirs(self.memory_dir, exist_ok=True)
                 return self._create_new_index()
         else:
-            print("🧠 No existing memory found. Creating a new one.")
+            logger.info("No existing memory found. Creating new index...")
             return self._create_new_index()
+    
+    def _create_new_index(self) -> FAISS:
+        """
+        Creates a fresh, empty FAISS index.
+        
+        Returns:
+            New FAISS vector store instance
+        """
+        os.makedirs(self.memory_dir, exist_ok=True)
+        
+        # Create with a minimal initialization document
+        init_doc = Document(
+            page_content="System initialized.",
+            metadata={
+                "type": "system",
+                "timestamp": datetime.now().isoformat(),
+                "importance": 0.0
+            }
+        )
+        
+        try:
+            vector_store = FAISS.from_documents([init_doc], self.embeddings)
+            vector_store.save_local(folder_path=self.memory_dir, index_name=INDEX_NAME)
+            logger.info("Created new memory index")
+            return vector_store
+            
+        except Exception as e:
+            logger.error(f"Failed to create new index: {e}")
+            raise
+    
+    def _compute_hash(self, text: str) -> str:
+        """Compute hash of text to detect duplicates."""
+        return hashlib.md5(text.encode('utf-8')).hexdigest()
+    
+    def add_to_memory(
+        self, 
+        text_to_add: str, 
+        metadata: Optional[Dict[str, Any]] = None,
+        importance: float = 0.5,
+        check_duplicate: bool = True
+    ) -> bool:
+        """
+        Add new information to memory with metadata.
+        
+        Args:
+            text_to_add: Text content to store
+            metadata: Additional metadata for the memory
+            importance: Importance score (0.0 to 1.0)
+            check_duplicate: Whether to check for duplicate memories
+            
+        Returns:
+            True if memory was added, False if duplicate detected
+        """
+        if not text_to_add or not text_to_add.strip():
+            logger.warning("Attempted to add empty memory")
+            return False
+        
+        # Check for duplicates
+        text_hash = self._compute_hash(text_to_add)
+        if check_duplicate:
+            if "memory_hashes" not in self.metadata:
+                self.metadata["memory_hashes"] = set()
+            elif not isinstance(self.metadata["memory_hashes"], set):
+                self.metadata["memory_hashes"] = set(self.metadata["memory_hashes"])
+            
+            if text_hash in self.metadata["memory_hashes"]:
+                logger.debug(f"Duplicate memory detected, skipping: {text_to_add[:50]}...")
+                return False
+        
+        # Prepare metadata
+        if metadata is None:
+            metadata = {}
+        
+        metadata.update({
+            "timestamp": datetime.now().isoformat(),
+            "importance": max(0.0, min(1.0, importance)),  # Clamp between 0 and 1
+            "hash": text_hash,
+            "length": len(text_to_add)
+        })
+        
+        # Add to vector store
+        try:
+            logger.info(f"Adding new memory: {text_to_add[:100]}...")
+            doc = Document(page_content=text_to_add, metadata=metadata)
+            self.vector_store.add_documents([doc])
+            self.vector_store.save_local(folder_path=self.memory_dir, index_name=INDEX_NAME)
+            
+            # Update metadata
+            if "memory_hashes" not in self.metadata:
+                self.metadata["memory_hashes"] = set()
+            elif not isinstance(self.metadata["memory_hashes"], set):
+                self.metadata["memory_hashes"] = set(self.metadata["memory_hashes"])
+            
+            self.metadata["memory_hashes"].add(text_hash)
+            self.metadata["total_memories"] = self.metadata.get("total_memories", 0) + 1
+            self.metadata["last_updated"] = datetime.now().isoformat()
+            self._save_metadata()
+            
+            logger.info("Memory added successfully")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to add memory: {e}")
+            return False
+    
+    def retrieve_relevant_memories(
+        self, 
+        query: str, 
+        k: int = 5,
+        score_threshold: Optional[float] = None,
+        filter_metadata: Optional[Dict[str, Any]] = None
+    ) -> List[Document]:
+        """
+        Retrieve memories relevant to a query using semantic search.
+        
+        Args:
+            query: Search query
+            k: Number of results to return
+            score_threshold: Minimum similarity score (if supported by vector store)
+            filter_metadata: Metadata filters to apply
+            
+        Returns:
+            List of relevant documents
+        """
+        if not query or not query.strip():
+            logger.warning("Empty query for memory retrieval")
+            return []
+        
+        try:
+            logger.info(f"Searching memory for: {query[:50]}...")
+            
+            # Perform similarity search
+            if filter_metadata:
+                # Note: FAISS doesn't natively support metadata filtering
+                # This would need custom implementation
+                results = self.vector_store.similarity_search(query, k=k*2)  # Get extra to filter
+                
+                # Manual filtering
+                filtered = []
+                for doc in results:
+                    match = all(
+                        doc.metadata.get(key) == value 
+                        for key, value in filter_metadata.items()
+                    )
+                    if match:
+                        filtered.append(doc)
+                        if len(filtered) >= k:
+                            break
+                results = filtered
+            else:
+                results = self.vector_store.similarity_search(query, k=k)
+            
+            logger.info(f"Retrieved {len(results)} relevant memories")
+            return results
+            
+        except Exception as e:
+            logger.error(f"Memory retrieval failed: {e}")
+            return []
+    
+    def get_memory_stats(self) -> Dict[str, Any]:
+        """
+        Get statistics about the memory store.
+        
+        Returns:
+            Dictionary with memory statistics
+        """
+        stats = {
+            "total_memories": self.metadata.get("total_memories", 0),
+            "created_at": self.metadata.get("created_at"),
+            "last_updated": self.metadata.get("last_updated"),
+            "memory_dir_size_mb": 0.0
+        }
+        
+        # Calculate directory size
+        if os.path.exists(self.memory_dir):
+            total_size = 0
+            for dirpath, _, filenames in os.walk(self.memory_dir):
+                for filename in filenames:
+                    filepath = os.path.join(dirpath, filename)
+                    try:
+                        total_size += os.path.getsize(filepath)
+                    except Exception:
+                        pass
+            stats["memory_dir_size_mb"] = round(total_size / (1024 * 1024), 2)
+        
+        return stats
+    
+    def cleanup_old_memories(self, days_to_keep: int = 30) -> int:
+        """
+        Remove memories older than specified days (not implemented for FAISS).
+        
+        Args:
+            days_to_keep: Number of days to keep memories
+            
+        Returns:
+            Number of memories removed
+        """
+        # FAISS doesn't support selective deletion easily
+        # This would require rebuilding the entire index
+        logger.warning("Memory cleanup not implemented for FAISS backend")
+        return 0
 
-    def _create_new_index(self):
-        """Creates a fresh, empty FAISS index."""
-        dummy_doc = [Document(page_content="Initial memory entry.")]
-        # Note: If memory needs to be truly empty, use a small, persistent dummy doc
-        # or handle an empty index creation if FAISS allows it. Keeping dummy for robustness.
-        vs = FAISS.from_documents(dummy_doc, self.embeddings)
-        vs.save_local(folder_path=MEMORY_DIR, index_name=INDEX_NAME)
-        return vs
-
-    def add_to_memory(self, text_to_add: str, metadata: dict):
-        print(f"📝 Adding new memory: {text_to_add[:100]}...")
-        doc = Document(page_content=text_to_add, metadata=metadata)
-        self.vector_store.add_documents([doc])
-        self.vector_store.save_local(folder_path=MEMORY_DIR, index_name=INDEX_NAME)
-
-    def retrieve_relevant_memories(self, query: str, k: int = 5) -> list[Document]:
-        print(f"🔍 Searching memory for: {query[:50]}...")
-        return self.vector_store.similarity_search(query, k=k)
+# Create singleton instance for import compatibility
+memory_manager = MemoryManager()
 
-# --- FIX: Instantiate the class globally to satisfy 'from memory_manager import memory_manager' ---
-# The name of the instance must match the name being imported by graph.py
-memory_manager = MemoryManager()
+# Export the instance for backward compatibility
+__all__ = ['MemoryManager', 'memory_manager']