Spaces:
Paused
Paused
File size: 2,964 Bytes
fcb16d0 414b83b c3b0bb2 906d397 fcb16d0 906d397 fcb16d0 6c52d80 fcb16d0 6c52d80 fcb16d0 6c52d80 fcb16d0 6c52d80 fcb16d0 6c52d80 fcb16d0 6c52d80 fcb16d0 6c52d80 fcb16d0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
# memory_manager.py - Cleaned version
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.")
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):
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
)
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)
# --- Instantiate the class globally to satisfy 'from memory_manager import memory_manager' ---
memory_manager = MemoryManager() |