From 18c0f6c5e81f661737c6e01cb280c7348bc4bc4e Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 9 Jul 2026 12:49:42 +0300 Subject: [PATCH] Add --no-expand to the visualize command for chunk-only grounding --- CHANGELOG.md | 1 + docs/cli.md | 2 + haiku_rag_slim/haiku/rag/app.py | 4 +- haiku_rag_slim/haiku/rag/cli.py | 7 ++- haiku_rag_slim/haiku/rag/client/__init__.py | 7 ++- haiku_rag_slim/haiku/rag/client/search.py | 11 +++-- tests/test_client.py | 55 +++++++++++++++++++++ 7 files changed, 79 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 75a1fb02..2da99610 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/cli.md b/docs/cli.md index b2b45f15..f54b9879 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -248,6 +248,8 @@ haiku-rag visualize 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. diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index e6a68fd0..b3eba176 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/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]" diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 9578628c..7b201dd4 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/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") diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index d8c55438..c817180f 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 7d74e84f..49bd75f0 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/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 = [ diff --git a/tests/test_client.py b/tests/test_client.py index ccf1be3e..40884c0a 100644 --- a/tests/test_client.py +++ b/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 # =============================================================================