From 0f90882d693367552a706d42e634b1e47d54de28 Mon Sep 17 00:00:00 2001 From: The Butcher Date: Fri, 21 Aug 2026 13:37:35 -0400 Subject: [PATCH] Add attachment-to-reference image bridge tool --- .../attachment_references/__init__.py | 1 + .../register_callbacks.py | 22 +++ .../attachment_references/tool.py | 137 +++++++++++++++++ pyproject.toml | 1 + tests/test_attachment_references.py | 143 ++++++++++++++++++ 5 files changed, 304 insertions(+) create mode 100644 code_puppy_core_plugins/attachment_references/__init__.py create mode 100644 code_puppy_core_plugins/attachment_references/register_callbacks.py create mode 100644 code_puppy_core_plugins/attachment_references/tool.py create mode 100644 tests/test_attachment_references.py diff --git a/code_puppy_core_plugins/attachment_references/__init__.py b/code_puppy_core_plugins/attachment_references/__init__.py new file mode 100644 index 0000000..c6178ed --- /dev/null +++ b/code_puppy_core_plugins/attachment_references/__init__.py @@ -0,0 +1 @@ +"""Bridge user-attached images into on-disk reference files.""" diff --git a/code_puppy_core_plugins/attachment_references/register_callbacks.py b/code_puppy_core_plugins/attachment_references/register_callbacks.py new file mode 100644 index 0000000..a202bf9 --- /dev/null +++ b/code_puppy_core_plugins/attachment_references/register_callbacks.py @@ -0,0 +1,22 @@ +"""Register the attachment-to-reference bridge tool with Code Puppy.""" + +from __future__ import annotations + +from typing import Any + +from code_puppy.callbacks import register_callback + +from .tool import TOOL_NAME, register_tools_callback + + +def _register_tools() -> list[dict[str, Any]]: + return register_tools_callback() + + +def _advertise_tool(agent_name: str | None = None) -> list[str]: + del agent_name + return [TOOL_NAME] + + +register_callback("register_tools", _register_tools) +register_callback("register_agent_tools", _advertise_tool) diff --git a/code_puppy_core_plugins/attachment_references/tool.py b/code_puppy_core_plugins/attachment_references/tool.py new file mode 100644 index 0000000..a437843 --- /dev/null +++ b/code_puppy_core_plugins/attachment_references/tool.py @@ -0,0 +1,137 @@ +"""Materialize the latest user-attached images into reference files. + +Pasted or dragged images arrive as ``BinaryContent`` stapled onto the user +message; they never touch the filesystem. Tools such as ``codex_imagegen`` +condition on ``reference_images`` *paths*, so there is otherwise no way to point +image generation at something the user just pasted. This tool bridges that gap +by writing the latest turn's image attachments to disk and returning their +paths -- nothing more (SRP/YAGNI). +""" + +from __future__ import annotations + +import asyncio +import mimetypes +import os +import uuid +from pathlib import Path +from typing import Any, List, Sequence + +from pydantic_ai import BinaryContent, RunContext +from pydantic_ai.messages import UserPromptPart + +from code_puppy import config + +TOOL_NAME = "save_attachments_as_references" + + +def _is_image(item: Any) -> bool: + return isinstance(item, BinaryContent) and str( + getattr(item, "media_type", "") or "" + ).startswith("image/") + + +def _extract_latest_user_images(messages: Sequence[Any] | None) -> List[BinaryContent]: + """Return image attachments from the most recent user message only. + + Walks history backwards to the last request carrying a ``UserPromptPart`` + (an actual user turn, not a tool return) and collects its image parts in + order. Returns an empty list when the latest user turn has no images. + """ + if not messages: + return [] + for message in reversed(list(messages)): + parts = getattr(message, "parts", None) + if not parts: + continue + user_parts = [part for part in parts if isinstance(part, UserPromptPart)] + if not user_parts: + continue + images: List[BinaryContent] = [] + for part in user_parts: + content = part.content + if isinstance(content, str): + continue + for item in content: + if _is_image(item): + images.append(item) + return images + return [] + + +def _extension_for(media_type: str | None) -> str: + guessed = mimetypes.guess_extension(media_type or "") if media_type else None + if guessed in {".jpe", ".jpeg"}: + return ".jpg" + return guessed or ".png" + + +def _save_reference_images(images: Sequence[BinaryContent]) -> List[Path]: + """Write each image to the Code Puppy data dir and return the paths.""" + if not images: + return [] + output_dir = Path(config.DATA_DIR) / "attachment_references" + output_dir.mkdir(parents=True, exist_ok=True, mode=0o700) + saved: List[Path] = [] + for image in images: + extension = _extension_for(getattr(image, "media_type", None)) + output_path = output_dir / f"attachment-{uuid.uuid4().hex}{extension}" + temporary_path = output_dir / f"{output_path.name}.tmp" + try: + temporary_path.write_bytes(image.data) + os.chmod(temporary_path, 0o600) + temporary_path.replace(output_path) + except OSError: + temporary_path.unlink(missing_ok=True) + raise + saved.append(output_path) + return saved + + +def register_save_attachments_as_references(agent: Any) -> None: + """Register the attachment-to-reference bridge on a pydantic-ai agent.""" + + @agent.tool + async def save_attachments_as_references(context: RunContext) -> dict[str, Any]: + """Save images the user attached in their latest message to disk. + + Use this to turn a pasted or dragged image into a real file path so + other tools can act on it. In particular, pass the returned paths to + ``codex_imagegen`` as ``reference_images`` to generate a new image that + preserves the subject or style of what the user just shared. + + Only images from the most recent user message are saved. Returns the + saved file paths; the paths are stable on disk under the Code Puppy + data directory. + """ + images = _extract_latest_user_images(getattr(context, "messages", None)) + if not images: + return { + "success": True, + "paths": [], + "count": 0, + "message": ( + "No image attachments were found in the latest user message." + ), + } + try: + paths = await asyncio.to_thread(_save_reference_images, images) + except OSError as exc: + return { + "success": False, + "error": f"Could not save reference images: {exc}", + } + return { + "success": True, + "paths": [str(path) for path in paths], + "count": len(paths), + } + + +def register_tools_callback() -> list[dict[str, Any]]: + return [ + { + "name": TOOL_NAME, + "register_func": register_save_attachments_as_references, + } + ] diff --git a/pyproject.toml b/pyproject.toml index 3376980..9dfa3af 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,7 @@ Homepage = "https://github.com/mpfaffenberger/code_puppy_core_plugins" acp = "code_puppy_core_plugins.acp.register_callbacks" agent_skills = "code_puppy_core_plugins.agent_skills.register_callbacks" agent_creator_skill = "code_puppy_core_plugins.agent_creator_skill.register_callbacks" +attachment_references = "code_puppy_core_plugins.attachment_references.register_callbacks" aws_bedrock = "code_puppy_core_plugins.aws_bedrock.register_callbacks" azure_foundry = "code_puppy_core_plugins.azure_foundry.register_callbacks" btw = "code_puppy_core_plugins.btw.register_callbacks" diff --git a/tests/test_attachment_references.py b/tests/test_attachment_references.py new file mode 100644 index 0000000..0873697 --- /dev/null +++ b/tests/test_attachment_references.py @@ -0,0 +1,143 @@ +"""Tests for the attachment-to-reference bridge tool.""" + +import asyncio +from types import SimpleNamespace + +from pydantic_ai import BinaryContent +from pydantic_ai.messages import ( + ModelRequest, + ModelResponse, + TextPart, + ToolReturnPart, + UserPromptPart, +) + +from code_puppy import config +from code_puppy_core_plugins.attachment_references import register_callbacks +from code_puppy_core_plugins.attachment_references import tool as attachment_tool + + +def _png(marker: bytes) -> BinaryContent: + return BinaryContent(data=b"\x89PNG" + marker, media_type="image/png") + + +def test_extract_returns_only_latest_user_turn_images(): + first = _png(b"first") + second = _png(b"second") + history = [ + ModelRequest(parts=[UserPromptPart(content=["earlier", first])]), + ModelResponse(parts=[TextPart(content="ok")]), + ModelRequest(parts=[UserPromptPart(content=["latest", second])]), + ] + + images = attachment_tool._extract_latest_user_images(history) + + assert images == [second] + + +def test_extract_skips_tool_returns_and_non_image_binaries(): + image = _png(b"keep") + document = BinaryContent(data=b"%PDF-1.4", media_type="application/pdf") + history = [ + ModelRequest(parts=[UserPromptPart(content=["look", image, document])]), + ModelResponse(parts=[TextPart(content="thinking")]), + ModelRequest( + parts=[ + ToolReturnPart( + tool_name="grep", + content="no images here", + tool_call_id="call-1", + ) + ] + ), + ] + + images = attachment_tool._extract_latest_user_images(history) + + assert images == [image] + + +def test_extract_handles_empty_and_text_only_history(): + assert attachment_tool._extract_latest_user_images(None) == [] + assert attachment_tool._extract_latest_user_images([]) == [] + text_only = [ModelRequest(parts=[UserPromptPart(content="just words")])] + assert attachment_tool._extract_latest_user_images(text_only) == [] + + +def test_extension_for_maps_common_image_types(): + assert attachment_tool._extension_for("image/png") == ".png" + assert attachment_tool._extension_for("image/jpeg") == ".jpg" + assert attachment_tool._extension_for(None) == ".png" + assert attachment_tool._extension_for("image/unknownxyz") == ".png" + + +def test_save_reference_images_writes_files(tmp_path, monkeypatch): + monkeypatch.setattr(config, "DATA_DIR", str(tmp_path)) + images = [_png(b"a"), BinaryContent(data=b"jpegbytes", media_type="image/jpeg")] + + paths = attachment_tool._save_reference_images(images) + + assert len(paths) == 2 + assert paths[0].read_bytes() == b"\x89PNGa" + assert paths[0].suffix == ".png" + assert paths[1].suffix == ".jpg" + assert all(path.parent == tmp_path / "attachment_references" for path in paths) + assert not any(path.name.endswith(".tmp") for path in paths) + + +def _register_tool(): + registered = {} + + class FakeAgent: + def tool(self, function): + registered[function.__name__] = function + return function + + attachment_tool.register_save_attachments_as_references(FakeAgent()) + return registered["save_attachments_as_references"] + + +def test_tool_returns_paths_for_latest_attachment(tmp_path, monkeypatch): + monkeypatch.setattr(config, "DATA_DIR", str(tmp_path)) + save = _register_tool() + context = SimpleNamespace( + messages=[ModelRequest(parts=[UserPromptPart(content=["hi", _png(b"z")])])] + ) + + result = asyncio.run(save(context)) + + assert result["success"] is True + assert result["count"] == 1 + assert len(result["paths"]) == 1 + assert result["paths"][0].endswith(".png") + + +def test_tool_reports_when_no_images_present(tmp_path, monkeypatch): + monkeypatch.setattr(config, "DATA_DIR", str(tmp_path)) + save = _register_tool() + context = SimpleNamespace( + messages=[ModelRequest(parts=[UserPromptPart(content="text only")])] + ) + + result = asyncio.run(save(context)) + + assert result == { + "success": True, + "paths": [], + "count": 0, + "message": "No image attachments were found in the latest user message.", + } + + +def test_registration_contract(): + tools = register_callbacks._register_tools() + assert tools == [ + { + "name": "save_attachments_as_references", + "register_func": attachment_tool.register_save_attachments_as_references, + } + ] + assert register_callbacks._advertise_tool("code-puppy") == [ + "save_attachments_as_references" + ] + assert register_callbacks._advertise_tool() == ["save_attachments_as_references"]