393 lines
15 KiB
Python
393 lines
15 KiB
Python
import asyncio
|
|
import json
|
|
import logging
|
|
from datetime import timedelta
|
|
from importlib import metadata
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
import lancedb
|
|
from lancedb.pydantic import LanceModel, Vector
|
|
from pydantic import Field
|
|
|
|
from haiku.rag.config import AppConfig, Config
|
|
from haiku.rag.embeddings import get_embedder
|
|
|
|
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="{}")
|
|
created_at: str = Field(default_factory=lambda: "")
|
|
updated_at: str = Field(default_factory=lambda: "")
|
|
|
|
|
|
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
|
|
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,
|
|
read_only: bool = False,
|
|
):
|
|
self.db_path: Path = db_path
|
|
self._config = config
|
|
self.embedder = get_embedder(config=self._config)
|
|
self._vacuum_lock = asyncio.Lock()
|
|
self._read_only = read_only
|
|
|
|
# Create the ChunkRecord model with the correct vector dimension
|
|
self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
|
|
|
|
# Local filesystem handling for DB directory
|
|
if not self._has_cloud_config():
|
|
if read_only:
|
|
# Read operations should not create the database
|
|
if not db_path.exists():
|
|
raise FileNotFoundError(
|
|
f"Database does not exist: {db_path}. Use a write operation (add, add-src) to create it."
|
|
)
|
|
else:
|
|
# Write operations - ensure parent directories exist
|
|
if not db_path.parent.exists():
|
|
Path.mkdir(db_path.parent, parents=True)
|
|
|
|
# Connect to LanceDB
|
|
self.db = self._connect_to_lancedb(db_path)
|
|
|
|
# Initialize tables
|
|
self.create_or_update_db()
|
|
|
|
# Validate config compatibility after connection is established
|
|
if not skip_validation:
|
|
self._validate_configuration()
|
|
|
|
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.
|
|
"""
|
|
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() # type: ignore[assignment]
|
|
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() # type: ignore[assignment]
|
|
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 create_or_update_db(self):
|
|
"""Create the database tables."""
|
|
|
|
# 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=DocumentRecord
|
|
)
|
|
|
|
# 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 the new table with phrase query support
|
|
self.chunks_table.create_fts_index(
|
|
"content", 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))]
|
|
)
|
|
|
|
# Run pending upgrades based on stored version and package version
|
|
# Skip in read-only mode to avoid modifying the database
|
|
if not self._read_only:
|
|
try:
|
|
from haiku.rag.store.upgrades import run_pending_upgrades
|
|
|
|
current_version = metadata.version("haiku.rag-slim")
|
|
db_version = self.get_haiku_version()
|
|
|
|
if db_version != "0.0.0":
|
|
run_pending_upgrades(self, db_version, current_version)
|
|
|
|
# After upgrades complete (or if none), set stored version
|
|
# to the greater of the installed package version and the
|
|
# highest available upgrade step version in code.
|
|
try:
|
|
from packaging.version import parse as _v
|
|
|
|
from haiku.rag.store.upgrades import upgrades as _steps
|
|
|
|
highest_step = max((_v(u.version) for u in _steps), default=None)
|
|
effective_version = (
|
|
str(max(_v(current_version), highest_step))
|
|
if highest_step is not None
|
|
else current_version
|
|
)
|
|
except Exception:
|
|
effective_version = current_version
|
|
|
|
self.set_haiku_version(effective_version)
|
|
except Exception as e:
|
|
# Avoid hard failure on initial connection; log and continue so CLI remains usable.
|
|
logger.warning(
|
|
"Skipping upgrade due to error (db=%s -> pkg=%s): %s",
|
|
self.get_haiku_version(),
|
|
metadata.version("haiku.rag-slim"),
|
|
e,
|
|
)
|
|
|
|
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."""
|
|
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."""
|
|
# 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 the new table with phrase query support
|
|
self.chunks_table.create_fts_index(
|
|
"content", 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."""
|
|
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
|