Rebuild database in client & cli
This commit is contained in:
parent
954e96370a
commit
11d9f2701e
5 changed files with 103 additions and 0 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
46
tests/test_rebuild.py
Normal 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()
|
||||
Loading…
Reference in a new issue