From 2207b7d1ed726f66e1a29737d368e63aaebbf09d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 10 Sep 2025 11:24:51 +0300 Subject: [PATCH 1/3] vacuum() in store, client, cli. Removes old versions and reduces disk usage --- src/haiku/rag/app.py | 9 ++++++++ src/haiku/rag/cli.py | 12 ++++++++++ src/haiku/rag/client.py | 10 ++++++++ src/haiku/rag/migration.py | 6 ++--- src/haiku/rag/store/engine.py | 24 +++++++++++++++----- src/haiku/rag/store/repositories/settings.py | 11 ++++++--- tests/generate_benchmark_db.py | 3 +++ 7 files changed, 63 insertions(+), 12 deletions(-) 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(): From 455d90684b4adee535fc5e8a7a3ad8a70260ddc6 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 10 Sep 2025 11:35:56 +0300 Subject: [PATCH 2/3] Update docs --- docs/cli.md | 123 +++++++++++++++++++++++++++---------------------- docs/index.md | 1 - docs/python.md | 10 ++++ 3 files changed, 79 insertions(+), 55 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index 65f5eb91..0d20fa95 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -2,22 +2,17 @@ The `haiku-rag` CLI provides complete document management functionality. -## Shell Autocompletion +!!! note + All commands support: -Enable shell autocompletion for faster, error‑free usage. + - `--db` - Specify custom database path + - `-h` - Show help for specific command -- Temporary (current shell only): - ```bash - eval "$(haiku-rag --show-completion)" - ``` -- Permanent installation: - ```bash - haiku-rag --install-completion - ``` - -What’s completed: -- `get` and `delete`/`rm`: Document IDs from the selected database (respects `--db`). -- `add-src`: Local filesystem paths (URLs can still be typed manually). + Example: + ```bash + haiku-rag list --db /path/to/custom.db + haiku-rag add -h + ``` ## Document Management @@ -40,6 +35,10 @@ haiku-rag add-src /path/to/document.pdf haiku-rag add-src https://example.com/article.html ``` +!!! note + As you add documents to `haiku.rag` the database keeps growing. By default, `lanceDB` supports versioning + of your data. You can optimize and compact the database by running the [vaccum](#vacuum-optimize-and-cleanup) command. + ### Get Document ```bash @@ -55,33 +54,8 @@ haiku-rag delete haiku-rag rm # alias ``` -### Rebuild Database - -Rebuild the database by deleting all chunks & embeddings and re-indexing all documents: - -```bash -haiku-rag rebuild -``` - Use this when you want to change things like the embedding model or chunk size for example. -## Migration - -### Migrate from SQLite to LanceDB - -Migrate an existing SQLite database to LanceDB: - -```bash -haiku-rag migrate /path/to/old_database.sqlite -``` - -This will: -- Read all documents, chunks, embeddings, and settings from the SQLite database -- Create a new LanceDB database with the same data in the same directory -- Optimize the new database for best performance - -The original SQLite database remains unchanged, so you can safely migrate without risk of data loss. - ## Search Basic search: @@ -108,13 +82,6 @@ haiku-rag ask "Who is the author of haiku.rag?" --cite The QA agent will search your documents for relevant information and provide a comprehensive answer. With `--cite`, responses include citations showing which documents were used. -## Configuration - -View current configuration settings: -```bash -haiku-rag settings -``` - ## Server Start the MCP server: @@ -129,14 +96,62 @@ haiku-rag serve --stdio haiku-rag serve --sse ``` -## Options +## Settings -All commands support: -- `--db` - Specify custom database path -- `-h` - Show help for specific command - -Example: +View current configuration settings: ```bash -haiku-rag list --db /path/to/custom.db -haiku-rag add -h +haiku-rag settings ``` + +## Maintenance + +### Vacuum (Optimize and Cleanup) + +Reduce disk usage by optimizing and pruning old table versions across all tables: + +```bash +haiku-rag vacuum +``` + +### Rebuild Database + +Rebuild the database by deleting all chunks & embeddings and re-indexing all documents. This is useful +when want to switch embeddings provider or model: + +```bash +haiku-rag rebuild +``` + +## Migration + +### Migrate from SQLite to LanceDB + +Migrate an existing SQLite database to LanceDB: + +```bash +haiku-rag migrate /path/to/old_database.sqlite +``` + +This will: +- Read all documents, chunks, embeddings, and settings from the SQLite database +- Create a new LanceDB database with the same data in the same directory +- Optimize the new database for best performance + +The original SQLite database remains unchanged, so you can safely migrate without risk of data loss. + +## Shell Autocompletion + +Enable shell autocompletion for faster, error‑free usage. + +- Temporary (current shell only): + ```bash + eval "$(haiku-rag --show-completion)" + ``` +- Permanent installation: + ```bash + haiku-rag --install-completion + ``` + +What’s completed: +- `get` and `delete`/`rm`: Document IDs from the selected database (respects `--db`). +- `add-src`: Local filesystem paths (URLs can still be typed manually). diff --git a/docs/index.md b/docs/index.md index a806bbab..17b6d404 100644 --- a/docs/index.md +++ b/docs/index.md @@ -52,7 +52,6 @@ haiku-rag migrate old_database.sqlite # Migrate from SQLite - [Installation](installation.md) - Install haiku.rag with different providers - [Configuration](configuration.md) - Environment variables and settings - [CLI](cli.md) - Command line interface usage -- [Question Answering](qa.md) - QA agents and natural language queries - [Server](server.md) - File monitoring and server mode - [MCP](mcp.md) - Model Context Protocol integration - [Python](python.md) - Python API reference diff --git a/docs/python.md b/docs/python.md index c1f9af5a..547de767 100644 --- a/docs/python.md +++ b/docs/python.md @@ -99,6 +99,16 @@ async for doc_id in client.rebuild_database(): print(f"Processed document {doc_id}") ``` +## Maintenance + +Run maintenance to optimize storage and prune old table versions: + +```python +await client.vacuum() +``` + +This compacts tables and removes historical versions to keep disk usage in check. It’s safe to run anytime, for example after bulk imports or periodically in long‑running apps. + ## Searching Documents The search method performs native hybrid search (vector + full-text) using LanceDB with optional reranking for improved relevance: From ddaf1cb4d7005a01a3582ad059f96f2dae1a74fc Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 10 Sep 2025 13:36:12 +0300 Subject: [PATCH 3/3] Suppress warnings in cli --- src/haiku/rag/logging.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/haiku/rag/logging.py b/src/haiku/rag/logging.py index 3c5202e9..d727f84e 100644 --- a/src/haiku/rag/logging.py +++ b/src/haiku/rag/logging.py @@ -1,4 +1,5 @@ import logging +import warnings from rich.console import Console from rich.logging import RichHandler @@ -50,4 +51,6 @@ def configure_cli_logging(level: int = logging.INFO) -> logging.Logger: logger = get_logger() logger.setLevel(level) logger.propagate = False + + warnings.filterwarnings("ignore") return logger