danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Make styles consistent

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-09-23 12:24:09 +03:00
parent 52ea1ebb70
commit 5ea5de274c
No known key found for this signature in database
3 changed files with 24 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
src/haiku/rag/app.py
View file

@ -121,7 +121,7 @@ class HaikuRAGApp:
doc = await self.client.create_document(text)
self._rich_print_document(doc, truncate=True)
self.console.print(
f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]"
f"[bold green]Document {doc.id} added successfully.[/bold green]"
)
async def add_document_from_source(self, source: str, title: str | None = None):
@ -129,7 +129,7 @@ class HaikuRAGApp:
doc = await self.client.create_document_from_source(source, title=title)
self._rich_print_document(doc, truncate=True)
self.console.print(
f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]"
f"[bold green]Document {doc.id} added successfully.[/bold green]"
)
async def get_document(self, doc_id: str):
@ -144,7 +144,9 @@ class HaikuRAGApp:
async with HaikuRAG(db_path=self.db_path) as self.client:
deleted = await self.client.delete_document(doc_id)
if deleted:
self.console.print(f"[b]Document {doc_id} deleted successfully.[/b]")
self.console.print(
f"[bold green]Document {doc_id} deleted successfully.[/bold green]"
)
else:
self.console.print(
f"[yellow]Document with id {doc_id} not found.[/yellow]"
@ -154,7 +156,7 @@ class HaikuRAGApp:
async with HaikuRAG(db_path=self.db_path) as self.client:
results = await self.client.search(query, limit=limit)
if not results:
self.console.print("[red]No results found.[/red]")
self.console.print("[yellow]No results found.[/yellow]")
return
for chunk, score in results:
self._rich_print_search_result(chunk, score)
@ -287,14 +289,16 @@ class HaikuRAGApp:
return
self.console.print(
f"[b]Rebuilding database with {total_docs} documents...[/b]"
f"[bold cyan]Rebuilding database with {total_docs} documents...[/bold cyan]"
)
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]")
self.console.print(
"[bold green]Database rebuild completed successfully.[/bold green]"
)
except Exception as e:
self.console.print(f"[red]Error rebuilding database: {e}[/red]")
@ -303,7 +307,9 @@ class HaikuRAGApp:
try:
async with HaikuRAG(db_path=self.db_path, skip_validation=True) as client:
await client.vacuum()
self.console.print("[b]Vacuum completed successfully.[/b]")
self.console.print(
"[bold green]Vacuum completed successfully.[/bold green]"
)
except Exception as e:
self.console.print(f"[red]Error during vacuum: {e}[/red]")
@ -325,7 +331,9 @@ class HaikuRAGApp:
else:
display_value = field_value
self.console.print(f" [cyan]{field_name}[/cyan]: {display_value}")
self.console.print(
f" [repr.attrib_name]{field_name}[/repr.attrib_name]: {display_value}"
)
def _rich_print_document(self, doc: Document, truncate: bool = False):
"""Format a document for display."""

4
src/haiku/rag/migration.py
View file

@ -51,7 +51,7 @@ class SQLiteToLanceDBMigrator:
sqlite_conn.enable_load_extension(True)
sqlite_vec.load(sqlite_conn)
self.console.print("[blue]Loaded sqlite-vec extension[/blue]")
self.console.print("[cyan]Loaded sqlite-vec extension[/cyan]")
except Exception as e:
self.console.print(
f"[yellow]Warning: Could not load sqlite-vec extension: {e}[/yellow]"
@ -92,7 +92,7 @@ class SQLiteToLanceDBMigrator:
sqlite_conn.close()
# Optimize and cleanup using centralized vacuum
self.console.print("[blue]Optimizing LanceDB...[/blue]")
self.console.print("[cyan]Optimizing LanceDB...[/cyan]")
try:
lance_store.vacuum()
self.console.print("[green]✅ Optimization completed[/green]")

10
tests/test_app.py
View file

@ -57,7 +57,7 @@ async def test_add_document_from_text(app: HaikuRAGApp, monkeypatch):
mock_client.create_document.assert_called_once_with("test document")
mock_rich_print.assert_called_once_with(mock_doc, truncate=True)
mock_print.assert_called_once_with(
"[b]Document with id [cyan]1[/cyan] added successfully.[/b]"
"[bold green]Document 1 added successfully.[/bold green]"
)
@ -83,7 +83,7 @@ async def test_add_document_from_source(app: HaikuRAGApp, monkeypatch):
)
mock_rich_print.assert_called_once_with(mock_doc, truncate=True)
mock_print.assert_called_once_with(
"[b]Document with id [cyan]1[/cyan] added successfully.[/b]"
"[bold green]Document 1 added successfully.[/bold green]"
)
@ -135,7 +135,9 @@ async def test_delete_document(app: HaikuRAGApp, monkeypatch):
await app.delete_document("1")
mock_client.delete_document.assert_called_once_with("1")
mock_print.assert_called_once_with("[b]Document 1 deleted successfully.[/b]")
mock_print.assert_called_once_with(
"[bold green]Document 1 deleted successfully.[/bold green]"
)
@pytest.mark.asyncio
@ -170,7 +172,7 @@ async def test_search_no_results(app: HaikuRAGApp, monkeypatch):
await app.search("query")
mock_client.search.assert_called_once_with("query", limit=5)
mock_print.assert_called_once_with("[red]No results found.[/red]")
mock_print.assert_called_once_with("[yellow]No results found.[/yellow]")
@pytest.mark.asyncio