vector_store.py

# 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)