danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #591 from ggozad/fix/fts-index-coverage

Browse source 
Guard, prevent, and repair an FTS index that covers no rows
This commit is contained in:
Yiorgis Gozadinos 2026-08-31 11:05:30 -05:00 committed by GitHub
commit b0818f4d3d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 687 additions and 44 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
CHANGELOG.md
View file

@ -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.

3
docs/cli.md
View file

@ -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

2
docs/configuration/storage.md
View file

@ -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

4
haiku_rag_slim/haiku/rag/client/rebuild.py
View file

@ -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

55
haiku_rag_slim/haiku/rag/doctor.py
View file

@ -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))

58
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -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"

26
haiku_rag_slim/haiku/rag/store/schema.py
View file

@ -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

2
haiku_rag_slim/pyproject.toml
View file

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

220
tests/store/test_indexes.py
View file

@ -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"]

14
tests/store/test_v0_75_0_migration.py
View file

@ -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)

198
tests/test_chunk.py
View file

@ -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()

14
tests/test_client.py
View file

@ -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

44
tests/test_doctor.py
View file

@ -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)

3
tests/test_filter.py
View file

@ -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'")

2
tests/test_mcp.py
View file

@ -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:

64
tests/test_rebuild.py
View file

@ -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

12
uv.lock
View file

@ -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]]