Merge pull request #595 from ggozad/feat/vector-nprobes

Expose search.vector_nprobes

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Added
+
+- `search.vector_nprobes` (default 20) sets the IVF partitions each vector query searches.
+
 ### Documentation
 
 - `docs/configuration/storage.md` "Vector Indexing" carries the measured with/without IVF_PQ retrieval comparison; `docs/benchmarks.md` states the published numbers are measured without a vector index.

docs/cli.md

@@ -402,15 +402,16 @@ haiku-rag create-index [--db /path/to/your.lancedb]
 - Creates an IVF_PQ index using the configured `search.vector_index_metric` (cosine/l2)
 
 **When to use:**
-- After ingesting documents (indexes are not created automatically)
-- After adding significant new data to rebuild the index
-- Use `haiku-rag info` to check index status and see how many chunks are indexed/unindexed
+- On a collection over 100,000 chunks (below that, brute-force kNN is exact and fast enough)
+- After substantial corpus growth, to retrain the centroids
+- Use `haiku-rag info` to check index status and see how many chunks are indexed/unindexed, or `haiku-rag doctor` for the same as a health check
+
+See [Vector Indexing](configuration/storage.md#vector-indexing) for the measured accuracy and build cost.
 
 **Search behavior:**
 - Without index: Brute-force kNN search (exact nearest neighbors, slower for large datasets)
-- With index: Fast ANN (approximate nearest neighbors) using IVF_PQ
-- With stale index: LanceDB combines indexed results (fast ANN) + brute-force kNN on unindexed rows
-- Performance degrades as more unindexed data accumulates
+- With index: ANN (approximate nearest neighbors) using IVF_PQ, tuned by `search.vector_nprobes`
+- Between a write and the next `optimize()`: LanceDB combines ANN over indexed rows with brute-force kNN over the remainder
 
 ### Rebuild Database
 

docs/configuration/storage.md

@@ -300,6 +300,7 @@ Configure vector search settings:
 search:
   vector_index_metric: cosine  # cosine or l2
   vector_refine_factor: 30     # Re-ranking factor for accuracy
+  vector_nprobes: 20           # IVF partitions searched per query
 ```
 
 For search behavior settings (`limit`, `max_context_chars`), see [Search and Question Answering](qa.md#search-settings).
@@ -309,6 +310,8 @@ For search behavior settings (`limit`, `max_context_chars`), see [Search and Que
   - `l2`: Euclidean distance
 - **vector_refine_factor**: Improves accuracy when using a vector index by retrieving `refine_factor * limit` candidates (using approximate search) and re-ranking them with exact distances. Higher values increase accuracy but slow down queries. Default: 30
   - **Only applies with a vector index** - has no effect on brute-force search, which already returns exact results
+- **vector_nprobes**: How many IVF partitions each query searches. Higher values increase recall and latency. A larger corpus holds more partitions, so the same value covers a smaller fraction of it. Default: 20
+  - **Only applies with a vector index** - ignored by brute-force search
 
 !!! note
     Vector indexes are only necessary for large datasets with over 100,000 chunks. For smaller datasets, LanceDB's brute-force kNN search provides exact results with good performance. Only create an index if you notice search performance degradation on large datasets.
@@ -321,7 +324,7 @@ Retrieval MAP with and without an index, measured on copies of the benchmark dat
 | `orb_multimodal_nemotron` | 121,168 | 2048 | 0.9799 | 0.9800 | +0.0001 | 25.8 s | 3.38 GB |
 | `frames` | 425,940 | 2560 | 0.5431 | 0.5387 | -0.0044 | 34.1 s | 4.02 GB |
 
-An index costs no accuracy at 70k and 121k chunks and 0.0044 MAP at 426k. A larger corpus holds more IVF partitions, so the default number of probes covers a smaller fraction of the space, and `vector_refine_factor` can only re-score what those probes returned. Build cost is near-flat in row count because training samples the data rather than scanning it, and vector dimension drives it more than corpus size.
+An index costs no accuracy at 70k and 121k chunks and 0.0044 MAP at 426k. A larger corpus holds more IVF partitions, so the default number of probes covers a smaller fraction of the space, and `vector_refine_factor` can only re-score what those probes returned. Raise `vector_nprobes` to trade latency for recall on a large corpus. Build cost is near-flat in row count because training samples the data rather than scanning it, and vector dimension drives it more than corpus size.
 
 **Index creation:**
 

haiku_rag_slim/haiku/rag/config/models.py

@@ -306,6 +306,7 @@ class SearchConfig(ConfigModel):
     max_context_chars: int = Field(default=5000, gt=0)
     vector_index_metric: Literal["cosine", "l2"] = "cosine"
     vector_refine_factor: int = Field(default=30, gt=0)
+    vector_nprobes: int = Field(default=20, gt=0)
 
 
 class OllamaConfig(ConfigModel):

haiku_rag_slim/haiku/rag/store/repositories/chunk.py

@@ -295,6 +295,7 @@ class ChunkRepository:
                 .column("vector")
                 .distance_type(self.store._config.search.vector_index_metric)
                 .refine_factor(self.store._config.search.vector_refine_factor)
+                .nprobes(self.store._config.search.vector_nprobes)
             )
             # An image query has no text to match, so it stays vector-only.
             if search_type != "vector" and query.strip():

tests/test_chunk.py

@@ -863,3 +863,32 @@ async def test_process_search_results_rejects_unknown_score_column(temp_db_path)
 
         with pytest.raises(ValueError, match="Unknown search result format"):
             await client.chunk_repository._process_search_results(_Frame())
+
+
+@pytest.mark.parametrize("search_type", ["vector", "hybrid"])
+async def test_search_applies_configured_nprobes(
+    temp_db_path, monkeypatch, search_type
+):
+    """The configured probe count reaches the vector query."""
+    from lancedb.query import AsyncVectorQueryBase
+
+    probed: list[int] = []
+    original = AsyncVectorQueryBase.nprobes
+
+    def record(self, nprobes):
+        probed.append(nprobes)
+        return original(self, nprobes)
+
+    monkeypatch.setattr(AsyncVectorQueryBase, "nprobes", record)
+
+    config = get_config()
+    config.search.vector_nprobes = 7
+    async with HaikuRAG(db_path=temp_db_path, config=config, create=True) as client:
+        await _import_one(client)
+        await client.chunk_repository.search(
+            "gardens",
+            search_type=search_type,
+            query_vector=[0.1] * config.embeddings.model.vector_dim,
+        )
+
+    assert probed == [7]