surface picture image bytes in SearchResult and emit multimodal ToolReturn from the agent search tool

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/<filename>` URIs back into `data:<mime>;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
 

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:

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,

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
 
 

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:

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 []
 

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(

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)

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