From 11d9f2701e3378112ddbe0c92436face57a5a8ed Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 2 Jul 2025 10:55:06 +0300 Subject: [PATCH 1/3] 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() From 85443ee963836654a388ed8c1a1f589d31ce93d4 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 2 Jul 2025 11:10:31 +0300 Subject: [PATCH 2/3] 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 From fb837ede699b1cab0281cb71f0e024ac1417fcbf Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 2 Jul 2025 11:17:16 +0300 Subject: [PATCH 3/3] Documentation --- README.md | 3 +++ docs/cli.md | 10 ++++++++++ docs/python.md | 7 +++++++ 3 files changed, 20 insertions(+) diff --git a/README.md b/README.md index cf86ae99..bc3a9567 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,9 @@ haiku-rag search "query" # Ask questions haiku-rag ask "Who is the author of haiku.rag?" +# Rebuild database (re-chunk and re-embed all documents) +haiku-rag rebuild + # 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 fae3db8a..efedc679 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -35,6 +35,16 @@ haiku-rag get 1 haiku-rag delete 1 ``` +### Rebuild Database + +Rebuild the database by deleting all chunks & embeddings and re-indexing all documents: + +```bash +haiku-rag rebuild +``` + +Use this when you want to change things like the embedding model or chunk size for example. + ## Search Basic search: diff --git a/docs/python.md b/docs/python.md index ebc87f4c..8ad47f4e 100644 --- a/docs/python.md +++ b/docs/python.md @@ -67,6 +67,13 @@ await client.update_document(doc) await client.delete_document(doc.id) ``` +### Rebuilding the Database + +```python +async for doc_id in client.rebuild_database(): + print(f"Processed document {doc_id}") +``` + ## Searching Documents Basic search: