From 494f7744731757cccf8b07ae590a3833408ba2a1 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 8 Jul 2026 15:26:27 +0300 Subject: [PATCH] Visualize all constituent chunks of a merged citation --- CHANGELOG.md | 1 + app/backend/main.py | 18 +++-- app/frontend/components/CitationBlock.tsx | 17 ++++- app/frontend/lib/sessionStorage.ts | 1 + haiku_rag_slim/haiku/rag/chat/app.py | 11 ++- haiku_rag_slim/haiku/rag/client/__init__.py | 4 +- haiku_rag_slim/haiku/rag/client/search.py | 52 ++++++++----- .../rag/inspector/widgets/visual_modal.py | 8 +- haiku_rag_slim/haiku/rag/skills/_tools.py | 13 +++- tests/test_client.py | 76 +++++++++++++++++++ 10 files changed, 161 insertions(+), 40 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 516db165..b9c7254f 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/app/backend/main.py b/app/backend/main.py index ca110bbb..b8a62207 100644 --- a/app/backend/main.py +++ b/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) - if not chunk: + chunks = [] + 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, } ) diff --git a/app/frontend/components/CitationBlock.tsx b/app/frontend/components/CitationBlock.tsx index 3221e0ec..177a8f13 100644 --- a/app/frontend/components/CitationBlock.tsx +++ b/app/frontend/components/CitationBlock.tsx @@ -55,7 +55,13 @@ function CitationItem({ @@ -96,9 +102,12 @@ export default function CitationBlock({ citations }: CitationBlockProps) { }); try { - const response = await fetch(`/api/visualize/${chunkId}`, { - signal: controller.signal, - }); + const response = await fetch( + `/api/visualize/${encodeURIComponent(chunkId)}`, + { + signal: controller.signal, + }, + ); const data = await response.json(); if (controller.signal.aborted) return; diff --git a/app/frontend/lib/sessionStorage.ts b/app/frontend/lib/sessionStorage.ts index 70b0b01a..26827fb1 100644 --- a/app/frontend/lib/sessionStorage.ts +++ b/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[]; diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 34a1e53d..ad7e6e85 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/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) - if not chunk: + chunk_ids = citation.chunk_ids or [citation.chunk_id] + 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.""" diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 910ef6a0..fdd63e2a 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/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) diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 06e6f2f6..cec8c59c 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/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: - """Render page images with bounding box highlights for a chunk. +async def visualize_chunk(client: "HaikuRAG", chunk: "Chunk | Sequence[Chunk]") -> list: + """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() - if chunk_meta.doc_item_refs: - search_result = SearchResult( - content=chunk.content, + search_results = [ + SearchResult( + content=c.content, score=1.0, - chunk_id=chunk.id, - document_id=chunk.document_id, - doc_item_refs=chunk_meta.doc_item_refs, - page_numbers=chunk_meta.page_numbers, + chunk_id=c.id, + document_id=c.document_id, + doc_item_refs=meta.doc_item_refs, + page_numbers=meta.page_numbers, ) - expanded = await expand_context(client, [search_result]) - refs = expanded[0].doc_item_refs if expanded else chunk_meta.doc_item_refs + for c in chunks + 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())) 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 66fddbc6..d2daa3fc 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py +++ b/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: diff --git a/haiku_rag_slim/haiku/rag/skills/_tools.py b/haiku_rag_slim/haiku/rag/skills/_tools.py index 71132b54..470239bd 100644 --- a/haiku_rag_slim/haiku/rag/skills/_tools.py +++ b/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) - if chunk is None: + chunks = [] + 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, diff --git a/tests/test_client.py b/tests/test_client.py index a0147de6..cf166029 100644 --- a/tests/test_client.py +++ b/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 # =============================================================================