Spaces:
Running
Running
| # FAISS vector store module | |
| import faiss | |
| import pickle | |
| import numpy as np | |
| from config.settings import Settings | |
| from core.llm import embed_text | |
| class FaissStore: | |
| def __init__(self): | |
| self.index = faiss.IndexFlatL2(Settings.EMBED_DIM) | |
| self.meta = [] | |
| def add(self, emb, meta): | |
| self.index.add(np.array(emb, dtype="float32")) | |
| self.meta.extend(meta) | |
| def clear(self): | |
| """Clear all vectors and metadata - for retraining""" | |
| self.index = faiss.IndexFlatL2(Settings.EMBED_DIM) | |
| self.meta = [] | |
| print("🗑️ FAISS store cleared for fresh retraining") | |
| def save(self, user_id: str = "user_001"): | |
| """Save FAISS index to per-user directory""" | |
| from pathlib import Path | |
| # Use per-user directory | |
| if user_id: | |
| user_faiss_dir = Settings.STORAGE / "users" / user_id / "faiss" | |
| user_faiss_dir.mkdir(parents=True, exist_ok=True) | |
| else: | |
| user_faiss_dir = Settings.FAISS_DIR | |
| user_faiss_dir.mkdir(parents=True, exist_ok=True) | |
| idx_path = user_faiss_dir / "index.faiss" | |
| meta_path = user_faiss_dir / "meta.pkl" | |
| print(f"💾 Saving FAISS to: {idx_path}") | |
| faiss.write_index(self.index, str(idx_path)) | |
| with open(meta_path, "wb") as f: | |
| pickle.dump(self.meta, f) | |
| print(f"✅ FAISS saved: {self.index.ntotal} vectors, {len(self.meta)} metadata entries") | |
| def load_or_create(user_id: str = "user_001", fresh: bool = False): | |
| """ | |
| Load or create FAISS store for specific user. | |
| Args: | |
| user_id: User identifier | |
| fresh: If True, create fresh store ignoring existing data (for retraining) | |
| """ | |
| from pathlib import Path | |
| store = FaissStore() | |
| # If fresh=True, return empty store for clean retraining | |
| if fresh: | |
| print(f"🆕 Creating fresh FAISS store for user {user_id}") | |
| return store | |
| # Use per-user directory | |
| if user_id: | |
| user_faiss_dir = Settings.STORAGE / "users" / user_id / "faiss" | |
| idx = user_faiss_dir / "index.faiss" | |
| meta = user_faiss_dir / "meta.pkl" | |
| else: | |
| idx = Settings.FAISS_DIR / "index.faiss" | |
| meta = Settings.FAISS_DIR / "meta.pkl" | |
| print(f"📂 Loading FAISS from: {idx}") | |
| if idx.exists(): | |
| store.index = faiss.read_index(str(idx)) | |
| print(f"✅ FAISS loaded: {store.index.ntotal} vectors") | |
| else: | |
| print(f"⚠️ FAISS index not found at {idx}, creating new") | |
| if meta.exists(): | |
| store.meta = pickle.load(open(meta, "rb")) | |
| print(f"✅ Metadata loaded: {len(store.meta)} entries") | |
| else: | |
| print(f"⚠️ Metadata not found at {meta}") | |
| return store | |
| def delete_index(user_id: str = "user_001"): | |
| """Delete FAISS index files for user - for clean retraining""" | |
| from pathlib import Path | |
| import shutil | |
| if user_id: | |
| user_faiss_dir = Settings.STORAGE / "users" / user_id / "faiss" | |
| else: | |
| user_faiss_dir = Settings.FAISS_DIR | |
| if user_faiss_dir.exists(): | |
| shutil.rmtree(user_faiss_dir) | |
| user_faiss_dir.mkdir(parents=True, exist_ok=True) | |
| print(f"🗑️ Deleted FAISS index for user {user_id}") | |
| def search(self, query, k=5): | |
| """Search with automatic query embedding""" | |
| if self.index.ntotal == 0: | |
| print("⚠️ FAISS index is empty - no vectors to search") | |
| return [] | |
| # Embed query text if it's a string | |
| if isinstance(query, str): | |
| query_vector = embed_text(query) | |
| if query_vector is None: | |
| print("⚠️ Failed to embed query") | |
| return [] | |
| else: | |
| query_vector = query | |
| query_vector = np.array([query_vector], dtype="float32") | |
| # Limit k to available vectors | |
| actual_k = min(k, self.index.ntotal) | |
| _, ids = self.index.search(query_vector, actual_k) | |
| results = [] | |
| for i in ids[0]: | |
| if 0 <= i < len(self.meta): | |
| meta = self.meta[i] | |
| results.append({ | |
| "text": meta.get("text", ""), | |
| "metadata": meta | |
| }) | |
| return results | |