diff --git a/CHANGELOG.md b/CHANGELOG.md index 92e3d249..75a1fb02 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/app/backend/main.py b/app/backend/main.py index b8a62207..cd157ccf 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -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"}) 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 177a8f13..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); @@ -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(); diff --git a/app/frontend/lib/sessionStorage.ts b/app/frontend/lib/sessionStorage.ts index 26827fb1..a3369f0a 100644 --- a/app/frontend/lib/sessionStorage.ts +++ b/app/frontend/lib/sessionStorage.ts @@ -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 diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index ad7e6e85..98015ed4 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -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.""" diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index fdd63e2a..d8c55438 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index f2f8f06c..7d74e84f 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -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 ) 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 d2daa3fc..978a21cd 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py @@ -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: diff --git a/haiku_rag_slim/haiku/rag/store/models/citation.py b/haiku_rag_slim/haiku/rag/store/models/citation.py index 94e53e27..95510b63 100644 --- a/haiku_rag_slim/haiku/rag/store/models/citation.py +++ b/haiku_rag_slim/haiku/rag/store/models/citation.py @@ -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, ) ) diff --git a/tests/test_client.py b/tests/test_client.py index d9a2669c..ccf1be3e 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -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 # ============================================================================= 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"]) == []