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

Visualize all constituent chunks of a merged citation

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-07-08 15:26:27 +03:00
parent 0bcf34363a
commit 494f774473
No known key found for this signature in database
10 changed files with 161 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
CHANGELOG.md
View file

@ -17,6 +17,7 @@
- `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.
- 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

18
app/backend/main.py
View file

@ -203,7 +203,11 @@ 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 so a merged citation
can render the union of its constituent chunks' expansions.
"""
import base64
from io import BytesIO
@ -214,11 +218,15 @@ async def visualize_chunk(request: Request) -> JSONResponse:
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)
if not images:
return JSONResponse({"images": [], "message": "No visual grounding available"})
@ -233,7 +241,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,
}
)

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

@ -55,7 +55,13 @@ function CitationItem({
<button
type="button"
className="citation-view-btn"
onClick={() => onViewInDocument(citation.chunk_id)}
onClick={() =>
onViewInDocument(
citation.chunk_ids?.length
? citation.chunk_ids.join(",")
: citation.chunk_id,
)
}
>
View in Document
</button>
@ -96,9 +102,12 @@ export default function CitationBlock({ citations }: CitationBlockProps) {
});
try {
const response = await fetch(`/api/visualize/${chunkId}`, {
signal: controller.signal,
});
const response = await fetch(
`/api/visualize/${encodeURIComponent(chunkId)}`,
{
signal: controller.signal,
},
);
const data = await response.json();
if (controller.signal.aborted) return;

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

@ -2,6 +2,7 @@ export interface Citation {
index: number;
document_id: string;
chunk_id: string;
chunk_ids?: string[];
document_uri: string;
document_title: string | null;
page_numbers: number[];

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

@ -405,13 +405,18 @@ 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))
async def action_show_info(self) -> None:
"""Show database info modal."""

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

@ -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,7 +504,7 @@ 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]) -> list:
from haiku.rag.client.search import visualize_chunk
return await visualize_chunk(self, chunk)

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

@ -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,15 @@ 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]") -> 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
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.
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.
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 +243,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 []
@ -251,21 +260,28 @@ async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list:
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,
search_results = [
SearchResult(
content=c.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,
chunk_id=c.id,
document_id=c.document_id,
doc_item_refs=meta.doc_item_refs,
page_numbers=meta.page_numbers,
)
expanded = await expand_context(client, [search_result])
refs = expanded[0].doc_item_refs if expanded else chunk_meta.doc_item_refs
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]
meta = ChunkMetadata(doc_item_refs=refs)
else:
meta = chunk_meta
meta = chunks[0].get_chunk_metadata()
bounding_boxes = meta.resolve_bounding_boxes(docling_doc)
if not bounding_boxes:
return []
@ -278,7 +294,7 @@ async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list:
boxes_by_page[bbox.page_no].append(bbox)
# 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()))

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

@ -57,14 +57,14 @@ class VisualGroundingModal(Screen):
def __init__(
self,
chunk: "Chunk",
chunk: "Chunk | list[Chunk]",
client: "HaikuRAG",
document_uri: 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.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 +81,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)
await self._render_current_page()
async def _render_current_page(self) -> None:

13
haiku_rag_slim/haiku/rag/skills/_tools.py
View file

@ -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,

76
tests/test_client.py
View file

@ -1590,6 +1590,82 @@ 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
),
)
# 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):
docling_doc.add_text(
label=DocItemLabel.PARAGRAPH,
text=f"Paragraph {i}. " + "x" * 4000,
prov=ProvenanceItem(
page_no=page_no,
bbox=BoundingBox(l=50, t=700 - i * 100, r=550, b=650 - i * 100),
charspan=(0, 20),
),
)
chunks = [
Chunk(
content="Paragraph 0. " + "x" * 4000,
metadata={
"doc_item_refs": ["#/texts/0"],
"page_numbers": [1],
"labels": ["paragraph"],
},
order=0,
embedding=[0.1] * 2560,
),
Chunk(
content="Paragraph 3. " + "x" * 4000,
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
# =============================================================================
# convert() method tests
# =============================================================================