| """FAISS vector store with parallel chunk metadata. |
| |
| Public entry: `FAISSStore(dim)`. Vectors are L2-normalized on add and |
| search so inner-product == cosine similarity. Chunks are arbitrary dicts; |
| `text` and `source` keys are recommended but not enforced. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
| from typing import Any |
|
|
| import faiss |
| import numpy as np |
|
|
|
|
| class FAISSStore: |
| """Inner-product (cosine after L2-norm) FAISS store with chunk metadata.""" |
|
|
| def __init__(self, dim: int) -> None: |
| self.dim = dim |
| self._index: faiss.Index = faiss.IndexFlatIP(dim) |
| self._chunks: list[dict[str, Any]] = [] |
|
|
| def __len__(self) -> int: |
| return len(self._chunks) |
|
|
| def add(self, vectors: np.ndarray, chunks: list[dict[str, Any]]) -> None: |
| if vectors.shape[0] != len(chunks): |
| raise ValueError( |
| f"size mismatch: {vectors.shape[0]} vectors vs {len(chunks)} chunks" |
| ) |
| if vectors.shape[0] == 0: |
| return |
| v = np.array(vectors, dtype=np.float32, copy=True) |
| faiss.normalize_L2(v) |
| self._index.add(v) |
| self._chunks.extend(chunks) |
|
|
| def search(self, query: np.ndarray, k: int = 5) -> list[tuple[dict[str, Any], float]]: |
| if len(self._chunks) == 0: |
| return [] |
| q = np.array(query, dtype=np.float32, copy=True) |
| if q.ndim == 1: |
| q = q[np.newaxis, :] |
| faiss.normalize_L2(q) |
| k = min(k, len(self._chunks)) |
| scores, idx = self._index.search(q, k) |
| out: list[tuple[dict[str, Any], float]] = [] |
| for i, s in zip(idx[0], scores[0]): |
| if i == -1: |
| continue |
| out.append((self._chunks[int(i)], float(s))) |
| return out |
|
|
| def save(self, dir_path: Path) -> None: |
| dir_path.mkdir(parents=True, exist_ok=True) |
| faiss.write_index(self._index, str(dir_path / "index.bin")) |
| (dir_path / "chunks.json").write_text(json.dumps(self._chunks, indent=2)) |
|
|
| @classmethod |
| def load(cls, dir_path: Path, dim: int) -> "FAISSStore": |
| store = cls(dim=dim) |
| store._index = faiss.read_index(str(dir_path / "index.bin")) |
| store._chunks = json.loads((dir_path / "chunks.json").read_text()) |
| return store |
|
|
