305 lines
11 KiB
Python
305 lines
11 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
|
|
):
|
|
self.db_path: Path = db_path
|
|
self._config = config
|
|
self.embedder = get_embedder(config=self._config)
|
|
self._vacuum_lock = asyncio.Lock()
|
|
|
|
# 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 self._config.storage.disable_autocreate:
|
|
# LanceDB uses a directory path for local databases; enforce presence
|
|
if not db_path.exists():
|
|
raise FileNotFoundError(
|
|
f"LanceDB path does not exist: {db_path}. Auto-creation is disabled."
|
|
)
|
|
else:
|
|
# Ensure parent directories exist when autocreation allowed
|
|
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 _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
|
|
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
|