sufian
79e345097e
Streamline system prompts, rename functions for clarity, and improve search query generation. Fix memory query edge cases and enhance robustness when no indexed data exists. Minor wording adjustments and structure improvements for better maintainability.
42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
import faiss
|
|
import os
|
|
import pickle
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
# Embedder: small model is fast & good enough
|
|
embedder = SentenceTransformer("all-MiniLM-L6-v2") # Replaceable
|
|
|
|
class Memory:
|
|
def __init__(self, index_path="kshama.index", metadata_path="memory_meta.pkl"):
|
|
self.index_path = index_path
|
|
self.metadata_path = metadata_path
|
|
self.index = None
|
|
self.metadata = []
|
|
self._load()
|
|
|
|
def _load(self):
|
|
if os.path.exists(self.index_path):
|
|
self.index = faiss.read_index(self.index_path)
|
|
with open(self.metadata_path, "rb") as f:
|
|
self.metadata = pickle.load(f)
|
|
else:
|
|
self.index = faiss.IndexFlatL2(384) # Depends on embedder output dim
|
|
|
|
def add(self, text, tags=None):
|
|
vec = embedder.encode([text])
|
|
self.index.add(vec)
|
|
self.metadata.append({"text": text, "tags": tags or []})
|
|
self._save()
|
|
|
|
def query(self, text, top_k=5):
|
|
if self.index.ntotal == 0:
|
|
return []
|
|
vec = embedder.encode([text])
|
|
D, I = self.index.search(vec, top_k)
|
|
return [self.metadata[i]["text"] for i in I[0] if 0 <= i < len(self.metadata)]
|
|
|
|
def _save(self):
|
|
faiss.write_index(self.index, self.index_path)
|
|
with open(self.metadata_path, "wb") as f:
|
|
pickle.dump(self.metadata, f)
|