316 lines
11 KiB
Python
316 lines
11 KiB
Python
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
|
|
|
|
from haiku.rag.store.engine import Store
|
|
|
|
|
|
def deserialize_sqlite_embedding(data: bytes) -> list[float]:
|
|
"""Deserialize sqlite-vec embedding from bytes."""
|
|
if not data:
|
|
return []
|
|
# sqlite-vec stores embeddings as float32 arrays
|
|
num_floats = len(data) // 4
|
|
return list(struct.unpack(f"{num_floats}f", data))
|
|
|
|
|
|
class SQLiteToLanceDBMigrator:
|
|
"""Migrates data from SQLite to LanceDB."""
|
|
|
|
def __init__(self, sqlite_path: Path, lancedb_path: Path):
|
|
self.sqlite_path = sqlite_path
|
|
self.lancedb_path = lancedb_path
|
|
self.console = Console()
|
|
|
|
def migrate(self) -> bool:
|
|
"""Perform the migration."""
|
|
try:
|
|
self.console.print(
|
|
f"[blue]Starting migration from {self.sqlite_path} to {self.lancedb_path}[/blue]"
|
|
)
|
|
|
|
# Check if SQLite database exists
|
|
if not self.sqlite_path.exists():
|
|
self.console.print(
|
|
f"[red]SQLite database not found: {self.sqlite_path}[/red]"
|
|
)
|
|
return False
|
|
|
|
# Connect to SQLite database
|
|
sqlite_conn = sqlite3.connect(self.sqlite_path)
|
|
sqlite_conn.row_factory = sqlite3.Row
|
|
|
|
# Load the sqlite-vec extension
|
|
try:
|
|
import sqlite_vec
|
|
|
|
sqlite_conn.enable_load_extension(True)
|
|
sqlite_vec.load(sqlite_conn)
|
|
self.console.print("[blue]Loaded sqlite-vec extension[/blue]")
|
|
except Exception as e:
|
|
self.console.print(
|
|
f"[yellow]Warning: Could not load sqlite-vec extension: {e}[/yellow]"
|
|
)
|
|
self.console.print(
|
|
"[yellow]Install sqlite-vec with[/yellow]\n[green]uv pip install sqlite-vec [/green]"
|
|
)
|
|
exit(1)
|
|
|
|
# Create LanceDB store
|
|
lance_store = Store(self.lancedb_path, skip_validation=True)
|
|
|
|
with Progress() as progress:
|
|
# Migrate documents
|
|
doc_task = progress.add_task(
|
|
"[green]Migrating documents...", total=None
|
|
)
|
|
document_id_mapping = self._migrate_documents(
|
|
sqlite_conn, lance_store, progress, doc_task
|
|
)
|
|
|
|
# Migrate chunks and embeddings
|
|
chunk_task = progress.add_task(
|
|
"[yellow]Migrating chunks and embeddings...", total=None
|
|
)
|
|
self._migrate_chunks(
|
|
sqlite_conn, lance_store, progress, chunk_task, document_id_mapping
|
|
)
|
|
|
|
# Migrate settings
|
|
settings_task = progress.add_task(
|
|
"[blue]Migrating settings...", total=None
|
|
)
|
|
self._migrate_settings(
|
|
sqlite_conn, lance_store, progress, settings_task
|
|
)
|
|
|
|
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(document_id_mapping)} documents[/green]"
|
|
)
|
|
return True
|
|
|
|
except Exception as e:
|
|
self.console.print(f"[red]❌ Migration failed: {e}[/red]")
|
|
import traceback
|
|
|
|
self.console.print(f"[red]{traceback.format_exc()}[/red]")
|
|
return False
|
|
|
|
def _migrate_documents(
|
|
self,
|
|
sqlite_conn: sqlite3.Connection,
|
|
lance_store: Store,
|
|
progress: Progress,
|
|
task: TaskID,
|
|
) -> 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": new_uuid,
|
|
"content": row["content"],
|
|
"uri": row["uri"],
|
|
"metadata": json.loads(row["metadata"]) if row["metadata"] else {},
|
|
"created_at": row["created_at"],
|
|
"updated_at": row["updated_at"],
|
|
}
|
|
documents.append(doc_data)
|
|
|
|
# Batch insert documents to LanceDB
|
|
if documents:
|
|
from haiku.rag.store.engine import DocumentRecord
|
|
|
|
doc_records = [
|
|
DocumentRecord(
|
|
id=doc["id"],
|
|
content=doc["content"],
|
|
uri=doc["uri"],
|
|
metadata=json.dumps(doc["metadata"]),
|
|
created_at=doc["created_at"],
|
|
updated_at=doc["updated_at"],
|
|
)
|
|
for doc in documents
|
|
]
|
|
lance_store.documents_table.add(doc_records)
|
|
|
|
progress.update(task, completed=len(documents), total=len(documents))
|
|
return id_mapping
|
|
|
|
def _migrate_chunks(
|
|
self,
|
|
sqlite_conn: sqlite3.Connection,
|
|
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 first
|
|
cursor.execute("""
|
|
SELECT id, document_id, content, metadata
|
|
FROM chunks
|
|
ORDER BY id
|
|
""")
|
|
|
|
chunks_data = cursor.fetchall()
|
|
|
|
# Get embeddings using the sqlite-vec virtual table
|
|
embeddings_map = {}
|
|
try:
|
|
# Use the virtual table to get embeddings properly
|
|
cursor.execute("""
|
|
SELECT chunk_id, embedding
|
|
FROM chunk_embeddings
|
|
""")
|
|
|
|
for row in cursor.fetchall():
|
|
chunk_id = row[0]
|
|
embedding_blob = row[1]
|
|
if embedding_blob and chunk_id not in embeddings_map:
|
|
embeddings_map[chunk_id] = embedding_blob
|
|
|
|
except sqlite3.OperationalError as e:
|
|
self.console.print(
|
|
f"[yellow]Warning: Could not extract embeddings from virtual table: {e}[/yellow]"
|
|
)
|
|
|
|
chunks = []
|
|
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 = []
|
|
embedding_blob = embeddings_map.get(row["id"])
|
|
if embedding_blob:
|
|
try:
|
|
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": chunk_uuid,
|
|
"document_id": document_uuid,
|
|
"content": row["content"],
|
|
"metadata": json.loads(row["metadata"]) if row["metadata"] else {},
|
|
"vector": embedding,
|
|
}
|
|
chunks.append(chunk_data)
|
|
|
|
# Batch insert chunks to LanceDB
|
|
if chunks:
|
|
chunk_records = [
|
|
lance_store.ChunkRecord(
|
|
id=chunk["id"],
|
|
document_id=chunk["document_id"],
|
|
content=chunk["content"],
|
|
metadata=json.dumps(chunk["metadata"]),
|
|
vector=chunk["vector"],
|
|
)
|
|
for chunk in chunks
|
|
]
|
|
lance_store.chunks_table.add(chunk_records)
|
|
|
|
progress.update(task, completed=len(chunks), total=len(chunks))
|
|
|
|
def _migrate_settings(
|
|
self,
|
|
sqlite_conn: sqlite3.Connection,
|
|
lance_store: Store,
|
|
progress: Progress,
|
|
task: TaskID,
|
|
):
|
|
"""Migrate settings from SQLite to LanceDB."""
|
|
cursor = sqlite_conn.cursor()
|
|
|
|
try:
|
|
cursor.execute("SELECT id, settings FROM settings WHERE id = 1")
|
|
row = cursor.fetchone()
|
|
|
|
if row:
|
|
settings_data = json.loads(row["settings"]) if row["settings"] else {}
|
|
|
|
# Update the existing settings in LanceDB (use string ID)
|
|
lance_store.settings_table.update(
|
|
where="id = 'settings'",
|
|
values={"settings": json.dumps(settings_data)},
|
|
)
|
|
|
|
progress.update(task, completed=1, total=1)
|
|
else:
|
|
progress.update(task, completed=0, total=0)
|
|
|
|
except sqlite3.OperationalError:
|
|
# Settings table doesn't exist in old SQLite database
|
|
self.console.print(
|
|
"[yellow]No settings table found in SQLite database[/yellow]"
|
|
)
|
|
progress.update(task, completed=0, total=0)
|
|
|
|
|
|
async def migrate_sqlite_to_lancedb(
|
|
sqlite_path: Path, lancedb_path: Path | None = None
|
|
) -> bool:
|
|
"""
|
|
Migrate an existing SQLite database to LanceDB.
|
|
|
|
Args:
|
|
sqlite_path: Path to the existing SQLite database
|
|
lancedb_path: Path for the new LanceDB database (optional, will auto-generate if not provided)
|
|
|
|
Returns:
|
|
True if migration was successful, False otherwise
|
|
"""
|
|
if lancedb_path is None:
|
|
# Auto-generate LanceDB path
|
|
lancedb_path = sqlite_path.parent / (sqlite_path.stem + ".lancedb")
|
|
|
|
migrator = SQLiteToLanceDBMigrator(sqlite_path, lancedb_path)
|
|
return migrator.migrate()
|