From 221a00ab2f14edba5db4e30b7db5b9120ca3c05f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 3 Sep 2025 11:09:59 +0300 Subject: [PATCH] Turn migration into a cli command --- README.md | 9 ++- docs/cli.md | 17 ++++ docs/index.md | 3 + docs/mcp.md | 4 +- docs/server.md | 2 +- src/haiku/rag/cli.py | 16 ++++ src/haiku/rag/migration.py | 156 ++++++++++++++++++++----------------- 7 files changed, 131 insertions(+), 76 deletions(-) diff --git a/README.md b/README.md index 75ba3d1e..850f9df7 100644 --- a/README.md +++ b/README.md @@ -1,12 +1,14 @@ -# Haiku LanceDB RAG +# Haiku RAG Retrieval-Augmented Generation (RAG) library built on LanceDB. `haiku.rag` is a Retrieval-Augmented Generation (RAG) library built to work with LanceDB as a local vector database. It uses LanceDB for storing embeddings and performs semantic (vector) search as well as full-text search combined through native hybrid search with Reciprocal Rank Fusion. Both open-source (Ollama) as well as commercial (OpenAI, VoyageAI) embedding providers are supported. +> **Note**: Starting with version 0.7.0, haiku.rag uses LanceDB instead of SQLite. If you have an existing SQLite database, use `haiku-rag migrate old_database.sqlite` to migrate your data safely. + ## Features -- **Local LanceDB**: No external servers required +- **Local LanceDB**: No external servers required, supports also LanceDB cloud storage, S3, Google Cloud & Azure - **Multiple embedding providers**: Ollama, VoyageAI, OpenAI - **Multiple QA providers**: Any provider/model supported by Pydantic AI - **Native hybrid search**: Vector + full-text search with native LanceDB RRF reranking @@ -39,6 +41,9 @@ haiku-rag ask "Who is the author of haiku.rag?" --cite # Rebuild database (re-chunk and re-embed all documents) haiku-rag rebuild +# Migrate from SQLite to LanceDB +haiku-rag migrate old_database.sqlite + # Start server with file monitoring export MONITOR_DIRECTORIES="/path/to/docs" haiku-rag serve diff --git a/docs/cli.md b/docs/cli.md index 164d0efe..5b7d6bb6 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -45,6 +45,23 @@ 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: diff --git a/docs/index.md b/docs/index.md index 0b616ab9..a806bbab 100644 --- a/docs/index.md +++ b/docs/index.md @@ -2,6 +2,8 @@ `haiku.rag` is a Retrieval-Augmented Generation (RAG) library built to work with LanceDB as a local vector database. It uses LanceDB for storing embeddings and performs semantic (vector) search as well as full-text search combined through native hybrid search with Reciprocal Rank Fusion. Both open-source (Ollama, MixedBread AI) as well as commercial (OpenAI, VoyageAI) embedding providers are supported. +> **Note**: Starting with version 0.7.0, haiku.rag uses LanceDB instead of SQLite. If you have an existing SQLite database, use `haiku-rag migrate old_database.sqlite` to migrate your data safely. + ## Features - **Local LanceDB**: No need to run additional servers @@ -42,6 +44,7 @@ Or use the CLI: haiku-rag add "Your document content" haiku-rag search "query" haiku-rag ask "Who is the author of haiku.rag?" +haiku-rag migrate old_database.sqlite # Migrate from SQLite ``` ## Documentation diff --git a/docs/mcp.md b/docs/mcp.md index 22f0f667..f0bf8c93 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -19,10 +19,10 @@ The MCP server exposes `haiku.rag` as MCP tools for compatible MCP clients. ## Starting MCP Server -The MCP server starts automatically with the serve command and supports `Streamable HTTP`, `stdio` and `SSE` transports: +The MCP server starts automatically with the serve command and supports Streamable HTTP, stdio and SSE transports: ```bash -# Default HTTP transport +# Default streamable HTTP transport haiku-rag serve # stdio transport (for Claude Desktop) diff --git a/docs/server.md b/docs/server.md index d0126eb4..afd8a836 100644 --- a/docs/server.md +++ b/docs/server.md @@ -9,7 +9,7 @@ haiku-rag serve ``` Transport options: -- `--http` (default) - Streamable HTTP transport +- Default - Streamable HTTP transport - `--stdio` - Standard input/output transport - `--sse` - Server-sent events transport diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index 6594c923..efa5dbe5 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -8,6 +8,7 @@ from rich.console import Console from haiku.rag.app import HaikuRAGApp from haiku.rag.config import Config +from haiku.rag.migration import migrate_sqlite_to_lancedb from haiku.rag.utils import is_up_to_date if not Config.ENV == "development": @@ -222,5 +223,20 @@ def serve( asyncio.run(app.serve(transport=transport)) +@cli.command("migrate", help="Migrate an SQLite database to LanceDB") +def migrate( + sqlite_path: Path = typer.Argument( + help="Path to the SQLite database file to migrate", + ), +): + # Generate LanceDB path in same parent directory + lancedb_path = sqlite_path.parent / (sqlite_path.stem + ".lancedb") + + success = asyncio.run(migrate_sqlite_to_lancedb(sqlite_path, lancedb_path)) + + if not success: + raise typer.Exit(1) + + if __name__ == "__main__": cli() diff --git a/src/haiku/rag/migration.py b/src/haiku/rag/migration.py index 1a05efd3..834cb98d 100644 --- a/src/haiku/rag/migration.py +++ b/src/haiku/rag/migration.py @@ -8,11 +8,11 @@ This script will: 3. Preserve all documents, chunks, embeddings, and settings """ -import asyncio import json import sqlite3 import struct from pathlib import Path +from uuid import uuid4 from rich.console import Console from rich.progress import Progress, TaskID @@ -63,7 +63,7 @@ class SQLiteToLanceDBMigrator: doc_task = progress.add_task( "[green]Migrating documents...", total=None ) - documents = self._migrate_documents( + document_id_mapping = self._migrate_documents( sqlite_conn, lance_store, progress, doc_task ) @@ -71,7 +71,9 @@ class SQLiteToLanceDBMigrator: chunk_task = progress.add_task( "[yellow]Migrating chunks and embeddings...", total=None ) - self._migrate_chunks(sqlite_conn, lance_store, progress, chunk_task) + self._migrate_chunks( + sqlite_conn, lance_store, progress, chunk_task, document_id_mapping + ) # Migrate settings settings_task = progress.add_task( @@ -82,10 +84,23 @@ class SQLiteToLanceDBMigrator: ) sqlite_conn.close() + + # Optimize the chunks table after migration + self.console.print("[blue]Optimizing LanceDB...[/blue]") + try: + lance_store.chunks_table.optimize() + self.console.print("[green]✅ Optimization completed[/green]") + except Exception as e: + self.console.print( + f"[yellow]Warning: Optimization failed: {e}[/yellow]" + ) + lance_store.close() self.console.print("[green]✅ Migration completed successfully![/green]") - self.console.print(f"[green]✅ Migrated {len(documents)} documents[/green]") + self.console.print( + f"[green]✅ Migrated {len(document_id_mapping)} documents[/green]" + ) return True except Exception as e: @@ -101,17 +116,22 @@ class SQLiteToLanceDBMigrator: lance_store: Store, progress: Progress, task: TaskID, - ) -> list[dict]: - """Migrate documents from SQLite to LanceDB.""" + ) -> dict[int, str]: + """Migrate documents from SQLite to LanceDB and return ID mapping.""" cursor = sqlite_conn.cursor() cursor.execute( "SELECT id, content, uri, metadata, created_at, updated_at FROM documents ORDER BY id" ) documents = [] + id_mapping = {} # Maps old integer ID to new UUID + for row in cursor.fetchall(): + new_uuid = str(uuid4()) + id_mapping[row["id"]] = new_uuid + doc_data = { - "id": row["id"], + "id": new_uuid, "content": row["content"], "uri": row["uri"], "metadata": json.loads(row["metadata"]) if row["metadata"] else {}, @@ -138,7 +158,7 @@ class SQLiteToLanceDBMigrator: lance_store.documents_table.add(doc_records) progress.update(task, completed=len(documents), total=len(documents)) - return documents + return id_mapping def _migrate_chunks( self, @@ -146,37 +166,78 @@ class SQLiteToLanceDBMigrator: lance_store: Store, progress: Progress, task: TaskID, + document_id_mapping: dict[int, str], ): """Migrate chunks and embeddings from SQLite to LanceDB.""" cursor = sqlite_conn.cursor() - # Get chunks with their embeddings + # Get chunks first cursor.execute(""" - SELECT - c.id, c.document_id, c.content, c.metadata, - ce.embedding - FROM chunks c - LEFT JOIN chunk_embeddings ce ON c.id = ce.chunk_id - ORDER BY c.id + SELECT id, document_id, content, metadata + FROM chunks + ORDER BY id """) + chunks_data = cursor.fetchall() + + # Get embeddings separately to avoid vec0 virtual table issues + embeddings_map = {} + try: + # Try to get embeddings from the vec0 tables directly + cursor.execute(""" + SELECT + r.chunk_id, + v.vectors + FROM chunk_embeddings_rowids r + JOIN chunk_embeddings_vector_chunks00 v ON r.rowid = v.rowid + """) + + for row in cursor.fetchall(): + chunk_id = row[0] + vectors_blob = row[1] + if vectors_blob and chunk_id not in embeddings_map: + embeddings_map[chunk_id] = vectors_blob + + except sqlite3.OperationalError as e: + self.console.print( + f"[yellow]Warning: Could not extract embeddings: {e}[/yellow]" + ) + self.console.print( + "[yellow]Continuing migration without embeddings...[/yellow]" + ) + chunks = [] - for row in cursor.fetchall(): - # Deserialize the embedding + for row in chunks_data: + # Generate new UUID for chunk + chunk_uuid = str(uuid4()) + + # Map the old document_id to new UUID + document_uuid = document_id_mapping.get(row["document_id"]) + if not document_uuid: + self.console.print( + f"[yellow]Warning: Document ID {row['document_id']} not found in mapping for chunk {row['id']}[/yellow]" + ) + continue + + # Get embedding for this chunk embedding = [] - if row["embedding"]: + embedding_blob = embeddings_map.get(row["id"]) + if embedding_blob: try: - embedding = deserialize_sqlite_embedding(row["embedding"]) + embedding = deserialize_sqlite_embedding(embedding_blob) except Exception as e: self.console.print( f"[yellow]Warning: Failed to deserialize embedding for chunk {row['id']}: {e}[/yellow]" ) # Generate a zero vector of the expected dimension embedding = [0.0] * lance_store.embedder._vector_dim + else: + # No embedding found, generate zero vector + embedding = [0.0] * lance_store.embedder._vector_dim chunk_data = { - "id": row["id"], - "document_id": row["document_id"], + "id": chunk_uuid, + "document_id": document_uuid, "content": row["content"], "metadata": json.loads(row["metadata"]) if row["metadata"] else {}, "vector": embedding, @@ -216,9 +277,10 @@ class SQLiteToLanceDBMigrator: if row: settings_data = json.loads(row["settings"]) if row["settings"] else {} - # Update the existing settings in LanceDB + # Update the existing settings in LanceDB (use string ID) lance_store.settings_table.update( - where="id = 1", values={"settings": json.dumps(settings_data)} + where="id = 'settings'", + values={"settings": json.dumps(settings_data)}, ) progress.update(task, completed=1, total=1) @@ -252,51 +314,3 @@ async def migrate_sqlite_to_lancedb( migrator = SQLiteToLanceDBMigrator(sqlite_path, lancedb_path) return migrator.migrate() - - -def main(): - """CLI entry point for the migration script.""" - import argparse - - parser = argparse.ArgumentParser(description="Migrate SQLite database to LanceDB") - parser.add_argument("sqlite_path", type=Path, help="Path to SQLite database file") - parser.add_argument( - "--lancedb-path", type=Path, help="Path for LanceDB database (optional)" - ) - parser.add_argument( - "--backup", - action="store_true", - help="Create backup of SQLite database before migration", - ) - - args = parser.parse_args() - - console = Console() - - # Create backup if requested - if args.backup: - backup_path = args.sqlite_path.with_suffix(args.sqlite_path.suffix + ".backup") - console.print(f"[blue]Creating backup: {backup_path}[/blue]") - import shutil - - shutil.copy2(args.sqlite_path, backup_path) - console.print("[green]✅ Backup created[/green]") - - # Run migration - success = asyncio.run( - migrate_sqlite_to_lancedb(args.sqlite_path, args.lancedb_path) - ) - - if success: - console.print("[green]🎉 Migration completed successfully![/green]") - console.print(f"[green]SQLite database: {args.sqlite_path}[/green]") - console.print( - f"[green]LanceDB database: {args.lancedb_path or args.sqlite_path.with_suffix('.lancedb')}[/green]" - ) - else: - console.print("[red]❌ Migration failed![/red]") - exit(1) - - -if __name__ == "__main__": - main()