Fix retrieval correctness and input validation

README.md

@@ -55,7 +55,7 @@ Each layer runs on its own; you only need the ones you care about.
 **1. The Rust engine (no Python, no Node):**
 
 ```sh
-cargo test                              # 22 tests, incl. seeded recall vs. brute force
+cargo test                              # 24 tests, incl. seeded recall vs. brute force
 cargo run --release --example bench     # the benchmark table below
 ```
 
@@ -66,7 +66,7 @@ cd bindings
 python3 -m venv .venv && . .venv/bin/activate
 pip install maturin pytest
 maturin develop --release               # builds the Rust extension into the venv
-pytest tests/                           # 16 tests through the FFI boundary
+pytest tests/                           # 21 tests through the FFI boundary
 ```
 
 ```python
@@ -91,6 +91,12 @@ RAG_SERVICE_URL=http://localhost:8000 npm run dev   # http://localhost:3000
 
 Set `ANTHROPIC_API_KEY` before starting the service to get real Claude-generated answers instead of the extractive mock. The retrieval layer (HNSW search + cited sources) is identical either way.
 
+Retrieval drops chunks whose relevance score is below `RAG_MIN_SCORE`
+(default: `0.09`). If no chunk clears the threshold, the service returns no
+sources and says the indexed documents do not address the question. Tune the
+threshold for a different embedding model or set it to `-1` to preserve every
+cosine-search hit.
+
 ## Roadmap
 
 - [x] **Phase 0 — Scaffolding**: workspace layout, README with architecture.
@@ -139,9 +145,9 @@ This section grows as the project does; each phase documents the trade-offs it m
 
 | Layer | Command | Count |
 |-------|---------|-------|
-| Rust engine | `cargo test` | 23 (unit + seeded recall vs. brute force + doctest) |
-| Python bindings + helpers | `pytest` in `bindings/` | 16 (FFI surface, recall vs. brute force, chunking, embeddings, E2E retrieval) |
-| RAG service | `PYTHONPATH=. pytest` in `service/` | 16 (keyless E2E via FastAPI TestClient, citation parsing, startup seeding) |
+| Rust engine | `cargo test` | 24 (unit + seeded recall vs. brute force + doctest) |
+| Python bindings + helpers | `pytest` in `bindings/` | 21 (FFI surface, atomic batches, validation, chunking, embeddings, E2E retrieval) |
+| RAG service | `PYTHONPATH=. pytest` in `service/` | 17 (keyless E2E, relevance filtering, citation parsing, startup seeding) |
 | Next.js app | `npm run build` | type-checked production build |
 
 CI (`.github/workflows/ci.yml`) runs all four on every push, in parallel jobs: `cargo fmt --check` + `cargo clippy -- -D warnings` + `cargo test`, the bindings suite, the service suite, and the app build.

bindings/python/hnsw_rag/chunking.py

@@ -32,9 +32,19 @@ def _split_paragraphs(text: str) -> list[str]:
     return [p.strip() for p in parts if p.strip()]
 
 
-def _split_long_paragraph(words: list[str], max_words: int, overlap: int) -> list[list[str]]:
-    step = max(1, max_words - overlap)
-    return [words[i : i + max_words] for i in range(0, len(words), step)]
+def _split_long_paragraph(
+    words: list[str], max_words: int, overlap: int
+) -> list[list[str]]:
+    pieces: list[list[str]] = []
+    start = 0
+    step = max_words - overlap
+    while start < len(words):
+        end = min(start + max_words, len(words))
+        pieces.append(words[start:end])
+        if end == len(words):
+            break
+        start += step
+    return pieces
 
 
 def chunk_text(
@@ -56,6 +66,10 @@ def chunk_text(
     Returns:
         Chunks in document order, each tagged with its running index.
     """
+    if max_words <= 0:
+        raise ValueError("max_words must be greater than zero")
+    if overlap < 0:
+        raise ValueError("overlap must be non-negative")
     if overlap >= max_words:
         raise ValueError("overlap must be smaller than max_words")
 
@@ -76,8 +90,10 @@ def flush() -> None:
                 chunks.append(" ".join(piece))
             continue
         if current and len(current) + len(words) > max_words:
-            # Carry an overlap tail from the chunk we're closing.
-            tail = current[-overlap:] if overlap else []
+            # Carry as much overlap as fits beside the next paragraph. A
+            # full-size paragraph leaves no room for an overlap tail.
+            tail_size = min(overlap, max_words - len(words))
+            tail = current[-tail_size:] if tail_size else []
             flush()
             current.extend(tail)
         current.extend(words)

bindings/src/lib.rs

@@ -75,6 +75,23 @@ impl Hnsw {
 
     /// Insert many vectors; returns their ids.
     fn insert_batch(&mut self, vectors: Vec<Vec<f32>>) -> PyResult<Vec<u32>> {
+        // Validate the complete batch before mutating the index. Otherwise a
+        // bad vector halfway through would leave earlier vectors inserted even
+        // though Python receives an exception for the overall operation.
+        for (batch_index, vector) in vectors.iter().enumerate() {
+            if vector.len() != self.inner.dim() {
+                return Err(PyValueError::new_err(format!(
+                    "batch vector {batch_index}: dimension mismatch: index holds {}-d vectors, got {}-d",
+                    self.inner.dim(),
+                    vector.len()
+                )));
+            }
+            if let Some(position) = vector.iter().position(|value| !value.is_finite()) {
+                return Err(PyValueError::new_err(format!(
+                    "batch vector {batch_index} contains a non-finite value at position {position}"
+                )));
+            }
+        }
         vectors.into_iter().map(|v| self.insert(v)).collect()
     }
 

bindings/tests/test_helpers.py

@@ -41,6 +41,25 @@ def test_chunk_splits_on_paragraphs():
 def test_chunk_overlap_validation():
     with pytest.raises(ValueError):
         chunk_text("x y z", max_words=10, overlap=10)
+    with pytest.raises(ValueError, match="greater than zero"):
+        chunk_text("x y z", max_words=0, overlap=0)
+    with pytest.raises(ValueError, match="non-negative"):
+        chunk_text("x y z", max_words=10, overlap=-1)
+
+
+def test_paragraph_overlap_never_exceeds_word_budget():
+    first = " ".join(f"first{i}" for i in range(100))
+    second = " ".join(f"second{i}" for i in range(100))
+    chunks = chunk_text(f"{first}\n\n{second}", max_words=100, overlap=20)
+    assert [len(chunk.text.split()) for chunk in chunks] == [100, 100]
+
+
+def test_long_paragraph_has_no_redundant_final_window():
+    words = " ".join(f"w{i}" for i in range(500))
+    chunks = chunk_text(words, max_words=100, overlap=20)
+    assert len(chunks) == 6
+    assert chunks[-1].text.split()[0] == "w400"
+    assert chunks[-1].text.split()[-1] == "w499"
 
 
 def test_hashed_embedder_is_deterministic_and_normalized():

bindings/tests/test_hnsw.py

@@ -58,6 +58,28 @@ def test_dimension_mismatch_raises():
         index.search([1.0, 2.0], k=1)
 
 
+def test_non_finite_values_are_rejected():
+    index = Hnsw(dim=2)
+    with pytest.raises(ValueError, match="non-finite"):
+        index.insert([float("nan"), 0.0])
+    with pytest.raises(ValueError, match="non-finite"):
+        index.search([0.0, float("inf")], k=1)
+
+
+def test_insert_batch_is_atomic_on_dimension_error():
+    index = Hnsw(dim=2)
+    with pytest.raises(ValueError, match="batch vector 1.*dimension mismatch"):
+        index.insert_batch([[1.0, 0.0], [1.0]])
+    assert len(index) == 0
+
+
+def test_insert_batch_is_atomic_on_non_finite_value():
+    index = Hnsw(dim=2)
+    with pytest.raises(ValueError, match="batch vector 1.*non-finite"):
+        index.insert_batch([[1.0, 0.0], [float("inf"), 1.0]])
+    assert len(index) == 0
+
+
 def test_basic_roundtrip():
     index = Hnsw(dim=3, metric="euclidean")
     ids = index.insert_batch([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 5.0, 0.0]])

engine/src/hnsw.rs

@@ -54,6 +54,7 @@ pub struct Neighbor {
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum Error {
     DimensionMismatch { expected: usize, got: usize },
+    NonFiniteValue { position: usize },
 }
 
 impl std::fmt::Display for Error {
@@ -65,6 +66,12 @@ impl std::fmt::Display for Error {
                     "dimension mismatch: index holds {expected}-d vectors, got {got}-d"
                 )
             }
+            Error::NonFiniteValue { position } => {
+                write!(
+                    f,
+                    "vector contains a non-finite value at position {position}"
+                )
+            }
         }
     }
 }
@@ -230,6 +237,9 @@ impl Hnsw {
                 got: vector.len(),
             });
         }
+        if let Some(position) = vector.iter().position(|value| !value.is_finite()) {
+            return Err(Error::NonFiniteValue { position });
+        }
         if self.store.metric == Metric::Cosine {
             distance::normalize(&mut vector);
         }
@@ -301,6 +311,9 @@ impl Hnsw {
                 got: query.len(),
             });
         }
+        if let Some(position) = query.iter().position(|value| !value.is_finite()) {
+            return Err(Error::NonFiniteValue { position });
+        }
         let Some(entry) = self.entry else {
             return Ok(Vec::new());
         };
@@ -498,6 +511,23 @@ mod tests {
         assert!(idx.search(&[0.0; 3], 1, 10).is_err());
     }
 
+    #[test]
+    fn rejects_non_finite_values_without_mutation() {
+        let mut idx = Hnsw::new(2, Metric::Cosine, HnswParams::default());
+        assert_eq!(
+            idx.insert(vec![f32::NAN, 0.0]),
+            Err(Error::NonFiniteValue { position: 0 })
+        );
+        assert!(idx.is_empty());
+
+        idx.insert(vec![1.0, 0.0]).unwrap();
+        assert_eq!(
+            idx.search(&[0.0, f32::INFINITY], 1, 10),
+            Err(Error::NonFiniteValue { position: 1 })
+        );
+        assert_eq!(idx.len(), 1);
+    }
+
     #[test]
     fn finds_exact_match_first() {
         let vectors: Vec<Vec<f32>> = vec![

service/rag_service/app.py

@@ -33,7 +33,10 @@
 # One in-memory store for the process. The embedder backend is chosen at
 # startup: real model if available, deterministic hashed fallback otherwise.
 _embedder = get_embedder(os.environ.get("RAG_EMBEDDER", "auto"))
-_store = DocumentStore(embedder=_embedder)
+_store = DocumentStore(
+    embedder=_embedder,
+    min_score=float(os.environ.get("RAG_MIN_SCORE", "0.09")),
+)
 
 SAMPLE_DOCS = pathlib.Path(__file__).resolve().parent.parent / "sample_docs"
 

service/rag_service/store.py

@@ -51,6 +51,7 @@ class RetrievedChunk:
 class DocumentStore:
     embedder: Embedder
     metric: str = "cosine"
+    min_score: Optional[float] = None
     m: int = 16
     ef_construction: int = 200
     seed: int = 0x5EED
@@ -121,6 +122,10 @@ def retrieve(self, query: str, k: int = 5, ef_search: int = 100) -> List[Retriev
             # Cosine distance is 1 - similarity; report similarity so higher
             # is more relevant, which is what a reader expects from a score.
             score = 1.0 - distance if self.metric == "cosine" else -distance
+            if self.min_score is not None and score < self.min_score:
+                # Hits are closest-first, so their relevance scores only
+                # decrease. Weak nearest neighbors are not useful grounding.
+                break
             out.append(
                 RetrievedChunk(
                     id=sc.id,
@@ -140,6 +145,7 @@ def stats(self) -> dict:
                 "chunks": len(self._chunks),
                 "dim": self.dim,
                 "metric": self.metric,
+                "min_score": self.min_score,
                 # Which embedding backend is actually live. Worth surfacing:
                 # HashedEmbedder matches on term overlap, not meaning, so a
                 # reader should not mistake it for semantic search.

service/tests/test_e2e.py

@@ -36,6 +36,7 @@ def test_stats_after_seed(client):
     assert stats["documents"] == 3
     assert stats["chunks"] >= 3
     assert stats["metric"] == "cosine"
+    assert stats["min_score"] == pytest.approx(0.09)
 
 
 def test_documents_listed(client):
@@ -71,10 +72,34 @@ def test_query_retrieves_relevant_document(client, question, expected_doc):
     assert body["answer"]
 
 
-def test_unrelated_question_still_responds_cleanly(client):
-    # An off-topic query returns whatever is nearest; the contract is only that
-    # the endpoint responds cleanly with a well-formed body.
+def test_unrelated_question_returns_no_sources(client):
     resp = client.post("/query", json={"question": "how do I prune roses", "k": 2})
     assert resp.status_code == 200
     body = resp.json()
-    assert "answer" in body and "sources" in body
+    assert body["sources"] == []
+    assert body["model"] == "mock"
+    assert body["answer"] == (
+        "I don't have any indexed documents that address that question."
+    )
+
+
+def test_store_filters_weak_matches_and_keeps_relevant_ones():
+    from rag_service.store import DocumentStore
+
+    class OrthogonalEmbedder:
+        dim = 2
+
+        def embed(self, texts):
+            return [
+                [1.0, 0.0] if text in {"known", "known query"} else [0.0, 1.0]
+                for text in texts
+            ]
+
+    store = DocumentStore(embedder=OrthogonalEmbedder(), min_score=0.15)
+    store.add_document("known document", "known", max_words=10, overlap=0)
+
+    assert store.retrieve("unrelated", k=1) == []
+    relevant = store.retrieve("known query", k=1)
+    assert len(relevant) == 1
+    assert relevant[0].doc_title == "known document"
+    assert relevant[0].score == pytest.approx(1.0)