From c184a25d6875b9e2d9cefb9ece427f93f49233fd Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 17 Aug 2026 14:54:01 +0300 Subject: [PATCH 01/11] Index every hot lookup key from one shared definition `_init_tables` left `chunks.id`, `chunks.document_id` and `documents.id` unindexed, so those lookups scanned the column. On object storage that is network I/O per query, on paths that run per document: citation lookup, delete-by-document, the re-ingest merge, and every dedup probe. Declare the index set per table in `index_specs()` and apply it through `ensure_indexes()`, which skips a column only when it is already indexed with the declared type. Both halves of that are load-bearing. Skipping is required because `create_index(replace=True)` rebuilds an identical index, writing a new index and a new table version and orphaning the old files until the next vacuum. Comparing the type is required because column coverage alone would let a wrong-typed index stand, and a BTree on `label` silently loses the low-cardinality equality lookup the Bitmap is there for. Columns not declared for a table are left alone, so an externally created index such as a vector index on `chunks` survives. `_init_tables`, `recreate_embeddings_table`, `ChunkRepository.delete_all` and `DocumentRepository.delete_all` now all route through it instead of repeating their own subsets. Also recreate `document_items` from `get_document_items_arrow_schema()` in `DocumentRepository.delete_all`, which was using the LanceModel and so returned `picture_data` as 32-bit `binary`. Existing databases are unchanged; the migration follows separately. --- CHANGELOG.md | 6 ++ haiku_rag_slim/haiku/rag/store/engine.py | 98 +++++++++++++------ .../haiku/rag/store/repositories/chunk.py | 9 +- .../haiku/rag/store/repositories/document.py | 25 ++--- tests/store/test_indexes.py | 98 +++++++++++++++++++ 5 files changed, 179 insertions(+), 57 deletions(-) create mode 100644 tests/store/test_indexes.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5bcc7be9..2b4aa52a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,10 +6,16 @@ - `evaluations run --filter/-f CLAUSE`: SQL `WHERE` clause over document columns, applied to the retrieval benchmark's searches and to every capability search during QA. Recorded as `document_filter` in experiment metadata. - `mtrag_clapnq` / `mtrag_clapnq_rewrite` / `mtrag_clapnq_live` / `mtrag_clapnq_live_uncompacted` evaluation datasets: multi-turn QA with gold-prefix and live-session conversation replay, live arms with and without `EvidenceCompactionCapability`, Recall@k/nDCG@k retrieval metrics, eligibility-aware citation scoring, refusal precision/recall, per-turn tool-traffic attributes, and `citation_status` / `turn_citation_status` eval attributes. +- BTree indexes on `chunks.id`, `chunks.document_id` and `documents.id`, and a Bitmap index on `document_items.label`. + ### Changed - `import_documents` embeds chunks across the whole batch in one pass instead of per document. +### Fixed + +- `DocumentRepository.delete_all` recreated `document_items` from `DocumentItemRecord` instead of `get_document_items_arrow_schema()`, returning `picture_data` as `binary` rather than `large_binary`. + ### Removed - `wix` evaluation dataset and its reference config `evaluations/configs/wix.yaml`. diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 50d4a73a..3ef990af 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -12,7 +12,7 @@ from uuid import uuid4 import lancedb import pyarrow as pa -from lancedb.index import FTS, BTree, IvfPq +from lancedb.index import FTS, Bitmap, BTree, IvfPq from lancedb.pydantic import LanceModel, Vector from packaging.version import parse from pydantic import BaseModel, Field @@ -176,6 +176,66 @@ def get_document_items_arrow_schema() -> pa.Schema: return pa.schema(fields) +def index_specs(table_name: str) -> list[tuple[str, Bitmap | BTree | FTS]]: + """The index set a haiku.rag table is expected to carry. + + Single source of truth for every path that creates a table: initialization, + migration, and the drop-and-recreate paths in the repositories. + + `label` gets a Bitmap rather than a BTree because it holds around ten + distinct values across every item of every document, and Bitmap is the + low-cardinality equality case. The FTS options are load-bearing: + `with_position` enables phrase queries and keeping stop words lets them match. + """ + match table_name: + case "documents": + return [("id", BTree())] + case "document_meta": + return [("id", BTree()), ("uri", BTree())] + case "chunks": + return [ + ("content_fts", FTS(with_position=True, remove_stop_words=False)), + ("id", BTree()), + ("document_id", BTree()), + ] + case "document_items": + return [ + ("document_id", BTree()), + ("position", BTree()), + ("self_ref", BTree()), + ("label", Bitmap()), + ] + case _: + return [] + + +async def ensure_indexes(table: lancedb.AsyncTable, table_name: str) -> None: + """Create the table's declared indexes, skipping those already correct. + + Correct means the column is indexed *and* carries the declared index type. + Matching on the column alone would let a wrong-typed index stand: a BTree on + `label` covers the column while losing the low-cardinality equality lookup a + Bitmap gives, and no amount of column coverage reveals that. + + Skipping matters as much as creating: `create_index(replace=True)` rebuilds + an identical index, writing a fresh index and a new table version rather + than no-oping, and leaves the previous index behind until the next vacuum. + On a large table over object storage that is a full column sort per pass. + + Columns not declared for the table are left untouched, so an externally + added index (a vector index on `chunks`, say) survives. + """ + indexed = { + column: index.index_type + for index in await table.list_indices() + for column in index.columns + } + for column, config in index_specs(table_name): + if indexed.get(column) == type(config).__name__: + continue + await table.create_index(column, config=config, replace=True) + + class SettingsRecord(LanceModel): id: str = Field(default="settings") settings: str = Field(default="{}") @@ -697,22 +757,17 @@ class Store: self.documents_table = await self.db.create_table( "documents", schema=get_documents_arrow_schema() ) + await ensure_indexes(self.documents_table, "documents") # Create or open document_meta table (mutable attributes kept out of the - # blob-bearing documents row). Indexed by document_id and uri — both are - # hot look-up keys (get_by_id, get_by_uri). + # blob-bearing documents row). if "document_meta" in existing_tables: self.document_meta_table = await self.db.open_table("document_meta") else: self.document_meta_table = await self.db.create_table( "document_meta", schema=DocumentMetaRecord ) - await self.document_meta_table.create_index( - "id", config=BTree(), replace=True - ) - await self.document_meta_table.create_index( - "uri", config=BTree(), replace=True - ) + await ensure_indexes(self.document_meta_table, "document_meta") # Create or open chunks table if "chunks" in existing_tables: @@ -721,12 +776,7 @@ class Store: self.chunks_table = await self.db.create_table( "chunks", schema=self.ChunkRecord ) - # Create FTS index on content_fts (contextualized content) for better search - await self.chunks_table.create_index( - "content_fts", - config=FTS(with_position=True, remove_stop_words=False), - replace=True, - ) + await ensure_indexes(self.chunks_table, "chunks") # Create or open document_items table if "document_items" in existing_tables: @@ -735,15 +785,7 @@ class Store: self.document_items_table = await self.db.create_table( "document_items", schema=get_document_items_arrow_schema() ) - await self.document_items_table.create_index( - "document_id", config=BTree(), replace=True - ) - await self.document_items_table.create_index( - "position", config=BTree(), replace=True - ) - await self.document_items_table.create_index( - "self_ref", config=BTree(), replace=True - ) + await ensure_indexes(self.document_items_table, "document_items") # Create or open settings table if "settings" in existing_tables: @@ -872,13 +914,7 @@ class Store: self.chunks_table = await self.db.create_table( "chunks", schema=self.ChunkRecord ) - - # Create FTS index on content_fts (contextualized content) for better search - await self.chunks_table.create_index( - "content_fts", - config=FTS(with_position=True, remove_stop_words=False), - replace=True, - ) + await ensure_indexes(self.chunks_table, "chunks") def close(self): """Close the database connection.""" diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 09cc3498..ecfe9f49 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -9,7 +9,7 @@ if TYPE_CHECKING: from lancedb.index import FTS from lancedb.rerankers import RRFReranker -from haiku.rag.store.engine import Store, query_to_pydantic +from haiku.rag.store.engine import Store, ensure_indexes, query_to_pydantic from haiku.rag.store.models.chunk import Chunk, SearchType from haiku.rag.utils import escape_sql_string @@ -187,12 +187,7 @@ class ChunkRepository: self.store.chunks_table = await self.store.db.create_table( "chunks", schema=self.store.ChunkRecord ) - # Create FTS index on content_fts (contextualized content) for better search - await self.store.chunks_table.create_index( - "content_fts", - config=FTS(with_position=True, remove_stop_words=False), - replace=True, - ) + await ensure_indexes(self.store.chunks_table, "chunks") async def delete_by_document_id(self, document_id: str) -> bool: """Delete all chunks for a document.""" diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 9bfd9b57..be242be3 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -3,12 +3,12 @@ from datetime import datetime from typing import overload from uuid import uuid4 -from lancedb.index import BTree - from haiku.rag.store.engine import ( DocumentMetaRecord, DocumentRecord, Store, + ensure_indexes, + get_document_items_arrow_schema, get_documents_arrow_schema, query_to_pydantic, ) @@ -407,23 +407,14 @@ class DocumentRepository: async def delete_all(self) -> None: """Delete all documents from the database.""" self.store._assert_writable() - from haiku.rag.store.engine import DocumentItemRecord # Delete all chunks and items first await self.chunk_repository.delete_all() await self.store.db.drop_table("document_items") self.store.document_items_table = await self.store.db.create_table( - "document_items", schema=DocumentItemRecord - ) - await self.store.document_items_table.create_index( - "document_id", config=BTree(), replace=True - ) - await self.store.document_items_table.create_index( - "position", config=BTree(), replace=True - ) - await self.store.document_items_table.create_index( - "self_ref", config=BTree(), replace=True + "document_items", schema=get_document_items_arrow_schema() ) + await ensure_indexes(self.store.document_items_table, "document_items") # Get count before deletion count = len( @@ -437,13 +428,9 @@ class DocumentRepository: self.store.documents_table = await self.store.db.create_table( "documents", schema=get_documents_arrow_schema() ) + await ensure_indexes(self.store.documents_table, "documents") await self.store.db.drop_table("document_meta") self.store.document_meta_table = await self.store.db.create_table( "document_meta", schema=DocumentMetaRecord ) - await self.store.document_meta_table.create_index( - "id", config=BTree(), replace=True - ) - await self.store.document_meta_table.create_index( - "uri", config=BTree(), replace=True - ) + await ensure_indexes(self.store.document_meta_table, "document_meta") diff --git a/tests/store/test_indexes.py b/tests/store/test_indexes.py new file mode 100644 index 00000000..a0f011c4 --- /dev/null +++ b/tests/store/test_indexes.py @@ -0,0 +1,98 @@ +import pyarrow as pa +import pytest +from lancedb.index import BTree + +from haiku.rag.store.engine import Store, ensure_indexes +from haiku.rag.store.models import Document +from haiku.rag.store.repositories.document import DocumentRepository + +EXPECTED_INDEXED_COLUMNS = { + "documents": {"id"}, + "document_meta": {"id", "uri"}, + "chunks": {"content_fts", "id", "document_id"}, + "document_items": {"document_id", "position", "self_ref", "label"}, +} + + +async def _indexed_columns(table) -> set[str]: + return {column for index in await table.list_indices() for column in index.columns} + + +async def _index_type(table, column: str) -> str | None: + for index in await table.list_indices(): + if column in index.columns: + return index.index_type + return None + + +@pytest.mark.asyncio +async def test_fresh_database_indexes_every_hot_lookup_key(temp_db_path): + """A new database carries the full index set, not a subset.""" + async with Store(temp_db_path, create=True) as store: + for name, table in store._tables().items(): + expected = EXPECTED_INDEXED_COLUMNS.get(name, set()) + assert await _indexed_columns(table) == expected, name + + +@pytest.mark.asyncio +async def test_ensure_indexes_skips_existing_instead_of_rebuilding(temp_db_path): + """`create_index(replace=True)` rebuilds an identical index and writes a new + table version, so a second pass must skip rather than replace.""" + async with Store(temp_db_path, create=True) as store: + table = store.chunks_table + version_before = await table.version() + + await ensure_indexes(table, "chunks") + + assert await table.version() == version_before + assert await _indexed_columns(table) == EXPECTED_INDEXED_COLUMNS["chunks"] + + +@pytest.mark.asyncio +async def test_ensure_indexes_corrects_an_index_of_the_wrong_type(temp_db_path): + """A column indexed with the wrong type must be re-indexed. `label` is the + live case: a BTree over ~ten distinct values loses the low-cardinality + equality lookup a Bitmap gives, and column coverage alone cannot see it. + """ + async with Store(temp_db_path, create=True) as store: + table = store.document_items_table + await table.create_index("label", config=BTree(), replace=True) + assert await _index_type(table, "label") == "BTree" + + await ensure_indexes(table, "document_items") + + assert await _index_type(table, "label") == "Bitmap" + assert ( + await _indexed_columns(table) == EXPECTED_INDEXED_COLUMNS["document_items"] + ) + + +@pytest.mark.asyncio +async def test_delete_all_restores_the_full_index_set(temp_db_path): + """delete_all drops and recreates tables; the recreated tables must come + back with the same indexes a fresh database gets.""" + async with Store(temp_db_path, create=True) as store: + repo = DocumentRepository(store) + await repo.create(Document(content="A document")) + + await repo.delete_all() + + for name, table in store._tables().items(): + expected = EXPECTED_INDEXED_COLUMNS.get(name, set()) + assert await _indexed_columns(table) == expected, name + + +@pytest.mark.asyncio +async def test_delete_all_keeps_picture_data_as_large_binary(temp_db_path): + """document_items must be recreated from the Arrow schema, which declares + picture_data as large_binary. The 32-bit `binary` type overflows its offsets + once a fragment holds enough embedded pictures. + """ + async with Store(temp_db_path, create=True) as store: + repo = DocumentRepository(store) + await repo.create(Document(content="A document")) + + await repo.delete_all() + + schema = await store.document_items_table.schema() + assert schema.field("picture_data").type == pa.large_binary() From 8e93b639bc62701e3afe404a6b3e2269574f7b81 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 17 Aug 2026 15:19:25 +0300 Subject: [PATCH 02/11] Migrate existing databases to the full index set Adds the 0.75.0 upgrade, which brings a pre-existing database up to the index set `_init_tables` now creates. It rewrites no table data, so unlike the earlier data migrations its cost is the index builds alone, each of which reads the column it indexes. `ensure_indexes` ensures an index of the declared *type* covers each declared column, rather than checking that the column is indexed at all. The distinction is what makes it safe to run against a database of unknown provenance: - A wrong-typed index no longer satisfies the check. A BTree on `label` covers the column while losing the low-cardinality equality lookup the Bitmap is for. - Nothing is dropped or converted away from. Two index types over one column can be deliberate, serving different query shapes, so an index this function did not declare survives even on a column it does. The one thing it overwrites is an index at LanceDB's default name, `{column}_idx`, which is the name it creates itself. - A column already carrying the declared type is skipped, so a database with the full set migrates instantly rather than re-sorting every indexed column. - Undeclared columns are untouched, so a vector index on `chunks` survives. It returns the columns it acted on, because a change is not always visible from outside: adding a Bitmap beside an existing BTree leaves the column indexed before and after. The version bump to 0.75.0 is required, not incidental: `_set_initial_version` stamps a new database with the installed package version, so a migration numbered above it would be pending the moment the database was created. `test_client_update_document_replaces_rows_with_bounded_versions` turns auto_vacuum off. Indexing `documents` means a background vacuum now has an index to maintain on that table, so `optimize()` writes a version where it previously had nothing to do, and it landed inside the window the test measures. The document update itself is still one version, so the bound stays exact. --- CHANGELOG.md | 2 +- app/backend/pyproject.toml | 2 +- docs/cli.md | 3 + evaluations/pyproject.toml | 2 +- haiku_rag_slim/haiku/rag/store/engine.py | 53 ++++++--- .../haiku/rag/store/upgrades/__init__.py | 4 + .../haiku/rag/store/upgrades/v0_75_0.py | 39 ++++++ haiku_rag_slim/pyproject.toml | 2 +- pyproject.toml | 10 +- tests/store/test_indexes.py | 47 ++++++++ tests/store/test_v0_75_0_migration.py | 111 ++++++++++++++++++ tests/test_client.py | 12 +- uv.lock | 6 +- 13 files changed, 260 insertions(+), 33 deletions(-) create mode 100644 haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py create mode 100644 tests/store/test_v0_75_0_migration.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 2b4aa52a..355eabd9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ - `evaluations run --filter/-f CLAUSE`: SQL `WHERE` clause over document columns, applied to the retrieval benchmark's searches and to every capability search during QA. Recorded as `document_filter` in experiment metadata. - `mtrag_clapnq` / `mtrag_clapnq_rewrite` / `mtrag_clapnq_live` / `mtrag_clapnq_live_uncompacted` evaluation datasets: multi-turn QA with gold-prefix and live-session conversation replay, live arms with and without `EvidenceCompactionCapability`, Recall@k/nDCG@k retrieval metrics, eligibility-aware citation scoring, refusal precision/recall, per-turn tool-traffic attributes, and `citation_status` / `turn_citation_status` eval attributes. -- BTree indexes on `chunks.id`, `chunks.document_id` and `documents.id`, and a Bitmap index on `document_items.label`. +- BTree indexes on `chunks.id`, `chunks.document_id` and `documents.id`, and a Bitmap index on `document_items.label`. Existing databases need `haiku-rag migrate`. ### Changed diff --git a/app/backend/pyproject.toml b/app/backend/pyproject.toml index bb0290b8..5e2b669d 100644 --- a/app/backend/pyproject.toml +++ b/app/backend/pyproject.toml @@ -8,7 +8,7 @@ dependencies = [ "uvicorn[standard]>=0.40.0", "pydantic-ai-slim[ag-ui,anthropic,openai]>=2.18.0,<3.0.0", "python-dotenv>=1.2.1", - "haiku.rag-slim>=0.74.0", + "haiku.rag-slim>=0.75.0", "logfire[pydantic-ai]>=3.17.0", ] diff --git a/docs/cli.md b/docs/cli.md index 06c2b726..9a3df4a4 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -359,6 +359,9 @@ Migration completed successfully. !!! tip Back up your database before running migrations. While migrations are designed to be safe, having a backup provides peace of mind for production databases. +!!! note "Upgrading to 0.75.0" + The 0.75.0 migration adds scalar indexes: BTree on `documents.id`, `chunks.id` and `chunks.document_id`, and Bitmap on `document_items.label`. It rewrites no table data, so it is far cheaper than the earlier data migrations, but building an index reads the whole column it indexes. On a large database, and particularly on object storage, budget for reading `chunks.id` and `chunks.document_id` in full. Indexes already present with the expected type are left alone, so a database that already carries the full set migrates instantly. + ### Download Models Download required runtime models: diff --git a/evaluations/pyproject.toml b/evaluations/pyproject.toml index 28deceb0..158a5362 100644 --- a/evaluations/pyproject.toml +++ b/evaluations/pyproject.toml @@ -2,7 +2,7 @@ name = "haiku.rag-evals" description = "Benchmarking and evaluation scripts for haiku.rag" -version = "0.74.0" +version = "0.75.0" authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] license = { text = "MIT" } requires-python = ">=3.12" diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 3ef990af..5464fc57 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -209,31 +209,48 @@ def index_specs(table_name: str) -> list[tuple[str, Bitmap | BTree | FTS]]: return [] -async def ensure_indexes(table: lancedb.AsyncTable, table_name: str) -> None: - """Create the table's declared indexes, skipping those already correct. +async def ensure_indexes(table: lancedb.AsyncTable, table_name: str) -> list[str]: + """Ensure an index of the declared type covers each declared column. - Correct means the column is indexed *and* carries the declared index type. - Matching on the column alone would let a wrong-typed index stand: a BTree on - `label` covers the column while losing the low-cardinality equality lookup a - Bitmap gives, and no amount of column coverage reveals that. + Returns the columns it indexed, so callers can report what changed. - Skipping matters as much as creating: `create_index(replace=True)` rebuilds - an identical index, writing a fresh index and a new table version rather - than no-oping, and leaves the previous index behind until the next vacuum. - On a large table over object storage that is a full column sort per pass. + The condition is the presence of the declared *type*, not that the column's + index happens to be that type. Checking coverage alone would let a wrong-typed + index satisfy the check: a BTree on `label` covers the column while losing the + low-cardinality equality lookup a Bitmap gives. - Columns not declared for the table are left untouched, so an externally - added index (a vector index on `chunks`, say) survives. + Nothing is ever dropped or converted away from. Two index types over one + column can be deliberate, since they serve different query shapes (BTree for + equality and range, `Fm` for `contains`), so an index this function did not + declare is left in place even on a column it does. Undeclared columns are + untouched entirely, so a vector index on `chunks` survives. The one thing it + will overwrite is an index at LanceDB's default name for a declared column, + `{column}_idx`, which is the name this function itself creates. + + Skipping matters as much as creating: `create_index(replace=True)` rebuilds an + identical index, writing a fresh index and a new table version rather than + no-oping, and leaves the previous index behind until the next vacuum. On a + large table over object storage that is a full column sort per pass. """ - indexed = { - column: index.index_type - for index in await table.list_indices() - for column in index.columns - } + covering: dict[str, set[str]] = {} + for index in await table.list_indices(): + for column in index.columns: + covering.setdefault(column, set()).add(index.index_type) + + applied: list[str] = [] for column, config in index_specs(table_name): - if indexed.get(column) == type(config).__name__: + declared = type(config).__name__ + present = covering.get(column, set()) + if declared in present: continue + if present: + logger.info( + f"Adding {declared} index on {table_name}.{column}, which carries " + f"{', '.join(sorted(present))}" + ) await table.create_index(column, config=config, replace=True) + applied.append(column) + return applied class SettingsRecord(LanceModel): diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py index 4280ff35..200f0f15 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py @@ -99,6 +99,9 @@ from haiku.rag.store.upgrades.v0_58_0 import ( from haiku.rag.store.upgrades.v0_64_0 import ( upgrade_rename_document_meta_id as upgrade_0_64_0_rename_document_meta_id, ) +from haiku.rag.store.upgrades.v0_75_0 import ( + upgrade_index_hot_lookup_keys as upgrade_0_75_0_index_hot_lookup_keys, +) upgrades.append(upgrade_0_20_0_docling) upgrades.append(upgrade_0_23_1_contextualize) @@ -110,3 +113,4 @@ upgrades.append(upgrade_0_48_0_heading_hierarchy) upgrades.append(upgrade_0_50_0_canonical_metadata_keys) upgrades.append(upgrade_0_58_0_split_document_meta) upgrades.append(upgrade_0_64_0_rename_document_meta_id) +upgrades.append(upgrade_0_75_0_index_hot_lookup_keys) diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py new file mode 100644 index 00000000..8aa2f4ab --- /dev/null +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py @@ -0,0 +1,39 @@ +import logging + +from haiku.rag.store.engine import Store, ensure_indexes +from haiku.rag.store.upgrades import Upgrade + +logger = logging.getLogger(__name__) + + +async def _apply_index_hot_lookup_keys(store: Store) -> None: + """Bring an existing database up to the declared index set. + + Adds what earlier versions never created: BTree on `documents.id`, + `chunks.id` and `chunks.document_id`, and a Bitmap on `document_items.label`. + Without them those lookups scan the column, which on object storage is + network I/O on paths that run per document. + + Does not materialize rows in Python or rewrite table data. The cost is the + index builds themselves, which read the indexed columns to sort them. + + `ensure_indexes` skips a column already indexed with the declared type, so a + database that already has the full set (a merged one, say) comes through + untouched rather than re-sorting every indexed column. Columns it does not + declare are left alone, so a vector index or an externally added index + survives. + """ + for table_name, table in store._tables().items(): + applied = await ensure_indexes(table, table_name) + if applied: + logger.info( + f"Indexed {table_name} (created or corrected): " + f"{', '.join(sorted(applied))}" + ) + + +upgrade_index_hot_lookup_keys = Upgrade( + version="0.75.0", + apply=_apply_index_hot_lookup_keys, + description="Index documents.id, chunks.id, chunks.document_id and document_items.label", +) diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 5977bf40..656a183a 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -2,7 +2,7 @@ name = "haiku.rag-slim" description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies" -version = "0.74.0" +version = "0.75.0" authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] license = { text = "MIT" } readme = { file = "README.md", content-type = "text/markdown" } diff --git a/pyproject.toml b/pyproject.toml index 194fc7d1..47d4349c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,7 +2,7 @@ name = "haiku.rag" description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling" -version = "0.74.0" +version = "0.75.0" authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] license = { text = "MIT" } readme = { file = "README.md", content-type = "text/markdown" } @@ -30,7 +30,7 @@ classifiers = [ ] dependencies = [ - "haiku.rag-slim[docling,voyageai,cohere,zeroentropy,tui,cross-encoder]==0.74.0", + "haiku.rag-slim[docling,voyageai,cohere,zeroentropy,tui,cross-encoder]==0.75.0", ] [project.scripts] @@ -38,9 +38,9 @@ haiku-rag = "haiku.rag.cli:cli" [project.optional-dependencies] tui = ["textual>=8.2.4"] -s3 = ["haiku.rag-slim[s3]==0.74.0"] -cross-encoder = ["haiku.rag-slim[cross-encoder]==0.74.0"] -ingester = ["haiku.rag-slim[ingester]==0.74.0"] +s3 = ["haiku.rag-slim[s3]==0.75.0"] +cross-encoder = ["haiku.rag-slim[cross-encoder]==0.75.0"] +ingester = ["haiku.rag-slim[ingester]==0.75.0"] [build-system] requires = ["hatchling"] diff --git a/tests/store/test_indexes.py b/tests/store/test_indexes.py index a0f011c4..5bbd3a7e 100644 --- a/tests/store/test_indexes.py +++ b/tests/store/test_indexes.py @@ -25,6 +25,15 @@ async def _index_type(table, column: str) -> str | None: return None +async def _covering(table, column: str) -> list[tuple[str, str]]: + """Every index over `column`, as (name, index_type).""" + return [ + (index.name, index.index_type) + for index in await table.list_indices() + if column in index.columns + ] + + @pytest.mark.asyncio async def test_fresh_database_indexes_every_hot_lookup_key(temp_db_path): """A new database carries the full index set, not a subset.""" @@ -67,6 +76,44 @@ async def test_ensure_indexes_corrects_an_index_of_the_wrong_type(temp_db_path): ) +@pytest.mark.asyncio +async def test_ensure_indexes_adds_the_declared_type_beside_a_custom_index( + temp_db_path, +): + """A wrong-typed index under a custom name must neither satisfy the check nor + be destroyed. `replace=True` replaces by name and the default name is derived + from the column, so a custom-named index is invisible to it either way. + """ + async with Store(temp_db_path, create=True) as store: + table = store.document_items_table + await table.drop_index("label_idx") + await table.create_index("label", config=BTree(), name="operator_label") + + await ensure_indexes(table, "document_items") + + covering = dict(await _covering(table, "label")) + assert covering["operator_label"] == "BTree" + assert "Bitmap" in covering.values() + + +@pytest.mark.asyncio +async def test_ensure_indexes_keeps_an_operator_index_on_a_declared_column( + temp_db_path, +): + """Two index types over one column can be deliberate, serving different query + shapes. An index this function did not declare is left alone even on a column + it does, so a migration never deletes an operator's index. + """ + async with Store(temp_db_path, create=True) as store: + table = store.document_items_table + await table.create_index("label", config=BTree(), name="operator_label") + + await ensure_indexes(table, "document_items") + + covering = dict(await _covering(table, "label")) + assert covering == {"label_idx": "Bitmap", "operator_label": "BTree"} + + @pytest.mark.asyncio async def test_delete_all_restores_the_full_index_set(temp_db_path): """delete_all drops and recreates tables; the recreated tables must come diff --git a/tests/store/test_v0_75_0_migration.py b/tests/store/test_v0_75_0_migration.py new file mode 100644 index 00000000..4c7e10e4 --- /dev/null +++ b/tests/store/test_v0_75_0_migration.py @@ -0,0 +1,111 @@ +import pytest +from lancedb.index import BTree + +from haiku.rag.store.engine import Store +from haiku.rag.store.models import Document +from haiku.rag.store.repositories.document import DocumentRepository +from haiku.rag.store.upgrades.v0_75_0 import _apply_index_hot_lookup_keys + +# What a pre-0.75.0 database carried: no index on `documents` at all, FTS only on +# `chunks`, and `document_items` without the `label` Bitmap. +LEGACY_DROPPED = { + "documents": ["id_idx"], + "chunks": ["id_idx", "document_id_idx"], + "document_items": ["label_idx"], +} + + +async def _indexed(table) -> dict[str, str]: + return { + column: index.index_type + for index in await table.list_indices() + for column in index.columns + } + + +async def _make_legacy(store: Store) -> None: + for table_name, indexes in LEGACY_DROPPED.items(): + table = store._tables()[table_name] + for index in indexes: + await table.drop_index(index) + + +@pytest.mark.asyncio +async def test_adds_the_missing_indexes(temp_db_path): + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _make_legacy(store) + + await _apply_index_hot_lookup_keys(store) + + assert await _indexed(store.documents_table) == {"id": "BTree"} + assert await _indexed(store.chunks_table) == { + "content_fts": "FTS", + "id": "BTree", + "document_id": "BTree", + } + assert await _indexed(store.document_items_table) == { + "document_id": "BTree", + "position": "BTree", + "self_ref": "BTree", + "label": "Bitmap", + } + + +@pytest.mark.asyncio +async def test_keeps_every_row(temp_db_path): + """Indexing must not touch data.""" + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + repo = DocumentRepository(store) + await repo.create(Document(content="Kept", uri="test://kept")) + await _make_legacy(store) + + await _apply_index_hot_lookup_keys(store) + + docs = await repo.list_all(include_content=True) + assert [d.content for d in docs] == ["Kept"] + + +@pytest.mark.asyncio +async def test_is_a_no_op_on_an_already_indexed_database(temp_db_path): + """An already-migrated database must not be re-indexed. `create_index` with + replace=True rebuilds, so a second run would re-sort every indexed column + and write a new version of every table. + """ + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _apply_index_hot_lookup_keys(store) + versions = {name: await t.version() for name, t in store._tables().items()} + + await _apply_index_hot_lookup_keys(store) + + assert {name: await t.version() for name, t in store._tables().items()} == ( + versions + ) + + +@pytest.mark.asyncio +async def test_replaces_a_wrong_typed_legacy_index(temp_db_path): + """A database that indexed `label` with a BTree gets the Bitmap it should + have, rather than being left alone because the column was covered.""" + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await store.document_items_table.create_index( + "label", config=BTree(), replace=True + ) + + await _apply_index_hot_lookup_keys(store) + + indexed = await _indexed(store.document_items_table) + assert indexed["label"] == "Bitmap" + + +@pytest.mark.asyncio +async def test_leaves_undeclared_indexes_alone(temp_db_path): + """A merged or externally built database may carry indexes haiku.rag never + creates. The migration must not drop them.""" + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await store.document_meta_table.create_index( + "title", config=BTree(), replace=True + ) + + await _apply_index_hot_lookup_keys(store) + + assert "title" in await _indexed(store.document_meta_table) diff --git a/tests/test_client.py b/tests/test_client.py index 939d1f07..94f240a5 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -993,10 +993,16 @@ async def test_client_import_documents_mixed_embeddings(temp_db_path): async def test_client_update_document_replaces_rows_with_bounded_versions( temp_db_path, ): - """Updating one document should replace stale rows with bounded versions.""" - dim = Config.embeddings.model.vector_dim + """Updating one document should replace stale rows with bounded versions. - async with HaikuRAG(temp_db_path, create=True) as client: + auto_vacuum is off: a background vacuum optimizes every table, which writes + versions of its own and would land inside the window being measured. + """ + dim = Config.embeddings.model.vector_dim + config = Config.model_copy(deep=True) + config.storage.auto_vacuum = False + + async with HaikuRAG(temp_db_path, config=config, create=True) as client: created = await client.import_document( _docling_doc("original", "Original body"), [Chunk(content="Original body", embedding=[0.1] * dim, order=0)], diff --git a/uv.lock b/uv.lock index c7c19729..99dac8fa 100644 --- a/uv.lock +++ b/uv.lock @@ -1577,7 +1577,7 @@ wheels = [ [[package]] name = "haiku-rag" -version = "0.74.0" +version = "0.75.0" source = { editable = "." } dependencies = [ { name = "haiku-rag-slim", extra = ["cohere", "cross-encoder", "docling", "tui", "voyageai", "zeroentropy"] }, @@ -1644,7 +1644,7 @@ dev = [ [[package]] name = "haiku-rag-evals" -version = "0.74.0" +version = "0.75.0" source = { editable = "evaluations" } dependencies = [ { name = "datasets" }, @@ -1667,7 +1667,7 @@ requires-dist = [ [[package]] name = "haiku-rag-slim" -version = "0.74.0" +version = "0.75.0" source = { editable = "haiku_rag_slim" } dependencies = [ { name = "docling-core" }, From 11644c7f4393d4531d7b332f12e1da583c253aac Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 12:38:02 +0300 Subject: [PATCH 03/11] Trim comments, docstrings and docs to what they need to say Also drops two things that were stale rather than merely verbose: the CLI docs note for 0.75.0, which was the only release-tagged note in the docs tree while the CHANGELOG already records that existing databases need `haiku-rag migrate`; and "(created or corrected)" from the migration log line, left over from the earlier behaviour that replaced wrong-typed indexes. --- docs/cli.md | 3 -- haiku_rag_slim/haiku/rag/store/engine.py | 36 ++++--------------- .../haiku/rag/store/upgrades/v0_75_0.py | 21 ++--------- tests/store/test_indexes.py | 28 ++++----------- tests/store/test_v0_75_0_migration.py | 14 +++----- tests/test_client.py | 3 +- 6 files changed, 21 insertions(+), 84 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index 9a3df4a4..06c2b726 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -359,9 +359,6 @@ Migration completed successfully. !!! tip Back up your database before running migrations. While migrations are designed to be safe, having a backup provides peace of mind for production databases. -!!! note "Upgrading to 0.75.0" - The 0.75.0 migration adds scalar indexes: BTree on `documents.id`, `chunks.id` and `chunks.document_id`, and Bitmap on `document_items.label`. It rewrites no table data, so it is far cheaper than the earlier data migrations, but building an index reads the whole column it indexes. On a large database, and particularly on object storage, budget for reading `chunks.id` and `chunks.document_id` in full. Indexes already present with the expected type are left alone, so a database that already carries the full set migrates instantly. - ### Download Models Download required runtime models: diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 5464fc57..002283a1 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -177,16 +177,7 @@ def get_document_items_arrow_schema() -> pa.Schema: def index_specs(table_name: str) -> list[tuple[str, Bitmap | BTree | FTS]]: - """The index set a haiku.rag table is expected to carry. - - Single source of truth for every path that creates a table: initialization, - migration, and the drop-and-recreate paths in the repositories. - - `label` gets a Bitmap rather than a BTree because it holds around ten - distinct values across every item of every document, and Bitmap is the - low-cardinality equality case. The FTS options are load-bearing: - `with_position` enables phrase queries and keeping stop words lets them match. - """ + """The index set each table carries.""" match table_name: case "documents": return [("id", BTree())] @@ -194,6 +185,7 @@ def index_specs(table_name: str) -> list[tuple[str, Bitmap | BTree | FTS]]: return [("id", BTree()), ("uri", BTree())] case "chunks": return [ + # Positions and stop words are required for phrase queries. ("content_fts", FTS(with_position=True, remove_stop_words=False)), ("id", BTree()), ("document_id", BTree()), @@ -210,27 +202,11 @@ def index_specs(table_name: str) -> list[tuple[str, Bitmap | BTree | FTS]]: async def ensure_indexes(table: lancedb.AsyncTable, table_name: str) -> list[str]: - """Ensure an index of the declared type covers each declared column. + """Create any declared index missing from a column. Returns the columns indexed. - Returns the columns it indexed, so callers can report what changed. - - The condition is the presence of the declared *type*, not that the column's - index happens to be that type. Checking coverage alone would let a wrong-typed - index satisfy the check: a BTree on `label` covers the column while losing the - low-cardinality equality lookup a Bitmap gives. - - Nothing is ever dropped or converted away from. Two index types over one - column can be deliberate, since they serve different query shapes (BTree for - equality and range, `Fm` for `contains`), so an index this function did not - declare is left in place even on a column it does. Undeclared columns are - untouched entirely, so a vector index on `chunks` survives. The one thing it - will overwrite is an index at LanceDB's default name for a declared column, - `{column}_idx`, which is the name this function itself creates. - - Skipping matters as much as creating: `create_index(replace=True)` rebuilds an - identical index, writing a fresh index and a new table version rather than - no-oping, and leaves the previous index behind until the next vacuum. On a - large table over object storage that is a full column sort per pass. + Matches on index type, not column coverage, so a BTree does not satisfy a + declared Bitmap. Never drops or converts an index it did not declare. + Re-creating is not free: `create_index(replace=True)` rebuilds. """ covering: dict[str, set[str]] = {} for index in await table.list_indices(): diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py index 8aa2f4ab..f71b7001 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py @@ -7,29 +7,14 @@ logger = logging.getLogger(__name__) async def _apply_index_hot_lookup_keys(store: Store) -> None: - """Bring an existing database up to the declared index set. + """Add the declared indexes to a database created before 0.75.0. - Adds what earlier versions never created: BTree on `documents.id`, - `chunks.id` and `chunks.document_id`, and a Bitmap on `document_items.label`. - Without them those lookups scan the column, which on object storage is - network I/O on paths that run per document. - - Does not materialize rows in Python or rewrite table data. The cost is the - index builds themselves, which read the indexed columns to sort them. - - `ensure_indexes` skips a column already indexed with the declared type, so a - database that already has the full set (a merged one, say) comes through - untouched rather than re-sorting every indexed column. Columns it does not - declare are left alone, so a vector index or an externally added index - survives. + Rewrites no rows. Each index build reads the column it indexes. """ for table_name, table in store._tables().items(): applied = await ensure_indexes(table, table_name) if applied: - logger.info( - f"Indexed {table_name} (created or corrected): " - f"{', '.join(sorted(applied))}" - ) + logger.info(f"Indexed {table_name}: {', '.join(sorted(applied))}") upgrade_index_hot_lookup_keys = Upgrade( diff --git a/tests/store/test_indexes.py b/tests/store/test_indexes.py index 5bbd3a7e..3fc2c542 100644 --- a/tests/store/test_indexes.py +++ b/tests/store/test_indexes.py @@ -36,7 +36,7 @@ async def _covering(table, column: str) -> list[tuple[str, str]]: @pytest.mark.asyncio async def test_fresh_database_indexes_every_hot_lookup_key(temp_db_path): - """A new database carries the full index set, not a subset.""" + """A new database carries the full index set.""" async with Store(temp_db_path, create=True) as store: for name, table in store._tables().items(): expected = EXPECTED_INDEXED_COLUMNS.get(name, set()) @@ -45,8 +45,7 @@ async def test_fresh_database_indexes_every_hot_lookup_key(temp_db_path): @pytest.mark.asyncio async def test_ensure_indexes_skips_existing_instead_of_rebuilding(temp_db_path): - """`create_index(replace=True)` rebuilds an identical index and writes a new - table version, so a second pass must skip rather than replace.""" + """A second pass must not rebuild: replace=True writes a new version.""" async with Store(temp_db_path, create=True) as store: table = store.chunks_table version_before = await table.version() @@ -59,10 +58,7 @@ async def test_ensure_indexes_skips_existing_instead_of_rebuilding(temp_db_path) @pytest.mark.asyncio async def test_ensure_indexes_corrects_an_index_of_the_wrong_type(temp_db_path): - """A column indexed with the wrong type must be re-indexed. `label` is the - live case: a BTree over ~ten distinct values loses the low-cardinality - equality lookup a Bitmap gives, and column coverage alone cannot see it. - """ + """A wrong-typed index does not satisfy the declared one.""" async with Store(temp_db_path, create=True) as store: table = store.document_items_table await table.create_index("label", config=BTree(), replace=True) @@ -80,10 +76,7 @@ async def test_ensure_indexes_corrects_an_index_of_the_wrong_type(temp_db_path): async def test_ensure_indexes_adds_the_declared_type_beside_a_custom_index( temp_db_path, ): - """A wrong-typed index under a custom name must neither satisfy the check nor - be destroyed. `replace=True` replaces by name and the default name is derived - from the column, so a custom-named index is invisible to it either way. - """ + """A custom-named index neither satisfies the check nor is destroyed.""" async with Store(temp_db_path, create=True) as store: table = store.document_items_table await table.drop_index("label_idx") @@ -100,10 +93,7 @@ async def test_ensure_indexes_adds_the_declared_type_beside_a_custom_index( async def test_ensure_indexes_keeps_an_operator_index_on_a_declared_column( temp_db_path, ): - """Two index types over one column can be deliberate, serving different query - shapes. An index this function did not declare is left alone even on a column - it does, so a migration never deletes an operator's index. - """ + """An index we did not declare survives, even on a declared column.""" async with Store(temp_db_path, create=True) as store: table = store.document_items_table await table.create_index("label", config=BTree(), name="operator_label") @@ -116,8 +106,7 @@ async def test_ensure_indexes_keeps_an_operator_index_on_a_declared_column( @pytest.mark.asyncio async def test_delete_all_restores_the_full_index_set(temp_db_path): - """delete_all drops and recreates tables; the recreated tables must come - back with the same indexes a fresh database gets.""" + """Recreated tables come back with the full index set.""" async with Store(temp_db_path, create=True) as store: repo = DocumentRepository(store) await repo.create(Document(content="A document")) @@ -131,10 +120,7 @@ async def test_delete_all_restores_the_full_index_set(temp_db_path): @pytest.mark.asyncio async def test_delete_all_keeps_picture_data_as_large_binary(temp_db_path): - """document_items must be recreated from the Arrow schema, which declares - picture_data as large_binary. The 32-bit `binary` type overflows its offsets - once a fragment holds enough embedded pictures. - """ + """picture_data must survive delete_all as large_binary, not binary.""" async with Store(temp_db_path, create=True) as store: repo = DocumentRepository(store) await repo.create(Document(content="A document")) diff --git a/tests/store/test_v0_75_0_migration.py b/tests/store/test_v0_75_0_migration.py index 4c7e10e4..3b996d51 100644 --- a/tests/store/test_v0_75_0_migration.py +++ b/tests/store/test_v0_75_0_migration.py @@ -6,8 +6,7 @@ from haiku.rag.store.models import Document from haiku.rag.store.repositories.document import DocumentRepository from haiku.rag.store.upgrades.v0_75_0 import _apply_index_hot_lookup_keys -# What a pre-0.75.0 database carried: no index on `documents` at all, FTS only on -# `chunks`, and `document_items` without the `label` Bitmap. +# The indexes a pre-0.75.0 database lacked. LEGACY_DROPPED = { "documents": ["id_idx"], "chunks": ["id_idx", "document_id_idx"], @@ -67,10 +66,7 @@ async def test_keeps_every_row(temp_db_path): @pytest.mark.asyncio async def test_is_a_no_op_on_an_already_indexed_database(temp_db_path): - """An already-migrated database must not be re-indexed. `create_index` with - replace=True rebuilds, so a second run would re-sort every indexed column - and write a new version of every table. - """ + """A second run must not re-index: replace=True rebuilds.""" async with Store(temp_db_path, create=True, skip_migration_check=True) as store: await _apply_index_hot_lookup_keys(store) versions = {name: await t.version() for name, t in store._tables().items()} @@ -84,8 +80,7 @@ async def test_is_a_no_op_on_an_already_indexed_database(temp_db_path): @pytest.mark.asyncio async def test_replaces_a_wrong_typed_legacy_index(temp_db_path): - """A database that indexed `label` with a BTree gets the Bitmap it should - have, rather than being left alone because the column was covered.""" + """A BTree on `label` does not satisfy the declared Bitmap.""" async with Store(temp_db_path, create=True, skip_migration_check=True) as store: await store.document_items_table.create_index( "label", config=BTree(), replace=True @@ -99,8 +94,7 @@ async def test_replaces_a_wrong_typed_legacy_index(temp_db_path): @pytest.mark.asyncio async def test_leaves_undeclared_indexes_alone(temp_db_path): - """A merged or externally built database may carry indexes haiku.rag never - creates. The migration must not drop them.""" + """Indexes haiku.rag never declared are not dropped.""" async with Store(temp_db_path, create=True, skip_migration_check=True) as store: await store.document_meta_table.create_index( "title", config=BTree(), replace=True diff --git a/tests/test_client.py b/tests/test_client.py index 94f240a5..e6e379a2 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -995,8 +995,7 @@ async def test_client_update_document_replaces_rows_with_bounded_versions( ): """Updating one document should replace stale rows with bounded versions. - auto_vacuum is off: a background vacuum optimizes every table, which writes - versions of its own and would land inside the window being measured. + auto_vacuum is off: its writes would land inside the measured window. """ dim = Config.embeddings.model.vector_dim config = Config.model_copy(deep=True) From bd7946178d0158feaeb886e98074cccb4a4b4628 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 14:28:14 +0300 Subject: [PATCH 04/11] Pin the eval judge to qwen3.8 --- CHANGELOG.md | 1 + docs/benchmarks.md | 14 ++++++++------ evaluations/configs/hotpotqa.yaml | 6 +++--- evaluations/configs/mtrag_clapnq.yaml | 6 +++--- evaluations/configs/orb_multimodal.yaml | 6 +++--- evaluations/configs/orb_multimodal_nemotron.yaml | 6 +++--- evaluations/configs/orb_text.yaml | 6 +++--- evaluations/evaluations/benchmark.py | 5 +++-- evaluations/tests/test_benchmark.py | 1 + evaluations/tests/test_reference_configs.py | 2 +- 10 files changed, 29 insertions(+), 24 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 355eabd9..de7fe3b8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ ### Changed +- Eval judge pinned to `qwen3.8`: `DEFAULT_JUDGE_MODEL` is `ollama:qwen3.8`, and the reference configs use `Inferact/Qwen3.8-27B-NVFP4` with `extra_body.chat_template_kwargs.reasoning_effort: low`. Results in `docs/benchmarks.md` were judged by `Qwen3.6-35B-A3B-NVFP4` and are not re-judged. - `import_documents` embeds chunks across the whole batch in one pass instead of per document. ### Fixed diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 68340bb0..3a2ce679 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -70,13 +70,13 @@ evaluations run hotpotqa --config /path/to/haiku.rag.yaml --db /path/to/custom.l If no config file is specified, the script searches standard locations: `./haiku.rag.yaml`, user config directory, then falls back to defaults. -To pin the LLM judge in YAML (rather than the default `ollama:qwen3.6`). These are the recommended settings: +To pin the LLM judge in YAML (rather than the default `ollama:qwen3.8`). These are the recommended settings: ```yaml evaluations: judge: provider: openai - name: RedHatAI/Qwen3.6-35B-A3B-NVFP4 + name: Inferact/Qwen3.8-27B-NVFP4 base_url: http://localhost:8000/v1 # optional, for OpenAI-compatible servers (vLLM, LM Studio, etc.) temperature: 0.6 max_tokens: 16384 @@ -85,7 +85,7 @@ evaluations: top_k: 20 min_p: 0 chat_template_kwargs: - enable_thinking: true + reasoning_effort: low # qwen3.8: low | medium | xhigh (default) ``` ### Restricting the corpus @@ -121,11 +121,13 @@ Filtering affects searches only — a run without `--skip-db` still populates th ### QA Accuracy -`pydantic-evals` coordinates an LLM judge to determine whether the capability's answer is correct. The default judge is `ollama:qwen3.6`, pinned so changes to the capability model don't change the judge underneath. Set `evaluations.judge` in `haiku.rag.yaml` to override (including a custom `base_url` for any OpenAI-compatible endpoint). Accuracy is the fraction of correctly answered questions. +`pydantic-evals` coordinates an LLM judge to determine whether the capability's answer is correct. The default judge is `ollama:qwen3.8`, pinned so changes to the capability model don't change the judge underneath. Set `evaluations.judge` in `haiku.rag.yaml` to override (including a custom `base_url` for any OpenAI-compatible endpoint). Accuracy is the fraction of correctly answered questions. A dataset that brings its own deterministic evaluator is scored by that evaluator instead, and no judge runs. T²-RAGBench is the only such dataset today, scored by `NumberMatchEvaluator`. -We picked `qwen3.6` over the previously-pinned `gpt-oss` after a 4-cell calibration (gpt-oss / qwen3.6 as both answerer and judge, with Claude Opus 4.7 as a reference). `qwen3.6` had κ ≥ 0.66 vs the reference on both same-family and cross-family answerers (vs ~0.39–0.55 for `gpt-oss`) and showed no measurable self-preference bias, while `gpt-oss` was ~10 pp more lenient on its own outputs. +`qwen3.8` replaced `qwen3.6` after a 120-case calibration on ORB, stratified 60 pass / 60 fail: agreement 0.950, Cohen's κ 0.900, and in all 6 disagreements it matched or beat `qwen3.6` (4 were `qwen3.6` failing answers that were equivalent in different notation). It emits no reasoning content, so it avoids the thinking spirals that made `qwen3.6` exceed its output budget and drop verdicts. `reasoning_effort` changes its verdicts in 1 case per 120, so the cheaper `low` is pinned. + +Before that, we picked `qwen3.6` over the previously-pinned `gpt-oss` after a 4-cell calibration (gpt-oss / qwen3.6 as both answerer and judge, with Claude Opus 4.7 as a reference). `qwen3.6` had κ ≥ 0.66 vs the reference on both same-family and cross-family answerers (vs ~0.39–0.55 for `gpt-oss`) and showed no measurable self-preference bias, while `gpt-oss` was ~10 pp more lenient on its own outputs. ### Citation Retrieval @@ -135,7 +137,7 @@ This is computed alongside QA accuracy from the same capability run, no extra in ## Current results -Numbers measured under the current pinned judge (`ollama:qwen3.6`) on a recent `haiku.rag` version. +Numbers below were measured under `Qwen3.6-35B-A3B-NVFP4` as judge, on a recent `haiku.rag` version. The pinned judge is now `qwen3.8`; rows are not re-judged, so compare rows to each other rather than to runs judged by `qwen3.8`. ### OpenRAG Bench (ORB) diff --git a/evaluations/configs/hotpotqa.yaml b/evaluations/configs/hotpotqa.yaml index 606ae86d..84aca56d 100644 --- a/evaluations/configs/hotpotqa.yaml +++ b/evaluations/configs/hotpotqa.yaml @@ -25,8 +25,8 @@ qa: evaluations: judge: provider: openai - name: RedHatAI/Qwen3.6-35B-A3B-NVFP4 - base_url: http://vllm:11430/v1 + name: Inferact/Qwen3.8-27B-NVFP4 + base_url: http://vllm:11439/v1 temperature: 0.6 max_tokens: 16384 extra_body: @@ -34,4 +34,4 @@ evaluations: top_k: 20 min_p: 0 chat_template_kwargs: - enable_thinking: true + reasoning_effort: low diff --git a/evaluations/configs/mtrag_clapnq.yaml b/evaluations/configs/mtrag_clapnq.yaml index 7fa34d84..5dbac876 100644 --- a/evaluations/configs/mtrag_clapnq.yaml +++ b/evaluations/configs/mtrag_clapnq.yaml @@ -44,8 +44,8 @@ qa: evaluations: judge: provider: openai - name: RedHatAI/Qwen3.6-35B-A3B-NVFP4 - base_url: http://vllm:11430/v1 + name: Inferact/Qwen3.8-27B-NVFP4 + base_url: http://vllm:11439/v1 temperature: 0.6 max_tokens: 16384 extra_body: @@ -53,4 +53,4 @@ evaluations: top_k: 20 min_p: 0 chat_template_kwargs: - enable_thinking: true + reasoning_effort: low diff --git a/evaluations/configs/orb_multimodal.yaml b/evaluations/configs/orb_multimodal.yaml index 2dad063e..b868fb3c 100644 --- a/evaluations/configs/orb_multimodal.yaml +++ b/evaluations/configs/orb_multimodal.yaml @@ -29,8 +29,8 @@ qa: evaluations: judge: provider: openai - name: RedHatAI/Qwen3.6-35B-A3B-NVFP4 - base_url: http://vllm:11430/v1 + name: Inferact/Qwen3.8-27B-NVFP4 + base_url: http://vllm:11439/v1 temperature: 0.6 max_tokens: 16384 extra_body: @@ -38,4 +38,4 @@ evaluations: top_k: 20 min_p: 0 chat_template_kwargs: - enable_thinking: true + reasoning_effort: low diff --git a/evaluations/configs/orb_multimodal_nemotron.yaml b/evaluations/configs/orb_multimodal_nemotron.yaml index eca56502..99e3899b 100644 --- a/evaluations/configs/orb_multimodal_nemotron.yaml +++ b/evaluations/configs/orb_multimodal_nemotron.yaml @@ -30,8 +30,8 @@ qa: evaluations: judge: provider: openai - name: RedHatAI/Qwen3.6-35B-A3B-NVFP4 - base_url: http://vllm:11430/v1 + name: Inferact/Qwen3.8-27B-NVFP4 + base_url: http://vllm:11439/v1 temperature: 0.6 max_tokens: 16384 extra_body: @@ -39,4 +39,4 @@ evaluations: top_k: 20 min_p: 0 chat_template_kwargs: - enable_thinking: true + reasoning_effort: low diff --git a/evaluations/configs/orb_text.yaml b/evaluations/configs/orb_text.yaml index 2475bac8..18f87cb9 100644 --- a/evaluations/configs/orb_text.yaml +++ b/evaluations/configs/orb_text.yaml @@ -31,8 +31,8 @@ qa: evaluations: judge: provider: openai - name: RedHatAI/Qwen3.6-35B-A3B-NVFP4 - base_url: http://vllm:11430/v1 + name: Inferact/Qwen3.8-27B-NVFP4 + base_url: http://vllm:11439/v1 temperature: 0.6 max_tokens: 16384 extra_body: @@ -40,4 +40,4 @@ evaluations: top_k: 20 min_p: 0 chat_template_kwargs: - enable_thinking: true + reasoning_effort: low diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index cb9f534c..0eb34507 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -47,10 +47,11 @@ TARGETS: tuple[Target, ...] = ("rag-capability", "analysis-capability") # Sampling follows Qwen's recommendation for thinking mode; its model cards # forbid greedy decoding. Only the keys ollama honours are set: it silently # ignores `top_k`, `min_p` and `chat_template_kwargs`. The vLLM reference -# configs under `evaluations/configs/` carry those too. +# configs under `evaluations/configs/` carry those too, plus +# `reasoning_effort`, which qwen3.8 reads from `chat_template_kwargs`. DEFAULT_JUDGE_MODEL = ModelConfig( provider="ollama", - name="qwen3.6", + name="qwen3.8", temperature=0.6, max_tokens=16384, extra_body={"top_p": 0.95}, diff --git a/evaluations/tests/test_benchmark.py b/evaluations/tests/test_benchmark.py index 16b43013..13a4e20a 100644 --- a/evaluations/tests/test_benchmark.py +++ b/evaluations/tests/test_benchmark.py @@ -677,6 +677,7 @@ class TestRunQaBenchmarkJudgeModel: from evaluations.benchmark import DEFAULT_JUDGE_MODEL assert DEFAULT_JUDGE_MODEL.temperature == 0.6 + assert DEFAULT_JUDGE_MODEL.name == "qwen3.8" assert DEFAULT_JUDGE_MODEL.max_tokens == 16384 assert DEFAULT_JUDGE_MODEL.extra_body == {"top_p": 0.95} diff --git a/evaluations/tests/test_reference_configs.py b/evaluations/tests/test_reference_configs.py index b25f762e..142ccb34 100644 --- a/evaluations/tests/test_reference_configs.py +++ b/evaluations/tests/test_reference_configs.py @@ -15,7 +15,7 @@ PINNED_JUDGE_SAMPLING = { "top_p": 0.95, "top_k": 20, "min_p": 0, - "chat_template_kwargs": {"enable_thinking": True}, + "chat_template_kwargs": {"reasoning_effort": "low"}, }, } From d28e2662e53e5f6d5714117fecc5bd4612e3b9c6 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 13:03:26 +0300 Subject: [PATCH 05/11] Fix the MCP registry entry and fill in package and docs metadata --- CHANGELOG.md | 11 +++++++---- docs/index.md | 34 ++++++++++++++++++++++++++++++++++ haiku_rag_slim/pyproject.toml | 26 ++++++++++++++++++++++++-- overrides/main.html | 4 ++-- pyproject.toml | 18 ++++++++++++++++-- server.json | 20 +++++--------------- zensical.toml | 2 +- 7 files changed, 89 insertions(+), 26 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index de7fe3b8..bbfe8308 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,15 +12,18 @@ - Eval judge pinned to `qwen3.8`: `DEFAULT_JUDGE_MODEL` is `ollama:qwen3.8`, and the reference configs use `Inferact/Qwen3.8-27B-NVFP4` with `extra_body.chat_template_kwargs.reasoning_effort: low`. Results in `docs/benchmarks.md` were judged by `Qwen3.6-35B-A3B-NVFP4` and are not re-judged. - `import_documents` embeds chunks across the whole batch in one pass instead of per document. - -### Fixed - -- `DocumentRepository.delete_all` recreated `document_items` from `DocumentItemRecord` instead of `get_document_items_arrow_schema()`, returning `picture_data` as `binary` rather than `large_binary`. +- `haiku.rag` and `haiku.rag-slim` summaries and keywords; both packages now publish `[project.urls]`. +- `server.json` declares `title` and `websiteUrl`, and drops the `keywords` and `license` keys, which are not in the server schema. ### Removed - `wix` evaluation dataset and its reference config `evaluations/configs/wix.yaml`. +### Fixed + +- `DocumentRepository.delete_all` recreated `document_items` from `DocumentItemRecord` instead of `get_document_items_arrow_schema()`, returning `picture_data` as `binary` rather than `large_binary`. +- `server.json` runtime arguments are `mcp --stdio`, was `serve --mcp`. + ## [0.74.0] - 2026-08-13 ### Added diff --git a/docs/index.md b/docs/index.md index d3e9886b..421df9c2 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,3 +1,37 @@ --- title: haiku.rag +description: Local-first agentic RAG. Index PDFs, web pages, and whole directories, then ask questions and get answers cited to page numbers and section headings. Hybrid search, reranking, and multimodal retrieval on embedded LanceDB. --- + +haiku.rag indexes PDFs, web pages, and whole directories, retrieves with hybrid search, and answers with citations down to the page number and section heading. It runs on an embedded database with open models, so your documents stay on your machine and there is no server to operate. + +```bash +uv pip install haiku.rag + +haiku-rag init +haiku-rag add-src ~/Documents/some-paper.pdf +haiku-rag ask "what does it conclude?" +``` + +[Quickstart](tutorial.md) covers provider setup and the first ingestion. + +## Why haiku.rag + +**Answers you can check.** Every answer carries citations with page numbers and section headings. Visual grounding shows the cited chunk highlighted on the original page image. Optional capabilities require an answer to declare what grounds it, including declaring that nothing does. + +**Local-first, no server.** Embedded [LanceDB](https://lancedb.com/) and open models through [Ollama](https://ollama.com/) by default. No database to run and no API keys required. The same code runs against S3, GCS, Azure, LanceDB Cloud, or any provider Pydantic AI supports. + +**Built for agents.** Native [Pydantic AI](https://ai.pydantic.dev/) capabilities compose into your own agents. An [MCP server](mcp.md) exposes the same database to Claude Desktop and other assistants. The analysis capability runs sandboxed Python across documents for questions that need computation rather than retrieval. + +**Measured, not asserted.** Retrieval and answer quality are tracked against public benchmarks with runnable configs. See [Benchmarks](benchmarks.md). + +## Start here + +- [Quickstart](tutorial.md): install, index, chat. +- [Overview](overview.md): what haiku.rag does, end to end. +- [Capabilities](capabilities/index.md): native RAG and analysis capabilities for Pydantic AI agents. +- [Python API](python.md): use haiku.rag from code. +- [MCP server](mcp.md): expose haiku.rag to Claude Desktop or other AI assistants. +- [Configuration](configuration/index.md): every setting. + +MIT licensed. Source on [GitHub](https://github.com/ggozad/haiku.rag). diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 656a183a..23757cd4 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -1,13 +1,28 @@ [project] name = "haiku.rag-slim" -description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies" +description = "Local-first agentic RAG with citations - hybrid search, reranking and multimodal retrieval over your own documents, no database server required - Minimal dependencies" version = "0.75.0" authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] license = { text = "MIT" } readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.12" -keywords = ["RAG", "lancedb", "vector-database", "ml", "mcp"] +keywords = [ + "RAG", + "agentic-rag", + "lancedb", + "vector-database", + "hybrid-search", + "reranking", + "multimodal-rag", + "embeddings", + "citations", + "document-ingestion", + "mcp", + "mcp-server", + "pydantic-ai", + "docling", +] classifiers = [ "Development Status :: 4 - Beta", "Environment :: Console", @@ -41,6 +56,13 @@ dependencies = [ "zstandard>=0.23.0; python_version<'3.14'", ] +[project.urls] +Homepage = "https://ggozad.github.io/haiku.rag/" +Documentation = "https://ggozad.github.io/haiku.rag/" +Repository = "https://github.com/ggozad/haiku.rag" +Issues = "https://github.com/ggozad/haiku.rag/issues" +Changelog = "https://ggozad.github.io/haiku.rag/changelog/" + [project.optional-dependencies] # Document processing docling = ["docling>=2.102.2,<3.0.0", "opencv-python-headless>=4.6.0.66,<5.0.0.0"] diff --git a/overrides/main.html b/overrides/main.html index 5bbf25f9..fc67d494 100644 --- a/overrides/main.html +++ b/overrides/main.html @@ -3,7 +3,7 @@ {% block extrahead %} - + @@ -17,7 +17,7 @@

haiku.rag

-

Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling. Runs locally, scales to production.

+

Ask questions about your own documents and get answers that cite their sources. Agentic RAG on LanceDB, Pydantic AI, and Docling. Runs locally, scales to production.

Get started Learn more diff --git a/pyproject.toml b/pyproject.toml index 47d4349c..e6836e02 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "haiku.rag" -description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling" +description = "Local-first agentic RAG with citations - hybrid search, reranking and multimodal retrieval over your own documents, no database server required" version = "0.75.0" authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] license = { text = "MIT" } @@ -9,10 +9,17 @@ readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.12" keywords = [ "RAG", + "agentic-rag", "lancedb", "vector-database", - "ml", + "hybrid-search", + "reranking", + "multimodal-rag", + "embeddings", + "citations", + "document-ingestion", "mcp", + "mcp-server", "pydantic-ai", "docling", ] @@ -33,6 +40,13 @@ dependencies = [ "haiku.rag-slim[docling,voyageai,cohere,zeroentropy,tui,cross-encoder]==0.75.0", ] +[project.urls] +Homepage = "https://ggozad.github.io/haiku.rag/" +Documentation = "https://ggozad.github.io/haiku.rag/" +Repository = "https://github.com/ggozad/haiku.rag" +Issues = "https://github.com/ggozad/haiku.rag/issues" +Changelog = "https://ggozad.github.io/haiku.rag/changelog/" + [project.scripts] haiku-rag = "haiku.rag.cli:cli" diff --git a/server.json b/server.json index 5c33f50a..23d187e9 100644 --- a/server.json +++ b/server.json @@ -1,24 +1,14 @@ { "$schema": "https://static.modelcontextprotocol.io/schemas/2025-12-11/server.schema.json", "name": "io.github.ggozad/haiku-rag", + "title": "haiku.rag", + "description": "Local-first agentic RAG with citations - hybrid search, reranking, multimodal document retrieval", "version": "{{VERSION}}", - "description": "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling", + "websiteUrl": "https://ggozad.github.io/haiku.rag/", "repository": { "url": "https://github.com/ggozad/haiku.rag", "source": "github" }, - "license": "MIT", - "keywords": [ - "rag", - "lancedb", - "vector-database", - "embeddings", - "search", - "qa", - "research", - "docling", - "pydantic-ai" - ], "packages": [ { "registryType": "pypi", @@ -29,11 +19,11 @@ "runtimeArguments": [ { "type": "positional", - "value": "serve" + "value": "mcp" }, { "type": "named", - "name": "--mcp" + "name": "--stdio" } ], "transport": { diff --git a/zensical.toml b/zensical.toml index fbf59e7f..0fa0fe4a 100644 --- a/zensical.toml +++ b/zensical.toml @@ -1,6 +1,6 @@ [project] site_name = "haiku.rag" -site_description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling." +site_description = "Local-first agentic RAG. Index your documents, then ask questions and get answers cited to page numbers and section headings. Hybrid search, reranking, and multimodal retrieval on embedded LanceDB." site_url = "https://ggozad.github.io/haiku.rag/" repo_url = "https://github.com/ggozad/haiku.rag" repo_name = "ggozad/haiku.rag" From 4bcfb9cc3b75f4a3d9b26c7e77fc62a287c03c34 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 13:11:37 +0300 Subject: [PATCH 06/11] Lead the README with what haiku.rag does --- README.md | 10 ++++++++-- evaluations/README.md | 2 +- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 869b3bfd..1cc489b0 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,15 @@ -# Haiku RAG +# haiku.rag +[![PyPI](https://img.shields.io/pypi/v/haiku.rag)](https://pypi.org/project/haiku.rag/) +[![Python](https://img.shields.io/pypi/pyversions/haiku.rag)](https://pypi.org/project/haiku.rag/) +[![Downloads](https://static.pepy.tech/badge/haiku-rag-slim/month)](https://pepy.tech/projects/haiku-rag-slim) +[![Docs](https://img.shields.io/badge/docs-ggozad.github.io-blue)](https://ggozad.github.io/haiku.rag/) [![Tests](https://github.com/ggozad/haiku.rag/actions/workflows/test.yml/badge.svg)](https://github.com/ggozad/haiku.rag/actions/workflows/test.yml) [![codecov](https://codecov.io/gh/ggozad/haiku.rag/graph/badge.svg)](https://codecov.io/gh/ggozad/haiku.rag) -Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.pydantic.dev/), and [Docling](https://docling-project.github.io/docling/). +Agentic RAG that answers questions about your own documents with citations to page numbers and section headings. Runs locally on an embedded database, no server required. + +Built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.pydantic.dev/), and [Docling](https://docling-project.github.io/docling/). Full documentation at [ggozad.github.io/haiku.rag](https://ggozad.github.io/haiku.rag/). > **New: vision and multimodal search.** Picture-aware ingestion captures embedded figure bytes; vision-capable QA models receive them alongside text. Multimodal embedders put picture vectors in the same space as text, enabling text-as-query → figure hits and image-as-query retrieval. diff --git a/evaluations/README.md b/evaluations/README.md index ff1c6f7d..341fb3a6 100644 --- a/evaluations/README.md +++ b/evaluations/README.md @@ -1,4 +1,4 @@ -# Haiku RAG - Evaluations +# haiku.rag - Evaluations Internal benchmarking and evaluation scripts for haiku.rag. From 198d9d7b7322a234199f4456991d4cfd90d28246 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 14:25:07 +0300 Subject: [PATCH 07/11] Lead the benchmarks page with results Current results moved above Methodology and Running Evaluations. A reader arriving from an external link met four screens of CLI flags and download instructions before any number. Benchmarks is promoted to a top-level nav entry, out of Reference, which holds Development and Changelog. The page move itself changes no wording or figures. --- docs/benchmarks.md | 266 ++++++++++++++++++++++----------------------- zensical.toml | 2 +- 2 files changed, 134 insertions(+), 134 deletions(-) diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 3a2ce679..c75d1040 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -2,139 +2,6 @@ We evaluate `haiku.rag` on a small set of datasets that exercise different parts of the pipeline. OpenRAG Bench (ORB), T²-RAGBench, HotpotQA, and MTRAG are the datasets we currently track. Retrieval, QA accuracy, and citation retrieval are scored end-to-end through the RAG and analysis capabilities. -## Running Evaluations - -You can run evaluations with the `evaluations` CLI: - -```bash -evaluations run hotpotqa -evaluations run orb_text -``` - -The evaluation flow is orchestrated with [`pydantic-evals`](https://github.com/pydantic/pydantic-ai/tree/main/libs/pydantic-evals), which we leverage for dataset management, scoring, and report generation. - -### Pre-built Databases - -Building evaluation databases from scratch can take a long time, especially for large datasets like OpenRAG Bench. Pre-built databases are available on HuggingFace: - -```bash -# Download a specific dataset -evaluations download hotpotqa - -# Download all datasets -evaluations download all - -# Force re-download (overwrite existing) -evaluations download hotpotqa --force -``` - -Active datasets: - -| Dataset | Size | -|---------|------| -| `orb_text` — OpenRAG Bench, text embedder (`qwen3-embedding:4b`) with VLM picture descriptions baked into chunk content | ~18 GB | -| `orb_multimodal` — OpenRAG Bench, multimodal embedder (`qwen3-vl-embedding-8b`); picture vectors live in the same space as text for cross-modal retrieval | ~16 GB | -| `orb_multimodal_nemotron` — OpenRAG Bench, multimodal embedder (`nvidia/llama-nemotron-embed-vl-1b-v2`), the embedder behind the published headline results | ~16 GB | -| `t2_finqa` — T²-RAGBench (FinQA) financial QA, text embedder (`qwen3-embedding:4b`); scored by exact numeric match, run with `--target analysis-capability` | ~2 GB | -| `hotpotqa` — HotpotQA multi-hop QA over Wikipedia paragraphs, text embedder (`qwen3-embedding:4b`) | ~1.5 GB | -| `mtrag_clapnq` — MTRAG multi-turn RAG, ClapNQ (Wikipedia) passages, text embedder (`qwen3-embedding:4b`); also serves the `mtrag_clapnq_rewrite`, `mtrag_clapnq_live` and `mtrag_clapnq_live_uncompacted` keys | ~2.8 GB | - -After downloading, run benchmarks with `--skip-db`. Each database is built with a specific embedder, so pass its reference config from `evaluations/configs/` (a database only opens against a config whose embedder matches): - -```bash -evaluations run orb_multimodal_nemotron --skip-db --config configs/orb_multimodal_nemotron.yaml -``` - -The configs use `vllm` as the model host. Point `base_url` at your own OpenAI-compatible endpoints to reproduce the numbers. - -### Configuration - -The benchmark script accepts several options: - -```bash -evaluations run hotpotqa --config /path/to/haiku.rag.yaml --db /path/to/custom.lancedb -``` - -**Options:** - -- `--config PATH` - Specify a custom `haiku.rag.yaml` configuration file -- `--db PATH` - Override the database path (default: platform-specific user data directory) -- `--skip-db` - Skip updating the evaluation database -- `--skip-retrieval` - Skip retrieval benchmark -- `--skip-qa` - Skip QA benchmark -- `--limit N` - Limit number of test cases -- `--name NAME` - Override the evaluation name -- `--target {rag-capability,analysis-capability}` - Choose which [capability](capabilities/index.md) to benchmark end-to-end (default: `rag-capability`). The target names remain stable dataset identifiers. -- `--capability-model PROVIDER:NAME` - Override the capability model independently from the judge (default: `config.qa.model`, or `config.analysis.model` when set for `--target analysis-capability`). -- `--filter CLAUSE` / `-f CLAUSE` - Restrict every benchmark search to a subset of the database (see [Restricting the corpus](#restricting-the-corpus)). - -If no config file is specified, the script searches standard locations: `./haiku.rag.yaml`, user config directory, then falls back to defaults. - -To pin the LLM judge in YAML (rather than the default `ollama:qwen3.8`). These are the recommended settings: - -```yaml -evaluations: - judge: - provider: openai - name: Inferact/Qwen3.8-27B-NVFP4 - base_url: http://localhost:8000/v1 # optional, for OpenAI-compatible servers (vLLM, LM Studio, etc.) - temperature: 0.6 - max_tokens: 16384 - extra_body: - top_p: 0.95 - top_k: 20 - min_p: 0 - chat_template_kwargs: - reasoning_effort: low # qwen3.8: low | medium | xhigh (default) -``` - -### Restricting the corpus - -When a database holds documents from several corpora — only some of which a dataset's questions are drawn from — `--filter` restricts every benchmark search to a subset. It takes the same SQL `WHERE` clause as `haiku-rag search --filter`, over document columns (`id`, `uri`, `title`, `created_at`, `updated_at`, `metadata`). Each dataset writes its own URIs: `orb_text` uses bare arXiv ids such as `2407.01528v3`, `hotpotqa` uses page titles. - -```bash -evaluations run orb_text --skip-db --config haiku.rag.s3.yaml \ - --filter "uri LIKE '2407%'" -``` - -If the corpora are distinguished by a tag rather than by URI, attach it at ingest time as document metadata and match it with `LIKE`. `metadata` is stored as a `json.dumps` string, so there is no JSON subfield access — match the serialized key/value, including the space after the colon: - -```bash -evaluations run orb_text --skip-db --filter "metadata LIKE '%\"corpus\": \"orb_text\"%'" -``` - -The clause applies to both benchmark phases — the retrieval benchmark's searches and every search the capability runs during QA — so the two score the same subset. It is recorded as `document_filter` in the run's experiment metadata, so a filtered run is never mistaken for an unfiltered one when comparing results. - -Filtering affects searches only — a run without `--skip-db` still populates the database with the dataset's full corpus. - -## Methodology - -### Retrieval Metrics - -**Mean Average Precision (MAP)** scores ranked retrieval results against the gold `expected_uris`. - -- For each relevant document at position k, calculate precision@k = (relevant docs in top k) / k -- Average Precision (AP) = sum of these precision values / total relevant documents -- MAP is the mean of AP scores across all queries -- Range: 0 to 1. Rewards ranking relevant documents higher -- For single-doc queries this collapses to `1/rank` (i.e. reciprocal rank) - -### QA Accuracy - -`pydantic-evals` coordinates an LLM judge to determine whether the capability's answer is correct. The default judge is `ollama:qwen3.8`, pinned so changes to the capability model don't change the judge underneath. Set `evaluations.judge` in `haiku.rag.yaml` to override (including a custom `base_url` for any OpenAI-compatible endpoint). Accuracy is the fraction of correctly answered questions. - -A dataset that brings its own deterministic evaluator is scored by that evaluator instead, and no judge runs. T²-RAGBench is the only such dataset today, scored by `NumberMatchEvaluator`. - -`qwen3.8` replaced `qwen3.6` after a 120-case calibration on ORB, stratified 60 pass / 60 fail: agreement 0.950, Cohen's κ 0.900, and in all 6 disagreements it matched or beat `qwen3.6` (4 were `qwen3.6` failing answers that were equivalent in different notation). It emits no reasoning content, so it avoids the thinking spirals that made `qwen3.6` exceed its output budget and drop verdicts. `reasoning_effort` changes its verdicts in 1 case per 120, so the cheaper `low` is pinned. - -Before that, we picked `qwen3.6` over the previously-pinned `gpt-oss` after a 4-cell calibration (gpt-oss / qwen3.6 as both answerer and judge, with Claude Opus 4.7 as a reference). `qwen3.6` had κ ≥ 0.66 vs the reference on both same-family and cross-family answerers (vs ~0.39–0.55 for `gpt-oss`) and showed no measurable self-preference bias, while `gpt-oss` was ~10 pp more lenient on its own outputs. - -### Citation Retrieval - -Alongside QA accuracy, a second metric scores the URIs the capability registered via the `cite` tool against each dataset's gold `expected_uris`, using the same MAP math as raw retrieval. The score key is `cited_map`. Console output also includes the cite rate (% of cases with at least one citation) and the mean number of citations per case. - -This is computed alongside QA accuracy from the same capability run, no extra invocations. The signal complements raw retrieval: where raw retrieval measures whether the retriever surfaced the gold document at any rank, citation retrieval measures whether the capability grounded its answer on it. - ## Current results Numbers below were measured under `Qwen3.6-35B-A3B-NVFP4` as judge, on a recent `haiku.rag` version. The pinned judge is now `qwen3.8`; rows are not re-judged, so compare rows to each other rather than to runs judged by `qwen3.8`. @@ -262,3 +129,136 @@ The two live arms replay the same 29 conversations (224 turns) and differ only i - Answer pass rate: 185/224 vs 175/224 turns. Of the 18 turns where the arms disagree, 14 pass only compacted and 4 only uncompacted. McNemar exact two-sided p = 0.031. The paired difference is +4.5pp with a Wald 95% CI of +0.8 to +8.1pp, so the honest claim is an improvement of roughly 1 to 8 points, not the point estimate. - Citation MAP, macro-averaged over conversations with 208 of 224 turns eligible (turns with gold passages) in each arm: 0.4174 compacted vs 0.4230 uncompacted. The gold-prefix 0.35 is over 208 of 224 eligible cases. - Refusal precision and recall against the answerability labels (16 UNANSWERABLE turns per arm): compacted 0.33 precision and 0.44 recall (21 refusals), uncompacted 0.23 and 0.31 (22 refusals). Gold-prefix: 0.24 and 0.44 (29 refusals). + +## Methodology + +### Retrieval Metrics + +**Mean Average Precision (MAP)** scores ranked retrieval results against the gold `expected_uris`. + +- For each relevant document at position k, calculate precision@k = (relevant docs in top k) / k +- Average Precision (AP) = sum of these precision values / total relevant documents +- MAP is the mean of AP scores across all queries +- Range: 0 to 1. Rewards ranking relevant documents higher +- For single-doc queries this collapses to `1/rank` (i.e. reciprocal rank) + +### QA Accuracy + +`pydantic-evals` coordinates an LLM judge to determine whether the capability's answer is correct. The default judge is `ollama:qwen3.8`, pinned so changes to the capability model don't change the judge underneath. Set `evaluations.judge` in `haiku.rag.yaml` to override (including a custom `base_url` for any OpenAI-compatible endpoint). Accuracy is the fraction of correctly answered questions. + +A dataset that brings its own deterministic evaluator is scored by that evaluator instead, and no judge runs. T²-RAGBench is the only such dataset today, scored by `NumberMatchEvaluator`. + +`qwen3.8` replaced `qwen3.6` after a 120-case calibration on ORB, stratified 60 pass / 60 fail: agreement 0.950, Cohen's κ 0.900, and in all 6 disagreements it matched or beat `qwen3.6` (4 were `qwen3.6` failing answers that were equivalent in different notation). It emits no reasoning content, so it avoids the thinking spirals that made `qwen3.6` exceed its output budget and drop verdicts. `reasoning_effort` changes its verdicts in 1 case per 120, so the cheaper `low` is pinned. + +Before that, we picked `qwen3.6` over the previously-pinned `gpt-oss` after a 4-cell calibration (gpt-oss / qwen3.6 as both answerer and judge, with Claude Opus 4.7 as a reference). `qwen3.6` had κ ≥ 0.66 vs the reference on both same-family and cross-family answerers (vs ~0.39–0.55 for `gpt-oss`) and showed no measurable self-preference bias, while `gpt-oss` was ~10 pp more lenient on its own outputs. + +### Citation Retrieval + +Alongside QA accuracy, a second metric scores the URIs the capability registered via the `cite` tool against each dataset's gold `expected_uris`, using the same MAP math as raw retrieval. The score key is `cited_map`. Console output also includes the cite rate (% of cases with at least one citation) and the mean number of citations per case. + +This is computed alongside QA accuracy from the same capability run, no extra invocations. The signal complements raw retrieval: where raw retrieval measures whether the retriever surfaced the gold document at any rank, citation retrieval measures whether the capability grounded its answer on it. + +## Running Evaluations + +You can run evaluations with the `evaluations` CLI: + +```bash +evaluations run hotpotqa +evaluations run orb_text +``` + +The evaluation flow is orchestrated with [`pydantic-evals`](https://github.com/pydantic/pydantic-ai/tree/main/libs/pydantic-evals), which we leverage for dataset management, scoring, and report generation. + +### Pre-built Databases + +Building evaluation databases from scratch can take a long time, especially for large datasets like OpenRAG Bench. Pre-built databases are available on HuggingFace: + +```bash +# Download a specific dataset +evaluations download hotpotqa + +# Download all datasets +evaluations download all + +# Force re-download (overwrite existing) +evaluations download hotpotqa --force +``` + +Active datasets: + +| Dataset | Size | +|---------|------| +| `orb_text` — OpenRAG Bench, text embedder (`qwen3-embedding:4b`) with VLM picture descriptions baked into chunk content | ~18 GB | +| `orb_multimodal` — OpenRAG Bench, multimodal embedder (`qwen3-vl-embedding-8b`); picture vectors live in the same space as text for cross-modal retrieval | ~16 GB | +| `orb_multimodal_nemotron` — OpenRAG Bench, multimodal embedder (`nvidia/llama-nemotron-embed-vl-1b-v2`), the embedder behind the published headline results | ~16 GB | +| `t2_finqa` — T²-RAGBench (FinQA) financial QA, text embedder (`qwen3-embedding:4b`); scored by exact numeric match, run with `--target analysis-capability` | ~2 GB | +| `hotpotqa` — HotpotQA multi-hop QA over Wikipedia paragraphs, text embedder (`qwen3-embedding:4b`) | ~1.5 GB | +| `mtrag_clapnq` — MTRAG multi-turn RAG, ClapNQ (Wikipedia) passages, text embedder (`qwen3-embedding:4b`); also serves the `mtrag_clapnq_rewrite`, `mtrag_clapnq_live` and `mtrag_clapnq_live_uncompacted` keys | ~2.8 GB | + +After downloading, run benchmarks with `--skip-db`. Each database is built with a specific embedder, so pass its reference config from `evaluations/configs/` (a database only opens against a config whose embedder matches): + +```bash +evaluations run orb_multimodal_nemotron --skip-db --config configs/orb_multimodal_nemotron.yaml +``` + +The configs use `vllm` as the model host. Point `base_url` at your own OpenAI-compatible endpoints to reproduce the numbers. + +### Configuration + +The benchmark script accepts several options: + +```bash +evaluations run hotpotqa --config /path/to/haiku.rag.yaml --db /path/to/custom.lancedb +``` + +**Options:** + +- `--config PATH` - Specify a custom `haiku.rag.yaml` configuration file +- `--db PATH` - Override the database path (default: platform-specific user data directory) +- `--skip-db` - Skip updating the evaluation database +- `--skip-retrieval` - Skip retrieval benchmark +- `--skip-qa` - Skip QA benchmark +- `--limit N` - Limit number of test cases +- `--name NAME` - Override the evaluation name +- `--target {rag-capability,analysis-capability}` - Choose which [capability](capabilities/index.md) to benchmark end-to-end (default: `rag-capability`). The target names remain stable dataset identifiers. +- `--capability-model PROVIDER:NAME` - Override the capability model independently from the judge (default: `config.qa.model`, or `config.analysis.model` when set for `--target analysis-capability`). +- `--filter CLAUSE` / `-f CLAUSE` - Restrict every benchmark search to a subset of the database (see [Restricting the corpus](#restricting-the-corpus)). + +If no config file is specified, the script searches standard locations: `./haiku.rag.yaml`, user config directory, then falls back to defaults. + +To pin the LLM judge in YAML (rather than the default `ollama:qwen3.8`). These are the recommended settings: + +```yaml +evaluations: + judge: + provider: openai + name: Inferact/Qwen3.8-27B-NVFP4 + base_url: http://localhost:8000/v1 # optional, for OpenAI-compatible servers (vLLM, LM Studio, etc.) + temperature: 0.6 + max_tokens: 16384 + extra_body: + top_p: 0.95 + top_k: 20 + min_p: 0 + chat_template_kwargs: + reasoning_effort: low # qwen3.8: low | medium | xhigh (default) +``` + +### Restricting the corpus + +When a database holds documents from several corpora — only some of which a dataset's questions are drawn from — `--filter` restricts every benchmark search to a subset. It takes the same SQL `WHERE` clause as `haiku-rag search --filter`, over document columns (`id`, `uri`, `title`, `created_at`, `updated_at`, `metadata`). Each dataset writes its own URIs: `orb_text` uses bare arXiv ids such as `2407.01528v3`, `hotpotqa` uses page titles. + +```bash +evaluations run orb_text --skip-db --config haiku.rag.s3.yaml \ + --filter "uri LIKE '2407%'" +``` + +If the corpora are distinguished by a tag rather than by URI, attach it at ingest time as document metadata and match it with `LIKE`. `metadata` is stored as a `json.dumps` string, so there is no JSON subfield access — match the serialized key/value, including the space after the colon: + +```bash +evaluations run orb_text --skip-db --filter "metadata LIKE '%\"corpus\": \"orb_text\"%'" +``` + +The clause applies to both benchmark phases — the retrieval benchmark's searches and every search the capability runs during QA — so the two score the same subset. It is recorded as `document_filter` in the run's experiment metadata, so a filtered run is never mistaken for an unfiltered one when comparing results. + +Filtering affects searches only — a run without `--skip-db` still populates the database with the dataset's full corpus. diff --git a/zensical.toml b/zensical.toml index 0fa0fe4a..033af248 100644 --- a/zensical.toml +++ b/zensical.toml @@ -46,8 +46,8 @@ nav = [ { Toolsets = "tools.md" }, { "Web app" = "apps.md" }, ] }, + { Benchmarks = "benchmarks.md" }, { Reference = [ - { Benchmarks = "benchmarks.md" }, { Development = "development.md" }, { Changelog = "changelog.md" }, ] }, From 6f976ef2a980853e8812a4cef9f4c387ce1bbe5b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 12:59:09 +0300 Subject: [PATCH 08/11] Share the LanceDB session across connections Every `Store` built its own connection with its own caches and discarded them on close, so the index a vector query loads was refetched by the next connection. On object storage that first fetch dominates: measured on a ~500k-chunk 2560-dim corpus over a ~200ms link, the first query cost ~41s and the second ~3s, and a new connection reusing the session cost ~7s instead of ~47s. `connect_lancedb` now passes a process-wide session, keyed on the configured cache sizes so a caller asking for different sizes gets its own. Also sets `read_consistency_interval`, defaulting to 30s. It was None, meaning a connection never re-checked for other processes' writes. Per-call connections hid that; a shared session makes connections long-lived enough for a reader to go stale against the ingester. All three settings reject negatives at the config boundary. A negative cache size raises OverflowError and a negative interval panics inside Lance, so neither is catchable further in. Zero stays valid for both: no cache, and check on every read. The routing tests now assert the kwargs they care about rather than the full call signature, since every connection carries the two new kwargs. --- CHANGELOG.md | 2 +- docs/configuration/storage.md | 12 ++ haiku_rag_slim/haiku/rag/config/models.py | 11 ++ haiku_rag_slim/haiku/rag/store/engine.py | 37 +++++- tests/test_lancedb_connection.py | 151 ++++++++++++++++++++-- 5 files changed, 196 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bbfe8308..c63d1446 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,8 +5,8 @@ - `evaluations run --filter/-f CLAUSE`: SQL `WHERE` clause over document columns, applied to the retrieval benchmark's searches and to every capability search during QA. Recorded as `document_filter` in experiment metadata. - `mtrag_clapnq` / `mtrag_clapnq_rewrite` / `mtrag_clapnq_live` / `mtrag_clapnq_live_uncompacted` evaluation datasets: multi-turn QA with gold-prefix and live-session conversation replay, live arms with and without `EvidenceCompactionCapability`, Recall@k/nDCG@k retrieval metrics, eligibility-aware citation scoring, refusal precision/recall, per-turn tool-traffic attributes, and `citation_status` / `turn_citation_status` eval attributes. - - BTree indexes on `chunks.id`, `chunks.document_id` and `documents.id`, and a Bitmap index on `document_items.label`. Existing databases need `haiku-rag migrate`. +- `lancedb.read_consistency_interval_seconds` (default 30), `lancedb.index_cache_size_bytes` and `lancedb.metadata_cache_size_bytes`. The LanceDB session is shared across connections in a process, so its index and metadata caches survive a connection being closed. ### Changed diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index 4d86b215..57ab3aa4 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -122,6 +122,18 @@ The `storage_options` keys are case-insensitive and passed directly to the under **Note:** Table optimization is automatically handled by LanceDB Cloud (`db://` URIs) and is disabled for better performance. For object storage backends (S3, Azure, GCS), optimization and vector indexing are still performed normally. +### Caching and Read Consistency + +```yaml +lancedb: + read_consistency_interval_seconds: 30 # null to never re-check + index_cache_size_bytes: 536870912 # null for the LanceDB default + metadata_cache_size_bytes: 268435456 +``` + +- **read_consistency_interval_seconds**: how often a connection checks for writes from another process. `null` never checks, so a long-lived reader never sees the ingester's writes. `0` checks on every read. +- **index_cache_size_bytes** / **metadata_cache_size_bytes**: sizes for the caches held by the LanceDB session, which is shared across every connection in the process. The first vector query loads the index into it, so on object storage the cache is what stops the next connection refetching it. Size it for the total set of indexes a process keeps warm, against the memory available to it. + ### Deployment Pattern: One Writer, Many Readers LanceDB on S3 supports **exactly one writer + N readers** per database URI. Multiple writers against the same URI can race on the manifest commit and corrupt state. This is a LanceDB property, not something `haiku.rag` enforces. diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 0bfbfec7..2fa58a8e 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -62,10 +62,21 @@ class StorageConfig(BaseModel): class LanceDBConfig(BaseModel): + """LanceDB connection settings. + + read_consistency_interval_seconds bounds how stale a reader may be. None + never re-checks, so a long-lived reader never sees another process's writes. + The cache sizes are per process, since the session is shared across + connections. + """ + uri: str = "" api_key: str = "" region: str = "" storage_options: dict[str, str] = Field(default_factory=dict) + read_consistency_interval_seconds: float | None = Field(default=30, ge=0) + index_cache_size_bytes: int | None = Field(default=None, ge=0) + metadata_cache_size_bytes: int | None = Field(default=None, ge=0) class EmbeddingsConfig(BaseModel): diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 002283a1..f47a5d82 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -55,25 +55,56 @@ class ConnectionMode(Enum): return ConnectionMode.OBJECT_STORAGE +_sessions: dict[tuple[int | None, int | None], lancedb.Session] = {} + + +def _session(config: AppConfig) -> lancedb.Session: + """The process's session for these cache sizes. + + Sessions hold the index and metadata caches. Sharing one across connections + is what keeps a cached index from being refetched per connection, which on + object storage is the dominant cost of the first query. + """ + key = ( + config.lancedb.index_cache_size_bytes, + config.lancedb.metadata_cache_size_bytes, + ) + if key not in _sessions: + kwargs = {} + if key[0] is not None: + kwargs["index_cache_size_bytes"] = key[0] + if key[1] is not None: + kwargs["metadata_cache_size_bytes"] = key[1] + _sessions[key] = lancedb.Session(**kwargs) + return _sessions[key] + + async def connect_lancedb( config: AppConfig, db_path: Path | None = None ) -> lancedb.AsyncConnection: + interval = config.lancedb.read_consistency_interval_seconds + kwargs: dict[str, Any] = { + "session": _session(config), + "read_consistency_interval": ( + timedelta(seconds=interval) if interval is not None else None + ), + } mode = ConnectionMode.from_config(config) if mode == ConnectionMode.CLOUD: return await lancedb.connect_async( uri=config.lancedb.uri, api_key=config.lancedb.api_key, region=config.lancedb.region, + **kwargs, ) elif mode == ConnectionMode.OBJECT_STORAGE: - kwargs: dict[str, Any] = {"uri": config.lancedb.uri} if config.lancedb.storage_options: kwargs["storage_options"] = config.lancedb.storage_options - return await lancedb.connect_async(**kwargs) + return await lancedb.connect_async(uri=config.lancedb.uri, **kwargs) else: if db_path is None: raise ValueError("No lancedb.uri configured and no db_path provided") - return await lancedb.connect_async(db_path.absolute()) + return await lancedb.connect_async(db_path.absolute(), **kwargs) class DocumentRecord(LanceModel): diff --git a/tests/test_lancedb_connection.py b/tests/test_lancedb_connection.py index 8e624250..33fa349c 100644 --- a/tests/test_lancedb_connection.py +++ b/tests/test_lancedb_connection.py @@ -1,6 +1,8 @@ +from datetime import timedelta from unittest.mock import AsyncMock, patch import pytest +from pydantic import ValidationError from haiku.rag.config import Config from haiku.rag.config.models import AppConfig, LanceDBConfig @@ -49,7 +51,8 @@ class TestConnectLancedb: "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: await connect_lancedb(config, db_path=temp_db_path) - mock_connect.assert_called_once_with(temp_db_path.absolute()) + mock_connect.assert_awaited_once() + assert mock_connect.call_args.args == (temp_db_path.absolute(),) @pytest.mark.asyncio async def test_local_resolves_relative_db_path(self, tmp_path, monkeypatch): @@ -62,7 +65,8 @@ class TestConnectLancedb: "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: await connect_lancedb(config, db_path=relative) - mock_connect.assert_called_once_with(relative.absolute()) + mock_connect.assert_awaited_once() + assert mock_connect.call_args.args == (relative.absolute(),) @pytest.mark.asyncio async def test_cloud_passes_uri_api_key_region(self): @@ -75,9 +79,11 @@ class TestConnectLancedb: "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: await connect_lancedb(config) - mock_connect.assert_called_once_with( - uri="db://my-database", api_key="test-key", region="us-west-2" - ) + mock_connect.assert_awaited_once() + kwargs = mock_connect.call_args.kwargs + assert kwargs["uri"] == "db://my-database" + assert kwargs["api_key"] == "test-key" + assert kwargs["region"] == "us-west-2" @pytest.mark.asyncio async def test_object_storage_passes_uri_and_storage_options(self): @@ -94,13 +100,13 @@ class TestConnectLancedb: "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: await connect_lancedb(config) - mock_connect.assert_called_once_with( - uri="s3://bucket/path", - storage_options={ - "endpoint": "http://minio:9000", - "region": "us-east-1", - }, - ) + mock_connect.assert_awaited_once() + kwargs = mock_connect.call_args.kwargs + assert kwargs["uri"] == "s3://bucket/path" + assert kwargs["storage_options"] == { + "endpoint": "http://minio:9000", + "region": "us-east-1", + } @pytest.mark.asyncio async def test_object_storage_without_storage_options(self): @@ -109,7 +115,10 @@ class TestConnectLancedb: "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: await connect_lancedb(config) - mock_connect.assert_called_once_with(uri="s3://bucket/path") + mock_connect.assert_awaited_once() + kwargs = mock_connect.call_args.kwargs + assert kwargs["uri"] == "s3://bucket/path" + assert "storage_options" not in kwargs @pytest.mark.asyncio async def test_local_without_db_path_raises(self): @@ -398,3 +407,119 @@ class TestStoreMiscellany: async with Store(temp_db_path, create=True) as store: with pytest.raises(ValueError, match="Unknown table"): await store.list_table_versions("not_a_table") + + +class TestSessionAndConsistency: + @pytest.mark.asyncio + async def test_session_is_shared_across_connections(self): + config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path")) + with patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ) as mock_connect: + await connect_lancedb(config) + await connect_lancedb(config) + + sessions = [c.kwargs["session"] for c in mock_connect.call_args_list] + assert sessions[0] is sessions[1] + + @pytest.mark.asyncio + async def test_cache_sizes_select_distinct_sessions(self): + small = AppConfig( + lancedb=LanceDBConfig( + uri="s3://bucket/path", index_cache_size_bytes=1 << 20 + ) + ) + large = AppConfig( + lancedb=LanceDBConfig( + uri="s3://bucket/path", index_cache_size_bytes=1 << 30 + ) + ) + with patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ) as mock_connect: + await connect_lancedb(small) + await connect_lancedb(large) + + sessions = [c.kwargs["session"] for c in mock_connect.call_args_list] + assert sessions[0] is not sessions[1] + + @pytest.mark.asyncio + async def test_both_cache_sizes_are_applied(self): + config = AppConfig( + lancedb=LanceDBConfig( + uri="s3://bucket/path", + index_cache_size_bytes=2 << 20, + metadata_cache_size_bytes=4 << 20, + ) + ) + with ( + patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ), + patch("haiku.rag.store.engine.lancedb.Session") as mock_session, + ): + await connect_lancedb(config) + + mock_session.assert_called_once_with( + index_cache_size_bytes=2 << 20, metadata_cache_size_bytes=4 << 20 + ) + + @pytest.mark.asyncio + async def test_read_consistency_interval_is_forwarded(self): + config = AppConfig( + lancedb=LanceDBConfig( + uri="s3://bucket/path", read_consistency_interval_seconds=5 + ) + ) + with patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ) as mock_connect: + await connect_lancedb(config) + + assert mock_connect.call_args.kwargs["read_consistency_interval"] == timedelta( + seconds=5 + ) + + @pytest.mark.asyncio + async def test_read_consistency_interval_omitted_when_disabled(self): + config = AppConfig( + lancedb=LanceDBConfig( + uri="s3://bucket/path", read_consistency_interval_seconds=None + ) + ) + with patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ) as mock_connect: + await connect_lancedb(config) + + assert mock_connect.call_args.kwargs["read_consistency_interval"] is None + + @pytest.mark.asyncio + async def test_local_connection_also_gets_session_and_consistency(self, tmp_path): + config = AppConfig() + with patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ) as mock_connect: + await connect_lancedb(config, tmp_path / "db.lancedb") + + assert mock_connect.call_args.kwargs["session"] is not None + assert mock_connect.call_args.kwargs["read_consistency_interval"] == timedelta( + seconds=30 + ) + + +class TestLanceDBConfigValidation: + def test_negative_values_are_rejected(self): + """Negatives overflow or panic inside Lance, so reject them here.""" + with pytest.raises(ValidationError): + LanceDBConfig(read_consistency_interval_seconds=-1) + with pytest.raises(ValidationError): + LanceDBConfig(index_cache_size_bytes=-1) + with pytest.raises(ValidationError): + LanceDBConfig(metadata_cache_size_bytes=-1) + + def test_zero_is_allowed(self): + config = LanceDBConfig( + read_consistency_interval_seconds=0, index_cache_size_bytes=0 + ) + assert config.read_consistency_interval_seconds == 0 From da207da106593dfa8ef33e682fa2885ac882abb1 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 13:25:19 +0300 Subject: [PATCH 09/11] Read the table list and settings once per open Opening a database ran `list_tables` three times, opened the settings table twice, and read and parsed the same settings row three times: once for the stored vector dimension, once for the version behind the migration check, and once for config validation. On object storage each of those is a round trip. `_initialize` now reads both once and threads them down. `_init_tables` and `_check_migrations` take what it read instead of fetching their own copy, and `validate_config_compatibility` accepts the settings it should compare against, still reading for itself when called directly. Passing the pre-init read to validation is equivalent: nothing between the read and the validation rewrites `embeddings`, which is all it compares. The settings read no longer swallows every exception. It did before, when the only consequence was falling back to the configured vector dimension; now the same empty result feeds the migration check, where it would read as version 0.0.0 and declare every migration pending. Only decode failures are tolerated, and a decoded non-object normalizes to {} rather than reaching callers that expect a mapping. --- haiku_rag_slim/haiku/rag/store/engine.py | 99 +++++++++---------- .../haiku/rag/store/repositories/settings.py | 7 +- tests/store/test_open_path.py | 71 +++++++++++++ tests/test_lancedb_connection.py | 4 +- 4 files changed, 125 insertions(+), 56 deletions(-) create mode 100644 tests/store/test_open_path.py diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index f47a5d82..08867fa6 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -207,6 +207,11 @@ def get_document_items_arrow_schema() -> pa.Schema: return pa.schema(fields) +def _stored_vector_dim(settings: dict) -> int | None: + """The vector dimension a database's chunks were written at.""" + return settings.get("embeddings", {}).get("model", {}).get("vector_dim") + + def index_specs(table_name: str) -> list[tuple[str, Bitmap | BTree | FTS]]: """The index set each table carries.""" match table_name: @@ -533,29 +538,27 @@ class Store: self._config, self.db_path ) - # For remote stores (and as a safety net for local paths that exist but - # have no tables — e.g. a previously failed init), detect new DB by - # checking whether any tables exist. - is_new_db = self._is_new_db - if not is_new_db: - existing_tables = (await self.db.list_tables()).tables - if not existing_tables: - is_new_db = True + # Read once and thread onward: on object storage each of these is a + # round trip. A local path that exists with no tables is a failed init, + # so treat it as new. + existing_tables = (await self.db.list_tables()).tables + is_new_db = self._is_new_db or not existing_tables - # For existing databases, read stored vector dimension to create ChunkRecord - # that can read existing chunks. For new databases, use config's dimension. - stored_vector_dim = None - if not is_new_db: - stored_vector_dim = await self._get_stored_vector_dim() + stored_settings: dict = {} + if not is_new_db and "settings" in existing_tables: + self.settings_table = await self.db.open_table("settings") + stored_settings = await self._read_stored_settings() - # Create ChunkRecord with stored dimension (for reading) or config dimension (for new DB) + # An existing database's chunks can only be read with the dimension they + # were written at. + stored_vector_dim = _stored_vector_dim(stored_settings) chunk_vector_dim = stored_vector_dim or self.embedder._vector_dim self.ChunkRecord: type[ChunkRecordBase] = create_chunk_model(chunk_vector_dim) # Initialize tables (creates them if they don't exist). For an existing # DB this raises MigrationRequiredError up front when migrations are # pending, before creating any newly-introduced table. - await self._init_tables(is_new_db) + await self._init_tables(is_new_db, existing_tables, stored_settings) # Set version for new databases. if is_new_db and not self._read_only: @@ -563,7 +566,7 @@ class Store: # Validate config compatibility after connection is established if not self._skip_validation: - await self._validate_configuration() + await self._validate_configuration(stored_settings) async def __aenter__(self): # If _initialize connects to LanceDB but then fails (e.g. migration @@ -585,33 +588,26 @@ class Store: """Whether the store is in read-only mode.""" return self._read_only - async def _get_stored_vector_dim(self) -> int | None: - """Read the stored vector dimension from the settings table. + async def _read_stored_settings(self) -> dict: + """The stored settings blob, or {} if it is absent or not a JSON object. - Returns: - The stored vector dimension, or None if not found. + Only decoding failures are tolerated. A storage failure must propagate: + read as empty settings it would look like version 0.0.0, and the + migration check would declare every migration pending. """ + rows = ( + await self.settings_table.query() + .where("id = 'settings'") + .limit(1) + .to_arrow() + ).to_pylist() + if not rows or not rows[0].get("settings"): + return {} try: - existing_tables = (await self.db.list_tables()).tables - if "settings" not in existing_tables: - return None - - settings_table = await self.db.open_table("settings") - rows = ( - await settings_table.query() - .where("id = 'settings'") - .limit(1) - .to_arrow() - ).to_pylist() - if not rows or not rows[0].get("settings"): - return None - - settings = json.loads(rows[0]["settings"]) - embeddings = settings.get("embeddings", {}) - model = embeddings.get("model", {}) - return model.get("vector_dim") - except Exception: - return None + decoded = json.loads(rows[0]["settings"]) + except (json.JSONDecodeError, TypeError): + return {} + return decoded if isinstance(decoded, dict) else {} def _assert_writable(self) -> None: """Raise ReadOnlyError if the store is in read-only mode.""" @@ -742,16 +738,19 @@ class Store: except Exception as e: logger.warning(f"Could not create vector index: {e}") - async def _validate_configuration(self) -> None: + async def _validate_configuration( + self, stored_settings: dict | None = None + ) -> None: """Validate that the configuration is compatible with the database.""" from haiku.rag.store.repositories.settings import SettingsRepository settings_repo = SettingsRepository(self) - await settings_repo.validate_config_compatibility() + await settings_repo.validate_config_compatibility(stored_settings) - async def _init_tables(self, is_new_db: bool): + async def _init_tables( + self, is_new_db: bool, existing_tables: list[str], stored_settings: dict + ): """Initialize database tables (create if they don't exist).""" - existing_tables = (await self.db.list_tables()).tables # Surface pending migrations BEFORE creating any newly-introduced table. # Otherwise opening a legacy DB would either mutate it (creating an empty @@ -763,8 +762,7 @@ class Store: and not self._skip_migration_check and "settings" in existing_tables ): - self.settings_table = await self.db.open_table("settings") - await self._check_migrations() + await self._check_migrations(stored_settings.get("version", "0.0.0")) missing_tables = set(REQUIRED_TABLES) - set(existing_tables) @@ -811,10 +809,8 @@ class Store: ) await ensure_indexes(self.document_items_table, "document_items") - # Create or open settings table - if "settings" in existing_tables: - self.settings_table = await self.db.open_table("settings") - else: + # _initialize opened the settings table when the database had one. + if "settings" not in existing_tables: self.settings_table = await self.db.create_table( "settings", schema=SettingsRecord ) @@ -828,7 +824,7 @@ class Store: """Set the initial version for a new database.""" await self.set_haiku_version(metadata.version("haiku.rag-slim")) - async def _check_migrations(self) -> None: + async def _check_migrations(self, db_version: str) -> None: """Raise if migrations are pending. Opening never writes the version. Raises: @@ -837,7 +833,6 @@ class Store: from haiku.rag.store.upgrades import get_pending_upgrades current_version = metadata.version("haiku.rag-slim") - db_version = await self.get_haiku_version() pending = get_pending_upgrades(db_version) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/settings.py b/haiku_rag_slim/haiku/rag/store/repositories/settings.py index d7b2986c..e38c17cf 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/settings.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/settings.py @@ -60,7 +60,9 @@ class SettingsRepository: ) await self.store.settings_table.add([settings_record]) - async def validate_config_compatibility(self) -> None: + async def validate_config_compatibility( + self, stored_settings: dict | None = None + ) -> None: """Validate the current configuration against stored settings without writing. Opening a database never modifies it. ``vector_dim`` mismatches raise — @@ -72,7 +74,8 @@ class SettingsRepository: while a read-only open continues. Stored settings are reconciled explicitly via ``haiku-rag rebuild --set-embedder``, never on open. """ - stored_settings = await self.get_current_settings() + if stored_settings is None: + stored_settings = await self.get_current_settings() # Nothing stored to validate against — never write on open. if not stored_settings: diff --git a/tests/store/test_open_path.py b/tests/store/test_open_path.py new file mode 100644 index 00000000..887db575 --- /dev/null +++ b/tests/store/test_open_path.py @@ -0,0 +1,71 @@ +import lancedb +import pytest + +from haiku.rag.store.engine import Store + + +@pytest.fixture +def counts(monkeypatch): + """Count the connection-level calls an open makes.""" + tally: dict[str, int] = {"list_tables": 0, "open_settings": 0, "settings_query": 0} + + list_tables = lancedb.AsyncConnection.list_tables + open_table = lancedb.AsyncConnection.open_table + query = lancedb.AsyncTable.query + + async def counted_list_tables(self, *args, **kwargs): + tally["list_tables"] += 1 + return await list_tables(self, *args, **kwargs) + + async def counted_open_table(self, name, *args, **kwargs): + if name == "settings": + tally["open_settings"] += 1 + return await open_table(self, name, *args, **kwargs) + + def counted_query(self): + if self.name == "settings": + tally["settings_query"] += 1 + return query(self) + + monkeypatch.setattr(lancedb.AsyncConnection, "list_tables", counted_list_tables) + monkeypatch.setattr(lancedb.AsyncConnection, "open_table", counted_open_table) + monkeypatch.setattr(lancedb.AsyncTable, "query", counted_query) + return tally + + +@pytest.mark.asyncio +async def test_reopening_reads_the_table_list_and_settings_once(temp_db_path, counts): + async with Store(temp_db_path, create=True): + pass + for key in counts: + counts[key] = 0 + + async with Store(temp_db_path): + pass + + assert counts["list_tables"] == 1 + assert counts["open_settings"] == 1 + assert counts["settings_query"] == 1 + + +@pytest.mark.asyncio +async def test_storage_failures_propagate(temp_db_path): + """A read failure must not read as empty settings: the migration check would + then see version 0.0.0 and declare every migration pending.""" + async with Store(temp_db_path, create=True) as store: + + def boom(): + raise RuntimeError("s3 is having a day") + + store.settings_table.query = boom + + with pytest.raises(RuntimeError, match="s3 is having a day"): + await store._read_stored_settings() + + +@pytest.mark.asyncio +async def test_non_dict_settings_read_as_empty(temp_db_path): + async with Store(temp_db_path, create=True) as store: + await store.settings_table.update({"settings": "[]"}, where="id = 'settings'") + + assert await store._read_stored_settings() == {} diff --git a/tests/test_lancedb_connection.py b/tests/test_lancedb_connection.py index 33fa349c..ed65a5f9 100644 --- a/tests/test_lancedb_connection.py +++ b/tests/test_lancedb_connection.py @@ -276,7 +276,7 @@ class TestInitFailureCleanup: async def fake_connect(*args, **kwargs): return mock_conn - async def failing_init_tables(self, is_new_db): + async def failing_init_tables(self, *args): raise RuntimeError("simulated table init failure") monkeypatch.setattr("haiku.rag.store.engine.connect_lancedb", fake_connect) @@ -390,7 +390,7 @@ class TestStoreMiscellany: {"settings": "not json at all"}, where="id = 'settings'" ) - assert await store._get_stored_vector_dim() is None + assert await store._read_stored_settings() == {} @pytest.mark.asyncio async def test_vacuum_skips_when_already_running(self, temp_db_path): From 8a4d488a726dbdbaa837e04292bffebefedd9119 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 14:20:45 +0300 Subject: [PATCH 10/11] Give the MCP server one client for its lifetime All ten tool bodies opened their own `HaikuRAG`, so every tool call paid a connection open and, on object storage, refetched the index the previous call had just cached. The client is now opened once, lazily so that calling a tool function directly still works, and eagerly from the lifespan so an unopenable database fails startup instead of every call. Teardown clears the cached client in a finally, since `_lifespan_manager` can be re-entered and would otherwise hand out a closed connection, including when the close itself fails. `delete_document` no longer opens its own connection with `skip_validation=True`. Keeping it separate broke consistency once connections became long-lived: the delete committed on one connection while reads served from another, which with a 30s consistency interval showed the deleted document as still present. A connection always sees its own writes, so sharing one is what makes delete visible to the next read. So the server no longer opts out of embedding-config validation. Drift that validation rejects now fails MCP startup, where before the server started and only `delete_document` worked while every read returned empty. Same-dimension identity drift still starts a read-only server, matching every other read verb. Delete under drift is now a CLI operation; CLAUDE.md and the CHANGELOG record it. --- CHANGELOG.md | 1 + haiku_rag_slim/haiku/rag/mcp.py | 142 +++++++++++++++++++------------ tests/test_mcp.py | 144 ++++++++++++++++++++++++++++++++ 3 files changed, 232 insertions(+), 55 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c63d1446..310f5ecd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ ### Changed - Eval judge pinned to `qwen3.8`: `DEFAULT_JUDGE_MODEL` is `ollama:qwen3.8`, and the reference configs use `Inferact/Qwen3.8-27B-NVFP4` with `extra_body.chat_template_kwargs.reasoning_effort: low`. Results in `docs/benchmarks.md` were judged by `Qwen3.6-35B-A3B-NVFP4` and are not re-judged. +- The MCP server opens one database client for its lifetime instead of one per tool call, and fails at startup if the database cannot be opened. `delete_document` no longer opens its own connection with `skip_validation=True`, so the server no longer opts out of embedding-config validation: drift that validation rejects (any `vector_dim` mismatch, or identity drift on a writable server) now fails MCP startup instead of serving a delete-only server. Use the CLI to delete under drift. - `import_documents` embeds chunks across the whole batch in one pass instead of per document. - `haiku.rag` and `haiku.rag-slim` summaries and keywords; both packages now publish `[project.urls]`. - `server.json` declares `title` and `websiteUrl`, and drops the `keywords` and `license` keys, which are not in the server schema. diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 30c2ed11..a8a44b50 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -1,3 +1,6 @@ +import asyncio +from collections.abc import AsyncIterator +from contextlib import AsyncExitStack, asynccontextmanager from pathlib import Path from typing import Any @@ -28,7 +31,42 @@ def create_mcp_server( config: Configuration to use. read_only: If True, write tools (add_document_*, delete_document) are not registered. """ - mcp = FastMCP("haiku-rag") + client: HaikuRAG | None = None + stack = AsyncExitStack() + client_lock = asyncio.Lock() + + async def _client() -> HaikuRAG: + """The server's client, opened once. + + Opening cost is per connection, and on object storage the first vector + query loads the index into the session cache, so a client per tool call + pays that repeatedly. + """ + nonlocal client + async with client_lock: + if client is None: + client = await stack.enter_async_context( + HaikuRAG(db_path, config=config, read_only=read_only) + ) + return client + + @asynccontextmanager + async def lifespan(_server: FastMCP) -> AsyncIterator[None]: + # Open eagerly so an unopenable database fails startup rather than + # every tool call. + nonlocal client + await _client() + try: + yield + finally: + # The lifespan can be re-entered; without the reset the next cycle + # hands out the closed client, including when aclose itself fails. + try: + await stack.aclose() + finally: + client = None + + mcp = FastMCP("haiku-rag", lifespan=lifespan) # Write tools - only registered when not in read-only mode if not read_only: @@ -41,14 +79,14 @@ def create_mcp_server( ) -> str | None: """Add a document to the RAG system from a file path.""" try: - async with HaikuRAG(db_path, config=config) as rag: - result = await rag.create_document_from_source( - Path(file_path), title=title, metadata=metadata or {} - ) - # Handle both single document and list of documents (directories) - if isinstance(result, list): - return result[0].id if result else None - return result.id + rag = await _client() + result = await rag.create_document_from_source( + Path(file_path), title=title, metadata=metadata or {} + ) + # Handle both single document and list of documents (directories) + if isinstance(result, list): + return result[0].id if result else None + return result.id except Exception: return None @@ -58,14 +96,14 @@ def create_mcp_server( ) -> str | None: """Add a document to the RAG system from a URL.""" try: - async with HaikuRAG(db_path, config=config) as rag: - result = await rag.create_document_from_source( - url, title=title, metadata=metadata or {} - ) - # Handle both single document and list of documents - if isinstance(result, list): - return result[0].id if result else None - return result.id + rag = await _client() + result = await rag.create_document_from_source( + url, title=title, metadata=metadata or {} + ) + # Handle both single document and list of documents + if isinstance(result, list): + return result[0].id if result else None + return result.id except Exception: return None @@ -78,11 +116,11 @@ def create_mcp_server( ) -> str | None: """Add a document to the RAG system from text content.""" try: - async with HaikuRAG(db_path, config=config) as rag: - document = await rag.create_document( - content, uri, title=title, metadata=metadata or {} - ) - return document.id + rag = await _client() + document = await rag.create_document( + content, uri, title=title, metadata=metadata or {} + ) + return document.id except Exception: return None @@ -90,10 +128,8 @@ def create_mcp_server( async def delete_document(document_id: str) -> bool: """Delete a document by its ID.""" try: - async with HaikuRAG( - db_path, config=config, skip_validation=True - ) as rag: - return await rag.delete_document(document_id) + rag = await _client() + return await rag.delete_document(document_id) except Exception: return False @@ -110,10 +146,8 @@ def create_mcp_server( response (smaller JSON payload for plain-text consumers). """ try: - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - return await rag.search( - query, limit=limit, include_images=include_images - ) + rag = await _client() + return await rag.search(query, limit=limit, include_images=include_images) except Exception: return [] @@ -145,10 +179,8 @@ def create_mcp_server( except Exception: return [] try: - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - return await rag.search( - raw, limit=limit, include_images=include_images - ) + rag = await _client() + return await rag.search(raw, limit=limit, include_images=include_images) except Exception: return [] @@ -156,8 +188,8 @@ def create_mcp_server( async def get_document(document_id: str) -> Document | None: """Get a document by its ID.""" try: - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - return await rag.get_document_by_id(document_id) + rag = await _client() + return await rag.get_document_by_id(document_id) except Exception: return None @@ -175,18 +207,18 @@ def create_mcp_server( filter: Optional SQL WHERE clause to filter documents. """ try: - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - documents = await rag.list_documents(limit, offset, filter) + rag = await _client() + documents = await rag.list_documents(limit, offset, filter) - return [ - DocumentInfo( - id=doc.id, - title=doc.title or "Untitled", - uri=doc.uri or "", - created=doc.created_at.strftime("%Y-%m-%d"), - ) - for doc in documents - ] + return [ + DocumentInfo( + id=doc.id, + title=doc.title or "Untitled", + uri=doc.uri or "", + created=doc.created_at.strftime("%Y-%m-%d"), + ) + for doc in documents + ] except Exception: return [] @@ -209,11 +241,11 @@ def create_mcp_server( """ try: images = _decode_images(images_base64) - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - answer, citations = await rag.ask(question, images=images) - if cite and citations: - answer += "\n\n" + format_citations(citations) - return answer + rag = await _client() + answer, citations = await rag.ask(question, images=images) + if cite and citations: + answer += "\n\n" + format_citations(citations) + return answer except Exception as e: return f"Error answering question: {e!s}" @@ -240,9 +272,9 @@ def create_mcp_server( """ try: images = _decode_images(images_base64) - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - result = await rag.analyze(question, filter=filter, images=images) - return result.answer + rag = await _client() + result = await rag.analyze(question, filter=filter, images=images) + return result.answer except Exception as e: return f"Error running analysis capability: {e!s}" diff --git a/tests/test_mcp.py b/tests/test_mcp.py index f0c36474..9077586e 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -481,3 +481,147 @@ class TestMCPToolsDegradeOnError: assert "AI Overview" in with_cite assert await ask(question="q", cite=False) == "the answer" + + +class TestMCPClientLifetime: + @pytest.mark.asyncio + async def test_tool_calls_share_one_database_open(self, mcp_db, monkeypatch): + from haiku.rag.store.engine import Store + + opens = 0 + initialize = Store._initialize + + async def counted(self): + nonlocal opens + opens += 1 + return await initialize(self) + + monkeypatch.setattr(Store, "_initialize", counted) + + mcp = create_mcp_server(mcp_db, read_only=True) + search = await _get_tool(mcp, "search_documents") + list_docs = await _get_tool(mcp, "list_documents") + await search(query="artificial intelligence") + await list_docs() + await search(query="machine learning") + + assert opens == 1 + + @pytest.mark.asyncio + async def test_concurrent_reads_share_one_open(self, mcp_db, monkeypatch): + import asyncio + + from haiku.rag.store.engine import Store + + opens = 0 + initialize = Store._initialize + + async def counted(self): + nonlocal opens + opens += 1 + return await initialize(self) + + monkeypatch.setattr(Store, "_initialize", counted) + + mcp = create_mcp_server(mcp_db, read_only=True) + list_docs = await _get_tool(mcp, "list_documents") + + results = await asyncio.gather(*(list_docs() for _ in range(5))) + + assert opens == 1 + assert all(len(r) == 2 for r in results) + + @pytest.mark.asyncio + async def test_a_write_is_visible_to_the_next_read(self, mcp_db): + """One connection sees its own writes, whatever the consistency interval.""" + mcp = create_mcp_server(mcp_db, read_only=False) + list_docs = await _get_tool(mcp, "list_documents") + delete_doc = await _get_tool(mcp, "delete_document") + + docs = await list_docs() + assert await delete_doc(document_id=docs[0].id) is True + + assert len(await list_docs()) == len(docs) - 1 + + @pytest.mark.asyncio + async def test_lifespan_opens_and_closes_once(self, mcp_db, monkeypatch): + from haiku.rag.store.engine import Store + + opens = 0 + initialize = Store._initialize + + async def counted(self): + nonlocal opens + opens += 1 + return await initialize(self) + + monkeypatch.setattr(Store, "_initialize", counted) + + mcp = create_mcp_server(mcp_db, read_only=True) + # _lifespan_manager is what every transport enters; the public + # lifespan() combines provider lifespans only. + async with mcp._lifespan_manager(): + assert opens == 1, "startup should open the database, not the first call" + search = await _get_tool(mcp, "search_documents") + await search(query="artificial intelligence") + assert opens == 1 + + assert opens == 1 + + @pytest.mark.asyncio + async def test_startup_fails_when_the_database_cannot_open(self, tmp_path): + mcp = create_mcp_server(tmp_path / "does-not-exist.lancedb", read_only=True) + + with pytest.raises(FileNotFoundError): + async with mcp._lifespan_manager(): + pass + + @pytest.mark.asyncio + async def test_a_second_lifespan_cycle_opens_a_fresh_client( + self, mcp_db, monkeypatch + ): + from haiku.rag.store.engine import Store + + opens = 0 + initialize = Store._initialize + + async def counted(self): + nonlocal opens + opens += 1 + return await initialize(self) + + monkeypatch.setattr(Store, "_initialize", counted) + + mcp = create_mcp_server(mcp_db, read_only=True) + search = await _get_tool(mcp, "search_documents") + + async with mcp._lifespan_manager(): + await search(query="artificial intelligence") + assert opens == 1 + + async with mcp._lifespan_manager(): + results = await search(query="artificial intelligence") + assert opens == 2 + assert len(results) > 0 + + @pytest.mark.asyncio + async def test_same_dim_drift_starts_read_only_but_not_writable(self, mcp_db): + """Validation is unchanged: same-dimension identity drift warns in + read-only mode and raises in writable mode. The MCP server no longer + opts out of it for deletion.""" + from haiku.rag.config import Config + from haiku.rag.store.repositories.settings import ConfigMismatchError + + drifted = Config.model_copy(deep=True) + drifted.embeddings.model.name = "a-different-model" + + async with create_mcp_server( + mcp_db, config=drifted, read_only=True + )._lifespan_manager(): + pass + + with pytest.raises(ConfigMismatchError): + async with create_mcp_server( + mcp_db, config=drifted, read_only=False + )._lifespan_manager(): + pass From 0882dc9fedb1f39dfe8e94d227c9d1a0f9cef74d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 14:41:11 +0300 Subject: [PATCH 11/11] Lend the caller's client to the capability `client.ask` and `client.analyze` built their capability from a db_path, so the capability opened a second connection to the database the client already had open, once per call. Ownership is now explicit rather than inferred. `rag` stays the connection the capability opened and must close; `borrowed_rag` is a caller's, which `_ensure_rag` prefers and `_close` never touches. Two fields rather than a flag, so closing a borrowed connection is not expressible. `for_run` still clears `rag` per run, since a run owns what it opens. It leaves `borrowed_rag` alone: that connection belongs to the caller and outlives the run. --- CHANGELOG.md | 1 + .../haiku/rag/capabilities/_base.py | 5 + .../haiku/rag/capabilities/analysis.py | 7 +- haiku_rag_slim/haiku/rag/capabilities/rag.py | 7 +- haiku_rag_slim/haiku/rag/client/agents.py | 2 + tests/capabilities/test_borrowed_client.py | 132 ++++++++++++++++++ 6 files changed, 152 insertions(+), 2 deletions(-) create mode 100644 tests/capabilities/test_borrowed_client.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 310f5ecd..bce01898 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ ### Changed - Eval judge pinned to `qwen3.8`: `DEFAULT_JUDGE_MODEL` is `ollama:qwen3.8`, and the reference configs use `Inferact/Qwen3.8-27B-NVFP4` with `extra_body.chat_template_kwargs.reasoning_effort: low`. Results in `docs/benchmarks.md` were judged by `Qwen3.6-35B-A3B-NVFP4` and are not re-judged. +- `create_capability(rag=...)` lends a capability an open client rather than having it open its own; `client.ask`/`client.analyze` now pass theirs. - The MCP server opens one database client for its lifetime instead of one per tool call, and fails at startup if the database cannot be opened. `delete_document` no longer opens its own connection with `skip_validation=True`, so the server no longer opts out of embedding-config validation: drift that validation rejects (any `vector_dim` mismatch, or identity drift on a writable server) now fails MCP startup instead of serving a delete-only server. Use the CLI to delete under drift. - `import_documents` embeds chunks across the whole batch in one pass instead of per document. - `haiku.rag` and `haiku.rag-slim` summaries and keywords; both packages now publish `[project.urls]`. diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index b9f0ff60..2f0bea64 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -129,6 +129,9 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): state: StateT | None = field(default=None, repr=False) outer_state: dict[str, Any] | None = field(default=None, repr=False) rag: HaikuRAG | None = field(default=None, repr=False) + """A connection this capability opened, and must close.""" + borrowed_rag: HaikuRAG | None = field(default=None, repr=False) + """A caller's connection, reused and never closed here.""" rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) resource_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) search_count: int = field(default=0, repr=False) @@ -347,6 +350,8 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): raise error async def _ensure_rag(self) -> HaikuRAG: + if self.borrowed_rag is not None: + return self.borrowed_rag if self.rag is None: async with self.resource_lock: if self.rag is None: diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index e2de8868..b75fa06c 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -1,13 +1,16 @@ from dataclasses import dataclass, field from functools import cache from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from pydantic import BaseModel, Field from pydantic_ai import RunContext, ToolFailed from pydantic_ai.messages import ToolReturn from pydantic_ai.toolsets import FunctionToolset +if TYPE_CHECKING: + from haiku.rag.client import HaikuRAG + from haiku.rag.capabilities._base import ( CodeExecutionEntry, RAGCapabilityBase, @@ -162,6 +165,7 @@ def create_capability( config: AppConfig | None = None, *, defer_loading: bool = True, + rag: "HaikuRAG | None" = None, request_limit: int | None = 30, vision: bool | None = None, ) -> AnalysisCapability: @@ -180,6 +184,7 @@ def create_capability( return AnalysisCapability( db_path=resolve_db_path(db_path, config), config=config, + borrowed_rag=rag, state_type=AnalysisState, state_namespace=STATE_NAMESPACE, instruction_text=instructions(), diff --git a/haiku_rag_slim/haiku/rag/capabilities/rag.py b/haiku_rag_slim/haiku/rag/capabilities/rag.py index 79216a87..f8ea44e9 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/rag.py +++ b/haiku_rag_slim/haiku/rag/capabilities/rag.py @@ -1,13 +1,16 @@ from dataclasses import dataclass from functools import cache from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from pydantic import BaseModel, Field from pydantic_ai import RunContext from pydantic_ai.messages import ToolReturn from pydantic_ai.toolsets import FunctionToolset +if TYPE_CHECKING: + from haiku.rag.client import HaikuRAG + from haiku.rag.capabilities._base import ( RAGCapabilityBase, resolve_db_path, @@ -72,6 +75,7 @@ def create_capability( config: AppConfig | None = None, *, defer_loading: bool = True, + rag: "HaikuRAG | None" = None, request_limit: int | None = 20, vision: bool | None = None, ) -> RAGCapability: @@ -88,6 +92,7 @@ def create_capability( return RAGCapability( db_path=resolve_db_path(db_path, config), config=config, + borrowed_rag=rag, state_type=RAGState, state_namespace=STATE_NAMESPACE, instruction_text=instructions(), diff --git a/haiku_rag_slim/haiku/rag/client/agents.py b/haiku_rag_slim/haiku/rag/client/agents.py index 9a739f74..c85fa1ab 100644 --- a/haiku_rag_slim/haiku/rag/client/agents.py +++ b/haiku_rag_slim/haiku/rag/client/agents.py @@ -63,6 +63,7 @@ async def ask( capability = create_capability( db_path=client.store.db_path, config=client._config, + rag=client, defer_loading=False, ) deps = _AgentDeps( @@ -115,6 +116,7 @@ async def analyze( capability = create_capability( db_path=client.store.db_path, config=client._config, + rag=client, defer_loading=False, ) deps = _AgentDeps( diff --git a/tests/capabilities/test_borrowed_client.py b/tests/capabilities/test_borrowed_client.py new file mode 100644 index 00000000..af1dcec6 --- /dev/null +++ b/tests/capabilities/test_borrowed_client.py @@ -0,0 +1,132 @@ +import pytest + +from haiku.rag.capabilities.rag import create_capability +from haiku.rag.client import HaikuRAG + + +@pytest.mark.asyncio +async def test_a_borrowed_client_is_reused_not_reopened(temp_db_path, monkeypatch): + """A capability handed a client must not open a second connection to the + same database.""" + from haiku.rag.store.engine import Store + + async with HaikuRAG(temp_db_path, create=True) as client: + opens = 0 + initialize = Store._initialize + + async def counted(self): + nonlocal opens + opens += 1 + return await initialize(self) + + monkeypatch.setattr(Store, "_initialize", counted) + + capability = create_capability( + db_path=client.store.db_path, config=client._config, rag=client + ) + + assert await capability._ensure_rag() is client + assert opens == 0 + + +@pytest.mark.asyncio +async def test_closing_never_closes_a_borrowed_client(temp_db_path): + """`_close` owns only what it opened. Closing the caller's client would be a + use-after-close for the caller.""" + async with HaikuRAG(temp_db_path, create=True) as client: + capability = create_capability( + db_path=client.store.db_path, config=client._config, rag=client + ) + await capability._ensure_rag() + + await capability._close() + + # Still usable by its owner. + assert await client.list_documents() == [] + + +@pytest.mark.asyncio +async def test_a_borrowed_client_survives_for_run(temp_db_path): + """for_run clears the owned connection per run; a borrowed one is the + caller's and carries into the run copy.""" + from tests.capabilities.test_capabilities import Deps, make_context + + async with HaikuRAG(temp_db_path, create=True) as client: + capability = create_capability( + db_path=client.store.db_path, config=client._config, rag=client + ) + + run_capability = await capability.for_run(make_context(Deps())) + + assert run_capability is not capability + assert run_capability.rag is None + assert run_capability.borrowed_rag is client + assert await run_capability._ensure_rag() is client + + +@pytest.mark.asyncio +async def test_ask_hands_its_client_to_the_capability(temp_db_path, monkeypatch): + """`ask` built the capability from a db_path alone, so the capability opened + its own connection to a database the client already had open.""" + from haiku.rag.capabilities import rag as rag_capability + from haiku.rag.store.engine import Store + + real = rag_capability.create_capability + built = {} + + def spy(**kwargs): + built["capability"] = real(**kwargs) + raise RuntimeError("stop before running the agent") + + async with HaikuRAG(temp_db_path, create=True) as client: + monkeypatch.setattr(rag_capability, "create_capability", spy) + + with pytest.raises(RuntimeError, match="stop before running the agent"): + await client.ask("anything") + + capability = built["capability"] + opens = 0 + initialize = Store._initialize + + async def counted(self): + nonlocal opens + opens += 1 + return await initialize(self) + + monkeypatch.setattr(Store, "_initialize", counted) + + assert await capability._ensure_rag() is client + assert opens == 0 + + +@pytest.mark.asyncio +async def test_analyze_hands_its_client_to_the_capability(temp_db_path, monkeypatch): + from haiku.rag.capabilities import analysis as analysis_capability + from haiku.rag.store.engine import Store + + real = analysis_capability.create_capability + built = {} + + def spy(**kwargs): + built["capability"] = real(**kwargs) + raise RuntimeError("stop before running the agent") + + async with HaikuRAG(temp_db_path, create=True) as client: + monkeypatch.setattr(analysis_capability, "create_capability", spy) + + with pytest.raises(RuntimeError, match="stop before running the agent"): + await client.analyze("anything") + + capability = built["capability"] + opens = 0 + initialize = Store._initialize + + async def counted(self): + nonlocal opens + opens += 1 + return await initialize(self) + + monkeypatch.setattr(Store, "_initialize", counted) + + assert await capability._ensure_rag() is client + assert opens == 0