Merge pull request #371 from ggozad/fix/image-results
Verify picture bytes before attaching to multimodal tool returns
This commit is contained in:
Yiorgis Gozadinos
GitHub
commit
bdab201c42
5 changed files with 161 additions and 42 deletions
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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] = []
|
||||
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)
|
||||
|
||||
binary_parts = build_binary_parts_from_results(results)
|
||||
if binary_parts:
|
||||
return ToolReturn(return_value=formatted, content=binary_parts)
|
||||
return formatted
|
||||
|
|
|
|||
|
|
@ -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] = []
|
||||
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)
|
||||
|
||||
binary_parts = build_binary_parts_from_results(results_list)
|
||||
if binary_parts:
|
||||
return ToolReturn(return_value=text, content=binary_parts)
|
||||
return text
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Reference in a new issue