import asyncio import json import logging from datetime import datetime, timedelta from importlib import metadata from pathlib import Path from typing import Any from uuid import uuid4 import lancedb import pyarrow as pa from lancedb.pydantic import LanceModel, Vector from pydantic import Field from haiku.rag.config import AppConfig, Config from haiku.rag.embeddings import get_embedder from haiku.rag.store.exceptions import MigrationRequiredError, ReadOnlyError logger = logging.getLogger(__name__) class DocumentRecord(LanceModel): id: str = Field(default_factory=lambda: str(uuid4())) content: str uri: str | None = None title: str | None = None metadata: str = Field(default="{}") docling_document: bytes | None = None docling_version: str | None = None 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() fields = [] for field in base_schema: if field.name == "docling_document": fields.append(pa.field("docling_document", pa.large_binary())) else: fields.append(field) return pa.schema(fields) def create_chunk_model(vector_dim: int): """Create a ChunkRecord model with the specified vector dimension. This creates a model with proper vector typing for LanceDB. """ class ChunkRecord(LanceModel): 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: Vector(vector_dim) = Field(default_factory=lambda: [0.0] * vector_dim) # type: ignore return ChunkRecord class SettingsRecord(LanceModel): id: str = Field(default="settings") settings: str = Field(default="{}") 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._vacuum_lock = asyncio.Lock() # Check if database exists (for local filesystem only) is_new_db = False if not self._has_cloud_config(): 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." ) is_new_db = True # Ensure parent directories exist for new databases if not db_path.parent.exists(): Path.mkdir(db_path.parent, parents=True) # Connect to LanceDB self.db = self._connect_to_lancedb(db_path) # 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 = self._get_stored_vector_dim() # Create embedder with config's dimension (for generating new embeddings) self.embedder = get_embedder(config=self._config) # 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 = create_chunk_model(chunk_vector_dim) # Initialize tables (creates them if they don't exist) self._init_tables() # Checkout tables to historical state if before is specified if before is not None: self._checkout_tables_before(before) # Set version for new databases, check migrations for existing ones if is_new_db: if not self._read_only: self._set_initial_version() elif not skip_migration_check: self._check_migrations() # Validate config compatibility after connection is established if not skip_validation: self._validate_configuration() @property def is_read_only(self) -> bool: """Whether the store is in read-only mode.""" return self._read_only 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 = self.db.table_names() if "settings" not in existing_tables: return None settings_table = self.db.open_table("settings") rows = ( settings_table.search() .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._has_cloud_config() and str(self._config.lancedb.uri).startswith( "db://" ): 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.chunks_table, self.settings_table, ]: 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}") def _connect_to_lancedb(self, db_path: Path): """Establish connection to LanceDB (local, cloud, or object storage).""" # Check if we have cloud configuration if self._has_cloud_config(): return lancedb.connect( uri=self._config.lancedb.uri, api_key=self._config.lancedb.api_key, region=self._config.lancedb.region, ) else: # Local file system connection return lancedb.connect(db_path) def _has_cloud_config(self) -> bool: """Check if cloud configuration is complete.""" return bool( self._config.lancedb.uri and self._config.lancedb.api_key and self._config.lancedb.region ) def get_stats(self) -> dict: """Get comprehensive table statistics. Returns: Dictionary with statistics for documents and chunks tables including: - Row counts - Storage sizes - Vector index status and statistics """ stats_dict: dict = { "documents": {"exists": False}, "chunks": {"exists": False}, } # Documents table stats doc_stats: dict = self.documents_table.stats() stats_dict["documents"] = { "exists": True, "num_rows": doc_stats.get("num_rows", 0), "total_bytes": doc_stats.get("total_bytes", 0), } # Chunks table stats chunk_stats: dict = self.chunks_table.stats() stats_dict["chunks"] = { "exists": True, "num_rows": chunk_stats.get("num_rows", 0), "total_bytes": chunk_stats.get("total_bytes", 0), } # Vector index stats indices = self.chunks_table.list_indices() has_vector_index = any("vector" in str(idx).lower() for idx in indices) stats_dict["chunks"]["has_vector_index"] = has_vector_index if has_vector_index: index_stats = self.chunks_table.index_stats("vector_idx") if index_stats is not None: stats_dict["chunks"]["num_indexed_rows"] = index_stats.num_indexed_rows stats_dict["chunks"]["num_unindexed_rows"] = ( index_stats.num_unindexed_rows ) return stats_dict 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._has_cloud_config(): return try: # Check if table has enough data (indexes require training data) row_count = 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...") self.chunks_table.create_index( metric=self._config.search.vector_index_metric, index_type="IVF_PQ", replace=True, # Explicit: replace existing index ) # Wait for index creation to complete # Index name is column_name + "_idx" 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}") 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) settings_repo.validate_config_compatibility() def _init_tables(self): """Initialize database tables (create if they don't exist).""" # Get list of existing tables existing_tables = self.db.table_names() # Create or get documents table if "documents" in existing_tables: self.documents_table = self.db.open_table("documents") else: self.documents_table = self.db.create_table( "documents", schema=get_documents_arrow_schema() ) # Create or get chunks table if "chunks" in existing_tables: self.chunks_table = self.db.open_table("chunks") else: self.chunks_table = self.db.create_table("chunks", schema=self.ChunkRecord) # Create FTS index on content_fts (contextualized content) for better search self.chunks_table.create_fts_index( "content_fts", replace=True, with_position=True, remove_stop_words=False ) # Create or get settings table if "settings" in existing_tables: self.settings_table = self.db.open_table("settings") else: self.settings_table = self.db.create_table( "settings", schema=SettingsRecord ) # Save current settings to the new database settings_data = self._config.model_dump(mode="json") self.settings_table.add( [SettingsRecord(id="settings", settings=json.dumps(settings_data))] ) def _set_initial_version(self): """Set the initial version for a new database.""" self.set_haiku_version(metadata.version("haiku.rag-slim")) def _check_migrations(self) -> None: """Check if migrations are pending and error or update version accordingly. Raises: MigrationRequiredError: If migrations are pending. """ from haiku.rag.store.upgrades import get_pending_upgrades current_version = metadata.version("haiku.rag-slim") db_version = 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." ) # No pending migrations - update version silently if needed (writable only) if not self._read_only and db_version != current_version: self.set_haiku_version(current_version) 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 = self.get_haiku_version() current_version = metadata.version("haiku.rag-slim") applied = run_pending_upgrades(self, db_version) # Update version after successful migration if applied or db_version != current_version: self.set_haiku_version(current_version) return applied def get_haiku_version(self) -> str: """Returns the user version stored in settings.""" settings_records = list( self.settings_table.search().limit(1).to_pydantic(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" 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 = list( self.settings_table.search().limit(1).to_pydantic(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 self.settings_table.update( where="id = 'settings'", values={"settings": json.dumps(current)}, ) else: # Create new settings record settings_data = Config.model_dump(mode="json") settings_data["version"] = version self.settings_table.add( [SettingsRecord(id="settings", settings=json.dumps(settings_data))] ) 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 try: self.db.drop_table("chunks") except Exception: pass # Update the ChunkRecord model with new vector dimension self.ChunkRecord = create_chunk_model(self.embedder._vector_dim) self.chunks_table = self.db.create_table("chunks", schema=self.ChunkRecord) # Create FTS index on content_fts (contextualized content) for better search self.chunks_table.create_fts_index( "content_fts", replace=True, with_position=True, remove_stop_words=False ) def close(self): """Close the database connection.""" # LanceDB connections are automatically managed pass def current_table_versions(self) -> dict[str, int]: """Capture current versions of key tables for rollback using LanceDB's API.""" return { "documents": int(self.documents_table.version), "chunks": int(self.chunks_table.version), "settings": int(self.settings_table.version), } 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() self.documents_table.restore(int(versions["documents"])) self.chunks_table.restore(int(versions["chunks"])) self.settings_table.restore(int(versions["settings"])) return True @property def _connection(self): """Compatibility property for repositories expecting _connection.""" return self 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), ("chunks", self.chunks_table), ("settings", self.settings_table), ] for table_name, table in tables: versions = 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 table.checkout(best_version) 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", "chunks", or "settings") Returns: List of version info dicts with "version" and "timestamp" keys """ table_map = { "documents": self.documents_table, "chunks": self.chunks_table, "settings": self.settings_table, } table = table_map.get(table_name) if table is None: raise ValueError(f"Unknown table: {table_name}") return list(table.list_versions())