Visualize all constituent chunks of a merged citation

CHANGELOG.md

@@ -17,6 +17,7 @@
 
 - `update_document` accepts a `uri` argument to change a document's URI.
 - `SearchResult.chunk_ids` and `Citation.chunk_ids` carry the chunk ids merged into an expanded result.
+- `visualize_chunk` accepts multiple chunks and reproduces merged-result expansion; chat, inspector, and app visualizations pass all cited chunks.
 - docling-serve requests fail over to another instance on transport/5xx errors and skip instances whose circuit breaker is open; tune via `providers.docling_serve.max_attempts` and `providers.docling_serve.circuit_breaker`.
 
 ### Fixed

app/backend/main.py

@@ -203,7 +203,11 @@ async def db_info(_: Request) -> JSONResponse:
 
 
 async def visualize_chunk(request: Request) -> JSONResponse:
-    """Return visual grounding images for a chunk as base64."""
+    """Return visual grounding images for one or more chunks as base64.
+
+    The path param accepts comma-separated chunk ids so a merged citation
+    can render the union of its constituent chunks' expansions.
+    """
     import base64
     from io import BytesIO
 
@@ -214,11 +218,15 @@ async def visualize_chunk(request: Request) -> JSONResponse:
 
     client = await get_client()
 
-    chunk = await client.chunk_repository.get_by_id(chunk_id)
+    chunks = []
-    if not chunk:
+    for cid in chunk_id.split(","):
+        chunk = await client.chunk_repository.get_by_id(cid)
+        if chunk:
+            chunks.append(chunk)
+    if not chunks:
         return JSONResponse({"error": "Chunk not found"}, status_code=404)
 
-    images = await client.visualize_chunk(chunk)
+    images = await client.visualize_chunk(chunks)
     if not images:
         return JSONResponse({"images": [], "message": "No visual grounding available"})
 
@@ -233,7 +241,7 @@ async def visualize_chunk(request: Request) -> JSONResponse:
         {
             "images": base64_images,
             "chunk_id": chunk_id,
-            "document_uri": chunk.document_uri,
+            "document_uri": chunks[0].document_uri,
         }
     )
 

app/frontend/components/CitationBlock.tsx

@@ -55,7 +55,13 @@ function CitationItem({
 					<button
 						type="button"
 						className="citation-view-btn"
-						onClick={() => onViewInDocument(citation.chunk_id)}
+						onClick={() =>
+							onViewInDocument(
+								citation.chunk_ids?.length
+									? citation.chunk_ids.join(",")
+									: citation.chunk_id,
+							)
+						}
 					>
 						View in Document
 					</button>
@@ -96,9 +102,12 @@ export default function CitationBlock({ citations }: CitationBlockProps) {
 		});
 
 		try {
-			const response = await fetch(`/api/visualize/${chunkId}`, {
+			const response = await fetch(
-				signal: controller.signal,
+				`/api/visualize/${encodeURIComponent(chunkId)}`,
-			});
+				{
+					signal: controller.signal,
+				},
+			);
 			const data = await response.json();
 
 			if (controller.signal.aborted) return;

app/frontend/lib/sessionStorage.ts

@@ -2,6 +2,7 @@ export interface Citation {
 	index: number;
 	document_id: string;
 	chunk_id: string;
+	chunk_ids?: string[];
 	document_uri: string;
 	document_title: string | null;
 	page_numbers: number[];

haiku_rag_slim/haiku/rag/chat/app.py

@@ -405,13 +405,18 @@ class ChatApp(App):
             return
 
         citation = selected_widgets[0].citation
-        chunk = await self.client.get_chunk_by_id(citation.chunk_id)
+        chunk_ids = citation.chunk_ids or [citation.chunk_id]
-        if not chunk:
+        chunks = []
+        for cid in chunk_ids:
+            chunk = await self.client.get_chunk_by_id(cid)
+            if chunk:
+                chunks.append(chunk)
+        if not chunks:
             return
 
         from haiku.rag.inspector.widgets.visual_modal import VisualGroundingModal
 
-        await self.push_screen(VisualGroundingModal(chunk=chunk, client=self.client))
+        await self.push_screen(VisualGroundingModal(chunk=chunks, client=self.client))
 
     async def action_show_info(self) -> None:
         """Show database info modal."""

haiku_rag_slim/haiku/rag/client/__init__.py

@@ -4,7 +4,7 @@ import json
 import logging
 import mimetypes
 import tempfile
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Sequence
 from datetime import datetime
 from enum import Enum
 from functools import cached_property
@@ -504,7 +504,7 @@ class HaikuRAG:
 
         return await analyze(self, question, filter)
 
-    async def visualize_chunk(self, chunk: Chunk) -> list:
+    async def visualize_chunk(self, chunk: Chunk | Sequence[Chunk]) -> list:
         from haiku.rag.client.search import visualize_chunk
 
         return await visualize_chunk(self, chunk)

haiku_rag_slim/haiku/rag/client/search.py

@@ -1,4 +1,5 @@
 import base64
+from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 from haiku.rag.store.models.chunk import Chunk, SearchResult, SearchType
@@ -223,12 +224,15 @@ async def expand_context(
     return expanded_results
 
 
-async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list:
+async def visualize_chunk(client: "HaikuRAG", chunk: "Chunk | Sequence[Chunk]") -> list:
-    """Render page images with bounding box highlights for a chunk.
+    """Render page images with bounding box highlights for one or more chunks.
 
-    Expands the chunk's context to find the full section, then resolves
+    Expands the chunks' context to find the full section, then resolves
     bounding boxes from all items in the expanded range. This ensures
-    visualization covers all pages the expanded content spans.
+    visualization covers all pages the expanded content spans. Passing all
+    constituent chunks of a merged search result (``SearchResult.chunk_ids``)
+    reproduces the merged expansion; chunks from a different document than
+    the first are ignored.
 
     Returns a list of PIL Image objects, one per page with bounding boxes.
     Empty list if no bounding boxes or page images available.
@@ -239,10 +243,15 @@ async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list:
 
     from haiku.rag.store.models.chunk import ChunkMetadata
 
-    if not chunk.document_id:
+    chunks = [chunk] if isinstance(chunk, Chunk) else list(chunk)
+    if not chunks:
         return []
+    document_id = chunks[0].document_id
+    if not document_id:
+        return []
+    chunks = [c for c in chunks if c.document_id == document_id]
 
-    doc = await client.document_repository.get_docling_data(chunk.document_id)
+    doc = await client.document_repository.get_docling_data(document_id)
     if not doc:
         return []
 
@@ -251,21 +260,28 @@ async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list:
         return []
 
     # Expand context to get all doc_item_refs in the section
-    chunk_meta = chunk.get_chunk_metadata()
+    search_results = [
-    if chunk_meta.doc_item_refs:
+        SearchResult(
-        search_result = SearchResult(
+            content=c.content,
-            content=chunk.content,
             score=1.0,
-            chunk_id=chunk.id,
+            chunk_id=c.id,
-            document_id=chunk.document_id,
+            document_id=c.document_id,
-            doc_item_refs=chunk_meta.doc_item_refs,
+            doc_item_refs=meta.doc_item_refs,
-            page_numbers=chunk_meta.page_numbers,
+            page_numbers=meta.page_numbers,
         )
-        expanded = await expand_context(client, [search_result])
+        for c in chunks
-        refs = expanded[0].doc_item_refs if expanded else chunk_meta.doc_item_refs
+        if (meta := c.get_chunk_metadata()).doc_item_refs
+    ]
+    if search_results:
+        expanded = await expand_context(client, search_results)
+        refs: list[str] = []
+        for result in expanded:
+            refs.extend(r for r in result.doc_item_refs if r not in refs)
+        if not refs:
+            refs = [r for sr in search_results for r in sr.doc_item_refs]
         meta = ChunkMetadata(doc_item_refs=refs)
     else:
-        meta = chunk_meta
+        meta = chunks[0].get_chunk_metadata()
     bounding_boxes = meta.resolve_bounding_boxes(docling_doc)
     if not bounding_boxes:
         return []
@@ -278,7 +294,7 @@ async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list:
         boxes_by_page[bbox.page_no].append(bbox)
 
     # Load only the needed page images
-    pages_doc = await client.document_repository.get_pages_data(chunk.document_id)
+    pages_doc = await client.document_repository.get_pages_data(document_id)
     if not pages_doc:
         return []
     page_images = pages_doc.get_page_images(list(boxes_by_page.keys()))

haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py

@@ -57,14 +57,14 @@ class VisualGroundingModal(Screen):
 
     def __init__(
         self,
-        chunk: "Chunk",
+        chunk: "Chunk | list[Chunk]",
         client: "HaikuRAG",
         document_uri: str | None = None,
     ):
         super().__init__()
-        self.chunk = chunk
+        self.chunks = chunk if isinstance(chunk, list) else [chunk]
         self.client = client
-        self.document_uri = document_uri or chunk.document_uri
+        self.document_uri = document_uri or self.chunks[0].document_uri
         self.images: list[PILImage] = []
         self.current_page_idx = 0
         self._image_widget: Widget = Static("Loading...", id="image-display")
@@ -81,7 +81,7 @@ class VisualGroundingModal(Screen):
 
     async def on_mount(self) -> None:
         """Load images and display the first page."""
-        self.images = await self.client.visualize_chunk(self.chunk)
+        self.images = await self.client.visualize_chunk(self.chunks)
         await self._render_current_page()
 
     async def _render_current_page(self) -> None:

haiku_rag_slim/haiku/rag/skills/_tools.py

@@ -131,14 +131,19 @@ def create_skill_extras(
     - 'visualize_chunk': returns visualizations for chunks in the database
     """
 
-    async def visualize_chunk(chunk_id: str) -> list:
+    async def visualize_chunk(chunk_id: str | list[str]) -> list:
         from haiku.rag.client import HaikuRAG
 
+        chunk_ids = [chunk_id] if isinstance(chunk_id, str) else chunk_id
         async with HaikuRAG(db_path, config=config, read_only=True) as rag:
-            chunk = await rag.get_chunk_by_id(chunk_id)
+            chunks = []
-            if chunk is None:
+            for cid in chunk_ids:
+                chunk = await rag.get_chunk_by_id(cid)
+                if chunk is not None:
+                    chunks.append(chunk)
+            if not chunks:
                 return []
-            return await rag.visualize_chunk(chunk)
+            return await rag.visualize_chunk(chunks)
 
     async def list_documents(
         limit: int | None = None,

tests/test_client.py

@@ -1590,6 +1590,82 @@ async def test_client_visualize_chunk_multi_page(temp_db_path):
             assert img.tobytes() != blank.tobytes()
 
 
+async def test_client_visualize_chunk_merged_chunks_union_pages(temp_db_path):
+    """Visualizing all chunks of a merged result covers the union of their
+    expansions, which a single constituent chunk alone does not reach."""
+    from docling_core.types.doc.base import BoundingBox, Size
+    from docling_core.types.doc.document import (
+        DoclingDocument,
+        ImageRef,
+        ProvenanceItem,
+    )
+    from docling_core.types.doc.labels import DocItemLabel
+    from PIL import Image as PilImageModule
+
+    docling_doc = DoclingDocument(name="merged-viz-test")
+    page_size = Size(width=612.0, height=792.0)
+    for page_no in (1, 2):
+        docling_doc.add_page(
+            page_no=page_no,
+            size=page_size,
+            image=ImageRef.from_pil(
+                PilImageModule.new("RGB", (612, 792), color="white"), dpi=72
+            ),
+        )
+
+    # Five large paragraphs: with max_context_chars=10000 each chunk's own
+    # outward expansion spans three items, so chunk one alone stays on page
+    # one while the merged ranges [0,2] and [2,4] union to cover page two.
+    pages = [1, 1, 1, 2, 2]
+    for i, page_no in enumerate(pages):
+        docling_doc.add_text(
+            label=DocItemLabel.PARAGRAPH,
+            text=f"Paragraph {i}. " + "x" * 4000,
+            prov=ProvenanceItem(
+                page_no=page_no,
+                bbox=BoundingBox(l=50, t=700 - i * 100, r=550, b=650 - i * 100),
+                charspan=(0, 20),
+            ),
+        )
+
+    chunks = [
+        Chunk(
+            content="Paragraph 0. " + "x" * 4000,
+            metadata={
+                "doc_item_refs": ["#/texts/0"],
+                "page_numbers": [1],
+                "labels": ["paragraph"],
+            },
+            order=0,
+            embedding=[0.1] * 2560,
+        ),
+        Chunk(
+            content="Paragraph 3. " + "x" * 4000,
+            metadata={
+                "doc_item_refs": ["#/texts/3"],
+                "page_numbers": [2],
+                "labels": ["paragraph"],
+            },
+            order=1,
+            embedding=[0.1] * 2560,
+        ),
+    ]
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        doc = await client.import_document(docling_doc, chunks, uri="test://merged")
+
+        stored_chunks = await client.chunk_repository.get_by_document_id(doc.id)
+        stored_chunks.sort(key=lambda c: c.order)
+        assert len(stored_chunks) == 2
+        c1, c2 = stored_chunks
+
+        solo_images = await client.visualize_chunk(c1)
+        assert len(solo_images) == 1
+
+        merged_images = await client.visualize_chunk([c1, c2])
+        assert len(merged_images) == 2
+
+
 # =============================================================================
 # convert() method tests
 # =============================================================================