danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Visualize the exact context the model saw via Citation.doc_item_refs

Browse source 
visualize_chunk re-expanded chunks from scratch to recover their refs,
which could not faithfully reproduce the original merge, scores, and
clip — so a visualization could highlight different pages than the
citation covered. Carry the cited items on Citation.doc_item_refs and
resolve bounding boxes from them directly; re-expansion remains only as
the fallback for callers with no stored context (CLI, inspector). Chat,
inspector, the app endpoint, and the frontend pass the refs through.
This commit is contained in:
Yiorgis Gozadinos 2026-07-09 10:50:04 +03:00
parent c424f51056
commit 995d081aa1
No known key found for this signature in database
12 changed files with 261 additions and 101 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
CHANGELOG.md
View file

@ -1,6 +1,12 @@
# 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.
### Changed
- Duplicate images within a document produce a single picture chunk.
@ -10,15 +16,14 @@
- 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
### Added
- `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

19
app/backend/main.py
View file

@ -205,14 +205,27 @@ async def db_info(_: Request) -> JSONResponse:
async def visualize_chunk(request: Request) -> JSONResponse:
"""Return visual grounding images for one or more chunks as base64.
The path param accepts comma-separated chunk ids so a merged citation
can render the union of its constituent chunks' expansions.
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)
@ -226,7 +239,7 @@ async def visualize_chunk(request: Request) -> JSONResponse:
if not chunks:
return JSONResponse({"error": "Chunk not found"}, status_code=404)
images = await client.visualize_chunk(chunks)
images = await client.visualize_chunk(chunks, refs)
if not images:
return JSONResponse({"images": [], "message": "No visual grounding available"})

8
app/frontend/app/api/visualize/[chunk_id]/route.ts
View file

@ -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) {

85
app/frontend/components/CitationBlock.tsx
View file

@ -20,7 +20,7 @@ function CitationItem({
onViewInDocument,
}: {
citation: Citation;
onViewInDocument: (chunkId: string) => void;
onViewInDocument: (chunkId: string, refs?: string[]) => void;
}) {
const [expanded, setExpanded] = useState(false);
@ -60,6 +60,7 @@ function CitationItem({
citation.chunk_ids?.length
? citation.chunk_ids.join(",")
: citation.chunk_id,
citation.doc_item_refs,
)
}
>
@ -88,48 +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,
});
setVisualGrounding({
isOpen: true,
chunkId,
images: [],
loading: true,
error: null,
});
try {
const response = await fetch(
`/api/visualize/${encodeURIComponent(chunkId)}`,
{
signal: controller.signal,
},
);
const data = await response.json();
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");
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();

1
app/frontend/lib/sessionStorage.ts
View file

@ -8,6 +8,7 @@ export interface Citation {
page_numbers: number[];
headings: string[] | null;
content: string;
doc_item_refs?: string[];
}
// Matches RAGState from the backend skill

8
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -416,7 +416,13 @@ class ChatApp(App):
from haiku.rag.inspector.widgets.visual_modal import VisualGroundingModal
await self.push_screen(VisualGroundingModal(chunk=chunks, 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."""

6
haiku_rag_slim/haiku/rag/client/__init__.py
View file

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

80
haiku_rag_slim/haiku/rag/client/search.py
View file

@ -224,19 +224,22 @@ async def expand_context(
return expanded_results
async def visualize_chunk(client: "HaikuRAG", chunk: "Chunk | Sequence[Chunk]") -> list:
async def visualize_chunk(
client: "HaikuRAG",
chunk: "Chunk | Sequence[Chunk]",
refs: list[str] | None = None,
) -> list:
"""Render page images with bounding box highlights for one or more chunks.
Expands the chunks' 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. Passing all
constituent chunks of a merged search result (``SearchResult.chunk_ids``)
reproduces the merged expansion; chunks from a different document than
the first are ignored.
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.
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.
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.
@ -263,35 +266,40 @@ async def visualize_chunk(client: "HaikuRAG", chunk: "Chunk | Sequence[Chunk]")
if not docling_doc:
return []
# Expand context to get all doc_item_refs in the 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)
refs: list[str] = []
for result in expanded:
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]
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)
else:
refs = chunks[0].get_chunk_metadata().doc_item_refs
# 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_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_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=sorted(matched_refs)
).resolve_bounding_boxes(docling_doc)
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
)

4
haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py
View file

@ -60,10 +60,12 @@ class VisualGroundingModal(Screen):
chunk: "Chunk | list[Chunk]",
client: "HaikuRAG",
document_uri: str | None = None,
refs: list[str] | None = None,
):
super().__init__()
self.chunks = chunk if isinstance(chunk, list) else [chunk]
self.client = client
self.refs = refs
self.document_uri = document_uri or self.chunks[0].document_uri
self.images: list[PILImage] = []
self.current_page_idx = 0
@ -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.chunks)
self.images = await self.client.visualize_chunk(self.chunks, self.refs)
await self._render_current_page()
async def _render_current_page(self) -> None:

11
haiku_rag_slim/haiku/rag/store/models/citation.py
View file

@ -20,9 +20,12 @@ class Citation(BaseModel):
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``). Visual grounding
passes them all to ``visualize_chunk`` so the rendered pages reproduce
the merged expansion.
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
@ -34,6 +37,7 @@ class Citation(BaseModel):
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)
@ -63,6 +67,7 @@ def resolve_citations(
page_numbers=r.page_numbers,
headings=r.headings,
content=r.content,
doc_item_refs=list(r.doc_item_refs),
picture_refs=picture_refs,
)
)

87
tests/test_client.py
View file

@ -1613,26 +1613,31 @@ async def test_client_visualize_chunk_merged_chunks_union_pages(temp_db_path):
),
)
# Five large paragraphs: with max_context_chars=10000 each chunk's own
# outward expansion spans three items, so chunk one alone stays on page
# one while the merged ranges [0,2] and [2,4] union to cover page two.
pages = [1, 1, 1, 2, 2]
for i, page_no in enumerate(pages):
# 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=DocItemLabel.PARAGRAPH,
text=f"Paragraph {i}. " + "x" * 4000,
label=label,
text=text,
prov=ProvenanceItem(
page_no=page_no,
bbox=BoundingBox(l=50, t=700 - i * 100, r=550, b=650 - i * 100),
bbox=BoundingBox(l=50, t=700 - (i % 2) * 100, r=550, b=650),
charspan=(0, 20),
),
)
chunks = [
Chunk(
content="Paragraph 0. " + "x" * 4000,
content="Page one body. " + "x" * 3000,
metadata={
"doc_item_refs": ["#/texts/0"],
"doc_item_refs": ["#/texts/1"],
"page_numbers": [1],
"labels": ["paragraph"],
},
@ -1640,7 +1645,7 @@ async def test_client_visualize_chunk_merged_chunks_union_pages(temp_db_path):
embedding=[0.1] * 2560,
),
Chunk(
content="Paragraph 3. " + "x" * 4000,
content="Page two body. " + "y" * 3000,
metadata={
"doc_item_refs": ["#/texts/3"],
"page_numbers": [2],
@ -1733,6 +1738,66 @@ async def test_client_visualize_chunk_two_tone_highlights(temp_db_path):
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
# =============================================================================
# convert() method tests
# =============================================================================

42
tests/test_skill_extras.py Normal file
View file

@ -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"]) == []