vacuum() in store, client, cli. Removes old versions and reduces disk usage
This commit is contained in:
parent
d1b62510d0
commit
2207b7d1ed
7 changed files with 63 additions and 12 deletions
|
|
@ -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]")
|
||||||
|
|
|
||||||
|
|
@ -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)"
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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(
|
||||||
|
|
|
||||||
|
|
@ -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")
|
||||||
|
|
|
||||||
|
|
@ -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(
|
||||||
|
|
|
||||||
|
|
@ -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():
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue