Turn rebuild into a generator, track progress in command line

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]")

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()

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