diff --git a/CHANGELOG.md b/CHANGELOG.md index bd166272..5971c43b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,10 +2,17 @@ ## [Unreleased] +### Changed + +- lancedb 0.37.1. +- An invalid search `filter` raises `ValueError`. + ### Added - `--full-citations` on `haiku-rag ask` and `haiku-rag analyze` renders citation text untruncated. `format_citations_rich` takes a `full` argument. +- `doctor` fails when the chunks FTS index covers no rows. +- FTS and hybrid searches log a warning when the FTS index covers no rows. ### Removed @@ -14,6 +21,9 @@ ### Fixed +- FTS and hybrid search on a database whose FTS index covers no rows. Chunk + writes now build the index and rebuild it if it covers none; `haiku-rag + vacuum` also repairs it. - Migration to 0.38.0 no longer fails with `UnicodeDecodeError` on a `docling_document` blob written as zstd. - Vector search sets `search.vector_index_metric` on every query. diff --git a/docs/cli.md b/docs/cli.md index e300a8ed..cd1fba8e 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -329,6 +329,7 @@ Checks include: - the configured embedding identity matches the stored settings - no database migrations are pending - the vector index covers all chunks +- the full-text index covers the chunks it searches - near-identical documents (by embedding-centroid similarity) are grouped and reported, with the largest member flagged as the likely one to keep (advisory only, never deleted, tuned via `doctor.duplicates` in config) - API keys are set for configured providers @@ -340,7 +341,7 @@ It also probes the external endpoints the config uses and reports them under a P SaaS providers (OpenAI, Anthropic, Cohere, Jina, ZeroEntropy, Voyage) are covered by the API-key check rather than a network probe. In-process local models (sentence-transformers, cross-encoder, jina-local) have no endpoint and are reported as such. -Each failure prints the command that fixes it (`rebuild`, `create-index`, `migrate`, `rebuild --set-embedder`). `doctor` makes no changes. It exits with status 1 when any check fails, so it can gate CI or monitoring. +Each failure prints the command that fixes it (`rebuild`, `create-index`, `vacuum`, `migrate`, `rebuild --set-embedder`). `doctor` makes no changes. It exits with status 1 when any check fails, so it can gate CI or monitoring. ### Migrate Database diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index e9e9df2b..a94b2b3c 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -64,6 +64,8 @@ If that number times six exceeds available RAM, use one of: - Set `generate_page_images: false` if visual grounding through `visualize_chunk()` is not needed. This removes page rasters entirely. - Set `auto_vacuum: false` and run `haiku-rag vacuum` manually when the machine is otherwise idle, so the peak does not land alongside ingestion. +Vacuum also folds new rows into the full-text index. Search stays correct without it but scans the uncovered rows on every query. `haiku-rag doctor` reports the coverage. + This is an upstream limitation rather than a `haiku.rag` setting. Compaction bounds itself by row count instead of bytes, and LanceDB's async API exposes no batch size or fragment target to override it. Tracked at [lancedb/lancedb#2325](https://github.com/lancedb/lancedb/issues/2325). The requirement above will drop once compaction batches by bytes. ### Changing the Default Database Path diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index be8dffe5..4de95b14 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -21,7 +21,7 @@ from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document from haiku.rag.store.models.document_item import extract_items from haiku.rag.store.repositories.settings import SettingsRepository -from haiku.rag.store.schema import ChunkRecordBase +from haiku.rag.store.schema import ChunkRecordBase, ensure_indexes if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument @@ -499,10 +499,12 @@ async def _rebuild_embed_only( if len(yielded_docs) % _REBUILD_BATCH_SIZE == 0 and pending_records: await session.store.chunks_table.add(pending_records) + await ensure_indexes(session.store.chunks_table, "chunks") pending_records = [] if pending_records: await session.store.chunks_table.add(pending_records) + await ensure_indexes(session.store.chunks_table, "chunks") # Phase 2 finished. Drop the recovery state — marker first so a crash # between the two drops leaves only staging behind, which the next diff --git a/haiku_rag_slim/haiku/rag/doctor.py b/haiku_rag_slim/haiku/rag/doctor.py index e6708456..570eceba 100644 --- a/haiku_rag_slim/haiku/rag/doctor.py +++ b/haiku_rag_slim/haiku/rag/doctor.py @@ -17,7 +17,7 @@ from haiku.rag.config.models import ( ) from haiku.rag.store.engine import Store, connect_lancedb from haiku.rag.store.info import get_database_stats -from haiku.rag.store.schema import REQUIRED_TABLES +from haiku.rag.store.schema import REQUIRED_TABLES, index_specs from haiku.rag.store.upgrades import get_pending_upgrades # Cap how many offending ids we collect per check; doctor is a summary, not a dump. @@ -225,6 +225,56 @@ def _classify_unchunked( return results +async def _check_fts_coverage(store: Store) -> CheckResult: + """An FTS index that covers no rows, and a populated table with no FTS + index at all, both make lance serve results unsorted by score with + matching rows dropped. optimize indexes the rows of an index that + exists; it never creates one that is absent.""" + from lancedb.index import FTS + + uncovered: list[str] = [] + missing: list[str] = [] + for table_name, table in store._tables().items(): + declared = [c for c, cfg in index_specs(table_name) if isinstance(cfg, FTS)] + if not declared: + continue + rows = await table.count_rows() + if not rows: + continue + indices = await table.list_indices() + for column in declared: + index = next( + (i for i in indices if column in i.columns and i.index_type == "FTS"), + None, + ) + if index is None: + missing.append(f"{table_name}.{column}: no index over {rows} rows") + continue + stats = await table.index_stats(index.name) + if stats is None or stats.num_indexed_rows == 0: + uncovered.append(f"{table_name}.{column}: 0 of {rows} rows indexed") + if missing or uncovered: + return CheckResult( + name="fts_index_coverage", + severity=Severity.FAIL, + message=( + "Full-text search index does not cover its rows; FTS and " + "hybrid results are unsorted and incomplete." + ), + remediation=( + "Run 'haiku-rag rebuild --embed-only' to build the index." + if missing + else "Run 'haiku-rag vacuum' to index the rows." + ), + details=missing + uncovered, + ) + return CheckResult( + name="fts_index_coverage", + severity=Severity.OK, + message="Full-text search indexes are present and cover rows.", + ) + + async def _column_values(table, column: str) -> list: rows = await table.query().select([column]).to_list() return [row[column] for row in rows] @@ -702,6 +752,9 @@ async def run_db_checks( self_refs_by_doc.setdefault(row["document_id"], set()).add(row["self_ref"]) labels_by_doc.setdefault(row["document_id"], set()).add(row["label"]) + notify("Checking index coverage") + results.append(await _check_fts_coverage(store)) + notify("Checking referential integrity") results.append(_check_document_meta_parity(doc_ids, meta_doc_ids)) results.append(_check_orphaned_chunks(chunk_doc_ids, doc_ids)) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 5ab1ba9c..3cac1864 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -6,7 +6,6 @@ from uuid import uuid4 if TYPE_CHECKING: from lancedb.query import AsyncQueryBase -from lancedb.index import FTS from lancedb.rerankers import RRFReranker from haiku.rag.store.engine import Store @@ -23,17 +22,49 @@ class ChunkRepository: def __init__(self, store: Store) -> None: self.store = store self.embedder = store.embedder + self._fts_coverage_checked = False - async def _ensure_fts_index(self) -> None: - """Ensure FTS index exists on the content_fts column.""" + async def _warn_if_fts_uncovered(self) -> None: + """An FTS index covering no rows makes lance serve a broken scan path: + results unsorted by score, with matching documents dropped. Checked + on the first search that uses the index, once per repository after + the table holds rows.""" + if self._fts_coverage_checked: + return try: - await self.store.chunks_table.create_index( - "content_fts", - config=FTS(with_position=True, remove_stop_words=False), - replace=True, + indices = await self.store.chunks_table.list_indices() + index = next( + ( + i + for i in indices + if "content_fts" in i.columns and i.index_type == "FTS" + ), + None, ) - except Exception as e: - logger.warning(f"FTS index build failed; full-text search degraded: {e}") + stats = ( + await self.store.chunks_table.index_stats(index.name) if index else None + ) + if stats is not None and stats.num_indexed_rows > 0: + self._fts_coverage_checked = True + return + # An empty table proves nothing; check again once it has rows. + if not await self.store.chunks_table.count_rows(): + return + except Exception: + self._fts_coverage_checked = True + logger.debug("FTS coverage check failed", exc_info=True) + return + self._fts_coverage_checked = True + if index is None: + logger.warning( + "No full-text search index; FTS and hybrid results are " + "degraded. Run 'haiku-rag rebuild --embed-only'." + ) + return + logger.warning( + "Full-text search index covers 0 rows; FTS and hybrid results " + "are degraded. Run 'haiku-rag vacuum'." + ) def _contextualize_content(self, chunk: Chunk) -> str: """Generate contextualized content for FTS by prepending headings.""" @@ -73,8 +104,9 @@ class ChunkRepository: chunk_record = self._to_record(entity, chunk_id) await self.store.chunks_table.add([chunk_record]) - entity.id = chunk_id + await ensure_indexes(self.store.chunks_table, "chunks") + return entity chunks = entity @@ -95,6 +127,7 @@ class ChunkRepository: chunk.id = chunk_id await self.store.chunks_table.add(chunk_records) + await ensure_indexes(self.store.chunks_table, "chunks") return chunks @@ -128,6 +161,7 @@ class ChunkRepository: .when_not_matched_by_source_delete(f"document_id = '{safe_id}'") .execute(records) ) + await ensure_indexes(self.store.chunks_table, "chunks") return chunks async def get_by_id(self, entity_id: str) -> Chunk | None: @@ -196,6 +230,7 @@ class ChunkRepository: return False await self.store.chunks_table.delete(f"document_id = '{document_id}'") + await ensure_indexes(self.store.chunks_table, "chunks") return True async def search( @@ -240,6 +275,9 @@ class ChunkRepository: id_list = ", ".join(f"'{d}'" for d in docs_df["id"]) chunk_filter = f"document_id IN ({id_list})" + if search_type != "vector" and query.strip(): + await self._warn_if_fts_uncovered() + if search_type == "fts": results = self.store.chunks_table.query().nearest_to_text( query, columns="content_fts" diff --git a/haiku_rag_slim/haiku/rag/store/schema.py b/haiku_rag_slim/haiku/rag/store/schema.py index e9832ad4..6c21f841 100644 --- a/haiku_rag_slim/haiku/rag/store/schema.py +++ b/haiku_rag_slim/haiku/rag/store/schema.py @@ -161,25 +161,47 @@ async def ensure_indexes(table: lancedb.AsyncTable, table_name: str) -> list[str 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. + Re-creating is not free: `create_index(replace=True)` rebuilds. A declared + FTS index covering no rows of a populated table is rebuilt in place: it + serves the same broken scan path as a missing one. """ covering: dict[str, set[str]] = {} + fts_names: dict[str, str] = {} for index in await table.list_indices(): for column in index.columns: covering.setdefault(column, set()).add(index.index_type) + if index.index_type == "FTS": + fts_names.setdefault(column, index.name) applied: list[str] = [] for column, config in index_specs(table_name): declared = type(config).__name__ present = covering.get(column, set()) if declared in present: + if isinstance(config, FTS): + stats = await table.index_stats(fts_names[column]) + if ( + stats is None or stats.num_indexed_rows == 0 + ) and await table.count_rows(): + await table.create_index( + column, config=config, replace=True, name=fts_names[column] + ) + applied.append(column) + continue + # lance indexes nothing when the table is empty and never catches an + # FTS index up on add: the scan path it then serves returns results + # unsorted by score, with matching rows dropped. + if isinstance(config, FTS) and not await table.count_rows(): 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) + # Replace matches on the index name. + await table.create_index( + column, config=config, replace=True, name=f"{column}_idx" + ) applied.append(column) return applied diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 9f883005..abd4130b 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -41,7 +41,7 @@ dependencies = [ "httpx>=0.28.1", "jinja2>=3.1.0", "fastmcp>=3.3.0", - "lancedb==0.34.0", + "lancedb==0.37.1", "pathspec>=1.0.4", "pydantic>=2.12.5", "pydantic-ai-slim[openai,logfire,ag-ui]>=2.18.0,<3.0.0", diff --git a/tests/store/test_indexes.py b/tests/store/test_indexes.py index 0f9451f6..bf993cf8 100644 --- a/tests/store/test_indexes.py +++ b/tests/store/test_indexes.py @@ -1,19 +1,44 @@ import pyarrow as pa import pytest -from lancedb.index import BTree +from lancedb.index import FTS, BTree from haiku.rag.store.engine import Store -from haiku.rag.store.models import Document +from haiku.rag.store.models import Chunk, Document +from haiku.rag.store.repositories.chunk import ChunkRepository from haiku.rag.store.repositories.document import DocumentRepository from haiku.rag.store.schema import ensure_indexes EXPECTED_INDEXED_COLUMNS = { "documents": {"id"}, "document_meta": {"id", "uri"}, - "chunks": {"content_fts", "id", "document_id"}, + "chunks": {"id", "document_id"}, "document_items": {"document_id", "position", "self_ref", "label"}, } +# content_fts joins the set once the table holds rows to index. +EXPECTED_POPULATED_CHUNK_COLUMNS = {"content_fts", "id", "document_id"} + + +async def _fts_indexed_rows(table) -> int | None: + """Rows the FTS index covers, or None when there is no FTS index.""" + for index in await table.list_indices(): + if index.index_type == "FTS": + return (await table.index_stats(index.name)).num_indexed_rows + return None + + +async def _add_chunk( + store, content: str = "a chunk about gardens", document_id: str = "doc-1" +) -> None: + await ChunkRepository(store).create( + Chunk( + document_id=document_id, + content=content, + embedding=[0.1] * store.embedder.vector_dim, + order=0, + ) + ) + async def _indexed_columns(table) -> set[str]: return {column for index in await table.list_indices() for column in index.columns} @@ -62,7 +87,9 @@ async def test_ensure_indexes_corrects_an_index_of_the_wrong_type(temp_db_path): """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) + await table.create_index( + "label", config=BTree(), replace=True, name="label_idx" + ) assert await _index_type(table, "label") == "BTree" await ensure_indexes(table, "document_items") @@ -130,3 +157,188 @@ async def test_delete_all_keeps_picture_data_as_large_binary(temp_db_path): schema = await store.document_items_table.schema() assert schema.field("picture_data").type == pa.large_binary() + + +@pytest.mark.asyncio +async def test_empty_chunks_table_carries_no_fts_index(temp_db_path): + """An FTS index over no rows indexes nothing and lance never catches it up, + so it waits for rows rather than being built with the table.""" + async with Store(temp_db_path, create=True) as store: + assert await _fts_indexed_rows(store.chunks_table) is None + + +@pytest.mark.asyncio +async def test_ensure_indexes_skips_fts_while_the_table_is_empty(temp_db_path): + async with Store(temp_db_path, create=True) as store: + await ensure_indexes(store.chunks_table, "chunks") + + assert await _fts_indexed_rows(store.chunks_table) is None + assert ( + await _indexed_columns(store.chunks_table) + == (EXPECTED_INDEXED_COLUMNS["chunks"]) + ) + + +@pytest.mark.asyncio +async def test_ensure_indexes_builds_fts_over_existing_rows(temp_db_path): + async with Store(temp_db_path, create=True) as store: + await store.chunks_table.add( + [ + store.ChunkRecord( + document_id="doc-1", + content="a chunk", + content_fts="a chunk", + metadata="{}", + order=0, + vector=[0.1] * store.embedder.vector_dim, + ) + ] + ) + + await ensure_indexes(store.chunks_table, "chunks") + + assert await _fts_indexed_rows(store.chunks_table) == 1 + assert await _indexed_columns(store.chunks_table) == ( + EXPECTED_POPULATED_CHUNK_COLUMNS + ) + + +@pytest.mark.asyncio +async def test_creating_chunks_builds_a_covering_fts_index(temp_db_path): + async with Store(temp_db_path, create=True) as store: + await _add_chunk(store) + + assert await _fts_indexed_rows(store.chunks_table) == 1 + + +@pytest.mark.asyncio +async def test_replacing_chunks_builds_a_covering_fts_index(temp_db_path): + """replace_for_document inserts where nothing matches, so it can be the + first write into a fresh table.""" + async with Store(temp_db_path, create=True) as store: + await ChunkRepository(store).replace_for_document( + "doc-1", + [ + Chunk( + document_id="doc-1", + content="a chunk about gardens", + embedding=[0.1] * store.embedder.vector_dim, + order=0, + ) + ], + ) + + assert await _fts_indexed_rows(store.chunks_table) == 1 + + +@pytest.mark.asyncio +async def test_a_second_write_leaves_the_covering_index_in_place(temp_db_path): + """One indexed row is enough: later rows merge as a scanned tail, so the + index is not rebuilt per write.""" + async with Store(temp_db_path, create=True) as store: + await _add_chunk(store, "first") + version_after_first = await store.chunks_table.version() + + await _add_chunk(store, "second") + + assert await _fts_indexed_rows(store.chunks_table) == 1 + assert await store.chunks_table.version() == version_after_first + 1 + + +@pytest.mark.asyncio +async def test_delete_all_then_write_rebuilds_a_covering_fts_index(temp_db_path): + """delete_all recreates the table empty, so the next write owns the index.""" + async with Store(temp_db_path, create=True) as store: + await _add_chunk(store) + await ChunkRepository(store).delete_all() + assert await _fts_indexed_rows(store.chunks_table) is None + + await _add_chunk(store) + + assert await _fts_indexed_rows(store.chunks_table) == 1 + + +@pytest.mark.asyncio +async def test_deleting_the_indexed_rows_rebuilds_the_fts_index(temp_db_path): + """Deleting every row the index covers, while unindexed rows remain, + reaches the zero-coverage scan path; the delete repairs it.""" + async with Store(temp_db_path, create=True) as store: + await _add_chunk(store, "first", document_id="doc-a") + await _add_chunk(store, "second", document_id="doc-b") + assert await _fts_indexed_rows(store.chunks_table) == 1 + + await ChunkRepository(store).delete_by_document_id("doc-a") + + assert await store.chunks_table.count_rows() == 1 + assert await _fts_indexed_rows(store.chunks_table) == 1 + + +@pytest.mark.asyncio +async def test_replacing_the_indexed_rows_rebuilds_the_fts_index(temp_db_path): + """Replacement rewrites rows, and rewritten rows are unindexed.""" + async with Store(temp_db_path, create=True) as store: + await _add_chunk(store, "first", document_id="doc-a") + await _add_chunk(store, "second", document_id="doc-b") + + await ChunkRepository(store).replace_for_document( + "doc-a", + [ + Chunk( + document_id="doc-a", + content="rewritten", + embedding=[0.1] * store.embedder.vector_dim, + order=0, + ) + ], + ) + + assert await _fts_indexed_rows(store.chunks_table) == 2 + + +@pytest.mark.asyncio +async def test_a_write_repairs_a_legacy_index_that_covers_no_rows(temp_db_path): + """A database whose FTS index predates its rows is repaired by the first + write that runs index maintenance.""" + async with Store(temp_db_path, create=True) as store: + await store.chunks_table.create_index( + "content_fts", config=FTS(with_position=True, remove_stop_words=False) + ) + await store.chunks_table.add( + [ + store.ChunkRecord( + document_id="doc-a", + content="a legacy chunk", + content_fts="a legacy chunk", + metadata="{}", + order=0, + vector=[0.1] * store.embedder.vector_dim, + ) + ] + ) + assert await _fts_indexed_rows(store.chunks_table) == 0 + + await _add_chunk(store, "second", document_id="doc-b") + + assert await _fts_indexed_rows(store.chunks_table) == 2 + + +@pytest.mark.asyncio +async def test_unavailable_index_stats_repair_matches_doctor(temp_db_path, monkeypatch): + """index_stats may return None; doctor treats that as uncovered, so the + write-path repair does too.""" + from lancedb.table import AsyncTable + + async with Store(temp_db_path, create=True) as store: + await _add_chunk(store) + + original = AsyncTable.index_stats + + async def no_stats(self, name): + if self.name == "chunks": + return None + return await original(self, name) + + monkeypatch.setattr(AsyncTable, "index_stats", no_stats) + applied = await ensure_indexes(store.chunks_table, "chunks") + + assert applied == ["content_fts"] diff --git a/tests/store/test_v0_75_0_migration.py b/tests/store/test_v0_75_0_migration.py index 3b996d51..15b3d3b3 100644 --- a/tests/store/test_v0_75_0_migration.py +++ b/tests/store/test_v0_75_0_migration.py @@ -32,6 +32,18 @@ async def _make_legacy(store: Store) -> None: @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 store.chunks_table.add( + [ + store.ChunkRecord( + document_id="doc-1", + content="a chunk", + content_fts="a chunk", + metadata="{}", + order=0, + vector=[0.1] * store.embedder.vector_dim, + ) + ] + ) await _make_legacy(store) await _apply_index_hot_lookup_keys(store) @@ -83,7 +95,7 @@ async def test_replaces_a_wrong_typed_legacy_index(temp_db_path): """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 + "label", config=BTree(), replace=True, name="label_idx" ) await _apply_index_hot_lookup_keys(store) diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 242ba26b..be15b96d 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -1,3 +1,5 @@ +from unittest.mock import patch + import pytest from haiku.rag.client import HaikuRAG @@ -477,8 +479,100 @@ async def test_chunk_content_fts(temp_db_path, metadata, content, expected_conte assert record["content_fts"] == expected_content_fts -async def test_ensure_fts_index_warns_on_failure(temp_db_path): - """A failed FTS index build is surfaced at WARNING, not swallowed silently.""" +async def _import_one(client) -> None: + from docling_core.types.doc.document import DoclingDocument + from docling_core.types.doc.labels import DocItemLabel + + doc = DoclingDocument(name="one") + doc.add_text(label=DocItemLabel.TEXT, text="a document about gardens") + await client.import_document( + doc, + [ + Chunk( + content="a document about gardens", + embedding=[0.1] * get_config().embeddings.model.vector_dim, + order=0, + ) + ], + uri="test://one", + ) + + +async def _add_legacy_row(client) -> None: + """A row written without index maintenance, as older releases wrote them.""" + await client.store.chunks_table.add( + [ + client.store.ChunkRecord( + document_id="doc-legacy", + content="a document about gardens", + content_fts="a document about gardens", + metadata="{}", + order=0, + vector=[0.1] * get_config().embeddings.model.vector_dim, + ) + ] + ) + + +async def test_fts_search_warns_when_index_covers_no_rows(temp_db_path): + """A database whose FTS index predates its rows covers none of them; + searching that state warns, once.""" + import logging + + from lancedb.index import FTS + + from haiku.rag.store.repositories import chunk as chunk_module + + async with HaikuRAG( + db_path=temp_db_path, config=get_config(), create=True + ) as client: + await client.store.chunks_table.create_index( + "content_fts", config=FTS(with_position=True, remove_stop_words=False) + ) + await _add_legacy_row(client) + + with capture_logs(chunk_module.logger, logging.WARNING) as records: + await client.chunk_repository.search("gardens", search_type="fts") + await client.chunk_repository.search("gardens", search_type="fts") + + warned = [r for r in records if "covers 0 rows" in r.getMessage()] + assert len(warned) == 1 + + +async def test_fts_search_warns_when_index_is_missing(temp_db_path, monkeypatch): + import logging + + from lancedb.table import AsyncTable + + from haiku.rag.store.repositories import chunk as chunk_module + + original = AsyncTable.list_indices + + async def no_chunk_indices(self): + if self.name == "chunks": + return [] + return await original(self) + + async with HaikuRAG( + db_path=temp_db_path, config=get_config(), create=True + ) as client: + await _import_one(client) + monkeypatch.setattr(AsyncTable, "list_indices", no_chunk_indices) + + with capture_logs(chunk_module.logger, logging.WARNING) as records: + await client.chunk_repository.search("gardens", search_type="fts") + + warned = [ + r.getMessage() + for r in records + if "No full-text search index" in r.getMessage() + ] + assert len(warned) == 1 + assert "rebuild --embed-only" in warned[0] + assert "vacuum" not in warned[0] + + +async def test_fts_search_does_not_warn_when_index_covers_rows(temp_db_path): import logging from haiku.rag.store.repositories import chunk as chunk_module @@ -486,18 +580,100 @@ async def test_ensure_fts_index_warns_on_failure(temp_db_path): async with HaikuRAG( db_path=temp_db_path, config=get_config(), create=True ) as client: - repo = client.chunk_repository - - async def _boom(*_args, **_kwargs): - raise RuntimeError("index build failed") - - repo.store.chunks_table.create_index = _boom + await _import_one(client) + await client.store.vacuum(retention_seconds=0) with capture_logs(chunk_module.logger, logging.WARNING) as records: - await repo._ensure_fts_index() + results = await client.chunk_repository.search("gardens", search_type="fts") - assert [r for r in records if r.levelno == logging.WARNING] - assert any("index build failed" in r.getMessage() for r in records) + assert results + assert not records + + +async def test_fts_coverage_check_failure_does_not_break_search(temp_db_path): + """The coverage check is a diagnostic: a metadata failure must not take + the search down with it.""" + from lancedb.table import AsyncTable + + async def boom(self): + raise RuntimeError("metadata unavailable") + + async with HaikuRAG( + db_path=temp_db_path, config=get_config(), create=True + ) as client: + await _import_one(client) + + with patch.object(AsyncTable, "list_indices", boom): + results = await client.chunk_repository.search("gardens", search_type="fts") + + assert results + + +async def test_create_assigns_the_id_when_index_maintenance_fails(temp_db_path): + """The row is committed before the index is ensured, so the chunk carries + the id it was written with even when that ensure fails.""" + from haiku.rag.store.repositories import chunk as chunk_module + + async with HaikuRAG( + db_path=temp_db_path, config=get_config(), create=True + ) as client: + repo = client.chunk_repository + + async def boom(*_args, **_kwargs): + raise RuntimeError("index build failed") + + chunk = Chunk( + document_id="doc-1", + content="a chunk about gardens", + embedding=[0.1] * get_config().embeddings.model.vector_dim, + order=0, + ) + with patch.object(chunk_module, "ensure_indexes", boom): + with pytest.raises(RuntimeError, match="index build failed"): + await repo.create(chunk) + + assert chunk.id is not None + assert await client.store.chunks_table.count_rows() == 1 + + +async def test_fts_search_on_an_empty_table_does_not_suppress_later_warnings( + temp_db_path, +): + """An empty table proves nothing about coverage, so searching it must not + spend the once-per-repository check.""" + import logging + + from lancedb.index import FTS + + from haiku.rag.store.repositories import chunk as chunk_module + + async with HaikuRAG( + db_path=temp_db_path, config=get_config(), create=True + ) as client: + await client.store.chunks_table.create_index( + "content_fts", config=FTS(with_position=True, remove_stop_words=False) + ) + await client.chunk_repository.search("gardens", search_type="fts") + await _add_legacy_row(client) + + with capture_logs(chunk_module.logger, logging.WARNING) as records: + await client.chunk_repository.search("gardens", search_type="fts") + + assert any("covers 0 rows" in r.getMessage() for r in records) + + +async def test_fts_search_does_not_warn_on_an_empty_table(temp_db_path): + import logging + + from haiku.rag.store.repositories import chunk as chunk_module + + async with HaikuRAG( + db_path=temp_db_path, config=get_config(), create=True + ) as client: + with capture_logs(chunk_module.logger, logging.WARNING) as records: + await client.chunk_repository.search("gardens", search_type="fts") + + assert not records @pytest.mark.vcr() diff --git a/tests/test_client.py b/tests/test_client.py index ef1bdfad..99fdd429 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -895,8 +895,20 @@ async def test_client_import_documents_single_version_per_table(temp_db_path): assert all(d.id is not None for d in docs) assert len({d.id for d in docs}) == 3 + assert after["documents"] - before["documents"] == 1 + assert after["document_items"] - before["document_items"] == 1 + # Two on chunks: the batch, plus building the FTS index over the first + # rows the table has ever held. + assert after["chunks"] - before["chunks"] == 2 + + # With the index in place, a further batch is one version per table. + again = await client.import_documents( + [_import("d", "Delta document body", uri="mem://d", title="Delta")] + ) + assert [d.title for d in again] == ["Delta"] + latest = await client.store.current_table_versions() for table in ("documents", "chunks", "document_items"): - assert after[table] - before[table] == 1, table + assert latest[table] - after[table] == 1, table for doc, expected in zip(docs, ("Alpha", "Beta", "Gamma")): assert doc.id is not None diff --git a/tests/test_doctor.py b/tests/test_doctor.py index 90282b38..eaa59ccb 100644 --- a/tests/test_doctor.py +++ b/tests/test_doctor.py @@ -81,18 +81,25 @@ async def _build_db( name: str = "test", vector_dim: int = VECTOR_DIM, stored_vector_dim: int | None = None, + fts_index: str | None = "covering", ): """Create a consistent single-document database without touching an embedder. ``stored_vector_dim`` records a different dimension in settings than the chunks table actually uses, to exercise the vector-dimension check. + ``fts_index`` places the chunks FTS index: "covering" builds it over the + rows, "empty" builds it before them, None never builds it. """ + from lancedb.index import FTS + db = await lancedb.connect_async(path) settings_tbl = await db.create_table("settings", schema=SettingsRecord) docs_tbl = await db.create_table("documents", schema=DocumentRecord) meta_tbl = await db.create_table("document_meta", schema=DocumentMetaRecord) chunks_tbl = await db.create_table("chunks", schema=create_chunk_model(vector_dim)) items_tbl = await db.create_table("document_items", schema=DocumentItemRecord) + if fts_index == "empty": + await chunks_tbl.create_index("content_fts", config=FTS()) await settings_tbl.add( [ @@ -134,6 +141,8 @@ async def _build_db( ) ] ) + if fts_index == "covering": + await chunks_tbl.create_index("content_fts", config=FTS()) return db @@ -201,6 +210,41 @@ async def test_missing_table_fails_without_opening_store(temp_db_path): assert "documents" in tables.details +@pytest.mark.asyncio +async def test_missing_fts_index_fails(temp_db_path): + """optimize indexes the rows of an index that exists; it never creates a + missing one, so the remediation has to build it.""" + await _build_db(temp_db_path, fts_index=None) + report = await run_doctor(_config(), temp_db_path, {}) + result = _result(report, "fts_index_coverage") + assert result.severity is Severity.FAIL + assert "chunks.content_fts" in result.details[0] + assert "no index over 1 rows" in result.details[0] + assert "rebuild" in (result.remediation or "") + assert "vacuum" not in (result.remediation or "") + assert report.failed + + +@pytest.mark.asyncio +async def test_fts_index_covering_no_rows_fails(temp_db_path): + """The state a bulk write without a closing vacuum leaves behind.""" + await _build_db(temp_db_path, fts_index="empty") + report = await run_doctor(_config(), temp_db_path, {}) + result = _result(report, "fts_index_coverage") + assert result.severity is Severity.FAIL + assert "0 of 1 rows indexed" in result.details[0] + assert "vacuum" in (result.remediation or "") + + +@pytest.mark.asyncio +async def test_fts_coverage_passes_an_empty_table(temp_db_path): + db = await _build_db(temp_db_path) + chunks_tbl = await db.open_table("chunks") + await chunks_tbl.delete("id = 'c1'") + report = await run_doctor(_config(), temp_db_path, {}) + assert _result(report, "fts_index_coverage").severity is Severity.OK + + @pytest.mark.asyncio async def test_orphaned_chunk_fails(temp_db_path): db = await _build_db(temp_db_path) diff --git a/tests/test_filter.py b/tests/test_filter.py index 6e2e7b88..4b746bc6 100644 --- a/tests/test_filter.py +++ b/tests/test_filter.py @@ -157,8 +157,7 @@ async def test_search_with_invalid_filter(temp_db_path): title="Test Document", ) - # Invalid filter should raise RuntimeError - with pytest.raises(RuntimeError, match="No field named invalid"): + with pytest.raises(ValueError, match="No field named invalid"): await client.search("test", limit=5, filter="invalid = 'value'") diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 3ee9e237..16177122 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -91,7 +91,7 @@ class TestMCPReadTools: embedding=embedding, ) ) - await rag.chunk_repository._ensure_fts_index() + await rag.store.chunks_table.optimize() mcp = create_mcp_server(mcp_db, read_only=True) async with Client(mcp) as client: diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index 1cc7f985..7a8d55ef 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -101,8 +101,9 @@ async def test_rebuild_embed_only_multi_doc_streams_via_staging( the chunks table, then streams doc-by-doc. This test verifies: - chunks survive across multiple documents (correctness), - - the staging table is dropped at the end (no leak), and - - the rebuild yields every document with chunks. + - the staging table is dropped at the end (no leak), + - the rebuild yields every document with chunks, and + - the recreated chunks table carries an FTS index covering its rows. """ async with HaikuRAG(temp_db_path, create=True) as client: doc1 = await client.create_document(content=qa_corpus[0]["document_extracted"]) @@ -136,6 +137,16 @@ async def test_rebuild_embed_only_multi_doc_streams_via_staging( tables = (await client.store.db.list_tables()).tables assert "chunks_rebuild_staging" not in tables + # Phase 2 recreates the chunks table, so it owns the FTS index. + fts = [ + index + for index in await client.store.chunks_table.list_indices() + if index.index_type == "FTS" + ] + assert len(fts) == 1 + stats = await client.store.chunks_table.index_stats(fts[0].name) + assert stats.num_indexed_rows > 0 + @pytest.mark.vcr() async def test_rebuild_drops_leftover_staging_table( @@ -330,6 +341,55 @@ async def test_rebuild_non_embed_mode_drops_staging_recovery_state( assert "chunks_rebuild_marker" not in tables +async def test_rebuild_embed_only_covers_the_index_from_the_first_flush( + temp_db_path, monkeypatch +): + """A rebuild abandoned after a flushed batch leaves the FTS index + covering the flushed rows.""" + from docling_core.types.doc.document import DoclingDocument + from docling_core.types.doc.labels import DocItemLabel + + from haiku.rag import embeddings as embeddings_module + from haiku.rag.client import rebuild as rebuild_module + from haiku.rag.config import get_config + from haiku.rag.store.models.chunk import Chunk + + monkeypatch.setattr(rebuild_module, "_REBUILD_BATCH_SIZE", 1) + + async def keep_embeddings(chunks, embedder, config=None): + for chunk in chunks: + if chunk.embedding is None: + chunk.embedding = [0.1] * embedder.vector_dim + return chunks + + monkeypatch.setattr(embeddings_module, "embed_chunks", keep_embeddings) + + async with HaikuRAG(temp_db_path, create=True) as client: + dim = get_config().embeddings.model.vector_dim + for name in ("one", "two"): + doc = DoclingDocument(name=name) + doc.add_text(label=DocItemLabel.TEXT, text=f"document {name}") + await client.import_document( + doc, + [Chunk(content=f"document {name}", embedding=[0.1] * dim, order=0)], + uri=f"test://{name}", + ) + + rebuild = client.rebuild_database(mode=RebuildMode.EMBED_ONLY) + await anext(rebuild) + await anext(rebuild) + await rebuild.aclose() + + fts = [ + index + for index in await client.store.chunks_table.list_indices() + if index.index_type == "FTS" + ] + assert len(fts) == 1 + stats = await client.store.chunks_table.index_stats(fts[0].name) + assert stats.num_indexed_rows > 0 + + @pytest.mark.vcr() async def test_rebuild_embed_only_skips_unchanged( qa_corpus: list[dict[str, str]], temp_db_path diff --git a/uv.lock b/uv.lock index f4244487..5f01576c 100644 --- a/uv.lock +++ b/uv.lock @@ -1759,7 +1759,7 @@ requires-dist = [ { name = "haiku-rag-slim", extras = ["s3"], marker = "extra == 'ingester'", editable = "haiku_rag_slim" }, { name = "httpx", specifier = ">=0.28.1" }, { name = "jinja2", specifier = ">=3.1.0" }, - { name = "lancedb", specifier = "==0.34.0" }, + { name = "lancedb", specifier = "==0.37.1" }, { name = "obstore", marker = "extra == 's3'", specifier = ">=0.9,<0.10" }, { name = "opencv-python-headless", marker = "extra == 'docling'", specifier = ">=4.6.0.66,<5.0.0.0" }, { name = "pathspec", specifier = ">=1.0.4" }, @@ -2261,7 +2261,7 @@ wheels = [ [[package]] name = "lancedb" -version = "0.34.0" +version = "0.37.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "deprecation" }, @@ -2273,10 +2273,10 @@ dependencies = [ { name = "tqdm" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/df/f7/5262b9aa593f790757163c0165ab0da1dda054758901bea7e4f02c9cb633/lancedb-0.34.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:c462f2e6f933cad659fd0179394eaab578acbc9151fe2ef41bc29b36ecca5058", size = 52654213, upload-time = "2026-07-02T17:13:31.102Z" }, - { url = "https://files.pythonhosted.org/packages/69/99/05ea0d32229ebea695193ff20c15d6ecae25785ad82a9d4723d98832a284/lancedb-0.34.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:48829e88e708947d0520454ab9e4f8efa35f3e3626469eadd3a6e061b89cb223", size = 55434501, upload-time = "2026-07-02T17:13:34.81Z" }, - { url = "https://files.pythonhosted.org/packages/cd/4e/4325c13d5afa93c466428a5a0f168ad4d96f5eb4a77bbe7c5100d39c9897/lancedb-0.34.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:05ba8a5b58e064edfbe5be71b1abf2e411b4eaf295d1a173dcb1a55c5bfb5285", size = 58659359, upload-time = "2026-07-02T17:13:38.424Z" }, - { url = "https://files.pythonhosted.org/packages/d9/5d/8ca165f1386caf6c4d1c515afd52f345b66432264eecfdfb7fd33eefd9af/lancedb-0.34.0-cp39-abi3-win_amd64.whl", hash = "sha256:51cbc11808f9e3332819b9367c975b3a888541447a8e7bea09c57c852a279153", size = 63530726, upload-time = "2026-07-02T17:13:41.612Z" }, + { url = "https://files.pythonhosted.org/packages/23/2f/4ddcab82bb618c6c8de00725f3cf59585dcb9040964dde13cb0cae6ed3cd/lancedb-0.37.1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:c15c46f23cf6959c79fb93cdba2c76536cf784d3134386662da03dc6ccac3c26", size = 58474767, upload-time = "2026-08-10T10:41:45.675Z" }, + { url = "https://files.pythonhosted.org/packages/e5/5e/dac0fd9478a21685444f23e6ec937babf4ff48b1616b68e8137778d6ecb9/lancedb-0.37.1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:35d872d920cbfdc3771fbcd33f2c63bff6ff7203d6d2ba8fe4330e98eb859d12", size = 61666580, upload-time = "2026-08-10T10:41:49.535Z" }, + { url = "https://files.pythonhosted.org/packages/1f/fa/ee1cdb1e904872d75fa0ff44ad35cd23494e810299c33e68de0687de7a71/lancedb-0.37.1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:86597f4dbc51a33a07341550dc77d21a1ddd1f7539266eda5bb64c4f3bd11cca", size = 64802395, upload-time = "2026-08-10T10:41:52.969Z" }, + { url = "https://files.pythonhosted.org/packages/42/55/65c69307373b7f05dea38465b7d3836c86390f0ebc79b5c56d2c0919f229/lancedb-0.37.1-cp310-abi3-win_amd64.whl", hash = "sha256:488eca15361dfc34439500c9e2607c4fb2b8bf190fa1003bd54b1d6eb40e0316", size = 70994296, upload-time = "2026-08-10T10:41:56.609Z" }, ] [[package]]