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:
|
except Exception as e:
|
||||||
self.console.print(f"[red]Error: {e}[/red]")
|
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):
|
def _rich_print_document(self, doc: Document, truncate: bool = False):
|
||||||
"""Format a document for display."""
|
"""Format a document for display."""
|
||||||
if truncate:
|
if truncate:
|
||||||
|
|
|
||||||
|
|
@ -128,6 +128,21 @@ def ask(
|
||||||
event_loop.run_until_complete(app.ask(question=question))
|
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(
|
@cli.command(
|
||||||
"serve", help="Start the haiku.rag MCP server (by default in streamable HTTP mode)"
|
"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)
|
qa_agent = get_qa_agent(self)
|
||||||
return await qa_agent.answer(question)
|
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):
|
def close(self):
|
||||||
"""Close the underlying store connection."""
|
"""Close the underlying store connection."""
|
||||||
self.store.close()
|
self.store.close()
|
||||||
|
|
|
||||||
|
|
@ -208,6 +208,22 @@ class ChunkRepository(BaseRepository[Chunk]):
|
||||||
|
|
||||||
return created_chunks
|
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(
|
async def delete_by_document_id(
|
||||||
self, document_id: int, commit: bool = True
|
self, document_id: int, commit: bool = True
|
||||||
) -> bool:
|
) -> 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