From b01c649684e2644ac61843d077d2df883dc427df Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 30 Apr 2026 10:44:42 +0300 Subject: [PATCH] surface picture image bytes in SearchResult and emit multimodal ToolReturn from the agent search tool --- CHANGELOG.md | 2 + haiku_rag_slim/haiku/rag/agents/qa/prompts.py | 2 + haiku_rag_slim/haiku/rag/client/__init__.py | 3 +- haiku_rag_slim/haiku/rag/client/search.py | 56 +++- haiku_rag_slim/haiku/rag/context.py | 7 + haiku_rag_slim/haiku/rag/mcp.py | 14 +- .../haiku/rag/store/models/chunk.py | 12 +- haiku_rag_slim/haiku/rag/tools/search.py | 37 ++- tests/test_picture_in_context.py | 292 ++++++++++++++++++ 9 files changed, 416 insertions(+), 9 deletions(-) create mode 100644 tests/test_picture_in_context.py diff --git a/CHANGELOG.md b/CHANGELOG.md index f563c08c..dd8ce7ca 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,10 +5,12 @@ - **Storage column for embedded picture bytes.** `DocumentItemRecord` gains a `picture_data: bytes | None` column (Arrow `large_binary`) to hold per-`PictureItem` image bytes addressable by `(document_id, self_ref)`. New repository accessors `get_picture_bytes` and `get_pictures_for_chunk` expose them; the existing items read paths (`get_all_items`, `get_all_items_grouped`, `get_items_in_range`, `_record_to_item`) now project an explicit lightweight column set so context expansion and the analysis-sandbox `items.jsonl` build never pull picture bytes into memory. Existing databases pick up the column via the `0.45.0` migration alongside the picture-byte backfill (see below). Foundation for upcoming vision-in-context retrieval; not yet wired into ingestion or search. - **Embedded picture bytes captured at ingestion.** `extract_items` now decodes each `PictureItem.image.uri` data URI into raw bytes and writes them to `document_items.picture_data` so per-figure lookups don't require decompressing the full docling blob. The same path also surfaces VLM-generated picture descriptions (`meta.description.text`) into `DocumentItem.text` so picture-only chunks survive `expand_with_items`' text filter. The `0.45.0` migration adds the `picture_data` column to existing databases and backfills it by extracting bytes out of `docling_document`, stripping picture URIs from the structure blob in the process; `compress_docling_split` does the same for new ingests so the structure stays lean. Rebuild and update flows snapshot picture bytes via the new `DocumentItemRepository.get_all_picture_data` accessor before re-extraction so a re-chunk doesn't drop them. +- **Picture image bytes in search results and vision-capable QA.** `SearchResult` gains an `image_data: dict[str, str] | None` field carrying base64-encoded picture bytes keyed by `self_ref` for picture-labeled chunks. `client.search()` and the MCP `search_documents` tool gain an `include_images: bool = True` flag; set False to omit the bytes for plain-text consumers. `expand_context` now preserves picture self_refs with empty text so they aren't filtered out before reaching the image-data lookup. The agent-facing search tool (`tools/search.py`) returns `pydantic_ai.messages.ToolReturn(return_value=text, content=[BinaryContent(...), ...])` when picture data is present so a vision-capable QA model sees the figures alongside the text; otherwise it returns a plain string and non-vision flows are unchanged. ### Fixed - **docling-serve picture-image extraction.** docling-serve never populated `PictureItem.image` when called with `image_export_mode="embedded"` (per upstream issue [docling-project/docling-serve#576](https://github.com/docling-project/docling-serve/issues/576) — picture-image generation is only triggered when the server's `image_export_mode == "referenced"`). The converter now switches to `image_export_mode="referenced"` + `target_type="zip"` whenever picture images are requested, parses the returned zip, and rehydrates `artifacts/` URIs back into `data:;base64,...` URIs so downstream code sees the same shape as docling-local. The previously-`xfail`ed `test_convert_pdf_with_picture_images` integration test now passes. +- **`expand_context` now repopulates `image_data` after rebuilding `SearchResult` objects.** `expand_with_items` constructs fresh `SearchResult`s from the items table, so `image_data` attached upstream by `client.search()` was being dropped before reaching the agent's search tool — vision-capable QA models received the picture text/captions but never the picture bytes. The expansion path now re-runs the picture-bytes lookup against the post-expansion `doc_item_refs` (which can grow to include sibling pictures pulled in via section bounds). ## [0.44.0] - 2026-04-29 diff --git a/haiku_rag_slim/haiku/rag/agents/qa/prompts.py b/haiku_rag_slim/haiku/rag/agents/qa/prompts.py index 701e689c..5403ccaf 100644 --- a/haiku_rag_slim/haiku/rag/agents/qa/prompts.py +++ b/haiku_rag_slim/haiku/rag/agents/qa/prompts.py @@ -26,6 +26,8 @@ Each result includes: - Type: content type like paragraph, table, code, list_item (when available) - Content: the actual text +When a result is of `Type: picture` (a figure or diagram), the search tool may also attach the picture itself as image content alongside the text — use it for visual reasoning when answering. Reference it by its chunk_id like any other source. + IMPORTANT: You MUST include in cited_chunks the COMPLETE IDs of every chunk you reference. Copy the full ID string without brackets — e.g. "5ae52166-5329-42e9-b6a5-756fc0cb7200" not "[5ae52166]" or "5ae52166". Never truncate IDs. Never leave cited_chunks empty if you found relevant content. Guidelines: diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 2b9c4e9e..017454aa 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -322,10 +322,11 @@ class HaikuRAG: limit: int | None = None, search_type: str = "hybrid", filter: str | None = None, + include_images: bool = True, ) -> list[SearchResult]: from haiku.rag.client.search import search - return await search(self, query, limit, search_type, filter) + return await search(self, query, limit, search_type, filter, include_images) async def expand_context( self, diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 10af0621..9cdb8a3c 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -1,3 +1,4 @@ +import base64 from typing import TYPE_CHECKING from haiku.rag.reranking import get_reranker @@ -13,6 +14,7 @@ async def search( limit: int | None = None, search_type: str = "hybrid", filter: str | None = None, + include_images: bool = True, ) -> list[SearchResult]: """Search for relevant chunks with optional reranking. @@ -22,6 +24,10 @@ async def search( limit: Maximum number of results to return. Defaults to config.search.limit. search_type: Type of search - "vector", "fts", or "hybrid" (default). filter: Optional SQL WHERE clause to filter documents before searching chunks. + include_images: When True, populate ``SearchResult.image_data`` with + base64-encoded picture bytes for picture-labeled chunks. Set to + False to skip the lookup (e.g. for plain-text MCP consumers that + don't want the JSON bloat). Returns: List of SearchResult objects ordered by relevance. @@ -43,7 +49,52 @@ async def search( chunks = [chunk for chunk, _ in raw_results] chunk_results = await reranker.rerank(query, chunks, top_n=limit) - return [SearchResult.from_chunk(chunk, score) for chunk, score in chunk_results] + results = [SearchResult.from_chunk(chunk, score) for chunk, score in chunk_results] + + if include_images: + await _populate_image_data(client, results) + + return results + + +async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult]) -> None: + """Attach base64 picture bytes to ``SearchResult.image_data`` in-place. + + Groups results by document_id and batches one picture-bytes lookup per + document so a result set spanning N documents costs N reads, not one per + picture. Only refs starting with ``#/pictures/`` are queried. + """ + by_doc: dict[str, list[SearchResult]] = {} + for r in results: + if not r.document_id: + continue + if not any(ref.startswith("#/pictures/") for ref in r.doc_item_refs): + continue + by_doc.setdefault(r.document_id, []).append(r) + + for doc_id, doc_results in by_doc.items(): + wanted: list[str] = [] + seen: set[str] = set() + for r in doc_results: + for ref in r.doc_item_refs: + if ref.startswith("#/pictures/") and ref not in seen: + wanted.append(ref) + seen.add(ref) + if not wanted: + continue + bytes_by_ref = await client.document_item_repository.get_pictures_for_chunk( + doc_id, wanted + ) + if not bytes_by_ref: + continue + for r in doc_results: + attached: dict[str, str] = {} + for ref in r.doc_item_refs: + blob = bytes_by_ref.get(ref) + if blob: + attached[ref] = base64.b64encode(blob).decode("ascii") + if attached: + r.image_data = attached async def expand_context( @@ -92,6 +143,9 @@ async def expand_context( expanded_results.extend(expanded) expanded_results.sort(key=lambda r: r.score, reverse=True) + # expand_with_items rebuilds SearchResult objects, so attach picture bytes + # to the fresh set — picture self_refs may have grown via section expansion. + await _populate_image_data(client, expanded_results) return expanded_results diff --git a/haiku_rag_slim/haiku/rag/context.py b/haiku_rag_slim/haiku/rag/context.py index 2ee03ebd..f32b14e5 100644 --- a/haiku_rag_slim/haiku/rag/context.py +++ b/haiku_rag_slim/haiku/rag/context.py @@ -238,6 +238,13 @@ async def expand_with_items( if item.label: labels.add(item.label) pages.update(item.page_numbers) + 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) all_headings: list[str] = [] for r in original_results: diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 11871dd5..574bb711 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -91,12 +91,20 @@ def create_mcp_server( # Read tools - always registered @mcp.tool() async def search_documents( - query: str, limit: int | None = None + query: str, limit: int | None = None, include_images: bool = True ) -> list[SearchResult]: - """Search the RAG system for documents using hybrid search (vector similarity + full-text search).""" + """Search the RAG system for documents using hybrid search (vector similarity + full-text search). + + When include_images is True (default) and a picture-labeled chunk is + in the result set, ``SearchResult.image_data`` carries base64-encoded + PNG bytes keyed by self_ref. Set to False to omit the bytes from the + response (smaller JSON payload for plain-text consumers). + """ try: async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - return await rag.search(query, limit=limit) + return await rag.search( + query, limit=limit, include_images=include_images + ) except Exception: return [] diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index c11c4f50..7ee9a125 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -109,7 +109,14 @@ class Chunk(BaseModel): class SearchResult(BaseModel): - """Search result with optional provenance information for citations.""" + """Search result with optional provenance information for citations. + + ``image_data`` carries embedded picture bytes (base64-encoded PNG) keyed by + ``self_ref`` for picture-labeled chunks. Empty/None when no pictures or + 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. + """ content: str score: float @@ -122,12 +129,14 @@ class SearchResult(BaseModel): page_numbers: list[int] = [] headings: list[str] | None = None labels: list[str] = [] + image_data: dict[str, str] | None = None @classmethod def from_chunk( cls, chunk: "Chunk", score: float, + image_data: dict[str, str] | None = None, ) -> "SearchResult": """Create from a Chunk.""" meta = chunk.get_chunk_metadata() @@ -143,6 +152,7 @@ class SearchResult(BaseModel): page_numbers=meta.page_numbers, headings=meta.headings, labels=meta.labels, + image_data=image_data, ) def format_for_agent( diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index 63b96e75..1c76a3fd 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -1,6 +1,8 @@ +import base64 from collections.abc import Callable from pydantic_ai import FunctionToolset, RunContext +from pydantic_ai.messages import BinaryContent, ToolReturn from haiku.rag.config.models import AppConfig from haiku.rag.store.models import SearchResult @@ -40,7 +42,7 @@ def create_search_toolset( ctx: RunContext[RAGDeps], query: str, limit: int | None = None, - ) -> str: + ) -> str | ToolReturn: """Search the knowledge base for relevant documents. Args: @@ -48,7 +50,11 @@ def create_search_toolset( limit: Number of results to return (default: from config). Returns: - Formatted search results with content and metadata. + Formatted search results with content and metadata. When a + picture-labeled chunk is in the result set, returns a + ``pydantic_ai.messages.ToolReturn`` whose ``content`` carries the + corresponding ``BinaryContent`` parts so a vision-capable model + sees the figures alongside the text. """ rid = ctx.run_id or "" search_counts[rid] = search_counts.get(rid, 0) + 1 @@ -82,7 +88,32 @@ def create_search_toolset( r.format_for_agent(rank=i + 1, total=total) for i, r in enumerate(results_list) ] - return "\n\n".join(formatted) + text = "\n\n".join(formatted) + + binary_parts: list[BinaryContent] = [] + seen: set[str] = set() + for result in results_list: + if not result.image_data: + continue + for self_ref, b64 in result.image_data.items(): + if self_ref in seen: + continue + try: + raw = base64.b64decode(b64) + except (ValueError, TypeError): + continue + binary_parts.append( + BinaryContent( + data=raw, + media_type="image/png", + identifier=self_ref, + ) + ) + seen.add(self_ref) + + if binary_parts: + return ToolReturn(return_value=text, content=binary_parts) + return text toolset: FunctionToolset[RAGDeps] = FunctionToolset() toolset.add_function(search, name=tool_name, retries=3) diff --git a/tests/test_picture_in_context.py b/tests/test_picture_in_context.py new file mode 100644 index 00000000..6f1a7060 --- /dev/null +++ b/tests/test_picture_in_context.py @@ -0,0 +1,292 @@ +"""A3: SearchResult.image_data, expand_context preservation, multimodal ToolReturn.""" + +import base64 +from dataclasses import dataclass +from unittest.mock import AsyncMock + +import pytest +from pydantic_ai import RunContext +from pydantic_ai.messages import BinaryContent, ToolReturn +from pydantic_ai.models.test import TestModel +from pydantic_ai.usage import RunUsage + +from haiku.rag.client import HaikuRAG +from haiku.rag.client.search import _populate_image_data +from haiku.rag.config import Config +from haiku.rag.context import expand_with_items +from haiku.rag.store.models.chunk import SearchResult +from haiku.rag.store.models.document_item import DocumentItem +from haiku.rag.tools.search import create_search_toolset + +PICTURE_BYTES = b"\x89PNG\r\n\x1a\nfake-picture-bytes" +PICTURE_B64 = base64.b64encode(PICTURE_BYTES).decode("ascii") + + +@pytest.mark.asyncio +async def test_populate_image_data_attaches_base64(temp_db_path): + async with HaikuRAG(temp_db_path, create=True) as rag: + await rag.document_item_repository.create_items( + "doc-1", + [ + DocumentItem( + document_id="doc-1", + position=0, + self_ref="#/texts/0", + label="paragraph", + text="Some text", + ), + DocumentItem( + document_id="doc-1", + position=1, + self_ref="#/pictures/0", + label="picture", + text="", + picture_data=PICTURE_BYTES, + ), + ], + ) + + text_only = SearchResult( + content="Some text", + score=1.0, + document_id="doc-1", + doc_item_refs=["#/texts/0"], + labels=["paragraph"], + ) + with_picture = SearchResult( + content="A figure caption", + score=0.9, + document_id="doc-1", + doc_item_refs=["#/texts/0", "#/pictures/0"], + labels=["paragraph", "picture"], + ) + + await _populate_image_data(rag, [text_only, with_picture]) + + # Text-only result is unchanged + assert text_only.image_data is None + # Picture-bearing result has the bytes attached + assert with_picture.image_data == {"#/pictures/0": PICTURE_B64} + + +@pytest.mark.asyncio +async def test_client_search_include_images_false_skips_lookup(temp_db_path): + """include_images=False must short-circuit the picture-bytes lookup.""" + async with HaikuRAG(temp_db_path, create=True) as rag: + await rag.document_item_repository.create_items( + "doc-1", + [ + DocumentItem( + document_id="doc-1", + position=0, + self_ref="#/pictures/0", + label="picture", + picture_data=PICTURE_BYTES, + ), + ], + ) + # Spy that we never reach the picture-bytes accessor + rag.document_item_repository.get_pictures_for_chunk = AsyncMock( # type: ignore[method-assign] + wraps=rag.document_item_repository.get_pictures_for_chunk + ) + + from haiku.rag.client.search import search + + # Stub the chunk-search results so we don't depend on embeddings/FTS + async def fake_chunk_search(*args, **kwargs): + return [] + + rag.chunk_repository.search = fake_chunk_search # type: ignore[method-assign] + + await search(rag, "anything", include_images=False) + rag.document_item_repository.get_pictures_for_chunk.assert_not_called() + + +@pytest.mark.asyncio +async def test_expand_context_preserves_picture_refs_with_empty_text(temp_db_path): + """A picture item with empty text must keep its self_ref through expansion.""" + async with HaikuRAG(temp_db_path, create=True) as rag: + # Build an items table with a section header + a paragraph match + an + # adjacent picture row that has no text. expand_with_items used to + # filter the picture out via the `if item.text:` guard; A3 keeps it. + await rag.document_item_repository.create_items( + "doc-1", + [ + DocumentItem( + document_id="doc-1", + position=0, + self_ref="#/texts/0", + label="section_header", + text="Methods", + ), + DocumentItem( + document_id="doc-1", + position=1, + self_ref="#/texts/1", + label="paragraph", + text="The figure below shows the architecture.", + ), + DocumentItem( + document_id="doc-1", + position=2, + self_ref="#/pictures/0", + label="picture", + text="", + picture_data=PICTURE_BYTES, + ), + DocumentItem( + document_id="doc-1", + position=3, + self_ref="#/texts/2", + label="paragraph", + text="More commentary follows.", + ), + ], + ) + + # Match on the paragraph that mentions the figure. + seed = SearchResult( + content="The figure below shows the architecture.", + score=1.0, + document_id="doc-1", + doc_item_refs=["#/texts/1"], + labels=["paragraph"], + ) + expanded = await expand_with_items( + rag.document_item_repository, + "doc-1", + [seed], + max_chars=10_000, + ) + assert len(expanded) == 1 + out = expanded[0] + assert "#/pictures/0" in out.doc_item_refs, ( + "Picture self_ref should survive expansion even with empty text" + ) + assert "picture" in out.labels + + +@pytest.mark.asyncio +async def test_expand_context_repopulates_image_data(temp_db_path): + """expand_context rebuilds SearchResult objects via expand_with_items, so + it must re-attach picture bytes — otherwise vision flows downstream see + empty image_data after expansion.""" + async with HaikuRAG(temp_db_path, create=True) as rag: + await rag.document_item_repository.create_items( + "doc-1", + [ + DocumentItem( + document_id="doc-1", + position=0, + self_ref="#/texts/0", + label="section_header", + text="Methods", + ), + DocumentItem( + document_id="doc-1", + position=1, + self_ref="#/texts/1", + label="paragraph", + text="The figure below shows the architecture.", + ), + DocumentItem( + document_id="doc-1", + position=2, + self_ref="#/pictures/0", + label="picture", + text="", + picture_data=PICTURE_BYTES, + ), + ], + ) + seed = SearchResult( + content="The figure below shows the architecture.", + score=1.0, + document_id="doc-1", + doc_item_refs=["#/texts/1"], + labels=["paragraph"], + image_data=None, + ) + expanded = await rag.expand_context([seed]) + assert len(expanded) == 1 + out = expanded[0] + assert "#/pictures/0" in out.doc_item_refs + assert out.image_data == {"#/pictures/0": PICTURE_B64} + + +@dataclass +class _Deps: + client: object + + +@pytest.mark.asyncio +async def test_search_tool_returns_multimodal_when_picture_present(): + """The agent-facing search tool must wrap text + BinaryContent in ToolReturn + whenever a result carries picture image_data.""" + picture_result = SearchResult( + content="A diagram of the layout", + score=1.0, + chunk_id="chunk-1", + document_id="doc-1", + doc_item_refs=["#/pictures/0"], + labels=["picture"], + image_data={"#/pictures/0": PICTURE_B64}, + ) + + fake_client = AsyncMock() + fake_client.search = AsyncMock(return_value=[picture_result]) + fake_client.expand_context = AsyncMock(return_value=[picture_result]) + + toolset = create_search_toolset(Config, expand_context=False) + func = toolset.tools["search"].function + + ctx = RunContext( + deps=_Deps(client=fake_client), # type: ignore[arg-type] + model=TestModel(), + usage=RunUsage(), + run_id="run-1", + ) + result = await func(ctx, "anything") + + assert isinstance(result, ToolReturn) + assert isinstance(result.return_value, str) + assert "Type: picture" in result.return_value or "rank 1" in result.return_value + assert result.content is not None + assert len(result.content) == 1 + part = result.content[0] + assert isinstance(part, BinaryContent) + assert part.media_type == "image/png" + assert part.identifier == "#/pictures/0" + assert part.data == PICTURE_BYTES + + +@pytest.mark.asyncio +async def test_search_tool_returns_plain_string_when_no_pictures(): + """When no result carries image_data the tool returns a plain str (no + ToolReturn wrapper) so non-vision flows are unaffected.""" + text_result = SearchResult( + content="Some text", + score=1.0, + chunk_id="chunk-1", + document_id="doc-1", + doc_item_refs=["#/texts/0"], + labels=["paragraph"], + ) + + fake_client = AsyncMock() + fake_client.search = AsyncMock(return_value=[text_result]) + fake_client.expand_context = AsyncMock(return_value=[text_result]) + + toolset = create_search_toolset(Config, expand_context=False) + func = toolset.tools["search"].function + + ctx = RunContext( + deps=_Deps(client=fake_client), # type: ignore[arg-type] + model=TestModel(), + usage=RunUsage(), + run_id="run-1", + ) + result = await func(ctx, "anything") + + assert isinstance(result, str) + assert "rank 1" in result