Spaces:
Running
Running
| """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 | |
| 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 ---------------------------------------------------- | |
| 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), | |
| } | |