From 910d38274892ffca17e5a43681547a955a0d1029 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 3 Dec 2025 11:58:33 +0200 Subject: [PATCH] Introduce visualize_chunk(), refactor inspector to use it --- haiku_rag_slim/haiku/rag/client.py | 88 ++++++++++++++ haiku_rag_slim/haiku/rag/inspector/app.py | 23 +--- .../rag/inspector/widgets/search_modal.py | 38 ++---- .../rag/inspector/widgets/visual_modal.py | 114 ++++-------------- tests/test_client.py | 113 +++++++++++++++++ 5 files changed, 236 insertions(+), 140 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 064a8b0e..e4dde1b9 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -870,6 +870,94 @@ class HaikuRAG: qa_agent = get_qa_agent(self, config=self._config, system_prompt=system_prompt) return await qa_agent.answer(question) + async def visualize_chunk(self, chunk: Chunk) -> list: + """Render page images with bounding box highlights for a chunk. + + Gets the DoclingDocument from the chunk's document, resolves bounding boxes + from chunk metadata, and renders all pages that contain bounding boxes with + yellow/orange highlight overlays. + + Args: + chunk: The chunk to visualize. + + Returns: + List of PIL Image objects, one per page with bounding boxes. + Empty list if no bounding boxes or page images available. + """ + from copy import deepcopy + + from PIL import ImageDraw + + # Get the document + if not chunk.document_id: + return [] + + doc = await self.document_repository.get_by_id(chunk.document_id) + if not doc: + return [] + + # Get DoclingDocument + docling_doc = doc.get_docling_document() + if not docling_doc: + return [] + + # Resolve bounding boxes from chunk metadata + chunk_meta = chunk.get_chunk_metadata() + bounding_boxes = chunk_meta.resolve_bounding_boxes(docling_doc) + if not bounding_boxes: + return [] + + # Group bounding boxes by page + boxes_by_page: dict[int, list] = {} + for bbox in bounding_boxes: + if bbox.page_no not in boxes_by_page: + boxes_by_page[bbox.page_no] = [] + boxes_by_page[bbox.page_no].append(bbox) + + # Render each page with its bounding boxes + images = [] + for page_no in sorted(boxes_by_page.keys()): + if page_no not in docling_doc.pages: + continue + + page = docling_doc.pages[page_no] + if page.image is None or page.image.pil_image is None: + continue + + pil_image = page.image.pil_image + page_height = page.size.height + + # Calculate scale factor (image pixels vs document coordinates) + scale_x = pil_image.width / page.size.width + scale_y = pil_image.height / page.size.height + + # Draw bounding boxes + image = deepcopy(pil_image) + draw = ImageDraw.Draw(image, "RGBA") + + for bbox in boxes_by_page[page_no]: + # Convert from document coordinates to image coordinates + # Document coords are bottom-left origin, PIL uses top-left + x0 = bbox.left * scale_x + y0 = (page_height - bbox.top) * scale_y + x1 = bbox.right * scale_x + y1 = (page_height - bbox.bottom) * scale_y + + # Ensure proper ordering (y0 should be less than y1 for PIL) + if y0 > y1: + y0, y1 = y1, y0 + + # Draw filled rectangle with transparency + fill_color = (255, 255, 0, 80) # Yellow with transparency + outline_color = (255, 165, 0, 255) # Orange outline + + draw.rectangle([(x0, y0), (x1, y1)], fill=fill_color, outline=None) + draw.rectangle([(x0, y0), (x1, y1)], outline=outline_color, width=3) + + images.append(image) + + return images + async def rebuild_database( self, mode: RebuildMode = RebuildMode.FULL ) -> AsyncGenerator[str, None]: diff --git a/haiku_rag_slim/haiku/rag/inspector/app.py b/haiku_rag_slim/haiku/rag/inspector/app.py index e3101ed2..9ee0b101 100644 --- a/haiku_rag_slim/haiku/rag/inspector/app.py +++ b/haiku_rag_slim/haiku/rag/inspector/app.py @@ -187,31 +187,10 @@ class InspectorApp(App): # type: ignore[misc] return chunk = chunk_list.chunks[idx] - if not chunk.document_id: - return - - document = await self.client.get_document_by_id(chunk.document_id) - if not document: - return - - docling_doc = document.get_docling_document() - if not docling_doc: - self.notify("No DoclingDocument available", severity="warning") - return - - meta = chunk.get_chunk_metadata() - bounding_boxes = meta.resolve_bounding_boxes(docling_doc) from haiku.rag.inspector.widgets.visual_modal import VisualGroundingModal - await self.push_screen( - VisualGroundingModal( - docling_document=docling_doc, - bounding_boxes=bounding_boxes, - page_numbers=meta.page_numbers, - document_uri=chunk.document_uri, - ) - ) + await self.push_screen(VisualGroundingModal(chunk=chunk, client=self.client)) def run_inspector(db_path: Path | None = None) -> None: diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/search_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/search_modal.py index 1cdbc8a8..7f3df9aa 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/search_modal.py +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/search_modal.py @@ -166,39 +166,17 @@ class SearchModal(Screen): """Show visual grounding for the current chunk.""" list_view = self.query_one("#search-results", ListView) idx = list_view.index - if idx is None or not self.search_results: + if idx is None or not self.chunks: return - search_result = self.search_results[idx] - if not search_result.document_id or not search_result.chunk_id: - return - - status_label = self.query_one("#status-label", Static) - - document = await self.client.get_document_by_id(search_result.document_id) - if not document: - status_label.update("[yellow]Document not found[/yellow]") - return - - docling_doc = document.get_docling_document() - if not docling_doc: - status_label.update( - "[yellow]No DoclingDocument available for visual[/yellow]" - ) - return - - chunk = await self.client.chunk_repository.get_by_id(search_result.chunk_id) - bounding_boxes = [] - if chunk: - meta = chunk.get_chunk_metadata() - bounding_boxes = meta.resolve_bounding_boxes(docling_doc) + chunk = self.chunks[idx] from haiku.rag.inspector.widgets.visual_modal import VisualGroundingModal - modal = VisualGroundingModal( - docling_document=docling_doc, - bounding_boxes=bounding_boxes, - page_numbers=search_result.page_numbers, - document_uri=search_result.document_uri, + await self.app.push_screen( + VisualGroundingModal( + chunk=chunk, + client=self.client, + document_uri=chunk.document_uri, + ) ) - await self.app.push_screen(modal) diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py index b3078d15..fe4ceb85 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py @@ -1,4 +1,3 @@ -from copy import deepcopy from typing import TYPE_CHECKING from textual.app import ComposeResult @@ -9,12 +8,12 @@ from textual.widget import Widget from textual.widgets import Static from textual_image.widget import Image as TextualImage -from haiku.rag.store.models import BoundingBox - if TYPE_CHECKING: - from docling_core.types.doc.document import DoclingDocument from PIL.Image import Image as PILImage + from haiku.rag.client import HaikuRAG + from haiku.rag.store.models import Chunk + class VisualGroundingModal(Screen): """Modal screen for displaying visual grounding with bounding boxes.""" @@ -57,22 +56,21 @@ class VisualGroundingModal(Screen): def __init__( self, - docling_document: "DoclingDocument", - bounding_boxes: list[BoundingBox], - page_numbers: list[int], + chunk: "Chunk", + client: "HaikuRAG", document_uri: str | None = None, ): super().__init__() - self.docling_document = docling_document - self.document_uri = document_uri - self.bounding_boxes = bounding_boxes - self.page_numbers = sorted(set(page_numbers)) if page_numbers else [] + self.chunk = chunk + self.client = client + self.document_uri = document_uri or chunk.document_uri + self.images: list[PILImage] = [] self.current_page_idx = 0 self._image_widget: Widget = Static("Loading...", id="image-display") self._page_info = Static("", id="page-info") def compose(self) -> ComposeResult: - uri_display = self.document_uri or self.docling_document.name or "Document" + uri_display = self.document_uri or "Document" with Vertical(id="visual-header"): yield Static(f"[bold]Visual Grounding[/bold] - {uri_display}") with Horizontal(id="visual-content"): @@ -81,96 +79,36 @@ class VisualGroundingModal(Screen): yield self._page_info async def on_mount(self) -> None: - """Load and display the first page.""" + """Load images and display the first page.""" + self.images = await self.client.visualize_chunk(self.chunk) await self._render_current_page() async def _render_current_page(self) -> None: - """Render the current page with bounding boxes.""" - if not self.page_numbers: + """Render the current page.""" + if not self.images: if isinstance(self._image_widget, Static): - self._image_widget.update("[red]No page information available[/red]") + self._image_widget.update( + "[yellow]No page images available[/yellow]\n" + "This document was converted without page images." + ) self._page_info.update("") return - current_page = self.page_numbers[self.current_page_idx] self._page_info.update( - f"Page {current_page} " - f"({self.current_page_idx + 1}/{len(self.page_numbers)}) " - f"- Use ←/→ to navigate" + f"Page {self.current_page_idx + 1}/{len(self.images)} - Use ←/→ to navigate" ) try: - image = self._render_page_with_boxes(current_page) - if image: - new_widget = TextualImage(image, id="rendered-image") - await self._image_widget.remove() - content = self.query_one("#visual-content", Horizontal) - await content.mount(new_widget) - self._image_widget = new_widget - elif isinstance(self._image_widget, Static): - self._image_widget.update( - "[yellow]No page image available[/yellow]\n" - "This document was converted without page images." - ) + image = self.images[self.current_page_idx] + new_widget = TextualImage(image, id="rendered-image") + await self._image_widget.remove() + content = self.query_one("#visual-content", Horizontal) + await content.mount(new_widget) + self._image_widget = new_widget except Exception as e: if isinstance(self._image_widget, Static): self._image_widget.update(f"[red]Error: {e}[/red]") - def _render_page_with_boxes(self, page_no: int) -> "PILImage | None": - """Render a page from DoclingDocument with bounding boxes.""" - from PIL import ImageDraw - - # Get the page from DoclingDocument - if page_no not in self.docling_document.pages: - return None - - page = self.docling_document.pages[page_no] - if page.image is None: - return None - - pil_image = page.image.pil_image - if pil_image is None: - return None - - # Get page dimensions - page_height = page.size.height - - # Calculate scale factor (image pixels vs document coordinates) - scale_x = pil_image.width / page.size.width - scale_y = pil_image.height / page.size.height - - # Get bounding boxes for this page - page_boxes = [bb for bb in self.bounding_boxes if bb.page_no == page_no] - - if page_boxes: - # Draw bounding boxes - image = deepcopy(pil_image) - draw = ImageDraw.Draw(image, "RGBA") - - for bbox in page_boxes: - # Convert from document coordinates to image coordinates - # Document coords are typically bottom-left origin - # PIL uses top-left origin - x0 = bbox.left * scale_x - y0 = (page_height - bbox.top) * scale_y # Flip Y - x1 = bbox.right * scale_x - y1 = (page_height - bbox.bottom) * scale_y # Flip Y - - # Ensure proper ordering (y0 should be less than y1 for PIL) - if y0 > y1: - y0, y1 = y1, y0 - - # Draw filled rectangle with transparency - fill_color = (255, 255, 0, 80) # Yellow with transparency - outline_color = (255, 165, 0, 255) # Orange outline - - draw.rectangle([(x0, y0), (x1, y1)], fill=fill_color, outline=None) - draw.rectangle([(x0, y0), (x1, y1)], outline=outline_color, width=3) - - return image - - return pil_image - async def action_dismiss(self, result=None) -> None: self.app.pop_screen() @@ -182,6 +120,6 @@ class VisualGroundingModal(Screen): async def action_next_page(self) -> None: """Navigate to the next page.""" - if self.current_page_idx < len(self.page_numbers) - 1: + if self.current_page_idx < len(self.images) - 1: self.current_page_idx += 1 await self._render_current_page() diff --git a/tests/test_client.py b/tests/test_client.py index 96fdd091..583c76e3 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1282,3 +1282,116 @@ async def test_client_file_update_stores_docling_json(temp_db_path): assert doc2.docling_document_json is not None assert doc2.docling_document_json != original_json assert doc2.docling_version == original_version # Version stays same + + +@pytest.mark.asyncio +async def test_client_visualize_chunk_no_document(temp_db_path): + """Test visualize_chunk returns empty list when chunk has no document_id.""" + async with HaikuRAG(temp_db_path, create=True) as client: + chunk = Chunk(content="Orphan chunk", order=0) + images = await client.visualize_chunk(chunk) + assert images == [] + + +@pytest.mark.asyncio +async def test_client_visualize_chunk_no_docling_document(temp_db_path): + """Test visualize_chunk returns empty list when document has no DoclingDocument.""" + async with HaikuRAG(temp_db_path, create=True) as client: + # Create document with custom chunks (no DoclingDocument) + custom_chunks = [Chunk(content="Custom chunk", order=0)] + doc = await client.create_document(content="Test content", chunks=custom_chunks) + + assert doc.id is not None + chunks = await client.chunk_repository.get_by_document_id(doc.id) + assert len(chunks) == 1 + + images = await client.visualize_chunk(chunks[0]) + assert images == [] + + +@pytest.mark.asyncio +async def test_client_visualize_chunk_no_bounding_boxes(temp_db_path): + """Test visualize_chunk returns empty list when chunk has no bounding boxes.""" + async with HaikuRAG(temp_db_path, create=True) as client: + # Create document from text (will have DoclingDocument but no page images) + doc = await client.create_document( + content="Simple text content without structure", + uri="test://simple", + ) + + assert doc.id is not None + assert doc.docling_document_json is not None + + chunks = await client.chunk_repository.get_by_document_id(doc.id) + assert len(chunks) >= 1 + + # Text documents converted via markdown won't have page images + # so visualize_chunk should return empty list + images = await client.visualize_chunk(chunks[0]) + assert images == [] + + +@pytest.mark.asyncio +async def test_client_visualize_chunk_returns_list(temp_db_path): + """Test visualize_chunk returns a list (empty or with images).""" + async with HaikuRAG(temp_db_path, create=True) as client: + # Create a structured document + markdown_content = """# Chapter 1 + +This is paragraph one about topic A. + +This is paragraph two about topic A continued. + +# Chapter 2 + +This is paragraph four about topic C. +""" + doc = await client.create_document( + content=markdown_content, + uri="test://structured", + ) + + assert doc.id is not None + chunks = await client.chunk_repository.get_by_document_id(doc.id) + + # Find a chunk with doc_item_refs + chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs] + + if chunks_with_refs: + # visualize_chunk should return a list (possibly empty if no page images) + images = await client.visualize_chunk(chunks_with_refs[0]) + assert isinstance(images, list) + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_client_visualize_chunk_with_pdf(temp_db_path): + """Test visualize_chunk returns images with bounding boxes for PDF documents.""" + from PIL.Image import Image as PILImage + + pdf_path = Path("tests/data/doclaynet.pdf") + if not pdf_path.exists(): + pytest.skip("doclaynet.pdf not found") + + async with HaikuRAG(temp_db_path, create=True) as client: + doc = await client.create_document_from_source(pdf_path) + assert isinstance(doc, Document) + assert doc.id is not None + assert doc.docling_document_json is not None + + chunks = await client.chunk_repository.get_by_document_id(doc.id) + assert len(chunks) > 0 + + # Find a chunk with doc_item_refs (bounding box info) + chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs] + assert len(chunks_with_refs) > 0, "PDF should have chunks with doc_item_refs" + + # Visualize a chunk - should return images with bounding boxes drawn + images = await client.visualize_chunk(chunks_with_refs[0]) + + assert isinstance(images, list) + assert len(images) > 0, "PDF with page images should return visualizations" + + # Verify returned objects are PIL Images + for img in images: + assert isinstance(img, PILImage)