vacuum() in store, client, cli. Removes old versions and reduces disk usage

This commit is contained in:
Yiorgis Gozadinos 2025-09-10 11:24:51 +03:00
parent d1b62510d0
commit 2207b7d1ed
No known key found for this signature in database
7 changed files with 63 additions and 12 deletions

View file

@ -102,6 +102,15 @@ class HaikuRAGApp:
except Exception as e: except Exception as e:
self.console.print(f"[red]Error rebuilding database: {e}[/red]") self.console.print(f"[red]Error rebuilding database: {e}[/red]")
async def vacuum(self):
"""Run database maintenance: optimize and cleanup table history."""
try:
async with HaikuRAG(db_path=self.db_path, skip_validation=True) as client:
await client.vacuum()
self.console.print("[b]Vacuum completed successfully.[/b]")
except Exception as e:
self.console.print(f"[red]Error during vacuum: {e}[/red]")
def show_settings(self): def show_settings(self):
"""Display current configuration settings.""" """Display current configuration settings."""
self.console.print("[bold]haiku.rag configuration[/bold]") self.console.print("[bold]haiku.rag configuration[/bold]")

View file

@ -256,6 +256,18 @@ def rebuild(
asyncio.run(app.rebuild()) asyncio.run(app.rebuild())
@cli.command("vacuum", help="Optimize and clean up all tables to reduce disk usage")
def vacuum(
db: Path = typer.Option(
Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
"--db",
help="Path to the LanceDB database file",
),
):
app = HaikuRAGApp(db_path=db)
asyncio.run(app.vacuum())
@cli.command( @cli.command(
"serve", help="Start the haiku.rag MCP server (by default in streamable HTTP mode)" "serve", help="Start the haiku.rag MCP server (by default in streamable HTTP mode)"
) )

View file

@ -550,6 +550,16 @@ class HaikuRAG:
) )
yield doc.id yield doc.id
# Final maintenance: centralized vacuum to curb disk usage
try:
self.store.vacuum()
except Exception:
pass
async def vacuum(self) -> None:
"""Optimize and clean up old versions across all tables."""
self.store.vacuum()
def close(self): def close(self):
"""Close the underlying store connection.""" """Close the underlying store connection."""
self.store.close() self.store.close()

View file

@ -47,7 +47,7 @@ class SQLiteToLanceDBMigrator:
# Load the sqlite-vec extension # Load the sqlite-vec extension
try: try:
import sqlite_vec import sqlite_vec # type: ignore
sqlite_conn.enable_load_extension(True) sqlite_conn.enable_load_extension(True)
sqlite_vec.load(sqlite_conn) sqlite_vec.load(sqlite_conn)
@ -91,10 +91,10 @@ class SQLiteToLanceDBMigrator:
sqlite_conn.close() sqlite_conn.close()
# Optimize the chunks table after migration # Optimize and cleanup using centralized vacuum
self.console.print("[blue]Optimizing LanceDB...[/blue]") self.console.print("[blue]Optimizing LanceDB...[/blue]")
try: try:
lance_store.chunks_table.optimize() lance_store.vacuum()
self.console.print("[green]✅ Optimization completed[/green]") self.console.print("[green]✅ Optimization completed[/green]")
except Exception as e: except Exception as e:
self.console.print( self.console.print(

View file

@ -1,5 +1,6 @@
import json import json
import logging import logging
from datetime import timedelta
from importlib import metadata from importlib import metadata
from pathlib import Path from pathlib import Path
from uuid import uuid4 from uuid import uuid4
@ -62,6 +63,15 @@ class Store:
if not skip_validation: if not skip_validation:
self._validate_configuration() self._validate_configuration()
def vacuum(self) -> None:
"""Optimize and clean up old versions across all tables to reduce disk usage."""
if self._has_cloud_config() and str(Config.LANCEDB_URI).startswith("db://"):
return
# Perform maintenance per table using optimize() with cleanup_older_than 0
for table in [self.documents_table, self.chunks_table, self.settings_table]:
table.optimize(cleanup_older_than=timedelta(0))
def _connect_to_lancedb(self, db_path: Path): def _connect_to_lancedb(self, db_path: Path):
"""Establish connection to LanceDB (local, cloud, or object storage).""" """Establish connection to LanceDB (local, cloud, or object storage)."""
# Check if we have cloud configuration # Check if we have cloud configuration
@ -159,16 +169,18 @@ class Store:
self.settings_table.search().limit(1).to_pydantic(SettingsRecord) self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
) )
if settings_records: if settings_records:
settings = ( # Only write if version actually changes to avoid creating new table versions
current = (
json.loads(settings_records[0].settings) json.loads(settings_records[0].settings)
if settings_records[0].settings if settings_records[0].settings
else {} else {}
) )
settings["version"] = version if current.get("version") != version:
# Update the record current["version"] = version
self.settings_table.update( self.settings_table.update(
where="id = 'settings'", values={"settings": json.dumps(settings)} where="id = 'settings'",
) values={"settings": json.dumps(current)},
)
else: else:
# Create new settings record # Create new settings record
settings_data = Config.model_dump(mode="json") settings_data = Config.model_dump(mode="json")

View file

@ -84,10 +84,15 @@ class SettingsRepository:
) )
if existing: if existing:
# Update existing settings # Only update when configuration actually changed to avoid needless new versions
self.store.settings_table.update( existing_payload = (
where="id = 'settings'", values={"settings": json.dumps(current_config)} json.loads(existing[0].settings) if existing[0].settings else {}
) )
if existing_payload != current_config:
self.store.settings_table.update(
where="id = 'settings'",
values={"settings": json.dumps(current_config)},
)
else: else:
# Create new settings # Create new settings
settings_record = SettingsRecord( settings_record = SettingsRecord(

View file

@ -8,8 +8,10 @@ from rich.progress import Progress
from haiku.rag import logging # noqa from haiku.rag import logging # noqa
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.logging import configure_cli_logging
from haiku.rag.qa import get_qa_agent from haiku.rag.qa import get_qa_agent
configure_cli_logging()
console = Console() console = Console()
db_path = Path(__file__).parent / "data" / "benchmark.lancedb" db_path = Path(__file__).parent / "data" / "benchmark.lancedb"
@ -35,6 +37,7 @@ async def populate_db():
uri=uri, uri=uri,
) )
progress.advance(task) progress.advance(task)
rag.store.vacuum()
async def run_match_benchmark(): async def run_match_benchmark():