From 11d9f2701e3378112ddbe0c92436face57a5a8ed Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Wed, 2 Jul 2025 10:55:06 +0300
Subject: [PATCH] Rebuild database in client & cli

---
 src/haiku/rag/app.py                      |  8 ++++
 src/haiku/rag/cli.py                      | 15 ++++++++
 src/haiku/rag/client.py                   | 18 +++++++++
 src/haiku/rag/store/repositories/chunk.py | 16 ++++++++
 tests/test_rebuild.py                     | 46 +++++++++++++++++++++++
 5 files changed, 103 insertions(+)
 create mode 100644 tests/test_rebuild.py

diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py
index 1e81a14f..dca36151 100644
--- a/src/haiku/rag/app.py
+++ b/src/haiku/rag/app.py
@@ -72,6 +72,14 @@ class HaikuRAGApp:
             except Exception as e:
                 self.console.print(f"[red]Error: {e}[/red]")
 
+    async def rebuild(self):
+        async with HaikuRAG(db_path=self.db_path) as client:
+            try:
+                await client.rebuild_database()
+                self.console.print("[b]Database rebuild completed successfully.[/b]")
+            except Exception as e:
+                self.console.print(f"[red]Error rebuilding database: {e}[/red]")
+
     def _rich_print_document(self, doc: Document, truncate: bool = False):
         """Format a document for display."""
         if truncate:
diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py
index 2e012cf1..03f653a6 100644
--- a/src/haiku/rag/cli.py
+++ b/src/haiku/rag/cli.py
@@ -128,6 +128,21 @@ def ask(
     event_loop.run_until_complete(app.ask(question=question))
 
 
+@cli.command(
+    "rebuild",
+    help="Rebuild the database by deleting all chunks and re-indexing all documents",
+)
+def rebuild(
+    db: Path = typer.Option(
+        get_default_data_dir() / "haiku.rag.sqlite",
+        "--db",
+        help="Path to the SQLite database file",
+    ),
+):
+    app = HaikuRAGApp(db_path=db)
+    event_loop.run_until_complete(app.rebuild())
+
+
 @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 0f24b3b9..396a9fde 100644
--- a/src/haiku/rag/client.py
+++ b/src/haiku/rag/client.py
@@ -270,6 +270,24 @@ class HaikuRAG:
         qa_agent = get_qa_agent(self)
         return await qa_agent.answer(question)
 
+    async def rebuild_database(self) -> None:
+        """Rebuild the database by deleting all chunks and re-indexing all documents."""
+        documents = await self.list_documents()
+
+        if not documents:
+            return
+
+        await self.chunk_repository.delete_all()
+
+        for doc in documents:
+            if doc.id is not None:
+                await self.chunk_repository.create_chunks_for_document(
+                    doc.id, doc.content, commit=False
+                )
+
+        if self.store._connection:
+            self.store._connection.commit()
+
     def close(self):
         """Close the underlying store connection."""
         self.store.close()
diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py
index 1f33ec68..4261cfeb 100644
--- a/src/haiku/rag/store/repositories/chunk.py
+++ b/src/haiku/rag/store/repositories/chunk.py
@@ -208,6 +208,22 @@ class ChunkRepository(BaseRepository[Chunk]):
 
         return created_chunks
 
+    async def delete_all(self, commit: bool = True) -> bool:
+        """Delete all chunks from the database."""
+        if self.store._connection is None:
+            raise ValueError("Store connection is not available")
+
+        cursor = self.store._connection.cursor()
+
+        cursor.execute("DELETE FROM chunks_fts")
+        cursor.execute("DELETE FROM chunk_embeddings")
+        cursor.execute("DELETE FROM chunks")
+
+        deleted = cursor.rowcount > 0
+        if commit:
+            self.store._connection.commit()
+        return deleted
+
     async def delete_by_document_id(
         self, document_id: int, commit: bool = True
     ) -> bool:
diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py
new file mode 100644
index 00000000..1192e7d1
--- /dev/null
+++ b/tests/test_rebuild.py
@@ -0,0 +1,46 @@
+import pytest
+from datasets import Dataset
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.store.models.document import Document
+
+
+@pytest.mark.asyncio
+async def test_rebuild_database(qa_corpus: Dataset):
+    """Test rebuild functionality with existing documents."""
+    client = HaikuRAG(":memory:")
+
+    created_docs: list[Document] = []
+    for content in qa_corpus["document_extracted"][:3]:
+        doc = await client.create_document(
+            content=content,
+        )
+        created_docs.append(doc)
+
+    documents_before = await client.list_documents()
+    assert len(documents_before) == 3
+
+    chunks_before = []
+    for doc in created_docs:
+        assert doc.id is not None
+        doc_chunks = await client.chunk_repository.get_by_document_id(doc.id)
+        chunks_before.extend(doc_chunks)
+
+    assert len(chunks_before) > 0
+
+    # Perform rebuild
+    await client.rebuild_database()
+
+    documents_after = await client.list_documents()
+    assert len(documents_after) == 3
+
+    # Verify chunks were recreated
+    chunks_after = []
+    for doc in documents_after:
+        if doc.id is not None:
+            doc_chunks = await client.chunk_repository.get_by_document_id(doc.id)
+            chunks_after.extend(doc_chunks)
+
+    assert len(chunks_after) > 0
+
+    client.close()