Add attachment-to-reference image bridge tool

code_puppy_core_plugins/attachment_references/__init__.py

@@ -0,0 +1 @@
+"""Bridge user-attached images into on-disk reference files."""

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)

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,
+        }
+    ]

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"

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"]