Blame for core/vector_store.py

"""FAISS-backed vector store for the T.I.A. RAG engine.



Persists a FAISS index plus a JSONL sidecar of chunk metadata so the index

can be rebuilt offline (via ``data/master_harvest/ingest.py``) and loaded by

the API at startup.



Designed to fail soft: if ``faiss`` or ``sentence-transformers`` are not

installed, ``FaissVectorStore.is_available()`` returns ``False`` and callers

should fall back to the in-memory TF-IDF index in :mod:`core.rag_engine`.

"""

from __future__ import annotations

import json
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

DEFAULT_INDEX_DIR = Path("data") / "master_harvest" / ".index"
DEFAULT_INDEX_FILE = "harvest.faiss"
DEFAULT_META_FILE = "harvest.meta.jsonl"
DEFAULT_EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"


def _try_import_faiss():
    try:
        import faiss  # type: ignore

        return faiss
    except Exception:
        return None


def _try_import_st():
    try:
        from sentence_transformers import SentenceTransformer  # type: ignore

        return SentenceTransformer
    except Exception:
        return None


def _try_import_numpy():
    try:
        import numpy as np  # type: ignore

        return np
    except Exception:
        return None


@dataclass
class VectorHit:
    source: str
    index: int
    text: str
    score: float

    def to_dict(self) -> dict:
        return {
            "source": self.source,
            "index": self.index,
            "text": self.text,
            "score": round(float(self.score), 4),
        }


class FaissVectorStore:
    """Thin wrapper around a FAISS ``IndexFlatIP`` over normalized embeddings.



    Cosine similarity is computed by L2-normalizing inputs and using inner

    product, which is the standard and most-portable FAISS recipe.

    """

    def __init__(

        self,

        index_dir: str | Path = DEFAULT_INDEX_DIR,

        embed_model: str = DEFAULT_EMBED_MODEL,

    ) -> None:
        self.index_dir = Path(index_dir)
        self.embed_model_name = embed_model
        self._index = None  # type: ignore[assignment]
        self._meta: list[dict] = []
        self._encoder = None  # type: ignore[assignment]
        self._lock = threading.Lock()

    # -- availability ----------------------------------------------------
    @staticmethod
    def is_available() -> bool:
        return (
            _try_import_faiss() is not None
            and _try_import_st() is not None
            and _try_import_numpy() is not None
        )

    # -- encoder lazy-init ----------------------------------------------
    def _get_encoder(self):
        if self._encoder is not None:
            return self._encoder
        SentenceTransformer = _try_import_st()
        if SentenceTransformer is None:
            raise RuntimeError(
                "sentence-transformers is required for embeddings; "
                "install with `pip install sentence-transformers`."
            )
        self._encoder = SentenceTransformer(self.embed_model_name)
        return self._encoder

    def embed(self, texts: list[str]):
        np = _try_import_numpy()
        if np is None:
            raise RuntimeError("numpy is required for embeddings")
        encoder = self._get_encoder()
        vectors = encoder.encode(
            texts,
            convert_to_numpy=True,
            normalize_embeddings=True,
            show_progress_bar=False,
        )
        return vectors.astype(np.float32, copy=False)

    # -- build / persist -------------------------------------------------
    def build(self, chunks: Iterable[dict]) -> dict:
        """Build the FAISS index from an iterable of ``{source, index, text}`` dicts.



        Persists the index + JSONL metadata to ``index_dir``.

        """
        faiss = _try_import_faiss()
        np = _try_import_numpy()
        if faiss is None or np is None:
            raise RuntimeError(
                "faiss and numpy are required to build a FaissVectorStore."
            )
        chunk_list = [c for c in chunks if c.get("text", "").strip()]
        if not chunk_list:
            with self._lock:
                self._index = None
                self._meta = []
            return {"chunks_indexed": 0, "dimension": 0, "index_dir": str(self.index_dir)}
        texts = [c["text"] for c in chunk_list]
        vectors = self.embed(texts)
        dimension = int(vectors.shape[1])
        index = faiss.IndexFlatIP(dimension)
        index.add(vectors)
        with self._lock:
            self._index = index
            self._meta = chunk_list
        self._persist()
        return {
            "chunks_indexed": len(chunk_list),
            "dimension": dimension,
            "index_dir": str(self.index_dir),
        }

    def _persist(self) -> None:
        faiss = _try_import_faiss()
        if faiss is None:
            return
        self.index_dir.mkdir(parents=True, exist_ok=True)
        index_path = self.index_dir / DEFAULT_INDEX_FILE
        meta_path = self.index_dir / DEFAULT_META_FILE
        if self._index is not None:
            faiss.write_index(self._index, str(index_path))
        with meta_path.open("w", encoding="utf-8") as fh:
            for record in self._meta:
                fh.write(json.dumps(record, ensure_ascii=False) + "\n")

    # -- load ------------------------------------------------------------
    def load(self) -> bool:
        """Load a previously-persisted index. Returns True on success."""
        faiss = _try_import_faiss()
        if faiss is None:
            return False
        index_path = self.index_dir / DEFAULT_INDEX_FILE
        meta_path = self.index_dir / DEFAULT_META_FILE
        if not index_path.exists() or not meta_path.exists():
            return False
        try:
            index = faiss.read_index(str(index_path))
            meta: list[dict] = []
            with meta_path.open("r", encoding="utf-8") as fh:
                for line in fh:
                    line = line.strip()
                    if line:
                        meta.append(json.loads(line))
        except Exception:
            return False
        with self._lock:
            self._index = index
            self._meta = meta
        return True

    # -- search ----------------------------------------------------------
    def search(self, query: str, top_k: int = 5) -> list[VectorHit]:
        if top_k <= 0 or not query.strip():
            return []
        with self._lock:
            index = self._index
            meta = list(self._meta)
        if index is None or not meta:
            return []
        query_vec = self.embed([query])
        scores, ids = index.search(query_vec, min(top_k, len(meta)))
        hits: list[VectorHit] = []
        for score, idx in zip(scores[0].tolist(), ids[0].tolist()):
            if idx < 0 or idx >= len(meta):
                continue
            record = meta[idx]
            hits.append(
                VectorHit(
                    source=record.get("source", ""),
                    index=int(record.get("index", 0)),
                    text=record.get("text", ""),
                    score=float(score),
                )
            )
        return hits

    # -- introspection ---------------------------------------------------
    def stats(self) -> dict:
        with self._lock:
            count = len(self._meta)
            loaded = self._index is not None
        return {
            "loaded": loaded,
            "chunks_indexed": count,
            "embed_model": self.embed_model_name,
            "index_dir": str(self.index_dir),
        }