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

Merge pull request #566 from ggozad/refactor/store-split

Browse source 
Split the store module by responsibility
This commit is contained in:
Yiorgis Gozadinos 2026-08-20 12:26:45 +03:00 committed by GitHub
commit 3dcd463792
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
31 changed files with 429 additions and 382 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
CHANGELOG.md
View file

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

2
app/backend/main.py
View file

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

2
haiku_rag_slim/haiku/rag/app.py
View file

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

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

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

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

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

2
haiku_rag_slim/haiku/rag/ingester/api/routes/database.py
View file

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

7
haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py
View file

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

352
haiku_rag_slim/haiku/rag/store/engine.py
View file

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

166
haiku_rag_slim/haiku/rag/store/info.py Normal file
View file

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

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

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

6
haiku_rag_slim/haiku/rag/store/repositories/document.py
View file

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

3
haiku_rag_slim/haiku/rag/store/repositories/document_item.py
View file

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

3
haiku_rag_slim/haiku/rag/store/repositories/settings.py
View file

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

198
haiku_rag_slim/haiku/rag/store/schema.py Normal file
View file

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

3
haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py
View file

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

3
haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py
View file

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

3
haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py
View file

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

3
tests/store/test_database_info.py
View file

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

6
tests/store/test_document_items.py
View file

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

2
tests/store/test_document_meta_split.py
View file

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

3
tests/store/test_engine.py
View file

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

3
tests/store/test_indexes.py
View file

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

3
tests/store/test_tags.py
View file

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

3
tests/store/test_v0_40_0_migration.py
View file

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

3
tests/store/test_v0_45_0_migration.py
View file

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

3
tests/store/test_v0_48_0_migration.py
View file

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

2
tests/store/test_v0_58_0_migration.py
View file

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

2
tests/test_doctor.py
View file

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

4
tests/test_info.py
View file

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

2
tests/test_rebuild.py
View file

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

5
tests/test_s3_integration.py
View file

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