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" },