Spaces:
Running
Running
File size: 4,563 Bytes
09801ca | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | # 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")
@staticmethod
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
@staticmethod
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
|