Rebuild database in client & cli

This commit is contained in:
Yiorgis Gozadinos 2025-07-02 10:55:06 +03:00
parent 954e96370a
commit 11d9f2701e
No known key found for this signature in database
5 changed files with 103 additions and 0 deletions

View file

@ -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:

View file

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

View file

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

View file

@ -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:

46
tests/test_rebuild.py Normal file
View file

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