diff --git a/CHANGELOG.md b/CHANGELOG.md index 5bcc7be9..355eabd9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,10 +6,16 @@ - `evaluations run --filter/-f CLAUSE`: SQL `WHERE` clause over document columns, applied to the retrieval benchmark's searches and to every capability search during QA. Recorded as `document_filter` in experiment metadata. - `mtrag_clapnq` / `mtrag_clapnq_rewrite` / `mtrag_clapnq_live` / `mtrag_clapnq_live_uncompacted` evaluation datasets: multi-turn QA with gold-prefix and live-session conversation replay, live arms with and without `EvidenceCompactionCapability`, Recall@k/nDCG@k retrieval metrics, eligibility-aware citation scoring, refusal precision/recall, per-turn tool-traffic attributes, and `citation_status` / `turn_citation_status` eval attributes. +- BTree indexes on `chunks.id`, `chunks.document_id` and `documents.id`, and a Bitmap index on `document_items.label`. Existing databases need `haiku-rag migrate`. + ### Changed - `import_documents` embeds chunks across the whole batch in one pass instead of per document. +### Fixed + +- `DocumentRepository.delete_all` recreated `document_items` from `DocumentItemRecord` instead of `get_document_items_arrow_schema()`, returning `picture_data` as `binary` rather than `large_binary`. + ### Removed - `wix` evaluation dataset and its reference config `evaluations/configs/wix.yaml`. diff --git a/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/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 50d4a73a..002283a1 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -12,7 +12,7 @@ from uuid import uuid4 import lancedb import pyarrow as pa -from lancedb.index import FTS, BTree, IvfPq +from lancedb.index import FTS, Bitmap, BTree, IvfPq from lancedb.pydantic import LanceModel, Vector from packaging.version import parse from pydantic import BaseModel, Field @@ -176,6 +176,59 @@ def get_document_items_arrow_schema() -> pa.Schema: return pa.schema(fields) +def index_specs(table_name: str) -> list[tuple[str, Bitmap | BTree | FTS]]: + """The index set each table carries.""" + match table_name: + case "documents": + return [("id", BTree())] + case "document_meta": + return [("id", BTree()), ("uri", BTree())] + case "chunks": + return [ + # Positions and stop words are required for phrase queries. + ("content_fts", FTS(with_position=True, remove_stop_words=False)), + ("id", BTree()), + ("document_id", BTree()), + ] + case "document_items": + return [ + ("document_id", BTree()), + ("position", BTree()), + ("self_ref", BTree()), + ("label", Bitmap()), + ] + case _: + return [] + + +async def ensure_indexes(table: lancedb.AsyncTable, table_name: str) -> list[str]: + """Create any declared index missing from a column. Returns the columns indexed. + + Matches on index type, not column coverage, so a BTree does not satisfy a + declared Bitmap. Never drops or converts an index it did not declare. + Re-creating is not free: `create_index(replace=True)` rebuilds. + """ + covering: dict[str, set[str]] = {} + for index in await table.list_indices(): + for column in index.columns: + covering.setdefault(column, set()).add(index.index_type) + + applied: list[str] = [] + for column, config in index_specs(table_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): id: str = Field(default="settings") settings: str = Field(default="{}") @@ -697,22 +750,17 @@ class Store: self.documents_table = await self.db.create_table( "documents", schema=get_documents_arrow_schema() ) + await ensure_indexes(self.documents_table, "documents") # Create or open document_meta table (mutable attributes kept out of the - # blob-bearing documents row). Indexed by document_id and uri — both are - # hot look-up keys (get_by_id, get_by_uri). + # blob-bearing documents row). if "document_meta" in existing_tables: self.document_meta_table = await self.db.open_table("document_meta") else: self.document_meta_table = await self.db.create_table( "document_meta", schema=DocumentMetaRecord ) - await self.document_meta_table.create_index( - "id", config=BTree(), replace=True - ) - await self.document_meta_table.create_index( - "uri", config=BTree(), replace=True - ) + await ensure_indexes(self.document_meta_table, "document_meta") # Create or open chunks table if "chunks" in existing_tables: @@ -721,12 +769,7 @@ class Store: self.chunks_table = await self.db.create_table( "chunks", schema=self.ChunkRecord ) - # Create FTS index on content_fts (contextualized content) for better search - await self.chunks_table.create_index( - "content_fts", - config=FTS(with_position=True, remove_stop_words=False), - replace=True, - ) + await ensure_indexes(self.chunks_table, "chunks") # Create or open document_items table if "document_items" in existing_tables: @@ -735,15 +778,7 @@ class Store: self.document_items_table = await self.db.create_table( "document_items", schema=get_document_items_arrow_schema() ) - await self.document_items_table.create_index( - "document_id", config=BTree(), replace=True - ) - await self.document_items_table.create_index( - "position", config=BTree(), replace=True - ) - await self.document_items_table.create_index( - "self_ref", config=BTree(), replace=True - ) + await ensure_indexes(self.document_items_table, "document_items") # Create or open settings table if "settings" in existing_tables: @@ -872,13 +907,7 @@ class Store: self.chunks_table = await self.db.create_table( "chunks", schema=self.ChunkRecord ) - - # Create FTS index on content_fts (contextualized content) for better search - await self.chunks_table.create_index( - "content_fts", - config=FTS(with_position=True, remove_stop_words=False), - replace=True, - ) + await ensure_indexes(self.chunks_table, "chunks") def close(self): """Close the database connection.""" diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 09cc3498..ecfe9f49 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -9,7 +9,7 @@ if TYPE_CHECKING: from lancedb.index import FTS from lancedb.rerankers import RRFReranker -from haiku.rag.store.engine import Store, query_to_pydantic +from haiku.rag.store.engine import Store, ensure_indexes, query_to_pydantic from haiku.rag.store.models.chunk import Chunk, SearchType from haiku.rag.utils import escape_sql_string @@ -187,12 +187,7 @@ class ChunkRepository: self.store.chunks_table = await self.store.db.create_table( "chunks", schema=self.store.ChunkRecord ) - # Create FTS index on content_fts (contextualized content) for better search - await self.store.chunks_table.create_index( - "content_fts", - config=FTS(with_position=True, remove_stop_words=False), - replace=True, - ) + await ensure_indexes(self.store.chunks_table, "chunks") async def delete_by_document_id(self, document_id: str) -> bool: """Delete all chunks for a document.""" diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 9bfd9b57..be242be3 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -3,12 +3,12 @@ from datetime import datetime from typing import overload from uuid import uuid4 -from lancedb.index import BTree - from haiku.rag.store.engine import ( DocumentMetaRecord, DocumentRecord, Store, + ensure_indexes, + get_document_items_arrow_schema, get_documents_arrow_schema, query_to_pydantic, ) @@ -407,23 +407,14 @@ class DocumentRepository: async def delete_all(self) -> None: """Delete all documents from the database.""" self.store._assert_writable() - from haiku.rag.store.engine import DocumentItemRecord # Delete all chunks and items first await self.chunk_repository.delete_all() await self.store.db.drop_table("document_items") self.store.document_items_table = await self.store.db.create_table( - "document_items", schema=DocumentItemRecord - ) - await self.store.document_items_table.create_index( - "document_id", config=BTree(), replace=True - ) - await self.store.document_items_table.create_index( - "position", config=BTree(), replace=True - ) - await self.store.document_items_table.create_index( - "self_ref", config=BTree(), replace=True + "document_items", schema=get_document_items_arrow_schema() ) + await ensure_indexes(self.store.document_items_table, "document_items") # Get count before deletion count = len( @@ -437,13 +428,9 @@ class DocumentRepository: self.store.documents_table = await self.store.db.create_table( "documents", schema=get_documents_arrow_schema() ) + await ensure_indexes(self.store.documents_table, "documents") await self.store.db.drop_table("document_meta") self.store.document_meta_table = await self.store.db.create_table( "document_meta", schema=DocumentMetaRecord ) - await self.store.document_meta_table.create_index( - "id", config=BTree(), replace=True - ) - await self.store.document_meta_table.create_index( - "uri", config=BTree(), replace=True - ) + await ensure_indexes(self.store.document_meta_table, "document_meta") diff --git a/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..f71b7001 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py @@ -0,0 +1,24 @@ +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: + """Add the declared indexes to a database created before 0.75.0. + + Rewrites no rows. Each index build reads the column it indexes. + """ + for table_name, table in store._tables().items(): + applied = await ensure_indexes(table, table_name) + if applied: + logger.info(f"Indexed {table_name}: {', '.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 new file mode 100644 index 00000000..3fc2c542 --- /dev/null +++ b/tests/store/test_indexes.py @@ -0,0 +1,131 @@ +import pyarrow as pa +import pytest +from lancedb.index import BTree + +from haiku.rag.store.engine import Store, ensure_indexes +from haiku.rag.store.models import Document +from haiku.rag.store.repositories.document import DocumentRepository + +EXPECTED_INDEXED_COLUMNS = { + "documents": {"id"}, + "document_meta": {"id", "uri"}, + "chunks": {"content_fts", "id", "document_id"}, + "document_items": {"document_id", "position", "self_ref", "label"}, +} + + +async def _indexed_columns(table) -> set[str]: + return {column for index in await table.list_indices() for column in index.columns} + + +async def _index_type(table, column: str) -> str | None: + for index in await table.list_indices(): + if column in index.columns: + return index.index_type + return None + + +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.""" + async with Store(temp_db_path, create=True) as store: + for name, table in store._tables().items(): + expected = EXPECTED_INDEXED_COLUMNS.get(name, set()) + assert await _indexed_columns(table) == expected, name + + +@pytest.mark.asyncio +async def test_ensure_indexes_skips_existing_instead_of_rebuilding(temp_db_path): + """A second pass must not rebuild: replace=True writes a new version.""" + async with Store(temp_db_path, create=True) as store: + table = store.chunks_table + version_before = await table.version() + + await ensure_indexes(table, "chunks") + + assert await table.version() == version_before + assert await _indexed_columns(table) == EXPECTED_INDEXED_COLUMNS["chunks"] + + +@pytest.mark.asyncio +async def test_ensure_indexes_corrects_an_index_of_the_wrong_type(temp_db_path): + """A 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) + assert await _index_type(table, "label") == "BTree" + + await ensure_indexes(table, "document_items") + + assert await _index_type(table, "label") == "Bitmap" + assert ( + await _indexed_columns(table) == EXPECTED_INDEXED_COLUMNS["document_items"] + ) + + +@pytest.mark.asyncio +async def test_ensure_indexes_adds_the_declared_type_beside_a_custom_index( + temp_db_path, +): + """A custom-named index neither satisfies the check nor is destroyed.""" + async with Store(temp_db_path, create=True) as store: + table = store.document_items_table + await table.drop_index("label_idx") + 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, +): + """An index we did not declare survives, even on a declared column.""" + async with Store(temp_db_path, create=True) as store: + table = store.document_items_table + await table.create_index("label", config=BTree(), name="operator_label") + + 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): + """Recreated tables come back with the full index set.""" + async with Store(temp_db_path, create=True) as store: + repo = DocumentRepository(store) + await repo.create(Document(content="A document")) + + await repo.delete_all() + + for name, table in store._tables().items(): + expected = EXPECTED_INDEXED_COLUMNS.get(name, set()) + assert await _indexed_columns(table) == expected, name + + +@pytest.mark.asyncio +async def test_delete_all_keeps_picture_data_as_large_binary(temp_db_path): + """picture_data must survive delete_all as large_binary, not binary.""" + async with Store(temp_db_path, create=True) as store: + repo = DocumentRepository(store) + await repo.create(Document(content="A document")) + + await repo.delete_all() + + schema = await store.document_items_table.schema() + assert schema.field("picture_data").type == pa.large_binary() 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..3b996d51 --- /dev/null +++ b/tests/store/test_v0_75_0_migration.py @@ -0,0 +1,105 @@ +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 + +# The indexes a pre-0.75.0 database lacked. +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): + """A second run must not re-index: replace=True rebuilds.""" + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await _apply_index_hot_lookup_keys(store) + versions = {name: await t.version() for name, t in store._tables().items()} + + 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 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 + ) + + 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): + """Indexes haiku.rag never declared are not dropped.""" + async with Store(temp_db_path, create=True, skip_migration_check=True) as store: + await store.document_meta_table.create_index( + "title", config=BTree(), replace=True + ) + + 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..e6e379a2 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -993,10 +993,15 @@ 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: its writes would land inside the measured window. + """ + 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" },