diff --git a/CHANGELOG.md b/CHANGELOG.md index ce46cde0..495a23c1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ ### Changed +- `haiku.rag.store.engine` split: table records, Arrow schemas, `index_specs`, `ensure_indexes`, `REQUIRED_TABLES` and `query_to_pydantic` are in `haiku.rag.store.schema`; `gather_database_info`, `get_database_stats`, `DatabaseInfo` and its result models are in `haiku.rag.store.info`. `Store` keeps lifecycle, locks, migration coordination, vacuuming and tags. Update imports. - Source adapters moved from `haiku.rag.ingester.sources` to `haiku.rag.sources`: `FetchResult`, `SourceEvent`, `SourceEventKind`, `RevisionSnapshot`, `Source` and the `FSSource`/`HTTPSource`/`S3Source`/`WebDAVSource` adapters. One-shot client ingestion uses them too, so they were never ingester-only. Update imports; the `haiku.rag.sources` plugin entry-point group is unchanged. - `HaikuRAG.convert(url)` fetches through `HTTPSource`, the same adapter the ingester uses, instead of its own httpx client. `_write_fetch_body` moved from `client.documents` to `client.processing`. - Chunk embedding is owned by the persistence funnels: `create_document`, `update_document` and source ingestion no longer embed eagerly before handing chunks to a check that would embed them anyway. The `document.embed` span moved onto `ensure_chunks_embedded`, so every path is instrumented rather than only ingest. diff --git a/app/backend/main.py b/app/backend/main.py index ec67be0b..057af35f 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -172,7 +172,7 @@ async def db_info(_: Request) -> JSONResponse: } ) - from haiku.rag.store.engine import get_database_stats + from haiku.rag.store.info import get_database_stats client = await get_client() stats = await get_database_stats(client.store.db) diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 546f9e17..8f800cb6 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -64,7 +64,7 @@ class HaikuRAGApp: # pragma: no cover async def info(self): """Display read-only information about the database without modifying it.""" - from haiku.rag.store.engine import gather_database_info + from haiku.rag.store.info import gather_database_info # Basic: show path/URI self.console.print("[bold]haiku.rag database info[/bold]") diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index 9193123e..410aebd2 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -14,11 +14,11 @@ from haiku.rag.client.documents import ( ) from haiku.rag.converters import get_converter from haiku.rag.store.compression import compress_docling_split -from haiku.rag.store.engine import ChunkRecordBase 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 if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument @@ -519,7 +519,7 @@ async def _flush_rebuild_batch( document. Used by RECHUNK and FULL modes after the chunks table has been cleared. """ - from haiku.rag.store.engine import DocumentMetaRecord, DocumentRecord + from haiku.rag.store.schema import DocumentMetaRecord, DocumentRecord if not documents: return diff --git a/haiku_rag_slim/haiku/rag/doctor.py b/haiku_rag_slim/haiku/rag/doctor.py index 9355a26b..49d53328 100644 --- a/haiku_rag_slim/haiku/rag/doctor.py +++ b/haiku_rag_slim/haiku/rag/doctor.py @@ -11,13 +11,10 @@ from pydantic import BaseModel, Field from haiku.rag.config import AppConfig from haiku.rag.config.models import DuplicateDetectionConfig -from haiku.rag.store.engine import ( - REQUIRED_TABLES, - Store, - connect_lancedb, - get_database_stats, -) +from haiku.rag.store.engine import Store, connect_lancedb +from haiku.rag.store.info import get_database_stats from haiku.rag.store.repositories.settings import SettingsRepository +from haiku.rag.store.schema import REQUIRED_TABLES 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. diff --git a/haiku_rag_slim/haiku/rag/ingester/api/routes/database.py b/haiku_rag_slim/haiku/rag/ingester/api/routes/database.py index 886f742b..0750db46 100644 --- a/haiku_rag_slim/haiku/rag/ingester/api/routes/database.py +++ b/haiku_rag_slim/haiku/rag/ingester/api/routes/database.py @@ -1,7 +1,7 @@ from fastapi import APIRouter, Depends, HTTPException, status from haiku.rag.ingester.api.server import APIState, get_state -from haiku.rag.store.engine import DatabaseInfo, gather_database_info +from haiku.rag.store.info import DatabaseInfo, gather_database_info router = APIRouter(tags=["database"]) diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py index 539b21ca..4a00e379 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py @@ -64,11 +64,8 @@ class InfoModal(ModalScreen): async def on_mount(self) -> None: """Load and display database info.""" - from haiku.rag.store.engine import ( - ConnectionMode, - connect_lancedb, - get_database_stats, - ) + from haiku.rag.store.engine import ConnectionMode, connect_lancedb + from haiku.rag.store.info import get_database_stats lines: list[str] = [] diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 240cfe13..845fd948 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -8,39 +8,30 @@ from datetime import UTC, datetime, timedelta from enum import Enum from importlib import metadata from pathlib import Path -from typing import TYPE_CHECKING, Any, cast -from uuid import uuid4 +from typing import Any import lancedb -import pyarrow as pa -from lancedb.index import FTS, Bitmap, BTree, IvfPq -from lancedb.pydantic import LanceModel, Vector +from lancedb.index import IvfPq from packaging.version import parse -from pydantic import BaseModel, Field from haiku.rag.config import AppConfig, get_config from haiku.rag.embeddings import get_embedder from haiku.rag.store.exceptions import MigrationRequiredError, ReadOnlyError - -if TYPE_CHECKING: - from lancedb.query import AsyncQueryBase +from haiku.rag.store.schema import ( + REQUIRED_TABLES, + ChunkRecordBase, + DocumentMetaRecord, + SettingsRecord, + create_chunk_model, + ensure_indexes, + get_document_items_arrow_schema, + get_documents_arrow_schema, + query_to_pydantic, +) logger = logging.getLogger(__name__) -async def query_to_pydantic[T: LanceModel]( - query: "AsyncQueryBase", model: type[T] -) -> list[T]: - """Typed wrapper around AsyncQueryBase.to_pydantic. - - The upstream stub annotates `.to_pydantic()` as returning `list[LanceModel]` - regardless of the concrete model passed in. This helper narrows the return - type to the concrete model so attribute access on the results type-checks - at call sites without needing per-line cast / ignore comments. - """ - return cast("list[T]", await query.to_pydantic(model)) - - class ConnectionMode(Enum): LOCAL = "local" CLOUD = "cloud" @@ -108,177 +99,11 @@ async def connect_lancedb( return await lancedb.connect_async(db_path.absolute(), **kwargs) -class DocumentRecord(LanceModel): - id: str = Field(default_factory=lambda: str(uuid4())) - content: str - docling_document: bytes | None = None - docling_pages: bytes | None = None - docling_version: str | None = None - - -class DocumentMetaRecord(LanceModel): - """Mutable, lightweight document attributes, kept separate from the - write-once content/blobs in `documents`. Updating these (metadata, title, - source_revision) must not rewrite the multi-MB docling row.""" - - id: str - uri: str | None = None - title: str | None = None - metadata: str = Field(default="{}") - created_at: str = Field(default_factory=lambda: "") - updated_at: str = Field(default_factory=lambda: "") - - -def get_documents_arrow_schema() -> pa.Schema: - """Generate Arrow schema for documents table with large_binary for docling_document. - - LanceDB maps Python `bytes` to Arrow's `binary` type, which uses 32-bit offsets - and is limited to ~2GB per column in a fragment. When many large documents - (with embedded page images) are grouped in a single fragment, this limit is - exceeded, causing "byte array offset overflow" panics. - - This function overrides the default mapping to use `large_binary` instead, - which has 64-bit offsets and no practical size limit. - """ - base_schema = DocumentRecord.to_arrow_schema() - large_binary_columns = {"docling_document", "docling_pages"} - fields = [] - for field in base_schema: - if field.name in large_binary_columns: - fields.append(pa.field(field.name, pa.large_binary())) - else: - fields.append(field) - return pa.schema(fields) - - -class ChunkRecordBase(LanceModel): - """Static base for ChunkRecord — declares the fields so attribute access - and constructor calls type-check. The concrete `vector` field is overridden - by create_chunk_model() with a Vector(dim) whose fixed-size-list dimension - is only known at runtime. - """ - - id: str = Field(default_factory=lambda: str(uuid4())) - document_id: str - content: str - content_fts: str = Field(default="") - metadata: str = Field(default="{}") - order: int = Field(default=0) - vector: list[float] = Field(default_factory=list) - - -def create_chunk_model(vector_dim: int) -> type[ChunkRecordBase]: - """Create a ChunkRecord model with the specified vector dimension.""" - - class ChunkRecord(ChunkRecordBase): - vector: Vector(vector_dim) = Field(default_factory=lambda: [0.0] * vector_dim) # type: ignore - - return ChunkRecord - - -class DocumentItemRecord(LanceModel): - document_id: str - position: int - self_ref: str - label: str = Field(default="") - text: str = Field(default="") - page_numbers: str = Field(default="[]") - picture_data: bytes | None = None - heading_level: int = Field(default=0) - tree_depth: int = Field(default=0) - - -def get_document_items_arrow_schema() -> pa.Schema: - """Generate Arrow schema for document_items with large_binary for picture_data. - - LanceDB maps Python `bytes` to Arrow's `binary` type, which uses 32-bit offsets - and is limited to ~2GB per column in a fragment. Many embedded picture PNGs in - one fragment can exceed that limit. `large_binary` uses 64-bit offsets and has - no practical size limit — same reasoning as `docling_document` on the - documents table. - """ - base_schema = DocumentItemRecord.to_arrow_schema() - large_binary_columns = {"picture_data"} - fields = [] - for field in base_schema: - if field.name in large_binary_columns: - fields.append(pa.field(field.name, pa.large_binary())) - else: - fields.append(field) - return pa.schema(fields) - - def _stored_vector_dim(settings: dict) -> int | None: """The vector dimension a database's chunks were written at.""" return settings.get("embeddings", {}).get("model", {}).get("vector_dim") -def index_specs(table_name: str) -> list[tuple[str, Bitmap | BTree | FTS]]: - """The index set each table carries.""" - match table_name: - 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="{}") - - -REQUIRED_TABLES: tuple[str, ...] = ( - "documents", - "document_meta", - "chunks", - "document_items", - "settings", -) - # Keeps the vacuum cleanup cutoff safely older than the oldest tagged # version; guards against timestamp precision at the boundary. TAG_RETENTION_MARGIN = timedelta(seconds=1) @@ -342,157 +167,6 @@ class TagInfo: return not self.missing_tables -async def get_database_stats(db: lancedb.AsyncConnection) -> dict: - """Collect stats for every haiku.rag table on the connection. - - Missing tables return ``{"exists": False}``. Present tables include - ``num_rows``, ``total_bytes``, and ``num_versions``. The ``chunks`` - entry additionally reports vector index status and, when an index - exists, ``num_indexed_rows`` and ``num_unindexed_rows``. - """ - existing = set((await db.list_tables()).tables) - stats: dict = {} - tables: dict = {} - - for name in REQUIRED_TABLES: - if name not in existing: - stats[name] = {"exists": False} - continue - tbl = await db.open_table(name) - tables[name] = tbl - # lancedb's .stats() stub claims TableStatistics but returns a plain dict at runtime. - tbl_stats: dict = await tbl.stats() # type: ignore[assignment] # ty: ignore[invalid-assignment] - stats[name] = { - "exists": True, - "num_rows": tbl_stats.get("num_rows", 0), - "total_bytes": tbl_stats.get("total_bytes", 0), - "num_versions": len(await tbl.list_versions()), - } - - if stats["chunks"]["exists"]: - chunks_tbl = tables["chunks"] - indices = await chunks_tbl.list_indices() - has_vector_index = any("vector" in str(idx).lower() for idx in indices) - stats["chunks"]["has_vector_index"] = has_vector_index - if has_vector_index: - index_stats = await chunks_tbl.index_stats("vector_idx") - if index_stats is not None: - stats["chunks"]["num_indexed_rows"] = index_stats.num_indexed_rows - stats["chunks"]["num_unindexed_rows"] = index_stats.num_unindexed_rows - - return stats - - -class EmbeddingsInfo(BaseModel): - provider: str = "unknown" - name: str = "unknown" - vector_dim: int | None = None - - -class TableInfo(BaseModel): - name: str - exists: bool - num_rows: int = 0 - total_bytes: int = 0 - num_versions: int = 0 - - -class VectorIndexInfo(BaseModel): - exists: bool = False - indexed_rows: int = 0 - unindexed_rows: int = 0 - - -class PendingMigration(BaseModel): - version: str - description: str = "" - - -class DatabaseInfo(BaseModel): - """Structured snapshot of a haiku.rag database, shared by the `info` CLI - command and the ingester control plane. Read-only; gathered without - opening a Store.""" - - path: str - exists: bool - stored_version: str = "unknown" - embeddings: EmbeddingsInfo = Field(default_factory=EmbeddingsInfo) - tables: list[TableInfo] = Field(default_factory=list) - vector_index: VectorIndexInfo = Field(default_factory=VectorIndexInfo) - pending_migrations: list[PendingMigration] = Field(default_factory=list) - packages: dict[str, str] = Field(default_factory=dict) - - -async def gather_database_info(config: AppConfig, db_path: Path) -> DatabaseInfo: - """Collect read-only database state without going through Store, so a - database missing tables (e.g. pre-migration) still reports what it can.""" - from haiku.rag.store.upgrades import get_pending_upgrades - from haiku.rag.utils import get_package_versions - - display_path = config.lancedb.uri or str(db_path) - - db = await connect_lancedb(config, db_path) - stats = await get_database_stats(db) - - if not any(entry["exists"] for entry in stats.values()): - return DatabaseInfo(path=display_path, exists=False) - - stored_version = "unknown" - embeddings = EmbeddingsInfo() - if stats["settings"]["exists"]: - settings_tbl = await db.open_table("settings") - rows = ( - await settings_tbl.query().where("id = 'settings'").limit(1).to_arrow() - ).to_pylist() - if rows: - raw = rows[0].get("settings") or "{}" - data = json.loads(raw) if isinstance(raw, str) else (raw or {}) - stored_version = str(data.get("version", "unknown")) - model = data.get("embeddings", {}).get("model", {}) - embeddings = EmbeddingsInfo( - provider=model.get("provider", "unknown"), - name=model.get("name", "unknown"), - vector_dim=model.get("vector_dim"), - ) - - tables = [ - TableInfo( - name=name, - exists=stats[name]["exists"], - num_rows=stats[name].get("num_rows", 0), - total_bytes=stats[name].get("total_bytes", 0), - num_versions=stats[name].get("num_versions", 0), - ) - for name in ("documents", "document_meta", "chunks", "document_items") - ] - - vector_index = VectorIndexInfo() - if stats["chunks"]["exists"] and stats["chunks"].get("has_vector_index"): - vector_index = VectorIndexInfo( - exists=True, - indexed_rows=stats["chunks"].get("num_indexed_rows", 0), - unindexed_rows=stats["chunks"].get("num_unindexed_rows", 0), - ) - - pending = ( - get_pending_upgrades(stored_version) if stored_version != "unknown" else [] - ) - - return DatabaseInfo( - path=display_path, - exists=True, - stored_version=stored_version, - embeddings=embeddings, - tables=tables, - vector_index=vector_index, - pending_migrations=[ - PendingMigration(version=step.version, description=step.description or "") - for step in pending - ], - packages=get_package_versions(), - ) - - class Store: def __init__( self, diff --git a/haiku_rag_slim/haiku/rag/store/info.py b/haiku_rag_slim/haiku/rag/store/info.py new file mode 100644 index 00000000..3c7b1dc4 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/store/info.py @@ -0,0 +1,166 @@ +"""Read-only database information and its result models. + +Every function here only reads: it reports what a database contains and how it +is configured, and never writes a table or a version. +""" + +import json +from pathlib import Path + +import lancedb +from pydantic import BaseModel, Field + +from haiku.rag.config import AppConfig +from haiku.rag.store.engine import connect_lancedb +from haiku.rag.store.schema import REQUIRED_TABLES + + +async def get_database_stats(db: lancedb.AsyncConnection) -> dict: + """Collect stats for every haiku.rag table on the connection. + + Missing tables return ``{"exists": False}``. Present tables include + ``num_rows``, ``total_bytes``, and ``num_versions``. The ``chunks`` + entry additionally reports vector index status and, when an index + exists, ``num_indexed_rows`` and ``num_unindexed_rows``. + """ + existing = set((await db.list_tables()).tables) + stats: dict = {} + tables: dict = {} + + for name in REQUIRED_TABLES: + if name not in existing: + stats[name] = {"exists": False} + continue + tbl = await db.open_table(name) + tables[name] = tbl + # lancedb's .stats() stub claims TableStatistics but returns a plain dict at runtime. + tbl_stats: dict = await tbl.stats() # type: ignore[assignment] # ty: ignore[invalid-assignment] + stats[name] = { + "exists": True, + "num_rows": tbl_stats.get("num_rows", 0), + "total_bytes": tbl_stats.get("total_bytes", 0), + "num_versions": len(await tbl.list_versions()), + } + + if stats["chunks"]["exists"]: + chunks_tbl = tables["chunks"] + indices = await chunks_tbl.list_indices() + has_vector_index = any("vector" in str(idx).lower() for idx in indices) + stats["chunks"]["has_vector_index"] = has_vector_index + if has_vector_index: + index_stats = await chunks_tbl.index_stats("vector_idx") + if index_stats is not None: + stats["chunks"]["num_indexed_rows"] = index_stats.num_indexed_rows + stats["chunks"]["num_unindexed_rows"] = index_stats.num_unindexed_rows + + return stats + + +class EmbeddingsInfo(BaseModel): + provider: str = "unknown" + name: str = "unknown" + vector_dim: int | None = None + + +class TableInfo(BaseModel): + name: str + exists: bool + num_rows: int = 0 + total_bytes: int = 0 + num_versions: int = 0 + + +class VectorIndexInfo(BaseModel): + exists: bool = False + indexed_rows: int = 0 + unindexed_rows: int = 0 + + +class PendingMigration(BaseModel): + version: str + description: str = "" + + +class DatabaseInfo(BaseModel): + """Structured snapshot of a haiku.rag database, shared by the `info` CLI + command and the ingester control plane. Read-only; gathered without + opening a Store.""" + + path: str + exists: bool + stored_version: str = "unknown" + embeddings: EmbeddingsInfo = Field(default_factory=EmbeddingsInfo) + tables: list[TableInfo] = Field(default_factory=list) + vector_index: VectorIndexInfo = Field(default_factory=VectorIndexInfo) + pending_migrations: list[PendingMigration] = Field(default_factory=list) + packages: dict[str, str] = Field(default_factory=dict) + + +async def gather_database_info(config: AppConfig, db_path: Path) -> DatabaseInfo: + """Collect read-only database state without going through Store, so a + database missing tables (e.g. pre-migration) still reports what it can.""" + from haiku.rag.store.upgrades import get_pending_upgrades + from haiku.rag.utils import get_package_versions + + display_path = config.lancedb.uri or str(db_path) + + db = await connect_lancedb(config, db_path) + stats = await get_database_stats(db) + + if not any(entry["exists"] for entry in stats.values()): + return DatabaseInfo(path=display_path, exists=False) + + stored_version = "unknown" + embeddings = EmbeddingsInfo() + if stats["settings"]["exists"]: + settings_tbl = await db.open_table("settings") + rows = ( + await settings_tbl.query().where("id = 'settings'").limit(1).to_arrow() + ).to_pylist() + if rows: + raw = rows[0].get("settings") or "{}" + data = json.loads(raw) if isinstance(raw, str) else (raw or {}) + stored_version = str(data.get("version", "unknown")) + model = data.get("embeddings", {}).get("model", {}) + embeddings = EmbeddingsInfo( + provider=model.get("provider", "unknown"), + name=model.get("name", "unknown"), + vector_dim=model.get("vector_dim"), + ) + + tables = [ + TableInfo( + name=name, + exists=stats[name]["exists"], + num_rows=stats[name].get("num_rows", 0), + total_bytes=stats[name].get("total_bytes", 0), + num_versions=stats[name].get("num_versions", 0), + ) + for name in ("documents", "document_meta", "chunks", "document_items") + ] + + vector_index = VectorIndexInfo() + if stats["chunks"]["exists"] and stats["chunks"].get("has_vector_index"): + vector_index = VectorIndexInfo( + exists=True, + indexed_rows=stats["chunks"].get("num_indexed_rows", 0), + unindexed_rows=stats["chunks"].get("num_unindexed_rows", 0), + ) + + pending = ( + get_pending_upgrades(stored_version) if stored_version != "unknown" else [] + ) + + return DatabaseInfo( + path=display_path, + exists=True, + stored_version=stored_version, + embeddings=embeddings, + tables=tables, + vector_index=vector_index, + pending_migrations=[ + PendingMigration(version=step.version, description=step.description or "") + for step in pending + ], + packages=get_package_versions(), + ) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index d3005128..dbb7a425 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -9,8 +9,9 @@ if TYPE_CHECKING: from lancedb.index import FTS from lancedb.rerankers import RRFReranker -from haiku.rag.store.engine import Store, ensure_indexes, query_to_pydantic +from haiku.rag.store.engine import Store from haiku.rag.store.models.chunk import Chunk, SearchType +from haiku.rag.store.schema import ensure_indexes, query_to_pydantic from haiku.rag.utils import escape_sql_string logger = logging.getLogger(__name__) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index be242be3..4ee87c3b 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -3,16 +3,16 @@ from datetime import datetime from typing import overload from uuid import uuid4 -from haiku.rag.store.engine import ( +from haiku.rag.store.engine import Store +from haiku.rag.store.models.document import Document +from haiku.rag.store.schema import ( DocumentMetaRecord, DocumentRecord, - Store, ensure_indexes, get_document_items_arrow_schema, get_documents_arrow_schema, query_to_pydantic, ) -from haiku.rag.store.models.document import Document from haiku.rag.utils import escape_sql_string # Ids per `id IN (...)` content lookup. Keeps the filter string bounded on an diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py index 4f293498..676ed862 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py @@ -1,8 +1,9 @@ import json from collections.abc import Mapping, Sequence -from haiku.rag.store.engine import DocumentItemRecord, Store +from haiku.rag.store.engine import Store from haiku.rag.store.models.document_item import DocumentItem +from haiku.rag.store.schema import DocumentItemRecord from haiku.rag.utils import escape_sql_string # Per-item metadata columns. The payload column ``picture_data`` is fetched diff --git a/haiku_rag_slim/haiku/rag/store/repositories/settings.py b/haiku_rag_slim/haiku/rag/store/repositories/settings.py index e38c17cf..351f37f5 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/settings.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/settings.py @@ -1,7 +1,8 @@ import json import logging -from haiku.rag.store.engine import SettingsRecord, Store, query_to_pydantic +from haiku.rag.store.engine import Store +from haiku.rag.store.schema import SettingsRecord, query_to_pydantic logger = logging.getLogger(__name__) diff --git a/haiku_rag_slim/haiku/rag/store/schema.py b/haiku_rag_slim/haiku/rag/store/schema.py new file mode 100644 index 00000000..e9832ad4 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/store/schema.py @@ -0,0 +1,198 @@ +"""Table records, Arrow schemas and index specifications. + +Describes what the tables are; nothing here opens a connection, mutates a +table, or imports a client layer. +""" + +import logging +from typing import cast +from uuid import uuid4 + +import lancedb +import pyarrow as pa +from lancedb.index import FTS, Bitmap, BTree +from lancedb.pydantic import LanceModel, Vector +from lancedb.query import AsyncQueryBase +from pydantic import Field + +logger = logging.getLogger(__name__) + + +async def query_to_pydantic[T: LanceModel]( + query: "AsyncQueryBase", model: type[T] +) -> list[T]: + """Typed wrapper around AsyncQueryBase.to_pydantic. + + The upstream stub annotates `.to_pydantic()` as returning `list[LanceModel]` + regardless of the concrete model passed in. This helper narrows the return + type to the concrete model so attribute access on the results type-checks + at call sites without needing per-line cast / ignore comments. + """ + return cast("list[T]", await query.to_pydantic(model)) + + +class DocumentRecord(LanceModel): + id: str = Field(default_factory=lambda: str(uuid4())) + content: str + docling_document: bytes | None = None + docling_pages: bytes | None = None + docling_version: str | None = None + + +class DocumentMetaRecord(LanceModel): + """Mutable, lightweight document attributes, kept separate from the + write-once content/blobs in `documents`. Updating these (metadata, title, + source_revision) must not rewrite the multi-MB docling row.""" + + id: str + uri: str | None = None + title: str | None = None + metadata: str = Field(default="{}") + created_at: str = Field(default_factory=lambda: "") + updated_at: str = Field(default_factory=lambda: "") + + +def get_documents_arrow_schema() -> pa.Schema: + """Generate Arrow schema for documents table with large_binary for docling_document. + + LanceDB maps Python `bytes` to Arrow's `binary` type, which uses 32-bit offsets + and is limited to ~2GB per column in a fragment. When many large documents + (with embedded page images) are grouped in a single fragment, this limit is + exceeded, causing "byte array offset overflow" panics. + + This function overrides the default mapping to use `large_binary` instead, + which has 64-bit offsets and no practical size limit. + """ + base_schema = DocumentRecord.to_arrow_schema() + large_binary_columns = {"docling_document", "docling_pages"} + fields = [] + for field in base_schema: + if field.name in large_binary_columns: + fields.append(pa.field(field.name, pa.large_binary())) + else: + fields.append(field) + return pa.schema(fields) + + +class ChunkRecordBase(LanceModel): + """Static base for ChunkRecord — declares the fields so attribute access + and constructor calls type-check. The concrete `vector` field is overridden + by create_chunk_model() with a Vector(dim) whose fixed-size-list dimension + is only known at runtime. + """ + + id: str = Field(default_factory=lambda: str(uuid4())) + document_id: str + content: str + content_fts: str = Field(default="") + metadata: str = Field(default="{}") + order: int = Field(default=0) + vector: list[float] = Field(default_factory=list) + + +def create_chunk_model(vector_dim: int) -> type[ChunkRecordBase]: + """Create a ChunkRecord model with the specified vector dimension.""" + + class ChunkRecord(ChunkRecordBase): + vector: Vector(vector_dim) = Field(default_factory=lambda: [0.0] * vector_dim) # type: ignore + + return ChunkRecord + + +class DocumentItemRecord(LanceModel): + document_id: str + position: int + self_ref: str + label: str = Field(default="") + text: str = Field(default="") + page_numbers: str = Field(default="[]") + picture_data: bytes | None = None + heading_level: int = Field(default=0) + tree_depth: int = Field(default=0) + + +def get_document_items_arrow_schema() -> pa.Schema: + """Generate Arrow schema for document_items with large_binary for picture_data. + + LanceDB maps Python `bytes` to Arrow's `binary` type, which uses 32-bit offsets + and is limited to ~2GB per column in a fragment. Many embedded picture PNGs in + one fragment can exceed that limit. `large_binary` uses 64-bit offsets and has + no practical size limit — same reasoning as `docling_document` on the + documents table. + """ + base_schema = DocumentItemRecord.to_arrow_schema() + large_binary_columns = {"picture_data"} + fields = [] + for field in base_schema: + if field.name in large_binary_columns: + fields.append(pa.field(field.name, pa.large_binary())) + else: + fields.append(field) + 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="{}") + + +REQUIRED_TABLES: tuple[str, ...] = ( + "documents", + "document_meta", + "chunks", + "document_items", + "settings", +) diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py index c5df800f..6e2ffa91 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py @@ -3,7 +3,8 @@ import logging import pyarrow as pa from lancedb.index import BTree -from haiku.rag.store.engine import DocumentItemRecord, Store +from haiku.rag.store.engine import Store +from haiku.rag.store.schema import DocumentItemRecord from haiku.rag.store.upgrades import Upgrade from haiku.rag.utils import escape_sql_string diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py index f0de5aaa..b44251f3 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py @@ -1,7 +1,8 @@ import logging import shutil -from haiku.rag.store.engine import DocumentMetaRecord, Store +from haiku.rag.store.engine import Store +from haiku.rag.store.schema import DocumentMetaRecord from haiku.rag.store.upgrades import Upgrade logger = logging.getLogger(__name__) diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py index f71b7001..975332a5 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py @@ -1,6 +1,7 @@ import logging -from haiku.rag.store.engine import Store, ensure_indexes +from haiku.rag.store.engine import Store +from haiku.rag.store.schema import ensure_indexes from haiku.rag.store.upgrades import Upgrade logger = logging.getLogger(__name__) diff --git a/tests/store/test_database_info.py b/tests/store/test_database_info.py index 21bd7bb1..eecb2896 100644 --- a/tests/store/test_database_info.py +++ b/tests/store/test_database_info.py @@ -5,7 +5,8 @@ from lancedb.pydantic import LanceModel, Vector from pydantic import Field from haiku.rag.config.models import AppConfig -from haiku.rag.store.engine import DocumentItemRecord, gather_database_info +from haiku.rag.store.info import gather_database_info +from haiku.rag.store.schema import DocumentItemRecord class _SettingsRecord(LanceModel): diff --git a/tests/store/test_document_items.py b/tests/store/test_document_items.py index 6104f2fc..0f1412d4 100644 --- a/tests/store/test_document_items.py +++ b/tests/store/test_document_items.py @@ -483,7 +483,7 @@ class TestDocumentItemMigration: async def test_migration_populates_items_for_existing_documents(self, temp_db_path): """Test that the v0.40.0 migration populates items for pre-existing documents.""" from haiku.rag.store.compression import compress_docling_split - from haiku.rag.store.engine import DocumentRecord + from haiku.rag.store.schema import DocumentRecord from haiku.rag.store.upgrades.v0_40_0 import _apply_populate_document_items docling_doc = _make_docling_doc() @@ -529,7 +529,7 @@ class TestDocumentItemMigration: async def test_migration_skips_documents_without_docling(self, temp_db_path): """Test that migration handles documents without docling data.""" - from haiku.rag.store.engine import DocumentRecord + from haiku.rag.store.schema import DocumentRecord from haiku.rag.store.upgrades.v0_40_0 import _apply_populate_document_items async with Store(temp_db_path, create=True, skip_migration_check=True) as store: @@ -742,7 +742,7 @@ class TestPictureDataMigrationBackfill: import json from haiku.rag.store.compression import compress_json, decompress_json - from haiku.rag.store.engine import DocumentItemRecord, DocumentRecord + from haiku.rag.store.schema import DocumentItemRecord, DocumentRecord from haiku.rag.store.upgrades.v0_45_0 import _apply_extract_picture_bytes fake_png = b"\x89PNG\r\n\x1a\nlegacy-picture-bytes-for-test" diff --git a/tests/store/test_document_meta_split.py b/tests/store/test_document_meta_split.py index 4da4941a..51b61d85 100644 --- a/tests/store/test_document_meta_split.py +++ b/tests/store/test_document_meta_split.py @@ -151,7 +151,7 @@ async def test_batch_create_rolls_back_meta_on_failure(temp_db_path, monkeypatch async def test_get_by_uri_with_orphan_meta_returns_none(temp_db_path): """Defensive: a document_meta row whose documents row is missing (an invariant violation) resolves to None, not a half-hydrated document.""" - from haiku.rag.store.engine import DocumentMetaRecord + from haiku.rag.store.schema import DocumentMetaRecord async with Store(temp_db_path, create=True, skip_migration_check=True) as store: repo = DocumentRepository(store) diff --git a/tests/store/test_engine.py b/tests/store/test_engine.py index 89d3ef65..229c9e5b 100644 --- a/tests/store/test_engine.py +++ b/tests/store/test_engine.py @@ -1,6 +1,7 @@ import pytest -from haiku.rag.store.engine import Store, get_database_stats +from haiku.rag.store.engine import Store +from haiku.rag.store.info import get_database_stats class TestGetDatabaseStats: diff --git a/tests/store/test_indexes.py b/tests/store/test_indexes.py index 3fc2c542..0f9451f6 100644 --- a/tests/store/test_indexes.py +++ b/tests/store/test_indexes.py @@ -2,9 +2,10 @@ import pyarrow as pa import pytest from lancedb.index import BTree -from haiku.rag.store.engine import Store, ensure_indexes +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.schema import ensure_indexes EXPECTED_INDEXED_COLUMNS = { "documents": {"id"}, diff --git a/tests/store/test_tags.py b/tests/store/test_tags.py index 3b5f8854..c7241af0 100644 --- a/tests/store/test_tags.py +++ b/tests/store/test_tags.py @@ -4,9 +4,10 @@ import pytest from lancedb.table import AsyncTags from haiku.rag.store import ReadOnlyError -from haiku.rag.store.engine import REQUIRED_TABLES, Store +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.schema import REQUIRED_TABLES @pytest.mark.asyncio diff --git a/tests/store/test_v0_40_0_migration.py b/tests/store/test_v0_40_0_migration.py index 84e13605..e9e996eb 100644 --- a/tests/store/test_v0_40_0_migration.py +++ b/tests/store/test_v0_40_0_migration.py @@ -9,7 +9,8 @@ introduced by later migrations (``picture_data`` in v0.45.0, import pytest from haiku.rag.store.compression import compress_docling_split -from haiku.rag.store.engine import DocumentRecord, Store +from haiku.rag.store.engine import Store +from haiku.rag.store.schema import DocumentRecord from haiku.rag.store.upgrades.v0_40_0 import _apply_populate_document_items diff --git a/tests/store/test_v0_45_0_migration.py b/tests/store/test_v0_45_0_migration.py index d0cc662c..cc381574 100644 --- a/tests/store/test_v0_45_0_migration.py +++ b/tests/store/test_v0_45_0_migration.py @@ -13,7 +13,8 @@ import pyarrow as pa import pytest from haiku.rag.store.compression import compress_json, decompress_json -from haiku.rag.store.engine import DocumentItemRecord, DocumentRecord, Store +from haiku.rag.store.engine import Store +from haiku.rag.store.schema import DocumentItemRecord, DocumentRecord from haiku.rag.store.upgrades.v0_45_0 import _apply_extract_picture_bytes PNG_BYTES = b"\x89PNG\r\n\x1a\nfake-picture-bytes" diff --git a/tests/store/test_v0_48_0_migration.py b/tests/store/test_v0_48_0_migration.py index 4e73dbf1..6a94b086 100644 --- a/tests/store/test_v0_48_0_migration.py +++ b/tests/store/test_v0_48_0_migration.py @@ -3,7 +3,8 @@ import json import pytest from haiku.rag.store.compression import compress_docling_split -from haiku.rag.store.engine import DocumentItemRecord, DocumentRecord, Store +from haiku.rag.store.engine import Store +from haiku.rag.store.schema import DocumentItemRecord, DocumentRecord from haiku.rag.store.upgrades.v0_48_0 import _apply_backfill_heading_hierarchy diff --git a/tests/store/test_v0_58_0_migration.py b/tests/store/test_v0_58_0_migration.py index 30fe510e..05e8c8d3 100644 --- a/tests/store/test_v0_58_0_migration.py +++ b/tests/store/test_v0_58_0_migration.py @@ -103,7 +103,7 @@ class TestV0_58_0MigrationEdgeCases: async def test_resume_skips_already_migrated_rows(self, temp_db_path): """A half-finished prior run leaves some document_meta rows; re-running migrates only the rest and never duplicates.""" - from haiku.rag.store.engine import DocumentMetaRecord + from haiku.rag.store.schema import DocumentMetaRecord async with Store(temp_db_path, create=True, skip_migration_check=True) as store: await seed_legacy_documents( diff --git a/tests/test_doctor.py b/tests/test_doctor.py index 03771327..ce172238 100644 --- a/tests/test_doctor.py +++ b/tests/test_doctor.py @@ -40,7 +40,7 @@ from haiku.rag.doctor import ( run_doctor, run_provider_checks, ) -from haiku.rag.store.engine import ( +from haiku.rag.store.schema import ( DocumentItemRecord, DocumentMetaRecord, DocumentRecord, diff --git a/tests/test_info.py b/tests/test_info.py index a3b3f4c0..d6a48e48 100644 --- a/tests/test_info.py +++ b/tests/test_info.py @@ -5,7 +5,7 @@ import pytest from haiku.rag.app import HaikuRAGApp from haiku.rag.config.models import AppConfig, LanceDBConfig -from haiku.rag.store.engine import DocumentItemRecord +from haiku.rag.store.schema import DocumentItemRecord @pytest.mark.asyncio @@ -176,7 +176,7 @@ async def test_app_info_uses_connect_lancedb_for_remote(tmp_path): app = HaikuRAGApp(db_path=nonexistent, config=config) with patch( - "haiku.rag.store.engine.connect_lancedb", new_callable=AsyncMock + "haiku.rag.store.info.connect_lancedb", new_callable=AsyncMock ) as mock_connect: # Empty DB triggers the early-return path - enough to prove connect_lancedb was used mock_db = mock_connect.return_value diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index 278065d0..971f797b 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -261,7 +261,7 @@ def test_staging_chunk_record_mirrors_chunk_record_schema(): column on every crash-recovery cycle. This test fails loudly instead. """ from haiku.rag.client.rebuild import _StagingChunkRecord - from haiku.rag.store.engine import ChunkRecordBase + from haiku.rag.store.schema import ChunkRecordBase expected = set(ChunkRecordBase.model_fields) - {"content_fts", "vector"} assert set(_StagingChunkRecord.model_fields) == expected diff --git a/tests/test_s3_integration.py b/tests/test_s3_integration.py index 05185499..f916d07a 100644 --- a/tests/test_s3_integration.py +++ b/tests/test_s3_integration.py @@ -65,7 +65,7 @@ def config(): @pytest.mark.asyncio async def test_store_connect_and_create(tmp_path, config): - from haiku.rag.store.engine import get_database_stats + from haiku.rag.store.info import get_database_stats async with Store(tmp_path / "unused", config=config, create=True) as store: stats = await get_database_stats(store.db) @@ -81,7 +81,8 @@ async def test_store_vacuum(tmp_path, config): @pytest.mark.asyncio async def test_store_add_document(tmp_path, config): - from haiku.rag.store.engine import DocumentRecord, get_database_stats + from haiku.rag.store.info import get_database_stats + from haiku.rag.store.schema import DocumentRecord async with Store(tmp_path / "unused", config=config, create=True) as store: doc = DocumentRecord(content="The quick brown fox jumps over the lazy dog.")