Verify picture bytes before attaching to multimodal tool returns

This commit is contained in:
Yiorgis Gozadinos 2026-05-13 16:00:38 +03:00
parent 675c9030a4
commit a6e5664c78
No known key found for this signature in database
5 changed files with 161 additions and 42 deletions

View file

@ -17,6 +17,7 @@
### Fixed ### 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`. - **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. - **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). - **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).

View file

@ -1,16 +1,16 @@
import base64
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
from pydantic import BaseModel from pydantic import BaseModel
from pydantic_ai import RunContext 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.agents.research.models import Citation
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
from haiku.rag.skills._deps import AnalysisRunDeps, RAGRunDeps from haiku.rag.skills._deps import AnalysisRunDeps, RAGRunDeps
from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.tools.search import build_binary_parts_from_results
class CodeExecutionEntry(BaseModel): class CodeExecutionEntry(BaseModel):
@ -202,24 +202,7 @@ def create_skill_tools(
if not config.qa.model.vision: if not config.qa.model.vision:
return formatted 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: if binary_parts:
return ToolReturn(return_value=formatted, content=binary_parts) return ToolReturn(return_value=formatted, content=binary_parts)
return formatted return formatted

View file

@ -1,6 +1,8 @@
import base64 import base64
from collections.abc import Callable from collections.abc import Callable
from io import BytesIO
from PIL import Image
from pydantic_ai import FunctionToolset, RunContext from pydantic_ai import FunctionToolset, RunContext
from pydantic_ai.messages import BinaryContent, ToolReturn 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 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( def create_search_toolset(
config: AppConfig, config: AppConfig,
expand_context: bool = True, expand_context: bool = True,
@ -93,24 +133,7 @@ def create_search_toolset(
if not config.qa.model.vision: if not config.qa.model.vision:
return text 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: if binary_parts:
return ToolReturn(return_value=text, content=binary_parts) return ToolReturn(return_value=text, content=binary_parts)
return text return text

View file

@ -2,9 +2,11 @@
import base64 import base64
from dataclasses import dataclass from dataclasses import dataclass
from io import BytesIO
from unittest.mock import AsyncMock from unittest.mock import AsyncMock
import pytest import pytest
from PIL import Image as PILImageModule
from pydantic_ai import RunContext from pydantic_ai import RunContext
from pydantic_ai.messages import BinaryContent, ToolReturn from pydantic_ai.messages import BinaryContent, ToolReturn
from pydantic_ai.models.test import TestModel 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.store.models.document_item import DocumentItem
from haiku.rag.tools.search import create_search_toolset 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") 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 key on ``(document_id, self_ref)`` so each document's figure reaches
the model. Keying on ``self_ref`` alone silently drops the second the model. Keying on ``self_ref`` alone silently drops the second
document's bytes, leaving the model with text only for that result.""" 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") other_b64 = base64.b64encode(other_bytes).decode("ascii")
doc_a = SearchResult( 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) 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 @pytest.mark.asyncio
async def test_search_tool_returns_plain_string_when_no_pictures(): async def test_search_tool_returns_plain_string_when_no_pictures():
"""When no result carries image_data the tool returns a plain str (no """When no result carries image_data the tool returns a plain str (no

View file

@ -8,10 +8,12 @@ text only.
""" """
import base64 import base64
from io import BytesIO
from pathlib import Path from pathlib import Path
from unittest.mock import AsyncMock from unittest.mock import AsyncMock
import pytest import pytest
from PIL import Image as PILImageModule
from pydantic_ai import RunContext from pydantic_ai import RunContext
from pydantic_ai.messages import BinaryContent, ToolReturn from pydantic_ai.messages import BinaryContent, ToolReturn
from pydantic_ai.models.test import TestModel 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.skills.rag import RAGState
from haiku.rag.store.models.chunk import SearchResult 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") 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 """``#/pictures/0`` in document A and ``#/pictures/0`` in document B
are different figures. Dedup must key on ``(document_id, self_ref)``; are different figures. Dedup must key on ``(document_id, self_ref)``;
keying on ``self_ref`` alone would drop document B's bytes.""" 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") other_b64 = base64.b64encode(other_bytes).decode("ascii")
doc_a = SearchResult( doc_a = SearchResult(