Shashiguduri/github-code-explainer

1"""2backend/services/vector_store.py3----------------------------------4FAISS-based vector store: build, search, save, and load.5 6Index type: IndexFlatIP (Inner Product)7  - Exact nearest-neighbor search8  - Correct cosine similarity because embeddings are L2-normalised9  - Scales well for typical repository sizes (< 100k chunks)10"""11 12import os13import pickle14import logging15import numpy as np16import faiss17from langchain_core.documents import Document18 19logger = logging.getLogger(__name__)20 21INDEX_FILE = "faiss.index"22DOCS_FILE  = "faiss_docs.pkl"23DEFAULT_TOP_K = 524 25 26class VectorStore:27    """Wraps a FAISS index with the corresponding Document list."""28 29    def __init__(self):30        self.index: faiss.Index | None = None31        self.documents: list[Document] = []32 33    # ------------------------------------------------------------------34    # Building35    # ------------------------------------------------------------------36 37    def build(self, documents: list[Document], embeddings: np.ndarray) -> None:38        """39        Create a FAISS IndexFlatIP from pre-computed embeddings.40 41        Args:42            documents:  LangChain Documents — must be same order as embeddings.43            embeddings: float32 ndarray, shape (n, dim).44        """45        if len(documents) != len(embeddings):46            raise ValueError(47                f"Mismatch: {len(documents)} docs but {len(embeddings)} vectors."48            )49        if len(documents) == 0:50            raise ValueError("Cannot build an empty index.")51 52        n, dim = embeddings.shape53        self.documents = documents54        self.index = faiss.IndexFlatIP(dim)55        self.index.add(embeddings)56 57        logger.info("FAISS index built: %d vectors, %d dims.", n, dim)58 59    # ------------------------------------------------------------------60    # Searching61    # ------------------------------------------------------------------62 63    def search(self, query_vec: np.ndarray, top_k: int = DEFAULT_TOP_K) -> list[dict]:64        """65        Retrieve top-K documents most similar to the query vector.66 67        Returns:68            list[dict] with keys: document, score, source, snippet69        """70        if not self.is_ready:71            raise RuntimeError("VectorStore is empty. Call build() or load() first.")72 73        q = query_vec.reshape(1, -1).astype(np.float32)74        k = min(top_k, self.index.ntotal)75        scores, indices = self.index.search(q, k)76 77        results = []78        for score, idx in zip(scores[0], indices[0]):79            if idx < 0:80                continue81            doc = self.documents[idx]82            results.append({83                "document": doc,84                "score":    float(score),85                "source":   doc.metadata.get("source", "unknown"),86                "snippet":  doc.page_content[:300],87            })88        return results89 90    # ------------------------------------------------------------------91    # Persistence92    # ------------------------------------------------------------------93 94    def save(self, directory: str) -> None:95        """Persist index and docs to directory."""96        if not self.is_ready:97            raise RuntimeError("Nothing to save.")98        os.makedirs(directory, exist_ok=True)99        faiss.write_index(self.index, os.path.join(directory, INDEX_FILE))100        with open(os.path.join(directory, DOCS_FILE), "wb") as f:101            pickle.dump(self.documents, f)102        logger.info("VectorStore saved to '%s'.", directory)103 104    def load(self, directory: str) -> None:105        """Load a previously saved index and docs from directory."""106        ip = os.path.join(directory, INDEX_FILE)107        dp = os.path.join(directory, DOCS_FILE)108        if not (os.path.exists(ip) and os.path.exists(dp)):109            raise FileNotFoundError(f"No saved index found in '{directory}'.")110        self.index = faiss.read_index(ip)111        with open(dp, "rb") as f:112            self.documents = pickle.load(f)113        logger.info("VectorStore loaded: %d vectors.", self.index.ntotal)114 115    # ------------------------------------------------------------------116    # Properties117    # ------------------------------------------------------------------118 119    @property120    def is_ready(self) -> bool:121        return self.index is not None and self.index.ntotal > 0122 123    @property124    def total_chunks(self) -> int:125        return self.index.ntotal if self.index else 0126