feat(rag): FAISS inner-product store with chunk metadata + roundtrip

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

src/rag/store.py

@@ -0,0 +1,66 @@
+"""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.asarray(vectors, dtype=np.float32)
+        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.asarray(query, dtype=np.float32)
+        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

tests/rag/test_store.py

@@ -0,0 +1,54 @@
+"""Tests for src.rag.store — FAISS vector store with metadata."""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from src.rag.store import FAISSStore
+
+
+def _rand_vecs(n: int, d: int = 4, seed: int = 0) -> np.ndarray:
+    rng = np.random.default_rng(seed)
+    return rng.standard_normal((n, d), dtype=np.float32)
+
+
+class TestFAISSStore:
+    def test_add_then_search(self) -> None:
+        store = FAISSStore(dim=4)
+        vecs = _rand_vecs(3)
+        chunks = [{"text": f"chunk-{i}", "source": "test.md"} for i in range(3)]
+        store.add(vecs, chunks)
+        results = store.search(vecs[0], k=2)
+        assert len(results) == 2
+        # the closest hit is the chunk we used as the query (cosine ~1.0)
+        top_chunk, top_score = results[0]
+        assert top_chunk["text"] == "chunk-0"
+        assert top_score > 0.99
+
+    def test_add_size_mismatch_raises(self) -> None:
+        store = FAISSStore(dim=4)
+        with pytest.raises(ValueError, match="size mismatch"):
+            store.add(_rand_vecs(3), [{"text": "only-one"}])
+
+    def test_search_k_larger_than_corpus(self) -> None:
+        store = FAISSStore(dim=4)
+        store.add(_rand_vecs(2), [{"text": f"c{i}"} for i in range(2)])
+        results = store.search(_rand_vecs(1)[0], k=10)
+        assert len(results) == 2
+
+    def test_save_load_roundtrip(self, tmp_path: Path) -> None:
+        store = FAISSStore(dim=4)
+        vecs = _rand_vecs(3)
+        chunks = [{"text": f"chunk-{i}", "source": "test.md"} for i in range(3)]
+        store.add(vecs, chunks)
+        store.save(tmp_path / "idx")
+
+        restored = FAISSStore.load(tmp_path / "idx", dim=4)
+        results = restored.search(vecs[0], k=1)
+        assert results[0][0]["text"] == "chunk-0"
+
+    def test_search_on_empty_store_returns_empty(self) -> None:
+        store = FAISSStore(dim=4)
+        assert store.search(_rand_vecs(1)[0], k=5) == []