Update memory_manager.py

memory_manager.py

@@ -1,13 +1,7 @@
-# memory_manager.py - Production-ready memory management with FAISS
+# memory_manager.py - Cleaned version
 
 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
@@ -16,341 +10,59 @@ 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:
-    """
-    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
+    def __init__(self):
+        self.embeddings = HuggingFaceEmbeddings(model_name=MODEL_NAME)
         self.vector_store = self._load_or_create_vector_store()
-        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:
-                with open(self.metadata_path, 'r', encoding='utf-8') as f:
-                    return json.load(f)
-            except Exception as e:
-                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
+
+    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()
-            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")
-        
+
+    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):
             try:
-                logger.info("Loading existing memory from disk...")
-                vector_store = FAISS.load_local(
-                    folder_path=self.memory_dir,
+                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  # Required for loading
+                    allow_dangerous_deserialization=True
                 )
-                logger.info("Successfully loaded existing memory index")
-                return vector_store
-                
             except Exception as e:
-                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)
+                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:
-            logger.info("No existing memory found. Creating new index...")
+            print("🧠 No existing memory found. Creating a new one.")
             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
 
-# Create singleton instance for import compatibility
-memory_manager = MemoryManager()
+    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)
 
-# Export the instance for backward compatibility
-__all__ = ['MemoryManager', 'memory_manager']
+# --- Instantiate the class globally to satisfy 'from memory_manager import memory_manager' ---
+memory_manager = MemoryManager()