From 6d86237dd693b8924b8404dbcfd108155e45db36 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 8 Jul 2026 15:37:00 +0300 Subject: [PATCH] Draw matched content stronger than expanded context in visualizations --- CHANGELOG.md | 1 + docs/cli.md | 2 +- haiku_rag_slim/haiku/rag/client/search.py | 38 +++++++++---- tests/test_client.py | 67 +++++++++++++++++++++++ 4 files changed, 97 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b9c7254f..92e3d249 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,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. +- Chunk visualizations draw matched content in a stronger highlight than expanded context. - 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/docs/cli.md b/docs/cli.md index 92e25855..b2b45f15 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -246,7 +246,7 @@ Display visual grounding for a chunk - shows page images with highlighted boundi haiku-rag visualize ``` -This renders the source document pages with the chunk's location highlighted. Useful for verifying chunk boundaries and understanding document structure. +This renders the source document pages with the chunk's location highlighted. The chunk itself draws in a strong highlight, while surrounding context swept in by expansion draws fainter. Useful for verifying chunk boundaries and understanding document structure. !!! note Requires a terminal with image support (iTerm2, Kitty, WezTerm, etc.) and documents processed with docling that have page images stored. diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index cec8c59c..f2f8f06c 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -234,6 +234,10 @@ async def visualize_chunk(client: "HaikuRAG", chunk: "Chunk | Sequence[Chunk]") reproduces the merged expansion; chunks from a different document than the first are ignored. + The chunks' own items draw in a strong highlight; items swept in by + expansion draw fainter, so the matched content stands out from its + surrounding context. + Returns a list of PIL Image objects, one per page with bounding boxes. Empty list if no bounding boxes or page images available. """ @@ -279,19 +283,29 @@ async def visualize_chunk(client: "HaikuRAG", chunk: "Chunk | Sequence[Chunk]") 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 = chunks[0].get_chunk_metadata() - bounding_boxes = meta.resolve_bounding_boxes(docling_doc) - if not bounding_boxes: + refs = chunks[0].get_chunk_metadata().doc_item_refs + + matched_refs = {r for sr in search_results for r in sr.doc_item_refs} or set(refs) + swept_refs = [r for r in refs if r not in matched_refs] + + matched_boxes = ChunkMetadata( + doc_item_refs=sorted(matched_refs) + ).resolve_bounding_boxes(docling_doc) + swept_boxes = ChunkMetadata(doc_item_refs=swept_refs).resolve_bounding_boxes( + docling_doc + ) + if not matched_boxes and not swept_boxes: return [] - # Group bounding boxes by page + # Group bounding boxes by page; swept boxes first so matched draw on top boxes_by_page: dict[int, list] = {} - for bbox in bounding_boxes: + for bbox, is_matched in [(b, False) for b in swept_boxes] + [ + (b, True) for b in matched_boxes + ]: if bbox.page_no not in boxes_by_page: boxes_by_page[bbox.page_no] = [] - boxes_by_page[bbox.page_no].append(bbox) + boxes_by_page[bbox.page_no].append((bbox, is_matched)) # Load only the needed page images pages_doc = await client.document_repository.get_pages_data(document_id) @@ -319,7 +333,7 @@ async def visualize_chunk(client: "HaikuRAG", chunk: "Chunk | Sequence[Chunk]") image = deepcopy(pil_image) draw = ImageDraw.Draw(image, "RGBA") - for bbox in boxes_by_page[page_no]: + for bbox, is_matched in boxes_by_page[page_no]: # Document coords are bottom-left origin; PIL uses top-left x0 = bbox.left * scale_x y0 = (page_height - bbox.top) * scale_y @@ -329,8 +343,12 @@ async def visualize_chunk(client: "HaikuRAG", chunk: "Chunk | Sequence[Chunk]") if y0 > y1: y0, y1 = y1, y0 - fill_color = (255, 255, 0, 40) # Yellow with transparency - outline_color = (255, 165, 0, 100) # Orange outline + if is_matched: + fill_color = (255, 255, 0, 40) # Yellow with transparency + outline_color = (255, 165, 0, 100) # Orange outline + else: + fill_color = (255, 255, 0, 15) + outline_color = (255, 165, 0, 50) draw.rectangle([(x0, y0), (x1, y1)], fill=fill_color, outline=None) draw.rectangle([(x0, y0), (x1, y1)], outline=outline_color, width=1) diff --git a/tests/test_client.py b/tests/test_client.py index cf166029..d9a2669c 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1666,6 +1666,73 @@ async def test_client_visualize_chunk_merged_chunks_union_pages(temp_db_path): assert len(merged_images) == 2 +async def test_client_visualize_chunk_two_tone_highlights(temp_db_path): + """Matched content draws stronger than context swept in by expansion.""" + 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="two-tone-test") + page_size = Size(width=612.0, height=792.0) + docling_doc.add_page( + page_no=1, + size=page_size, + image=ImageRef.from_pil( + PilImageModule.new("RGB", (612, 792), color="white"), dpi=72 + ), + ) + + # Three small paragraphs; the chunk matches only the middle one, so + # expansion sweeps in its neighbors. + for i in range(3): + docling_doc.add_text( + label=DocItemLabel.PARAGRAPH, + text=f"Paragraph {i}.", + prov=ProvenanceItem( + page_no=1, + bbox=BoundingBox(l=50, t=700 - i * 100, r=550, b=650 - i * 100), + charspan=(0, 12), + ), + ) + + chunks = [ + Chunk( + content="Paragraph 1.", + metadata={ + "doc_item_refs": ["#/texts/1"], + "page_numbers": [1], + "labels": ["paragraph"], + }, + order=0, + embedding=[0.1] * 2560, + ) + ] + + async with HaikuRAG(temp_db_path, create=True) as client: + doc = await client.import_document(docling_doc, chunks, uri="test://two-tone") + stored_chunks = await client.chunk_repository.get_by_document_id(doc.id) + assert len(stored_chunks) == 1 + + images = await client.visualize_chunk(stored_chunks[0]) + assert len(images) == 1 + image = images[0] + + # Page dpi 72 == document coords, bottom-left origin flipped to + # top-left: item i's box spans y = 92 + i * 100 .. 142 + i * 100. + matched = image.getpixel((300, 217)) # inside #/texts/1 + swept = image.getpixel((300, 117)) # inside #/texts/0 + background = image.getpixel((300, 30)) # outside all boxes + + assert matched != background + assert swept != background + assert matched != swept + + # ============================================================================= # convert() method tests # =============================================================================