haiku.rag/haiku_rag_slim/haiku/rag/store/engine.py

658 lines
24 KiB
Python

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 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 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
def connect_lancedb(config: AppConfig, db_path: Path | None = None):
mode = ConnectionMode.from_config(config)
if mode == ConnectionMode.CLOUD:
return lancedb.connect(
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 lancedb.connect(**kwargs)
else:
if db_path is None:
raise ValueError("No lancedb.uri configured and no db_path provided")
return lancedb.connect(db_path)
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_pages: 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()
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)
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 DocumentItemRecord(LanceModel):
document_id: str
position: int
self_ref: str
label: str = Field(default="")
text: str = Field(default="")
page_numbers: str = Field(default="[]")
class SettingsRecord(LanceModel):
id: str = Field(default="settings")
settings: str = Field(default="{}")
REQUIRED_TABLES: tuple[str, ...] = ("documents", "chunks", "document_items", "settings")
def get_database_stats(db: lancedb.DBConnection) -> 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(db.list_tables().tables)
stats: dict = {}
tables: dict = {}
for name in REQUIRED_TABLES:
if name not in existing:
stats[name] = {"exists": False}
continue
tbl = db.open_table(name)
tables[name] = tbl
# lancedb's .stats() stub claims TableStatistics but returns a plain dict at runtime.
tbl_stats: dict = tbl.stats() # 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(list(tbl.list_versions())),
}
if stats["chunks"]["exists"]:
chunks_tbl = tables["chunks"]
indices = 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 = 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 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 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."
)
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 = connect_lancedb(self._config, db_path)
# For remote stores, detect new DB by checking if tables exist
if not is_new_db and self._connection_mode != ConnectionMode.LOCAL:
existing_tables = 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 = 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.list_tables().tables
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._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.chunks_table,
self.document_items_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}")
@property
def _connection_mode(self) -> ConnectionMode:
return ConnectionMode.from_config(self._config)
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 = 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)."""
existing_tables = self.db.list_tables().tables
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 = self.db.open_table("documents")
else:
self.documents_table = self.db.create_table(
"documents", schema=get_documents_arrow_schema()
)
# Create or open 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 open document_items table
if "document_items" in existing_tables:
self.document_items_table = self.db.open_table("document_items")
else:
self.document_items_table = self.db.create_table(
"document_items", schema=DocumentItemRecord
)
self.document_items_table.create_scalar_index(
"document_id", index_type="BTREE", replace=True
)
self.document_items_table.create_scalar_index(
"position", index_type="BTREE", replace=True
)
self.document_items_table.create_scalar_index(
"self_ref", index_type="BTREE", replace=True
)
# Create or open 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),
"document_items": int(self.document_items_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.document_items_table.restore(int(versions["document_items"]))
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),
("document_items", self.document_items_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,
"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(table.list_versions())