Merge pull request #371 from ggozad/fix/image-results

Verify picture bytes before attaching to multimodal tool returns

CHANGELOG.md

@@ -18,6 +18,7 @@
 
 ### Fixed
 
+- **Picture bytes attached to multimodal tool returns are PNG-verified via `PIL.Image.verify()`.** Bytes that fail verification are dropped.
 - **Conversion options now apply to non-PDF formats.** `DoclingLocalConverter` previously wired its `PdfPipelineOptions` only to `InputFormat.PDF`, so user settings (OCR knobs, `picture_description.enabled`, `images_scale`, etc.) silently no-op'd for HTML, Markdown, DOCX, PPTX, and IMAGE inputs. The converter now shares a single `PdfPipelineOptions` instance across PDF, IMAGE, HTML, MD, DOCX, and PPTX `FormatOption`s. SimplePipeline-backed formats ignore the PDF-specific fields; `ConvertPipelineOptions`-level enrichments (picture description / classification / chart extraction) now run uniformly. HTML and Markdown additionally receive `HTMLBackendOptions` / `MarkdownBackendOptions` gated on `fetch_remote_images`.
 - **HTML text ingest path picks up converter options.** `convert_text(format="html"/"md")` previously used a bare `DoclingDocConverter()` with zero format options — the wix corpus ingest path. It now uses the same shared `_build_format_options()` helper as the file path.
 - **Relative `<img>` paths resolve during URL ingest.** `HaikuRAG.convert()` and the converter `convert_file` / `convert_text` methods now thread a `source_uri` through to `HTMLBackendOptions.source_uri` / `MarkdownBackendOptions.source_uri`. URL ingest uses the originating URL; file ingest uses `file://`; raw text accepts an optional override. docling-serve accepts the kwarg as a no-op (its API has no equivalent option).

haiku_rag_slim/haiku/rag/skills/_tools.py

@@ -1,16 +1,16 @@
-import base64
 from pathlib import Path
 from typing import Any
 
 from pydantic import BaseModel
 from pydantic_ai import RunContext
-from pydantic_ai.messages import BinaryContent, ToolReturn
+from pydantic_ai.messages import ToolReturn
 
 from haiku.rag.agents.research.models import Citation
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import AppConfig
 from haiku.rag.skills._deps import AnalysisRunDeps, RAGRunDeps
 from haiku.rag.store.models.chunk import SearchResult
+from haiku.rag.tools.search import build_binary_parts_from_results
 
 
 class CodeExecutionEntry(BaseModel):
@@ -202,24 +202,7 @@ def create_skill_tools(
             if not config.qa.model.vision:
                 return formatted
 
-            binary_parts: list[BinaryContent] = []
+            binary_parts = build_binary_parts_from_results(results)
-            seen: set[tuple[str | None, str]] = set()
-            for result in results:
-                if not result.image_data:
-                    continue
-                for self_ref, b64 in result.image_data.items():
-                    key = (result.document_id, self_ref)
-                    if key in seen:
-                        continue
-                    binary_parts.append(
-                        BinaryContent(
-                            data=base64.b64decode(b64),
-                            media_type="image/png",
-                            identifier=self_ref,
-                        )
-                    )
-                    seen.add(key)
-
             if binary_parts:
                 return ToolReturn(return_value=formatted, content=binary_parts)
             return formatted

haiku_rag_slim/haiku/rag/tools/search.py

@@ -1,6 +1,8 @@
 import base64
 from collections.abc import Callable
+from io import BytesIO
 
+from PIL import Image
 from pydantic_ai import FunctionToolset, RunContext
 from pydantic_ai.messages import BinaryContent, ToolReturn
 
@@ -9,6 +11,44 @@ from haiku.rag.store.models import SearchResult
 from haiku.rag.tools.context import RAGDeps
 
 
+def build_binary_parts_from_results(
+    results: list[SearchResult],
+) -> list[BinaryContent]:
+    """Decode and validate picture bytes attached to search results.
+
+    Dedup keyed on ``(document_id, self_ref)`` so the same picture in
+    different chunks is sent once. Pictures that fail
+    ``PIL.Image.verify()`` are skipped — the model adapter renders one
+    vision placeholder per ``BinaryContent``, so emitting one for an
+    image the server can't decode leaves the processor with an
+    off-by-one count.
+    """
+    parts: list[BinaryContent] = []
+    seen: set[tuple[str | None, str]] = set()
+    for result in results:
+        if not result.image_data:
+            continue
+        for self_ref, b64 in result.image_data.items():
+            key = (result.document_id, self_ref)
+            if key in seen:
+                continue
+            data = base64.b64decode(b64)
+            try:
+                with Image.open(BytesIO(data)) as img:
+                    img.verify()
+            except Exception:
+                continue
+            parts.append(
+                BinaryContent(
+                    data=data,
+                    media_type="image/png",
+                    identifier=self_ref,
+                )
+            )
+            seen.add(key)
+    return parts
+
+
 def create_search_toolset(
     config: AppConfig,
     expand_context: bool = True,
@@ -93,24 +133,7 @@ def create_search_toolset(
         if not config.qa.model.vision:
             return text
 
-        binary_parts: list[BinaryContent] = []
+        binary_parts = build_binary_parts_from_results(results_list)
-        seen: set[tuple[str | None, str]] = set()
-        for result in results_list:
-            if not result.image_data:
-                continue
-            for self_ref, b64 in result.image_data.items():
-                key = (result.document_id, self_ref)
-                if key in seen:
-                    continue
-                binary_parts.append(
-                    BinaryContent(
-                        data=base64.b64decode(b64),
-                        media_type="image/png",
-                        identifier=self_ref,
-                    )
-                )
-                seen.add(key)
-
         if binary_parts:
             return ToolReturn(return_value=text, content=binary_parts)
         return text

tests/test_picture_in_context.py

@@ -2,9 +2,11 @@
 
 import base64
 from dataclasses import dataclass
+from io import BytesIO
 from unittest.mock import AsyncMock
 
 import pytest
+from PIL import Image as PILImageModule
 from pydantic_ai import RunContext
 from pydantic_ai.messages import BinaryContent, ToolReturn
 from pydantic_ai.models.test import TestModel
@@ -18,7 +20,14 @@ 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"
+
+def _make_png(color: str = "red", size: tuple[int, int] = (4, 4)) -> bytes:
+    buf = BytesIO()
+    PILImageModule.new("RGB", size, color).save(buf, "PNG")
+    return buf.getvalue()
+
+
+PICTURE_BYTES = _make_png("red")
 PICTURE_B64 = base64.b64encode(PICTURE_BYTES).decode("ascii")
 
 
@@ -296,7 +305,7 @@ async def test_search_tool_attaches_same_self_ref_from_different_documents():
     key on ``(document_id, self_ref)`` so each document's figure reaches
     the model. Keying on ``self_ref`` alone silently drops the second
     document's bytes, leaving the model with text only for that result."""
-    other_bytes = b"\x89PNG\r\n\x1a\nother-doc-bytes"
+    other_bytes = _make_png("blue")
     other_b64 = base64.b64encode(other_bytes).decode("ascii")
 
     doc_a = SearchResult(
@@ -628,6 +637,100 @@ async def test_search_tool_skips_binary_content_when_qa_model_is_text_only():
     assert isinstance(result, str)
 
 
+@pytest.mark.asyncio
+async def test_search_tool_drops_invalid_image_bytes():
+    """A picture whose bytes cannot be decoded by PIL must not produce a
+    BinaryContent part. Otherwise the model adapter emits a vision
+    placeholder for an image the server can't decode, leaving Qwen3-VL's
+    processor with an off-by-one count and a 400 from
+    ``Qwen3VLProcessor``."""
+    from io import BytesIO
+
+    from PIL import Image as PILImageModule
+
+    buf = BytesIO()
+    PILImageModule.new("RGB", (4, 4), "red").save(buf, "PNG")
+    valid_png = buf.getvalue()
+    valid_b64 = base64.b64encode(valid_png).decode("ascii")
+
+    # Truthy bytes (passes ``if blob`` guards) but not a decodable PNG.
+    invalid_bytes = b"\x89PNG\r\n\x1a\ngarbage"
+    invalid_b64 = base64.b64encode(invalid_bytes).decode("ascii")
+
+    picture_result = SearchResult(
+        content="Two figures",
+        score=1.0,
+        chunk_id="chunk-1",
+        document_id="doc-1",
+        doc_item_refs=["#/pictures/0", "#/pictures/1"],
+        labels=["picture"],
+        image_data={"#/pictures/0": valid_b64, "#/pictures/1": invalid_b64},
+    )
+
+    fake_client = AsyncMock()
+    fake_client.search = AsyncMock(return_value=[picture_result])
+    fake_client.expand_context = AsyncMock(return_value=[picture_result])
+
+    config = AppConfig()
+    config.qa.model.vision = True
+    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 result.content is not None
+    identifiers = {p.identifier for p in result.content}  # type: ignore[attr-defined]
+    assert identifiers == {"#/pictures/0"}, (
+        "Only the decodable PNG should reach the model — the corrupt "
+        "ref must be dropped so we don't emit a placeholder for an "
+        "image the server can't decode."
+    )
+
+
+@pytest.mark.asyncio
+async def test_search_tool_drops_all_invalid_returns_plain_text():
+    """If every picture in the result set fails decode, fall back to a
+    plain string return — there's nothing to attach, so wrapping in
+    ``ToolReturn`` with an empty ``content`` list would surface an empty
+    user message downstream."""
+    bad_b64 = base64.b64encode(b"\x89PNGnope").decode("ascii")
+    picture_result = SearchResult(
+        content="One broken figure",
+        score=1.0,
+        chunk_id="chunk-1",
+        document_id="doc-1",
+        doc_item_refs=["#/pictures/0"],
+        labels=["picture"],
+        image_data={"#/pictures/0": bad_b64},
+    )
+
+    fake_client = AsyncMock()
+    fake_client.search = AsyncMock(return_value=[picture_result])
+    fake_client.expand_context = AsyncMock(return_value=[picture_result])
+
+    config = AppConfig()
+    config.qa.model.vision = True
+    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)
+
+
 @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

tests/test_skill_tools.py

@@ -8,10 +8,12 @@ text only.
 """
 
 import base64
+from io import BytesIO
 from pathlib import Path
 from unittest.mock import AsyncMock
 
 import pytest
+from PIL import Image as PILImageModule
 from pydantic_ai import RunContext
 from pydantic_ai.messages import BinaryContent, ToolReturn
 from pydantic_ai.models.test import TestModel
@@ -23,7 +25,14 @@ from haiku.rag.skills._tools import create_skill_tools
 from haiku.rag.skills.rag import RAGState
 from haiku.rag.store.models.chunk import SearchResult
 
-PICTURE_BYTES = b"\x89PNG\r\n\x1a\nfake-picture-bytes"
+
+def _make_png(color: str = "red") -> bytes:
+    buf = BytesIO()
+    PILImageModule.new("RGB", (4, 4), color).save(buf, "PNG")
+    return buf.getvalue()
+
+
+PICTURE_BYTES = _make_png("red")
 PICTURE_B64 = base64.b64encode(PICTURE_BYTES).decode("ascii")
 
 
@@ -192,7 +201,7 @@ async def test_skill_search_keeps_same_self_ref_from_different_documents():
     """``#/pictures/0`` in document A and ``#/pictures/0`` in document B
     are different figures. Dedup must key on ``(document_id, self_ref)``;
     keying on ``self_ref`` alone would drop document B's bytes."""
-    other_bytes = b"\x89PNG\r\n\x1a\nother-doc-bytes"
+    other_bytes = _make_png("blue")
     other_b64 = base64.b64encode(other_bytes).decode("ascii")
 
     doc_a = SearchResult(