Create vector_store.py

vector_store.py

@@ -0,0 +1,39 @@
+# vector_store.py
+import faiss
+import numpy as np
+import pickle
+from typing import List
+
+class SimpleVectorStore:
+    def __init__(self, dim, index_path=None):
+        self.dim = dim
+        self.index = faiss.IndexFlatIP(dim)  # cosine if vectors normalized
+        self.metadata = []
+        self.index_path = index_path
+
+    def add(self, vec: np.ndarray, meta: dict):
+        if vec.ndim == 1:
+            vec = vec.reshape(1, -1)
+        self.index.add(vec.astype("float32"))
+        self.metadata.append(meta)
+
+    def search(self, query_vec: np.ndarray, k=5):
+        if query_vec.ndim == 1:
+            query_vec = query_vec.reshape(1, -1)
+        D, I = self.index.search(query_vec.astype("float32"), k)
+        results = []
+        for dist, idx in zip(D[0], I[0]):
+            if idx == -1 or idx >= len(self.metadata):
+                continue
+            results.append((float(dist), self.metadata[idx]))
+        return results
+
+    def save(self, path_prefix):
+        faiss.write_index(self.index, f"{path_prefix}.index")
+        with open(f"{path_prefix}.meta", "wb") as f:
+            pickle.dump(self.metadata, f)
+
+    def load(self, path_prefix):
+        self.index = faiss.read_index(f"{path_prefix}.index")
+        with open(f"{path_prefix}.meta", "rb") as f:
+            self.metadata = pickle.load(f)