Add --no-expand to the visualize command for chunk-only grounding

CHANGELOG.md

@@ -6,6 +6,7 @@
 - `SearchResult.chunk_ids` and `Citation.chunk_ids` carry the chunk ids merged into an expanded result.
 - `Citation.doc_item_refs` carries the items the model saw; `visualize_chunk` accepts a `refs` argument and resolves bounding boxes from them, so visualizations match the cited content instead of re-expanding.
 - Chunk visualizations draw matched content in a stronger highlight than surrounding context.
+- `haiku-rag visualize --no-expand` highlights only the chunk itself, without its expanded context (`visualize_chunk` gains an `expand` argument).
 
 ### Changed
 

docs/cli.md

@@ -248,6 +248,8 @@ haiku-rag visualize <chunk_id>
 
 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.
 
+Pass `--no-expand` to highlight only the chunk itself, without its expanded context.
+
 !!! 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/app.py

@@ -457,7 +457,7 @@ class HaikuRAGApp:  # pragma: no cover
             for result in results:
                 self._rich_print_search_result(result)
 
-    async def visualize_chunk(self, chunk_id: str):
+    async def visualize_chunk(self, chunk_id: str, expand: bool = True):
         """Display visual grounding images for a chunk."""
         from textual_image.renderable import Image as RichImage
 
@@ -472,7 +472,7 @@ class HaikuRAGApp:  # pragma: no cover
                 self.console.print(f"[red]Chunk with id {chunk_id} not found.[/red]")
                 return
 
-            images = await self.client.visualize_chunk(chunk)
+            images = await self.client.visualize_chunk(chunk, expand=expand)
             if not images:
                 self.console.print(
                     "[yellow]No visual grounding available for this chunk.[/yellow]"

haiku_rag_slim/haiku/rag/cli.py

@@ -346,9 +346,14 @@ def visualize(  # pragma: no cover
         "--db",
         help="Path to the LanceDB database file",
     ),
+    no_expand: bool = typer.Option(
+        False,
+        "--no-expand",
+        help="Highlight only the chunk itself, without its expanded context",
+    ),
 ):
     app = create_app(db)
-    asyncio.run(app.visualize_chunk(chunk_id=chunk_id))
+    asyncio.run(app.visualize_chunk(chunk_id=chunk_id, expand=not no_expand))
 
 
 @_cli.command("ask", help="Ask a question using the QA agent")

haiku_rag_slim/haiku/rag/client/__init__.py

@@ -505,11 +505,14 @@ class HaikuRAG:
         return await analyze(self, question, filter)
 
     async def visualize_chunk(
-        self, chunk: Chunk | Sequence[Chunk], refs: list[str] | None = None
+        self,
+        chunk: Chunk | Sequence[Chunk],
+        refs: list[str] | None = None,
+        expand: bool = True,
     ) -> list:
         from haiku.rag.client.search import visualize_chunk
 
-        return await visualize_chunk(self, chunk, refs)
+        return await visualize_chunk(self, chunk, refs, expand)
 
     async def rebuild_database(
         self, mode: RebuildMode = RebuildMode.FULL

haiku_rag_slim/haiku/rag/client/search.py

@@ -228,14 +228,16 @@ async def visualize_chunk(
     client: "HaikuRAG",
     chunk: "Chunk | Sequence[Chunk]",
     refs: list[str] | None = None,
+    expand: bool = True,
 ) -> list:
     """Render page images with bounding box highlights for one or more chunks.
 
     When ``refs`` is given (the ``doc_item_refs`` of the citation, i.e. the
     exact items the model saw), bounding boxes are resolved from them directly
-    so the visualization matches the cited context precisely. When ``refs`` is
-    ``None`` (e.g. the CLI or inspector, where there is no stored context), the
-    chunks' context is re-expanded to recover the surrounding section.
+    so the visualization matches the cited context precisely. Otherwise, with
+    ``expand=True`` (default) the chunks' context is re-expanded to recover the
+    surrounding section; with ``expand=False`` only the chunks' own items are
+    drawn, so the visualization shows just the retrieved chunk with no context.
 
     The chunks' own items draw in a strong highlight; the remaining items draw
     fainter, so the matched content stands out from its surrounding context.
@@ -270,6 +272,9 @@ async def visualize_chunk(
 
     if refs is not None:
         all_refs = list(refs)
+    elif not expand:
+        # Chunk-only: draw just the retrieved chunks' own items, no context.
+        all_refs = list(matched_refs)
     else:
         # No stored context: re-expand the chunks to recover their section.
         search_results = [

tests/test_client.py

@@ -1798,6 +1798,61 @@ async def test_client_visualize_chunk_uses_given_refs(temp_db_path):
         assert len(await client.visualize_chunk(chunk, refs=["#/texts/0"])) == 1
 
 
+async def test_client_visualize_chunk_no_expand_shows_only_chunk(temp_db_path):
+    """expand=False draws only the chunk's own items, not the expanded section."""
+    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="no-expand-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
+            ),
+        )
+    for page_no in (1, 2):
+        docling_doc.add_text(
+            label=DocItemLabel.PARAGRAPH,
+            text=f"Short paragraph on page {page_no}.",
+            prov=ProvenanceItem(
+                page_no=page_no,
+                bbox=BoundingBox(l=50, t=700, r=550, b=650),
+                charspan=(0, 20),
+            ),
+        )
+
+    chunks = [
+        Chunk(
+            content="Short paragraph on page 1.",
+            metadata={
+                "doc_item_refs": ["#/texts/0"],
+                "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://no-expand")
+        chunk = (await client.chunk_repository.get_by_document_id(doc.id))[0]
+
+        # Default expands the chunk's context outward → reaches page two.
+        assert len(await client.visualize_chunk(chunk)) == 2
+        # expand=False draws only the chunk's own page-one item.
+        assert len(await client.visualize_chunk(chunk, expand=False)) == 1
+
+
 # =============================================================================
 # convert() method tests
 # =============================================================================