import asyncio import json import logging from datetime import 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 import lancedb import pyarrow as pa from lancedb.index import FTS, BTree, IvfPq from lancedb.pydantic import LanceModel, Vector from packaging.version import parse from pydantic import BaseModel, Field from haiku.rag.config import AppConfig, Config from haiku.rag.embeddings import get_embedder from haiku.rag.store.exceptions import MigrationRequiredError, ReadOnlyError if TYPE_CHECKING: from lancedb.query import AsyncQueryBase 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" OBJECT_STORAGE = "object_storage" @staticmethod def from_config(config: AppConfig) -> "ConnectionMode": uri = config.lancedb.uri if not uri: return ConnectionMode.LOCAL if uri.startswith("db://"): return ConnectionMode.CLOUD return ConnectionMode.OBJECT_STORAGE async def connect_lancedb( config: AppConfig, db_path: Path | None = None ) -> lancedb.AsyncConnection: mode = ConnectionMode.from_config(config) if mode == ConnectionMode.CLOUD: return await lancedb.connect_async( uri=config.lancedb.uri, api_key=config.lancedb.api_key, region=config.lancedb.region, ) elif mode == ConnectionMode.OBJECT_STORAGE: kwargs: dict[str, Any] = {"uri": config.lancedb.uri} if config.lancedb.storage_options: kwargs["storage_options"] = config.lancedb.storage_options return await lancedb.connect_async(**kwargs) else: if db_path is None: raise ValueError("No lancedb.uri configured and no db_path provided") return await lancedb.connect_async(db_path.absolute()) 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.""" document_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) class SettingsRecord(LanceModel): id: str = Field(default="settings") settings: str = Field(default="{}") REQUIRED_TABLES: tuple[str, ...] = ( "documents", "document_meta", "chunks", "document_items", "settings", ) 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, db_path: Path, config: AppConfig = Config, skip_validation: bool = False, create: bool = False, read_only: bool = False, before: datetime | None = None, skip_migration_check: bool = False, ): self.db_path: Path = db_path self._config = config self._before = before # Time-travel mode is always read-only self._read_only = read_only or (before is not None) self._create = create self._skip_validation = skip_validation self._skip_migration_check = skip_migration_check self._vacuum_lock = asyncio.Lock() self._write_lock = asyncio.Lock() self._is_new_db = False # Check if database exists (for local filesystem only) if self._connection_mode == ConnectionMode.LOCAL: if not db_path.exists(): if not create: raise FileNotFoundError( f"Database does not exist at {self.db_path.absolute()}. " "Use 'haiku-rag init' to create a new database." ) self._is_new_db = True # Ensure parent directories exist for new databases if not db_path.parent.exists(): Path.mkdir(db_path.parent, parents=True) # Create embedder (sync — no LanceDB needed) self.embedder = get_embedder(config=self._config) async def _initialize(self): """Perform async initialization: connect to LanceDB, init tables, validate.""" # Connect to LanceDB self.db: lancedb.AsyncConnection = await connect_lancedb( self._config, self.db_path ) # For remote stores (and as a safety net for local paths that exist but # have no tables — e.g. a previously failed init), detect new DB by # checking whether any tables exist. is_new_db = self._is_new_db if not is_new_db: existing_tables = (await self.db.list_tables()).tables if not existing_tables: is_new_db = True # For existing databases, read stored vector dimension to create ChunkRecord # that can read existing chunks. For new databases, use config's dimension. stored_vector_dim = None if not is_new_db: stored_vector_dim = await self._get_stored_vector_dim() # Create ChunkRecord with stored dimension (for reading) or config dimension (for new DB) chunk_vector_dim = stored_vector_dim or self.embedder._vector_dim self.ChunkRecord: type[ChunkRecordBase] = create_chunk_model(chunk_vector_dim) # Initialize tables (creates them if they don't exist). For an existing # DB this raises MigrationRequiredError up front when migrations are # pending, before creating any newly-introduced table. await self._init_tables(is_new_db) # Checkout tables to historical state if before is specified if self._before is not None: await self._checkout_tables_before(self._before) # Set version for new databases. if is_new_db and not self._read_only: await self._set_initial_version() # Validate config compatibility after connection is established if not self._skip_validation: await self._validate_configuration() async def __aenter__(self): # If _initialize connects to LanceDB but then fails (e.g. migration # check, config validation), close the connection so it doesn't # leak — __aexit__ won't run because the `async with` never entered. try: await self._initialize() except BaseException: self.close() raise return self async def __aexit__(self, exc_type, exc_val, exc_tb): # noqa: ARG002 self.close() return False @property def is_read_only(self) -> bool: """Whether the store is in read-only mode.""" return self._read_only async def _get_stored_vector_dim(self) -> int | None: """Read the stored vector dimension from the settings table. Returns: The stored vector dimension, or None if not found. """ try: existing_tables = (await self.db.list_tables()).tables if "settings" not in existing_tables: return None settings_table = await self.db.open_table("settings") rows = ( await settings_table.query() .where("id = 'settings'") .limit(1) .to_arrow() ).to_pylist() if not rows or not rows[0].get("settings"): return None settings = json.loads(rows[0]["settings"]) embeddings = settings.get("embeddings", {}) model = embeddings.get("model", {}) return model.get("vector_dim") except Exception: return None def _assert_writable(self) -> None: """Raise ReadOnlyError if the store is in read-only mode.""" if self._read_only: raise ReadOnlyError("Cannot modify database in read-only mode") async def vacuum(self, retention_seconds: int | None = None) -> None: """Optimize and clean up old versions across all tables to reduce disk usage. Args: retention_seconds: Retention threshold in seconds. Only versions older than this will be removed. If None, uses config.storage.vacuum_retention_seconds. Note: If vacuum is already running, this method returns immediately without blocking. Use asyncio.create_task(store.vacuum()) for non-blocking background execution. Raises: ReadOnlyError: If the store is in read-only mode. """ self._assert_writable() if self._connection_mode == ConnectionMode.CLOUD: return # Skip if already running (non-blocking) if self._vacuum_lock.locked(): return async with self._vacuum_lock: try: # Evaluate config at runtime to allow dynamic changes if retention_seconds is None: retention_seconds = self._config.storage.vacuum_retention_seconds # Perform maintenance per table using optimize() with configurable retention retention = timedelta(seconds=retention_seconds) for table in [ self.documents_table, self.document_meta_table, self.chunks_table, self.document_items_table, self.settings_table, ]: await table.optimize(cleanup_older_than=retention) except (RuntimeError, OSError) as e: # Handle resource errors gracefully logger.debug(f"Vacuum skipped due to resource constraints: {e}") @property def _connection_mode(self) -> ConnectionMode: return ConnectionMode.from_config(self._config) async def _ensure_vector_index(self) -> None: """Create or rebuild vector index on chunks table. Cloud deployments auto-create indexes, so we skip for those. For self-hosted, creates an IVF_PQ index. If an index exists, it will be replaced (using replace=True parameter). Note: Index creation requires sufficient training data. """ if self._connection_mode == ConnectionMode.CLOUD: return try: # Check if table has enough data (indexes require training data) row_count = await self.chunks_table.count_rows() if row_count < 256: logger.debug( f"Skipping vector index creation: need at least 256 rows, have {row_count}" ) return # Create or replace index (replace=True is the default) logger.info("Creating vector index on chunks table...") await self.chunks_table.create_index( "vector", config=IvfPq( distance_type=self._config.search.vector_index_metric, ), replace=True, ) # Wait for index creation to complete # Index name is column_name + "_idx" await self.chunks_table.wait_for_index( ["vector_idx"], timeout=timedelta(hours=1) ) logger.info("Vector index created successfully") except Exception as e: logger.warning(f"Could not create vector index: {e}") async def _validate_configuration(self) -> None: """Validate that the configuration is compatible with the database.""" from haiku.rag.store.repositories.settings import SettingsRepository settings_repo = SettingsRepository(self) await settings_repo.validate_config_compatibility() async def _init_tables(self, is_new_db: bool): """Initialize database tables (create if they don't exist).""" existing_tables = (await self.db.list_tables()).tables # Surface pending migrations BEFORE creating any newly-introduced table. # Otherwise opening a legacy DB would either mutate it (creating an empty # document_meta on open) or raise the wrong ReadOnlyError instead of # telling the user to run `haiku-rag migrate`. The settings table exists # on any non-new DB, which is all _check_migrations needs. if ( not is_new_db and not self._skip_migration_check and "settings" in existing_tables ): self.settings_table = await self.db.open_table("settings") await self._check_migrations() missing_tables = set(REQUIRED_TABLES) - set(existing_tables) if missing_tables and self._read_only: raise ReadOnlyError( "Cannot create tables in read-only mode. " "Use 'haiku-rag init' to create a new database." ) # Create or open documents table if "documents" in existing_tables: self.documents_table = await self.db.open_table("documents") else: self.documents_table = await self.db.create_table( "documents", schema=get_documents_arrow_schema() ) # 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). 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( "document_id", config=BTree(), replace=True ) await self.document_meta_table.create_index( "uri", config=BTree(), replace=True ) # Create or open chunks table if "chunks" in existing_tables: self.chunks_table = await self.db.open_table("chunks") else: 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, ) # Create or open document_items table if "document_items" in existing_tables: self.document_items_table = await self.db.open_table("document_items") else: 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 ) # Create or open settings table if "settings" in existing_tables: self.settings_table = await self.db.open_table("settings") else: self.settings_table = await self.db.create_table( "settings", schema=SettingsRecord ) # Save current settings to the new database settings_data = self._config.model_dump(mode="json") await self.settings_table.add( [SettingsRecord(id="settings", settings=json.dumps(settings_data))] ) async def _set_initial_version(self): """Set the initial version for a new database.""" await self.set_haiku_version(metadata.version("haiku.rag-slim")) async def _check_migrations(self) -> None: """Raise if migrations are pending. Opening never writes the version. Raises: MigrationRequiredError: If migrations are pending. """ from haiku.rag.store.upgrades import get_pending_upgrades current_version = metadata.version("haiku.rag-slim") db_version = await self.get_haiku_version() pending = get_pending_upgrades(db_version) if pending: # Migrations are pending - require explicit migrate command raise MigrationRequiredError( f"Database requires migration from {db_version} to {current_version}. " f"{len(pending)} migration(s) pending. " "Run 'haiku-rag migrate' to upgrade." ) async def migrate(self) -> list[str]: """Run pending database migrations. Returns: List of descriptions of applied upgrades. Raises: ReadOnlyError: If the store is in read-only mode. """ self._assert_writable() from haiku.rag.store.upgrades import run_pending_upgrades db_version = await self.get_haiku_version() current_version = metadata.version("haiku.rag-slim") applied = await run_pending_upgrades(self, db_version) # Advance the schema marker only forward — never downgrade a database # opened with an older build than last stamped it. if parse(current_version) > parse(db_version): await self.set_haiku_version(current_version) return applied async def get_haiku_version(self) -> str: """Returns the user version stored in settings.""" settings_records = await query_to_pydantic( self.settings_table.query().limit(1), SettingsRecord ) if settings_records: settings = ( json.loads(settings_records[0].settings) if settings_records[0].settings else {} ) return settings.get("version", "0.0.0") return "0.0.0" async def set_haiku_version(self, version: str) -> None: """Updates the user version in settings. Raises: ReadOnlyError: If the store is in read-only mode. """ self._assert_writable() settings_records = await query_to_pydantic( self.settings_table.query().limit(1), SettingsRecord ) if settings_records: # Only write if version actually changes to avoid creating new table versions current = ( json.loads(settings_records[0].settings) if settings_records[0].settings else {} ) if current.get("version") != version: current["version"] = version await self.settings_table.update( {"settings": json.dumps(current)}, where="id = 'settings'", ) else: # Create new settings record settings_data = Config.model_dump(mode="json") settings_data["version"] = version await self.settings_table.add( [SettingsRecord(id="settings", settings=json.dumps(settings_data))] ) async def recreate_embeddings_table(self) -> None: """Recreate the chunks table with current vector dimensions. Raises: ReadOnlyError: If the store is in read-only mode. """ self._assert_writable() # Drop and recreate chunks table. Check existence first rather than # catching-and-swallowing drop_table's errors — a catch-all would # hide real failures (permissions, storage-backend errors) and then # the subsequent create_table would fail confusingly. if "chunks" in (await self.db.list_tables()).tables: await self.db.drop_table("chunks") # Update the ChunkRecord model with new vector dimension self.ChunkRecord = create_chunk_model(self.embedder._vector_dim) 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, ) def close(self): """Close the database connection.""" # AsyncConnection.close() is synchronous if hasattr(self, "db"): self.db.close() async def current_table_versions(self) -> dict[str, int]: """Capture current versions of key tables for rollback using LanceDB's API.""" return { "documents": await self.documents_table.version(), "document_meta": await self.document_meta_table.version(), "chunks": await self.chunks_table.version(), "document_items": await self.document_items_table.version(), "settings": await self.settings_table.version(), } async def restore_table_versions(self, versions: dict[str, int]) -> bool: """Restore tables to the provided versions using LanceDB's API. Raises: ReadOnlyError: If the store is in read-only mode. """ self._assert_writable() await self.documents_table.restore(int(versions["documents"])) await self.document_meta_table.restore(int(versions["document_meta"])) await self.chunks_table.restore(int(versions["chunks"])) await self.document_items_table.restore(int(versions["document_items"])) await self.settings_table.restore(int(versions["settings"])) return True async def _checkout_tables_before(self, before: datetime) -> None: """Checkout all tables to their state at or before the given datetime. Args: before: The datetime to checkout to Raises: ValueError: If no version exists before the given datetime """ # LanceDB stores timestamps as naive datetimes in local time. # Convert 'before' to naive local time for comparison. if before.tzinfo is not None: # Convert to local time and make naive before_local = before.astimezone().replace(tzinfo=None) else: # Already naive, assume local time before_local = before tables = [ ("documents", self.documents_table), ("document_meta", self.document_meta_table), ("chunks", self.chunks_table), ("document_items", self.document_items_table), ("settings", self.settings_table), ] for table_name, table in tables: versions = await table.list_versions() # Find the latest version at or before the target datetime # Versions are sorted by version number, not timestamp, so we need to check all best_version = None best_timestamp = None for v in versions: # LanceDB version timestamps are naive datetime objects in local time v_timestamp = v["timestamp"] # Make sure it's naive for comparison if v_timestamp.tzinfo is not None: v_timestamp = v_timestamp.replace(tzinfo=None) if v_timestamp <= before_local: if best_timestamp is None or v_timestamp > best_timestamp: best_version = v["version"] best_timestamp = v_timestamp if best_version is None: # Find the earliest version to report in error message if versions: earliest = min(versions, key=lambda v: v["timestamp"]) earliest_ts = earliest["timestamp"] raise ValueError( f"No data exists before {before}. " f"Database was created on {earliest_ts}" ) else: raise ValueError( f"No data exists before {before}. Table has no versions." ) # Checkout to the found version await table.checkout(best_version) async def list_table_versions(self, table_name: str) -> list[dict[str, Any]]: """List version history for a table. Args: table_name: Name of the table ("documents", "document_meta", "chunks", "document_items", or "settings") Returns: List of version info dicts with "version" and "timestamp" keys """ table_map = { "documents": self.documents_table, "document_meta": self.document_meta_table, "chunks": self.chunks_table, "document_items": self.document_items_table, "settings": self.settings_table, } table = table_map.get(table_name) if table is None: raise ValueError(f"Unknown table: {table_name}") return list(await table.list_versions())