vector_store.py

"""Build, save, and load the FAISS vector index."""

import os
import faiss
import numpy as np

from config import VECTOR_INDEX_PATH, CHUNKS_PATH, EMBEDDING_DIM
from utils import save_json, load_json, file_exists


def build_index(embeddings: np.ndarray) -> faiss.IndexFlatL2:
    """Create a flat L2 FAISS index and add all embeddings."""
    index = faiss.IndexFlatL2(EMBEDDING_DIM)
    index.add(embeddings)
    print(f"[vector_store] Index built with {index.ntotal} vectors.")
    return index


def save_index(index: faiss.IndexFlatL2, chunks: list[dict]) -> None:
    """Persist FAISS index and chunk metadata to disk."""
    os.makedirs(os.path.dirname(VECTOR_INDEX_PATH), exist_ok=True)
    faiss.write_index(index, VECTOR_INDEX_PATH)
    save_json(chunks, CHUNKS_PATH)
    print(f"[vector_store] Saved index to {VECTOR_INDEX_PATH}")


def load_index() -> tuple[faiss.IndexFlatL2, list[dict]]:
    """Load FAISS index and chunk metadata from disk."""
    if not file_exists(VECTOR_INDEX_PATH):
        raise FileNotFoundError(
            f"FAISS index not found at '{VECTOR_INDEX_PATH}'.\n"
            "Run the knowledge base builder first:\n"
            "  python -c \"from src.rag_pipeline import build_knowledge_base; build_knowledge_base()\""
        )
    index = faiss.read_index(VECTOR_INDEX_PATH)
    chunks = load_json(CHUNKS_PATH)
    print(f"[vector_store] Loaded index ({index.ntotal} vectors, {len(chunks)} chunks).")
    return index, chunks


def search_index(
    index: faiss.IndexFlatL2,
    chunks: list[dict],
    query_vec: np.ndarray,
    top_k: int,
) -> list[dict]:
    """Return top_k most similar chunks for a query vector."""
    distances, indices = index.search(query_vec, top_k)
    results = []
    for dist, idx in zip(distances[0], indices[0]):
        if idx < len(chunks):
            chunk = chunks[idx].copy()
            chunk["score"] = float(dist)
            results.append(chunk)
    return results