diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7e8786ad..2da99610 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,13 @@
# Changelog
## [Unreleased]
+### Added
+
+- `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
- Duplicate images within a document produce a single picture chunk.
@@ -9,6 +16,9 @@
### Fixed
- vLLM embedding and vLLM/Jina reranking reuse one HTTP client across requests instead of opening one per request.
+- Picture and table search hits expand within their section, never across section boundaries.
+- A merged search result anchors its `chunk_id` on the highest-scoring constituent chunk instead of the one earliest in the document.
+- A clipped citation's `page_numbers` and `doc_item_refs` reflect only the content that survived the budget, not the full expanded range.
## [0.64.0] - 2026-07-08
diff --git a/app/backend/main.py b/app/backend/main.py
index ca110bbb..cd157ccf 100644
--- a/app/backend/main.py
+++ b/app/backend/main.py
@@ -203,22 +203,43 @@ 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 (a merged citation's
+ constituent chunks). The optional ``refs`` query param is a JSON-encoded
+ list of the citation's ``doc_item_refs`` — the exact items the model saw —
+ so the highlight matches the cited content instead of re-expanding.
+ """
import base64
+ import json
from io import BytesIO
chunk_id = request.path_params["chunk_id"]
+ refs: list[str] | None = None
+ refs_param = request.query_params.get("refs")
+ if refs_param:
+ try:
+ parsed = json.loads(refs_param)
+ except ValueError:
+ parsed = None
+ if isinstance(parsed, list):
+ refs = [str(x) for x in parsed]
+
if not db_path.exists():
return JSONResponse({"error": "Database not found"}, status_code=404)
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, refs)
if not images:
return JSONResponse({"images": [], "message": "No visual grounding available"})
@@ -233,7 +254,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/app/api/visualize/[chunk_id]/route.ts b/app/frontend/app/api/visualize/[chunk_id]/route.ts
index 81f297ee..098c74f7 100644
--- a/app/frontend/app/api/visualize/[chunk_id]/route.ts
+++ b/app/frontend/app/api/visualize/[chunk_id]/route.ts
@@ -1,14 +1,18 @@
import { NextResponse } from "next/server";
export async function GET(
- _request: Request,
+ request: Request,
{ params }: { params: Promise<{ chunk_id: string }> },
) {
const { chunk_id } = await params;
const backendUrl = process.env.BACKEND_URL || "http://backend:8000";
+ const refs = new URL(request.url).searchParams.get("refs");
+ const query = refs ? `?refs=${encodeURIComponent(refs)}` : "";
try {
- const response = await fetch(`${backendUrl}/api/visualize/${chunk_id}`);
+ const response = await fetch(
+ `${backendUrl}/api/visualize/${chunk_id}${query}`,
+ );
const data = await response.json();
if (!response.ok) {
diff --git a/app/frontend/components/CitationBlock.tsx b/app/frontend/components/CitationBlock.tsx
index 3221e0ec..362ed113 100644
--- a/app/frontend/components/CitationBlock.tsx
+++ b/app/frontend/components/CitationBlock.tsx
@@ -20,7 +20,7 @@ function CitationItem({
onViewInDocument,
}: {
citation: Citation;
- onViewInDocument: (chunkId: string) => void;
+ onViewInDocument: (chunkId: string, refs?: string[]) => void;
}) {
const [expanded, setExpanded] = useState(false);
@@ -55,7 +55,14 @@ function CitationItem({
@@ -82,45 +89,54 @@ export default function CitationBlock({ citations }: CitationBlockProps) {
return () => abortRef.current?.abort();
}, []);
- const fetchVisualGrounding = useCallback(async (chunkId: string) => {
- abortRef.current?.abort();
- const controller = new AbortController();
- abortRef.current = controller;
+ const fetchVisualGrounding = useCallback(
+ async (chunkId: string, refs?: string[]) => {
+ abortRef.current?.abort();
+ const controller = new AbortController();
+ abortRef.current = controller;
- setVisualGrounding({
- isOpen: true,
- chunkId,
- images: [],
- loading: true,
- error: null,
- });
-
- try {
- const response = await fetch(`/api/visualize/${chunkId}`, {
- signal: controller.signal,
+ setVisualGrounding({
+ isOpen: true,
+ chunkId,
+ images: [],
+ loading: true,
+ error: null,
});
- const data = await response.json();
- if (controller.signal.aborted) return;
- if (!response.ok) {
- throw new Error(data.error || "Failed to fetch visual grounding");
+ const query = refs?.length
+ ? `?refs=${encodeURIComponent(JSON.stringify(refs))}`
+ : "";
+ try {
+ const response = await fetch(
+ `/api/visualize/${encodeURIComponent(chunkId)}${query}`,
+ {
+ signal: controller.signal,
+ },
+ );
+ const data = await response.json();
+
+ if (controller.signal.aborted) return;
+ if (!response.ok) {
+ throw new Error(data.error || "Failed to fetch visual grounding");
+ }
+
+ setVisualGrounding((prev) => ({
+ ...prev,
+ images: data.images || [],
+ loading: false,
+ error: data.images?.length === 0 ? data.message : null,
+ }));
+ } catch (err) {
+ if (controller.signal.aborted) return;
+ setVisualGrounding((prev) => ({
+ ...prev,
+ loading: false,
+ error: err instanceof Error ? err.message : "Unknown error",
+ }));
}
-
- setVisualGrounding((prev) => ({
- ...prev,
- images: data.images || [],
- loading: false,
- error: data.images?.length === 0 ? data.message : null,
- }));
- } catch (err) {
- if (controller.signal.aborted) return;
- setVisualGrounding((prev) => ({
- ...prev,
- loading: false,
- error: err instanceof Error ? err.message : "Unknown error",
- }));
- }
- }, []);
+ },
+ [],
+ );
const closeVisualGrounding = useCallback(() => {
abortRef.current?.abort();
diff --git a/app/frontend/lib/sessionStorage.ts b/app/frontend/lib/sessionStorage.ts
index 70b0b01a..a3369f0a 100644
--- a/app/frontend/lib/sessionStorage.ts
+++ b/app/frontend/lib/sessionStorage.ts
@@ -2,11 +2,13 @@ export interface Citation {
index: number;
document_id: string;
chunk_id: string;
+ chunk_ids?: string[];
document_uri: string;
document_title: string | null;
page_numbers: number[];
headings: string[] | null;
content: string;
+ doc_item_refs?: string[];
}
// Matches RAGState from the backend skill
diff --git a/docs/cli.md b/docs/cli.md
index 92e25855..f54b9879 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -246,7 +246,9 @@ 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.
+
+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/docs/configuration/qa.md b/docs/configuration/qa.md
index ec527a1e..ac7e61d9 100644
--- a/docs/configuration/qa.md
+++ b/docs/configuration/qa.md
@@ -13,7 +13,7 @@ search:
- **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, and QA. Default: 10
- **max_context_chars**: Hard limit on total characters in expanded content. Default: 10000.
-Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers). This naturally crosses into adjacent sections until the budget is filled. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
+Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers). This naturally crosses into adjacent sections until the budget is filled. Picture and table matches are exempt: they return their enclosing section as-is and never cross section boundaries. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
!!! note "Reranking behavior"
When a reranker is configured, search automatically retrieves 10x the requested limit, then reranks to return the final count. This improves result quality without requiring you to adjust `limit`.
diff --git a/docs/python.md b/docs/python.md
index 51322ad8..38308424 100644
--- a/docs/python.md
+++ b/docs/python.md
@@ -298,7 +298,7 @@ for result in expanded_results:
print(f"Expanded content: {result.content}")
```
-Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers). This naturally crosses into adjacent sections until the budget is filled. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
+Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers). This naturally crosses into adjacent sections until the budget is filled. Picture and table matches are exempt: they return their enclosing section as-is and never cross section boundaries. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
Configuration:
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/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py
index 34a1e53d..98015ed4 100644
--- a/haiku_rag_slim/haiku/rag/chat/app.py
+++ b/haiku_rag_slim/haiku/rag/chat/app.py
@@ -405,13 +405,24 @@ 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,
+ refs=citation.doc_item_refs or None,
+ )
+ )
async def action_show_info(self) -> None:
"""Show database info modal."""
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 910ef6a0..c817180f 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,10 +504,15 @@ 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],
+ refs: list[str] | None = None,
+ expand: bool = True,
+ ) -> list:
from haiku.rag.client.search import visualize_chunk
- return await visualize_chunk(self, chunk)
+ 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 06e6f2f6..ef814b9d 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,24 @@ 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]",
+ refs: list[str] | None = None,
+ expand: bool = True,
+) -> 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
- bounding boxes from all items in the expanded range. This ensures
- visualization covers all pages the expanded content spans.
+ 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. 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.
+ 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 +252,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 []
@@ -250,35 +268,60 @@ async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list:
if not docling_doc:
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,
- 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,
- )
- expanded = await expand_context(client, [search_result])
- refs = expanded[0].doc_item_refs if expanded else chunk_meta.doc_item_refs
- meta = ChunkMetadata(doc_item_refs=refs)
+ matched_refs = {r for c in chunks for r in c.get_chunk_metadata().doc_item_refs}
+
+ 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:
- meta = chunk_meta
- bounding_boxes = meta.resolve_bounding_boxes(docling_doc)
- if not bounding_boxes:
+ # No stored context: re-expand the chunks to recover their section.
+ search_results = [
+ SearchResult(
+ content=c.content,
+ score=1.0,
+ chunk_id=c.id,
+ document_id=c.document_id,
+ doc_item_refs=meta.doc_item_refs,
+ page_numbers=meta.page_numbers,
+ )
+ for c in chunks
+ if (meta := c.get_chunk_metadata()).doc_item_refs
+ ]
+ if search_results:
+ expanded = await expand_context(client, search_results)
+ all_refs = []
+ for result in expanded:
+ all_refs.extend(r for r in result.doc_item_refs if r not in all_refs)
+ if not all_refs:
+ all_refs = [r for sr in search_results for r in sr.doc_item_refs]
+ else:
+ all_refs = list(chunks[0].get_chunk_metadata().doc_item_refs)
+
+ matched_draw = [r for r in all_refs if r in matched_refs]
+ swept_refs = [r for r in all_refs if r not in matched_refs]
+
+ matched_boxes = ChunkMetadata(doc_item_refs=matched_draw).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(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()))
@@ -303,7 +346,7 @@ async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list:
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
@@ -313,8 +356,12 @@ async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list:
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, 150, 0, 55) # Orange, matched content
+ outline_color = (240, 130, 0, 150) # Orange outline
+ else:
+ fill_color = (255, 255, 0, 40) # Yellow, surrounding context
+ outline_color = (255, 165, 0, 100)
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/haiku_rag_slim/haiku/rag/context.py b/haiku_rag_slim/haiku/rag/context.py
index ee1ea311..0b6c37ef 100644
--- a/haiku_rag_slim/haiku/rag/context.py
+++ b/haiku_rag_slim/haiku/rag/context.py
@@ -12,7 +12,8 @@ For STRUCTURED documents (containing section_header or title labels):
5. If the section is too small (under 20% of max_context_chars), expand
item-by-item crossing into adjacent sections until the budget is filled.
This lets small sections (e.g., title+authors) grow into neighboring
- content.
+ content. Picture and table matches are exempt: they return their
+ enclosing section as-is, never crossing section boundaries.
6. Merge overlapping ranges from multiple results in the same document.
Adjacent but non-overlapping ranges stay separate to preserve section
independence.
@@ -36,6 +37,10 @@ from haiku.rag.store.repositories.document_item import DocumentItemRepository
_NOISE_LABELS = {"footnote", "page_header", "page_footer", "document_index"}
_SECTION_BOUNDARY_LABELS = {"section_header", "title"}
+# Labels whose pertinent unit is the item plus its own section: expansion
+# never crosses section boundaries for these matches.
+_TIGHT_LABELS = {"picture", "table"}
+
# Sections with fewer chars than this fraction of max_context_chars are
# considered too small — expansion falls through to item-by-item outward
# growth, which naturally crosses into adjacent sections.
@@ -99,16 +104,18 @@ def _evidence_anchors(content: str, max_chars: int) -> list[str]:
return anchors
-def _clip_to_budget(content: str, results: list[SearchResult], max_chars: int) -> str:
- """Clip expanded content to ``max_chars``, keeping the matched evidence.
+def _clip_window(
+ content: str, results: list[SearchResult], max_chars: int
+) -> tuple[int, int]:
+ """Return the ``[start, end)`` char window to keep when clipping ``content``.
Anchors on the first locatable result in ``results`` order (the primary
chunk that supplies the expanded result's identity) and returns a
- ``max_chars``-wide window centered on it. Falls back to a prefix cut only
+ ``max_chars``-wide window centered on it. Falls back to a prefix window only
when no anchor is locatable (heavy drift).
"""
if max_chars <= 0:
- return ""
+ return (0, 0)
evidence_start, evidence_len = -1, 0
for result in results:
for anchor in _evidence_anchors(result.content, max_chars):
@@ -119,14 +126,45 @@ def _clip_to_budget(content: str, results: list[SearchResult], max_chars: int) -
if evidence_start != -1:
break
if evidence_start == -1:
- return content[:max_chars]
+ return (0, min(len(content), max_chars))
center = evidence_start + evidence_len // 2
start = max(0, center - max_chars // 2)
end = min(len(content), start + max_chars)
start = max(0, end - max_chars)
+ return (start, end)
+
+
+def _clip_to_budget(content: str, results: list[SearchResult], max_chars: int) -> str:
+ """Clip expanded content to ``max_chars``, keeping the matched evidence."""
+ start, end = _clip_window(content, results, max_chars)
return content[start:end]
+def _collect_meta(
+ spans: list[tuple[int, int, DocumentItem]],
+) -> tuple[set[int], list[str], set[str]]:
+ """Union the page numbers, refs, and labels of the given item spans."""
+ pages: set[int] = set()
+ refs: list[str] = []
+ labels: set[str] = set()
+ for _start, _end, item in spans:
+ refs.append(item.self_ref)
+ if item.label:
+ labels.add(item.label)
+ pages.update(item.page_numbers)
+ return pages, refs, labels
+
+
+def _span_in_window(
+ span: tuple[int, int, DocumentItem], win_start: int, win_end: int
+) -> bool:
+ """Whether an item's char span overlaps the ``[win_start, win_end]`` clip window."""
+ start, end, _item = span
+ if start == end: # zero-width picture position
+ return win_start <= start <= win_end
+ return start < win_end and end > win_start
+
+
def _expand_outward(
items: list[DocumentItem],
center_idx: int,
@@ -222,6 +260,11 @@ def _find_expansion_range(
hi_bound=sec_end,
)
+ # Picture/table hits stay section-bounded: their pertinent unit is the
+ # figure or table plus its section, never neighboring sections.
+ if any(items[i].label in _TIGHT_LABELS for i in matched_indices):
+ return (items[sec_start].position, items[sec_end].position)
+
# Section too small (e.g., title+authors) — expand across boundaries
return _expand_outward(items, center_idx, max_chars, skip_noise=True)
@@ -280,9 +323,11 @@ async def expand_with_items(
final_results: list[SearchResult] = []
for range_start, range_end, original_results in merged:
content_parts: list[str] = []
- refs: list[str] = []
- pages: set[int] = set()
- labels: set[str] = set()
+ # Char span of each contributing item within the joined content, so
+ # metadata can be narrowed to whatever survives a budget clip.
+ item_spans: list[tuple[int, int, DocumentItem]] = []
+ cursor = 0
+ separator = "\n\n"
for pos in range(range_start, range_end + 1):
item = pos_to_item.get(pos)
@@ -291,56 +336,96 @@ async def expand_with_items(
if has_sections and item.label in _NOISE_LABELS:
continue
if item.text:
+ if content_parts:
+ cursor += len(separator)
+ start = cursor
content_parts.append(item.text)
- refs.append(item.self_ref)
- if item.label:
- labels.add(item.label)
- pages.update(item.page_numbers)
+ cursor += len(item.text)
+ item_spans.append((start, cursor, item))
elif item.label == "picture":
# Pictures may legitimately have empty text (no VLM
# description configured). Keep their self_ref so the
- # downstream image_data lookup can still attach bytes.
- refs.append(item.self_ref)
- labels.add(item.label)
- pages.update(item.page_numbers)
+ # downstream image_data lookup can still attach bytes. They
+ # occupy a zero-width position in reading order.
+ item_spans.append((cursor, cursor, item))
all_headings: list[str] = []
for r in original_results:
if r.headings:
all_headings.extend(h for h in r.headings if h not in all_headings)
- # Carry image_data and picture_captions through expansion so that
- # only pictures from the originally retrieved chunks get attached.
- # Pictures swept in by section expansion are referenced in `refs`
- # for cross-referencing but their bytes are not re-fetched —
- # otherwise a single search can balloon the response with adjacent
- # figures the model did not actually retrieve.
+ # Anchor identity (chunk_id, content/refs fallbacks) on the
+ # best-scoring constituent — the chunk that earned the result its
+ # rank — rather than whichever sits earliest in the document.
+ first = max(original_results, key=lambda r: r.score)
+
+ chunk_ids: list[str] = []
+ for r in original_results:
+ if r.chunk_id and r.chunk_id not in chunk_ids:
+ chunk_ids.append(r.chunk_id)
+
+ joined = separator.join(content_parts)
+
+ # Expansion should never return less content than the original chunk.
+ # This can happen when item texts are fragmented (e.g., docling splits
+ # formatted HTML list items into many small text nodes); fall back to
+ # the chunk's own content, described by the chunk's own metadata.
+ if len(joined) < len(first.content):
+ base_content, base_spans = first.content, None
+ else:
+ base_content, base_spans = joined, item_spans
+
+ if len(base_content) > max_chars:
+ # Clip to the budget, anchored on the primary (highest-scoring)
+ # chunk, and narrow metadata to whatever survives the window.
+ win_start, win_end = _clip_window(
+ base_content, [first, *original_results], max_chars
+ )
+ expanded_content = base_content[win_start:win_end]
+ if base_spans is None:
+ pages, refs, labels = (
+ set(first.page_numbers),
+ list(first.doc_item_refs),
+ set(first.labels),
+ )
+ else:
+ pages, refs, labels = _collect_meta(
+ [s for s in base_spans if _span_in_window(s, win_start, win_end)]
+ )
+ else:
+ expanded_content = base_content
+ if base_spans is None:
+ pages, refs, labels = (
+ set(first.page_numbers),
+ list(first.doc_item_refs),
+ set(first.labels),
+ )
+ else:
+ pages, refs, labels = _collect_meta(base_spans)
+
+ # Carry image_data and picture_captions from the originally retrieved
+ # chunks, but only for constituents whose refs survive the window — a
+ # chunk clipped out of the budget must not still ship its image to the
+ # model. Pictures swept in by section expansion are referenced in
+ # ``refs`` but their bytes are never re-fetched, so the multimodal
+ # payload stays bounded to what was actually retrieved and shown.
+ surviving_refs = set(refs)
merged_image_data: dict[str, str] = {}
merged_captions: dict[str, str] = {}
for r in original_results:
+ if r.doc_item_refs and not surviving_refs.intersection(r.doc_item_refs):
+ continue
if r.image_data:
merged_image_data.update(r.image_data)
if r.picture_captions:
merged_captions.update(r.picture_captions)
- first = original_results[0]
-
- # Expansion should never return less content than the original chunk.
- # This can happen when item texts are fragmented (e.g., docling splits
- # formatted HTML list items into many small text nodes).
- expanded_content = "\n\n".join(content_parts)
- if len(expanded_content) < len(first.content):
- expanded_content = first.content
- if len(expanded_content) > max_chars:
- expanded_content = _clip_to_budget(
- expanded_content, original_results, max_chars
- )
-
final_results.append(
SearchResult(
content=expanded_content,
score=max(r.score for r in original_results),
chunk_id=first.chunk_id,
+ chunk_ids=chunk_ids,
document_id=first.document_id,
document_uri=first.document_uri,
document_title=first.document_title,
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..978a21cd 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,16 @@ class VisualGroundingModal(Screen):
def __init__(
self,
- chunk: "Chunk",
+ chunk: "Chunk | list[Chunk]",
client: "HaikuRAG",
document_uri: str | None = None,
+ refs: list[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.refs = refs
+ 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 +83,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, self.refs)
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/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py
index d6d3ba6c..7eb1e57b 100644
--- a/haiku_rag_slim/haiku/rag/store/models/chunk.py
+++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py
@@ -124,11 +124,17 @@ class SearchResult(BaseModel):
when the caller asked to omit them via ``include_images=False`` on
``client.search``. Same shape is used everywhere — MCP, in-process search,
agent toolsets — so non-vision callers see ``None`` and pay nothing.
+
+ ``chunk_ids`` lists the ids of all chunks whose expansion ranges merged
+ into this result; empty means just ``chunk_id``. It lets citation
+ consumers (visual grounding) reproduce a merged expansion and is never
+ part of ``format_for_agent`` output.
"""
content: str
score: float
chunk_id: str | None = None
+ chunk_ids: list[str] = []
document_id: str | None = None
document_uri: str | None = None
document_title: str | None = None
diff --git a/haiku_rag_slim/haiku/rag/store/models/citation.py b/haiku_rag_slim/haiku/rag/store/models/citation.py
index 5a10c7b0..95510b63 100644
--- a/haiku_rag_slim/haiku/rag/store/models/citation.py
+++ b/haiku_rag_slim/haiku/rag/store/models/citation.py
@@ -18,16 +18,26 @@ class Citation(BaseModel):
cited chunk. Empty for text-only citations. UIs can fetch the picture
bytes via ``DocumentItemRepository.get_picture_bytes(document_id, ref)``
and render them alongside the text content.
+
+ ``chunk_ids`` lists the ids of all chunks whose expansion ranges merged
+ into the cited result (always includes ``chunk_id``).
+
+ ``doc_item_refs`` are the ``self_ref`` values of every item in the cited
+ content — the exact items the model saw. Visual grounding resolves bounding
+ boxes from them so the rendered pages match the citation precisely.
+ ``picture_refs`` is the picture-labeled subset.
"""
index: int | None = None
document_id: str
chunk_id: str
+ chunk_ids: list[str] = Field(default_factory=list)
document_uri: str
document_title: str | None = None
page_numbers: list[int] = Field(default_factory=list)
headings: list[str] | None = None
content: str
+ doc_item_refs: list[str] = Field(default_factory=list)
picture_refs: list[str] = Field(default_factory=list)
@@ -51,11 +61,13 @@ def resolve_citations(
Citation(
document_id=r.document_id or "",
chunk_id=chunk_id,
+ chunk_ids=r.chunk_ids or [chunk_id],
document_uri=r.document_uri or "",
document_title=r.document_title,
page_numbers=r.page_numbers,
headings=r.headings,
content=r.content,
+ doc_item_refs=list(r.doc_item_refs),
picture_refs=picture_refs,
)
)
diff --git a/tests/store/test_citation.py b/tests/store/test_citation.py
new file mode 100644
index 00000000..1b1d89db
--- /dev/null
+++ b/tests/store/test_citation.py
@@ -0,0 +1,41 @@
+from haiku.rag.store.models.chunk import SearchResult
+from haiku.rag.store.models.citation import resolve_citations
+
+
+def _result(chunk_id: str, chunk_ids: list[str] | None = None) -> SearchResult:
+ return SearchResult(
+ content="content",
+ score=0.9,
+ chunk_id=chunk_id,
+ chunk_ids=chunk_ids or [],
+ document_id="doc-1",
+ document_uri="test://doc",
+ )
+
+
+def test_resolve_citations_copies_merged_chunk_ids():
+ result = _result("c1", chunk_ids=["c1", "c2"])
+ citations = resolve_citations(["c1"], [result])
+ assert len(citations) == 1
+ assert citations[0].chunk_id == "c1"
+ assert citations[0].chunk_ids == ["c1", "c2"]
+
+
+def test_resolve_citations_falls_back_to_chunk_id():
+ result = _result("c1")
+ citations = resolve_citations(["c1"], [result])
+ assert len(citations) == 1
+ assert citations[0].chunk_ids == ["c1"]
+
+
+def test_resolve_citations_strips_brackets():
+ result = _result("c1")
+ citations = resolve_citations(["[c1]"], [result])
+ assert len(citations) == 1
+ assert citations[0].chunk_id == "c1"
+
+
+def test_resolve_citations_skips_unknown_ids():
+ result = _result("c1")
+ citations = resolve_citations(["c1", "missing"], [result])
+ assert len(citations) == 1
diff --git a/tests/test_client.py b/tests/test_client.py
index a0147de6..40884c0a 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -1590,6 +1590,269 @@ 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
+ ),
+ )
+
+ # Two sections, one per page, each large enough to be returned whole and
+ # under budget (so no cross-boundary expansion and no clip). A single
+ # chunk visualizes its own section's page; both chunks cover both pages.
+ layout = [
+ (DocItemLabel.SECTION_HEADER, "Section One", 1),
+ (DocItemLabel.PARAGRAPH, "Page one body. " + "x" * 3000, 1),
+ (DocItemLabel.SECTION_HEADER, "Section Two", 2),
+ (DocItemLabel.PARAGRAPH, "Page two body. " + "y" * 3000, 2),
+ ]
+ for i, (label, text, page_no) in enumerate(layout):
+ docling_doc.add_text(
+ label=label,
+ text=text,
+ prov=ProvenanceItem(
+ page_no=page_no,
+ bbox=BoundingBox(l=50, t=700 - (i % 2) * 100, r=550, b=650),
+ charspan=(0, 20),
+ ),
+ )
+
+ chunks = [
+ Chunk(
+ content="Page one body. " + "x" * 3000,
+ metadata={
+ "doc_item_refs": ["#/texts/1"],
+ "page_numbers": [1],
+ "labels": ["paragraph"],
+ },
+ order=0,
+ embedding=[0.1] * 2560,
+ ),
+ Chunk(
+ content="Page two body. " + "y" * 3000,
+ 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
+
+
+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
+
+
+async def test_client_visualize_chunk_uses_given_refs(temp_db_path):
+ """Explicit refs (the citation's doc_item_refs) restrict the visualization
+ to exactly those items, instead of re-expanding the chunk's context."""
+ 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="refs-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
+ ),
+ )
+ docling_doc.add_text(
+ label=DocItemLabel.PARAGRAPH,
+ text="Content on page one.",
+ prov=ProvenanceItem(
+ page_no=1, bbox=BoundingBox(l=50, t=700, r=550, b=650), charspan=(0, 20)
+ ),
+ )
+ docling_doc.add_text(
+ label=DocItemLabel.PARAGRAPH,
+ text="Content on page two.",
+ prov=ProvenanceItem(
+ page_no=2, bbox=BoundingBox(l=50, t=700, r=550, b=650), charspan=(0, 20)
+ ),
+ )
+
+ chunks = [
+ Chunk(
+ content="Content on page one.\nContent on page two.",
+ metadata={
+ "doc_item_refs": ["#/texts/0", "#/texts/1"],
+ "page_numbers": [1, 2],
+ "labels": ["paragraph", "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://refs")
+ chunk = (await client.chunk_repository.get_by_document_id(doc.id))[0]
+
+ # No refs: re-expands the chunk's own refs → both pages.
+ assert len(await client.visualize_chunk(chunk)) == 2
+ # Given only the page-one ref → only page one is rendered.
+ 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
# =============================================================================
diff --git a/tests/test_context.py b/tests/test_context.py
index 4c6a8239..5bc6dcf3 100644
--- a/tests/test_context.py
+++ b/tests/test_context.py
@@ -244,6 +244,46 @@ class TestFindExpansionRange:
# Should expand outward into the first section
assert hi >= 2
+ def test_picture_in_small_section_stays_section_bounded(self):
+ items = [
+ _item(0, label="section_header", text="Chapter 1"),
+ _item(1, text="Chapter 1 prose. " * 100),
+ _item(2, label="section_header", text="Figure heading"),
+ _item(3, label="picture", text="Diagram description."),
+ _item(4, label="caption", text="Figure 2-3. Balance arm."),
+ _item(5, label="section_header", text="Chapter 3"),
+ _item(6, text="Chapter 3 prose. " * 100),
+ ]
+ # Figure section (items 2-4) is far under 20% of 5000 chars.
+ lo, hi = _find_expansion_range(items, {3}, has_sections=True, max_chars=5000)
+ # Never crosses either header
+ assert (lo, hi) == (2, 4)
+
+ def test_table_in_small_section_stays_section_bounded(self):
+ items = [
+ _item(0, label="section_header", text="Chapter 1"),
+ _item(1, text="Chapter 1 prose. " * 100),
+ _item(2, label="section_header", text="Table heading"),
+ _item(3, label="table", text="Header | Value"),
+ _item(4, label="section_header", text="Chapter 3"),
+ _item(5, text="Chapter 3 prose. " * 100),
+ ]
+ lo, hi = _find_expansion_range(items, {3}, has_sections=True, max_chars=5000)
+ assert (lo, hi) == (2, 3)
+
+ def test_text_in_small_section_still_expands_outward(self):
+ items = [
+ _item(0, label="section_header", text="Chapter 1"),
+ _item(1, text="Chapter 1 prose. " * 100),
+ _item(2, label="section_header", text="Short note"),
+ _item(3, text="A brief remark."),
+ _item(4, label="section_header", text="Chapter 3"),
+ _item(5, text="Chapter 3 prose. " * 100),
+ ]
+ lo, hi = _find_expansion_range(items, {3}, has_sections=True, max_chars=5000)
+ # Text hits keep growing across section boundaries
+ assert lo < 2 or hi > 3
+
class TestEvidenceAnchors:
def test_empty_content(self):
@@ -369,6 +409,85 @@ class TestExpandWithItems:
# real content — so we get expanded content, not the fallback.
assert len(expanded[0].content) > 0
+ async def test_picture_expansion_stays_within_section_pages(self, temp_db_path):
+ from haiku.rag.client import HaikuRAG
+
+ async with HaikuRAG(temp_db_path, create=True) as rag:
+ items = [
+ DocumentItem(
+ document_id="doc-1",
+ position=0,
+ self_ref="#/texts/0",
+ label="section_header",
+ text="Chapter 1",
+ page_numbers=[10],
+ ),
+ DocumentItem(
+ document_id="doc-1",
+ position=1,
+ self_ref="#/texts/1",
+ label="text",
+ text="Chapter 1 prose. " * 200,
+ page_numbers=[10],
+ ),
+ DocumentItem(
+ document_id="doc-1",
+ position=2,
+ self_ref="#/texts/2",
+ label="section_header",
+ text="Figure heading",
+ page_numbers=[13],
+ ),
+ DocumentItem(
+ document_id="doc-1",
+ position=3,
+ self_ref="#/pictures/0",
+ label="picture",
+ text="Diagram description.",
+ page_numbers=[13],
+ ),
+ DocumentItem(
+ document_id="doc-1",
+ position=4,
+ self_ref="#/texts/3",
+ label="caption",
+ text="Figure 2-3. Balance arm.",
+ page_numbers=[13],
+ ),
+ DocumentItem(
+ document_id="doc-1",
+ position=5,
+ self_ref="#/texts/4",
+ label="section_header",
+ text="Chapter 3",
+ page_numbers=[15],
+ ),
+ DocumentItem(
+ document_id="doc-1",
+ position=6,
+ self_ref="#/texts/5",
+ label="text",
+ text="Chapter 3 prose. " * 200,
+ page_numbers=[15],
+ ),
+ ]
+ await rag.document_item_repository.create_items("doc-1", items)
+
+ result = SearchResult(
+ content="Diagram description.",
+ score=0.9,
+ document_id="doc-1",
+ doc_item_refs=["#/pictures/0"],
+ page_numbers=[13],
+ )
+ expanded = await expand_with_items(
+ rag.document_item_repository, "doc-1", [result], 5000
+ )
+ assert len(expanded) == 1
+ assert expanded[0].page_numbers == [13]
+ assert "Chapter 1 prose" not in expanded[0].content
+ assert "Chapter 3 prose" not in expanded[0].content
+
async def test_fragmented_items_preserve_chunk(self, temp_db_path):
"""When items are fragmented (e.g., list_item children), the original
chunk content is preserved if expansion produces less text."""
@@ -578,6 +697,187 @@ class TestExpandWithItems:
assert len(expanded[0].content) <= 5000
assert marker in expanded[0].content
+ async def test_solo_result_carries_own_chunk_id(self, temp_db_path):
+ from haiku.rag.client import HaikuRAG
+
+ async with HaikuRAG(temp_db_path, create=True) as rag:
+ items = [
+ DocumentItem(
+ document_id="doc-1",
+ position=i,
+ self_ref=f"#/texts/{i}",
+ label="text",
+ text=f"Paragraph {i}. " * 10,
+ )
+ for i in range(5)
+ ]
+ await rag.document_item_repository.create_items("doc-1", items)
+
+ result = SearchResult(
+ content="Paragraph 2.",
+ score=0.9,
+ chunk_id="c1",
+ document_id="doc-1",
+ doc_item_refs=["#/texts/2"],
+ )
+ expanded = await expand_with_items(
+ rag.document_item_repository, "doc-1", [result], 5000
+ )
+ assert len(expanded) == 1
+ assert expanded[0].chunk_ids == ["c1"]
+
+ async def test_merged_results_carry_all_chunk_ids(self, temp_db_path):
+ from haiku.rag.client import HaikuRAG
+
+ async with HaikuRAG(temp_db_path, create=True) as rag:
+ items = [
+ DocumentItem(
+ document_id="doc-1",
+ position=i,
+ self_ref=f"#/texts/{i}",
+ label="text",
+ text=f"Paragraph {i}. " * 10,
+ )
+ for i in range(5)
+ ]
+ await rag.document_item_repository.create_items("doc-1", items)
+
+ r1 = SearchResult(
+ content="Paragraph 1.",
+ score=0.9,
+ chunk_id="c1",
+ document_id="doc-1",
+ doc_item_refs=["#/texts/1"],
+ )
+ r2 = SearchResult(
+ content="Paragraph 3.",
+ score=0.85,
+ chunk_id="c2",
+ document_id="doc-1",
+ doc_item_refs=["#/texts/3"],
+ )
+ expanded = await expand_with_items(
+ rag.document_item_repository, "doc-1", [r1, r2], 5000
+ )
+ # Ranges around positions 1 and 3 overlap → one merged result.
+ assert len(expanded) == 1
+ assert expanded[0].chunk_id == "c1"
+ assert expanded[0].chunk_ids == ["c1", "c2"]
+ # Sibling chunk ids are plumbing for visualization, never shown
+ # to the model.
+ assert "c2" not in expanded[0].format_for_agent()
+
+ async def test_merged_anchor_is_highest_scoring_constituent(self, temp_db_path):
+ """A merged result's chunk_id anchors on the best-scoring constituent,
+ not whichever chunk sits earliest in the document."""
+ from haiku.rag.client import HaikuRAG
+
+ async with HaikuRAG(temp_db_path, create=True) as rag:
+ items = [
+ DocumentItem(
+ document_id="doc-1",
+ position=i,
+ self_ref=f"#/texts/{i}",
+ label="text",
+ text=f"Paragraph {i}. " * 10,
+ )
+ for i in range(5)
+ ]
+ await rag.document_item_repository.create_items("doc-1", items)
+
+ # earlier in the document, lower score
+ r_early = SearchResult(
+ content="Paragraph 1.",
+ score=0.40,
+ chunk_id="c-early",
+ document_id="doc-1",
+ doc_item_refs=["#/texts/1"],
+ )
+ # later in the document, higher score — the real hit
+ r_best = SearchResult(
+ content="Paragraph 3.",
+ score=0.95,
+ chunk_id="c-best",
+ document_id="doc-1",
+ doc_item_refs=["#/texts/3"],
+ )
+ expanded = await expand_with_items(
+ rag.document_item_repository, "doc-1", [r_early, r_best], 5000
+ )
+ assert len(expanded) == 1
+ assert expanded[0].chunk_id == "c-best"
+ assert expanded[0].score == 0.95
+ # provenance still lists both
+ assert set(expanded[0].chunk_ids) == {"c-early", "c-best"}
+
+ async def test_clipped_merged_result_keeps_anchor_evidence_and_pages(
+ self, temp_db_path
+ ):
+ """When a merged result is clipped to budget, the surviving window is
+ centered on the anchor (highest-scoring) chunk, and page_numbers reflect
+ only the content that survived — not the full merged range."""
+ from haiku.rag.client import HaikuRAG
+
+ async with HaikuRAG(temp_db_path, create=True) as rag:
+ items = [
+ DocumentItem(
+ document_id="doc-1",
+ position=0,
+ self_ref="#/texts/0",
+ label="text",
+ text="LOWMARK " + "a" * 400,
+ page_numbers=[1],
+ ),
+ DocumentItem(
+ document_id="doc-1",
+ position=1,
+ self_ref="#/texts/1",
+ label="text",
+ text="b" * 400,
+ page_numbers=[2],
+ ),
+ DocumentItem(
+ document_id="doc-1",
+ position=2,
+ self_ref="#/texts/2",
+ label="text",
+ text="c" * 400 + " HIGHMARK",
+ page_numbers=[3],
+ ),
+ ]
+ await rag.document_item_repository.create_items("doc-1", items)
+
+ r_low = SearchResult(
+ content="LOWMARK " + "a" * 400,
+ score=0.4,
+ chunk_id="c-low",
+ document_id="doc-1",
+ doc_item_refs=["#/texts/0"],
+ page_numbers=[1],
+ )
+ r_high = SearchResult(
+ content="c" * 400 + " HIGHMARK",
+ score=0.9,
+ chunk_id="c-high",
+ document_id="doc-1",
+ doc_item_refs=["#/texts/2"],
+ page_numbers=[3],
+ )
+ expanded = await expand_with_items(
+ rag.document_item_repository, "doc-1", [r_low, r_high], 500
+ )
+ assert len(expanded) == 1
+ e = expanded[0]
+ # anchor is the high-scoring chunk, and its evidence survives clipping
+ assert e.chunk_id == "c-high"
+ assert "HIGHMARK" in e.content
+ assert "LOWMARK" not in e.content
+ # page_numbers reflect only the surviving window, not the full range
+ assert 3 in e.page_numbers
+ assert 1 not in e.page_numbers
+ # refs likewise exclude the clipped-out item
+ assert "#/texts/0" not in e.doc_item_refs
+
async def test_fuzzy_match_preserves_central_marker(self, temp_db_path):
"""The chunk's text need not be verbatim in the joined item text: a clean
central marker is still located via the central-slice anchor."""
@@ -715,3 +1015,61 @@ class TestExpandWithItemsPictureBytes:
"#/pictures/1": "A",
"#/pictures/3": "B",
}
+
+ async def test_clipped_out_picture_bytes_dropped(self, temp_db_path):
+ """A lower-scoring picture chunk clipped out of the budget window no
+ longer contributes its image bytes — the model must not receive an
+ image the citation and visualization omit."""
+ from haiku.rag.client import HaikuRAG
+
+ async with HaikuRAG(temp_db_path, create=True) as rag:
+ items = [
+ DocumentItem(
+ document_id="doc-1",
+ position=0,
+ self_ref="#/pictures/0",
+ label="picture",
+ text="LOWPIC " + "a" * 400,
+ page_numbers=[1],
+ ),
+ DocumentItem(
+ document_id="doc-1",
+ position=1,
+ self_ref="#/texts/1",
+ label="text",
+ text="b" * 400,
+ page_numbers=[2],
+ ),
+ DocumentItem(
+ document_id="doc-1",
+ position=2,
+ self_ref="#/pictures/1",
+ label="picture",
+ text="c" * 400 + " HIGHPIC",
+ page_numbers=[3],
+ ),
+ ]
+ await rag.document_item_repository.create_items("doc-1", items)
+
+ r_low = SearchResult(
+ content="LOWPIC " + "a" * 400,
+ score=0.4,
+ chunk_id="c-low",
+ document_id="doc-1",
+ doc_item_refs=["#/pictures/0"],
+ image_data={"#/pictures/0": "LOWBYTES"},
+ )
+ r_high = SearchResult(
+ content="c" * 400 + " HIGHPIC",
+ score=0.9,
+ chunk_id="c-high",
+ document_id="doc-1",
+ doc_item_refs=["#/pictures/1"],
+ image_data={"#/pictures/1": "HIGHBYTES"},
+ )
+ expanded = await expand_with_items(
+ rag.document_item_repository, "doc-1", [r_low, r_high], 500
+ )
+ assert len(expanded) == 1
+ assert "#/pictures/0" not in expanded[0].doc_item_refs
+ assert expanded[0].image_data == {"#/pictures/1": "HIGHBYTES"}
diff --git a/tests/test_skill_extras.py b/tests/test_skill_extras.py
new file mode 100644
index 00000000..7a667010
--- /dev/null
+++ b/tests/test_skill_extras.py
@@ -0,0 +1,42 @@
+"""Tests for non-tool utilities from ``haiku.rag.skills._tools.create_skill_extras``."""
+
+from docling_core.types.doc.document import DoclingDocument
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.config import AppConfig
+from haiku.rag.skills._tools import create_skill_extras
+from haiku.rag.store.models.chunk import Chunk
+
+
+async def _seed_chunk(db_path) -> str:
+ docling_doc = DoclingDocument(name="extras-test")
+ chunk = Chunk(
+ content="Some content.",
+ metadata={"doc_item_refs": ["#/texts/0"], "page_numbers": [1]},
+ order=0,
+ embedding=[0.1] * 2560,
+ )
+ async with HaikuRAG(db_path, create=True) as client:
+ doc = await client.import_document(docling_doc, [chunk], uri="test://extras")
+ stored = await client.chunk_repository.get_by_document_id(doc.id)
+ return stored[0].id
+
+
+async def test_extras_visualize_chunk_accepts_str_and_list(temp_db_path):
+ chunk_id = await _seed_chunk(temp_db_path)
+ extras = create_skill_extras(temp_db_path, AppConfig())
+ visualize_chunk = extras["visualize_chunk"]
+
+ # A document imported without page images yields no visualizations, but the
+ # str and list inputs must both resolve the chunk and reach visualize_chunk.
+ assert await visualize_chunk(chunk_id) == []
+ assert await visualize_chunk([chunk_id]) == []
+
+
+async def test_extras_visualize_chunk_unknown_id_returns_empty(temp_db_path):
+ await _seed_chunk(temp_db_path)
+ extras = create_skill_extras(temp_db_path, AppConfig())
+ visualize_chunk = extras["visualize_chunk"]
+
+ assert await visualize_chunk("does-not-exist") == []
+ assert await visualize_chunk(["does-not-exist"]) == []