diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index ad591fe1..334f0cc2 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -102,6 +102,15 @@ class HaikuRAGApp: except Exception as e: 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): """Display current configuration settings.""" self.console.print("[bold]haiku.rag configuration[/bold]") diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index c9e8fde2..f64e63a5 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -256,6 +256,18 @@ def 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( "serve", help="Start the haiku.rag MCP server (by default in streamable HTTP mode)" ) diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index e46ce00e..7a360835 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -550,6 +550,16 @@ class HaikuRAG: ) 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): """Close the underlying store connection.""" self.store.close() diff --git a/src/haiku/rag/migration.py b/src/haiku/rag/migration.py index b5122c0e..d5b19198 100644 --- a/src/haiku/rag/migration.py +++ b/src/haiku/rag/migration.py @@ -47,7 +47,7 @@ class SQLiteToLanceDBMigrator: # Load the sqlite-vec extension try: - import sqlite_vec + import sqlite_vec # type: ignore sqlite_conn.enable_load_extension(True) sqlite_vec.load(sqlite_conn) @@ -91,10 +91,10 @@ class SQLiteToLanceDBMigrator: sqlite_conn.close() - # Optimize the chunks table after migration + # Optimize and cleanup using centralized vacuum self.console.print("[blue]Optimizing LanceDB...[/blue]") try: - lance_store.chunks_table.optimize() + lance_store.vacuum() self.console.print("[green]✅ Optimization completed[/green]") except Exception as e: self.console.print( diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py index db51f22e..036729b2 100644 --- a/src/haiku/rag/store/engine.py +++ b/src/haiku/rag/store/engine.py @@ -1,5 +1,6 @@ import json import logging +from datetime import timedelta from importlib import metadata from pathlib import Path from uuid import uuid4 @@ -62,6 +63,15 @@ class Store: if not skip_validation: 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): """Establish connection to LanceDB (local, cloud, or object storage).""" # Check if we have cloud configuration @@ -159,16 +169,18 @@ class Store: self.settings_table.search().limit(1).to_pydantic(SettingsRecord) ) if settings_records: - settings = ( + # 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 {} ) - settings["version"] = version - # Update the record - self.settings_table.update( - where="id = 'settings'", values={"settings": json.dumps(settings)} - ) + 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") diff --git a/src/haiku/rag/store/repositories/settings.py b/src/haiku/rag/store/repositories/settings.py index 4b44fb3d..af752fda 100644 --- a/src/haiku/rag/store/repositories/settings.py +++ b/src/haiku/rag/store/repositories/settings.py @@ -84,10 +84,15 @@ class SettingsRepository: ) if existing: - # Update existing settings - self.store.settings_table.update( - where="id = 'settings'", values={"settings": json.dumps(current_config)} + # Only update when configuration actually changed to avoid needless new versions + existing_payload = ( + 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: # Create new settings settings_record = SettingsRecord( diff --git a/tests/generate_benchmark_db.py b/tests/generate_benchmark_db.py index 47baabd9..aa11a9bc 100644 --- a/tests/generate_benchmark_db.py +++ b/tests/generate_benchmark_db.py @@ -8,8 +8,10 @@ from rich.progress import Progress from haiku.rag import logging # noqa from haiku.rag.client import HaikuRAG +from haiku.rag.logging import configure_cli_logging from haiku.rag.qa import get_qa_agent +configure_cli_logging() console = Console() db_path = Path(__file__).parent / "data" / "benchmark.lancedb" @@ -35,6 +37,7 @@ async def populate_db(): uri=uri, ) progress.advance(task) + rag.store.vacuum() async def run_match_benchmark():