From 85443ee963836654a388ed8c1a1f589d31ce93d4 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Wed, 2 Jul 2025 11:10:31 +0300
Subject: [PATCH] Turn rebuild into a generator, track progress in command line

---
 src/haiku/rag/app.py    | 19 ++++++++++++++++++-
 src/haiku/rag/client.py | 10 ++++++++--
 tests/test_rebuild.py   |  8 +++++++-
 3 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py
index dca36151..e7d483e1 100644
--- a/src/haiku/rag/app.py
+++ b/src/haiku/rag/app.py
@@ -3,6 +3,7 @@ from pathlib import Path
 
 from rich.console import Console
 from rich.markdown import Markdown
+from rich.progress import Progress
 
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
@@ -75,7 +76,23 @@ class HaikuRAGApp:
     async def rebuild(self):
         async with HaikuRAG(db_path=self.db_path) as client:
             try:
-                await client.rebuild_database()
+                documents = await client.list_documents()
+                total_docs = len(documents)
+
+                if total_docs == 0:
+                    self.console.print(
+                        "[yellow]No documents found in database.[/yellow]"
+                    )
+                    return
+
+                self.console.print(
+                    f"[b]Rebuilding database with {total_docs} documents...[/b]"
+                )
+                with Progress() as progress:
+                    task = progress.add_task("Rebuilding...", total=total_docs)
+                    async for _ in client.rebuild_database():
+                        progress.update(task, advance=1)
+
                 self.console.print("[b]Database rebuild completed successfully.[/b]")
             except Exception as e:
                 self.console.print(f"[red]Error rebuilding database: {e}[/red]")
diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py
index 396a9fde..74478654 100644
--- a/src/haiku/rag/client.py
+++ b/src/haiku/rag/client.py
@@ -1,6 +1,7 @@
 import hashlib
 import mimetypes
 import tempfile
+from collections.abc import AsyncGenerator
 from pathlib import Path
 from typing import Literal
 from urllib.parse import urlparse
@@ -270,8 +271,12 @@ 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."""
+    async def rebuild_database(self) -> AsyncGenerator[int, None]:
+        """Rebuild the database by deleting all chunks and re-indexing all documents.
+
+        Yields:
+            int: The ID of the document currently being processed
+        """
         documents = await self.list_documents()
 
         if not documents:
@@ -284,6 +289,7 @@ class HaikuRAG:
                 await self.chunk_repository.create_chunks_for_document(
                     doc.id, doc.content, commit=False
                 )
+                yield doc.id
 
         if self.store._connection:
             self.store._connection.commit()
diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py
index 1192e7d1..3254ce1d 100644
--- a/tests/test_rebuild.py
+++ b/tests/test_rebuild.py
@@ -29,7 +29,13 @@ async def test_rebuild_database(qa_corpus: Dataset):
     assert len(chunks_before) > 0
 
     # Perform rebuild
-    await client.rebuild_database()
+    processed_doc_ids = []
+    async for doc_id in client.rebuild_database():
+        processed_doc_ids.append(doc_id)
+
+    # Verify all documents were processed
+    expected_doc_ids = [doc.id for doc in created_docs]
+    assert set(processed_doc_ids) == set(expected_doc_ids)
 
     documents_after = await client.list_documents()
     assert len(documents_after) == 3