Draw matched content stronger than expanded context in visualizations

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

docs/cli.md

@@ -246,7 +246,7 @@ Display visual grounding for a chunk - shows page images with highlighted boundi
 haiku-rag visualize <chunk_id>
 ```
 
-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.

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()
+        refs = chunks[0].get_chunk_metadata().doc_item_refs
-    bounding_boxes = meta.resolve_bounding_boxes(docling_doc)
+
-    if not bounding_boxes:
+    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
+            if is_matched:
-            outline_color = (255, 165, 0, 100)  # Orange outline
+                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)

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
 # =============================================================================