utils/retriever.py

# ─────────────────────────────────────────────────────────────
# utils/retriever.py
# FAISS vector store — build, save, load, search
# ─────────────────────────────────────────────────────────────

import faiss
import numpy as np
import os
import json


class FAISSRetriever:
    """
    Builds and searches a FAISS vector index
    from your embedded documents.
    """

    def __init__(self, index_path: str = None):
        self.index_path = index_path or os.getenv("FAISS_INDEX_PATH", "vector_store/index.faiss")
        self.docs_path  = self.index_path.replace(".faiss", "_docs.json")
        self.index      = None
        self.documents  = []
        self.dimension  = None

    def build(self, documents: list, embeddings: np.ndarray):
        """Build FAISS index from documents and their embeddings."""
        self.documents = documents
        self.dimension = embeddings.shape[1]
        self.index     = faiss.IndexFlatL2(self.dimension)
        self.index.add(embeddings)
        print(f"Index built — {len(documents)} documents, dimension {self.dimension}")

    def save(self):
        """Save FAISS index and documents to disk."""
        os.makedirs(os.path.dirname(self.index_path), exist_ok=True)
        faiss.write_index(self.index, self.index_path)
        with open(self.docs_path, "w") as f:
            json.dump(self.documents, f)
        print(f"Index saved to: {self.index_path}")

    def load(self):
        """Load FAISS index and documents from disk."""
        if not os.path.exists(self.index_path):
            raise FileNotFoundError(f"No index at {self.index_path} — build first!")
        self.index = faiss.read_index(self.index_path)
        with open(self.docs_path) as f:
            self.documents = json.load(f)
        print(f"Index loaded — {len(self.documents)} documents")

    def search(self, query_vector: np.ndarray, top_k: int = 3) -> list:
        """Return top_k most relevant documents for a query vector."""
        query = query_vector.reshape(1, -1).astype(np.float32)
        distances, indices = self.index.search(query, top_k)
        results = []
        for i, dist in zip(indices[0], distances[0]):
            if i < len(self.documents):
                results.append({
                    "text"     : self.documents[i],
                    "distance" : float(dist),
                    "index"    : int(i)
                })
        return results

    def add_documents(self, new_docs: list, new_embeddings: np.ndarray):
        """Add new documents to existing index."""
        self.documents.extend(new_docs)
        self.index.add(new_embeddings)
        print(f"Added {len(new_docs)} documents. Total: {len(self.documents)}")


# ── Quick test ────────────────────────────────────────────────
if __name__ == "__main__":
    retriever = FAISSRetriever()
    docs      = ["Hello world", "Goodbye world"]
    embeddings = np.random.rand(2, 384).astype(np.float32)
    retriever.build(docs, embeddings)
    retriever.save()
    print("Retriever test passed!")