diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 47c66bf8..330a3410 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -274,6 +274,38 @@ class HaikuRAGApp: for result in results: self._rich_print_search_result(result) + async def visualize_chunk(self, chunk_id: str): + """Display visual grounding images for a chunk.""" + from textual_image.renderable import Image as RichImage + + async with HaikuRAG(db_path=self.db_path, config=self.config) as self.client: + chunk = await self.client.chunk_repository.get_by_id(chunk_id) + if not chunk: + self.console.print(f"[red]Chunk with id {chunk_id} not found.[/red]") + return + + images = await self.client.visualize_chunk(chunk) + if not images: + self.console.print( + "[yellow]No visual grounding available for this chunk.[/yellow]" + ) + self.console.print( + "This may be because the document was converted without page images." + ) + return + + self.console.print(f"[bold]Visual grounding for chunk {chunk_id}[/bold]") + if chunk.document_uri: + self.console.print( + f"[repr.attrib_name]document[/repr.attrib_name]: {chunk.document_uri}" + ) + + for i, img in enumerate(images): + self.console.print( + f"\n[bold cyan]Page {i + 1}/{len(images)}[/bold cyan]" + ) + self.console.print(RichImage(img)) + async def ask( self, question: str, @@ -624,11 +656,13 @@ class HaikuRAGApp: content = Markdown(result.content) self.console.print( f"[repr.attrib_name]document_id[/repr.attrib_name]: {result.document_id} " + f"[repr.attrib_name]chunk_id[/repr.attrib_name]: {result.chunk_id} " f"[repr.attrib_name]score[/repr.attrib_name]: {result.score:.4f}" ) if result.document_uri: - self.console.print("[repr.attrib_name]document uri[/repr.attrib_name]:") - self.console.print(result.document_uri) + self.console.print( + f"[repr.attrib_name]document uri[/repr.attrib_name]: {result.document_uri}" + ) if result.document_title: self.console.print("[repr.attrib_name]document title[/repr.attrib_name]:") self.console.print(result.document_title) diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 8d996fc1..e02e2ff9 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -261,6 +261,21 @@ def search( asyncio.run(app.search(query=query, limit=limit, filter=filter)) +@cli.command("visualize", help="Show visual grounding for a chunk") +def visualize( + chunk_id: str = typer.Argument( + help="The ID of the chunk to visualize", + ), + db: Path | None = typer.Option( + None, + "--db", + help="Path to the LanceDB database file", + ), +): + app = create_app(db) + asyncio.run(app.visualize_chunk(chunk_id=chunk_id)) + + @cli.command("ask", help="Ask a question using the QA agent") def ask( question: str = typer.Argument( diff --git a/haiku_rag_slim/haiku/rag/graph/common/models.py b/haiku_rag_slim/haiku/rag/graph/common/models.py index 98f1eebf..8d1f7145 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/models.py +++ b/haiku_rag_slim/haiku/rag/graph/common/models.py @@ -31,6 +31,8 @@ class ResearchPlan(BaseModel): class Citation(BaseModel): """Resolved citation with full metadata for display/visual grounding.""" + document_id: str + chunk_id: str document_uri: str document_title: str | None = None page_numbers: list[int] = Field(default_factory=list) @@ -75,6 +77,8 @@ def resolve_citations( continue citations.append( Citation( + document_id=r.document_id or "", + chunk_id=chunk_id, document_uri=r.document_uri or "", document_title=r.document_title, page_numbers=r.page_numbers, diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index 899ababc..fc7fd41f 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -280,8 +280,11 @@ def format_citations(citations: "list[Citation]") -> str: return "" lines = ["## Citations\n"] for c in citations: - # Build citation header - parts = [f"- **{c.document_uri}**"] + # Build citation header with document_id and chunk_id + parts = [ + f"- document_id: `{c.document_id}` chunk_id: `{c.chunk_id}` " + f"uri: **{c.document_uri}**" + ] if c.document_title: parts.append(f' - "{c.document_title}"') location_parts = []