memory.py

from sentence_transformers import SentenceTransformer
import numpy as np
import faiss
import json
import os

model = SentenceTransformer("all-MiniLM-L6-v2")
MEMORY_LOG = "memory_log.json"

def init_memory():
    dim = 384
    index = faiss.IndexFlatL2(dim)
    memory = {"index": index, "texts": []}
    if os.path.exists(MEMORY_LOG):
        with open(MEMORY_LOG, "r") as f:
            memory["texts"] = json.load(f)
        vectors = np.array([model.encode([text])[0] for text in memory["texts"]])
        if len(vectors) > 0:
            memory["index"].add(vectors)
    return memory

def add_to_memory(text, memory):
    vec = model.encode([text])[0]
    memory["index"].add(np.array([vec]))
    memory["texts"].append(text)
    with open(MEMORY_LOG, "w") as f:
        json.dump(memory["texts"], f)

def search_memory(query, memory, k=5):
    vec = model.encode([query])[0]
    D, I = memory["index"].search(np.array([vec]), k)
    return [memory["texts"][i] for i in I[0] if i < len(memory["texts"])]