diff --git a/CHANGELOG.md b/CHANGELOG.md index ced4010b..d58e4213 100644 --- a/CHANGELOG.md +++ b/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 `` 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). diff --git a/haiku_rag_slim/haiku/rag/skills/_tools.py b/haiku_rag_slim/haiku/rag/skills/_tools.py index c5116c8c..5c8d5279 100644 --- a/haiku_rag_slim/haiku/rag/skills/_tools.py +++ b/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] = [] - 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 diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index 5715227c..acdf6425 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 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 diff --git a/tests/test_picture_in_context.py b/tests/test_picture_in_context.py index c675225d..fefa1d9a 100644 --- a/tests/test_picture_in_context.py +++ b/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 diff --git a/tests/test_skill_tools.py b/tests/test_skill_tools.py index b6403e1c..cfe6797b 100644 --- a/tests/test_skill_tools.py +++ b/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(