diff --git a/plugins/image_gen/openai-codex/__init__.py b/plugins/image_gen/openai-codex/__init__.py index 339e390be1f5..280d01c1e678 100644 --- a/plugins/image_gen/openai-codex/__init__.py +++ b/plugins/image_gen/openai-codex/__init__.py @@ -19,9 +19,12 @@ from __future__ import annotations +import base64 import json import logging -from typing import Any, Dict, List, Optional, Tuple +import mimetypes +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Tuple from agent.image_gen_provider import ( DEFAULT_ASPECT_RATIO, @@ -80,6 +83,16 @@ "using the image_generation tool when provided." ) +_MAX_REFERENCE_IMAGES = 8 +_MAX_REFERENCE_IMAGE_BYTES = 20 * 1024 * 1024 +_MAX_DATA_IMAGE_URL_LENGTH = 30 * 1024 * 1024 +_SUPPORTED_IMAGE_MIMES = { + "image/png", + "image/jpeg", + "image/webp", + "image/gif", +} + # --------------------------------------------------------------------------- # Config + auth helpers @@ -143,8 +156,97 @@ def _read_codex_access_token() -> Optional[str]: return None -def _build_responses_payload(*, prompt: str, size: str, quality: str) -> Dict[str, Any]: +def _detect_image_mime(raw: bytes) -> Optional[str]: + """Detect common web image formats from magic bytes.""" + if raw.startswith(b"\x89PNG\r\n\x1a\n"): + return "image/png" + if raw.startswith(b"\xff\xd8\xff"): + return "image/jpeg" + if raw.startswith((b"GIF87a", b"GIF89a")): + return "image/gif" + if len(raw) >= 12 and raw[:4] == b"RIFF" and raw[8:12] == b"WEBP": + return "image/webp" + return None + + +def _reference_image_to_url(ref: str) -> str: + """Return a Responses-compatible image URL for a local path, URL, or data URL.""" + value = str(ref).strip() + if not value: + raise ValueError("reference image path/URL must be non-empty") + if value.startswith(("http://", "https://")): + return value + if value.startswith("data:image/"): + if len(value) > _MAX_DATA_IMAGE_URL_LENGTH: + raise ValueError("reference image data URL is too large") + header, sep, data = value.partition(",") + if not sep or ";base64" not in header: + raise ValueError("reference image data URL must be base64-encoded") + mime = header.removeprefix("data:").split(";", 1)[0] + if mime not in _SUPPORTED_IMAGE_MIMES: + raise ValueError(f"unsupported reference image MIME type: {mime}") + try: + base64.b64decode(data, validate=True) + except Exception as exc: + raise ValueError("reference image data URL contains invalid base64") from exc + return value + + path = Path(value).expanduser() + if not path.is_file(): + raise ValueError(f"reference image not found: {path}") + try: + if path.stat().st_size > _MAX_REFERENCE_IMAGE_BYTES: + raise ValueError("local reference image is too large") + data = path.read_bytes() + except OSError as exc: + raise ValueError(f"could not read reference image: {path}") from exc + mime = _detect_image_mime(data) or mimetypes.guess_type(str(path))[0] + if mime not in _SUPPORTED_IMAGE_MIMES: + raise ValueError(f"unsupported or invalid reference image: {path}") + return f"data:{mime};base64,{base64.b64encode(data).decode('ascii')}" + + +def _normalize_reference_images(reference_images: Optional[Iterable[str]]) -> List[str]: + """Validate reference count and convert refs into Responses image URLs.""" + if not reference_images: + return [] + if isinstance(reference_images, str): + reference_images = [reference_images] + refs = [str(ref) for ref in reference_images if str(ref).strip()] + if len(refs) > _MAX_REFERENCE_IMAGES: + raise ValueError(f"at most {_MAX_REFERENCE_IMAGES} reference images are supported") + return [_reference_image_to_url(ref) for ref in refs] + + +def _build_input_content(prompt: str, reference_images: Optional[Iterable[str]]) -> List[Dict[str, str]]: + """Build Responses message content with optional image references.""" + content: List[Dict[str, str]] = [{"type": "input_text", "text": prompt}] + for image_url in _normalize_reference_images(reference_images): + content.append({"type": "input_image", "image_url": image_url}) + return content + + +def _build_responses_payload( + *, + prompt: str, + size: str, + quality: str, + reference_images: Optional[Iterable[str]] = None, +) -> Dict[str, Any]: """Build the Codex Responses request body for an image_generation call.""" + input_content = _build_input_content(prompt, reference_images) + tool_options = { + "type": "image_generation", + "model": API_MODEL, + "size": size, + "quality": quality, + "output_format": "png", + "background": "opaque", + "partial_images": 1, + } + if len(input_content) > 1: + tool_options["action"] = "edit" + return { "model": _CODEX_CHAT_MODEL, "store": False, @@ -152,17 +254,9 @@ def _build_responses_payload(*, prompt: str, size: str, quality: str) -> Dict[st "input": [{ "type": "message", "role": "user", - "content": [{"type": "input_text", "text": prompt}], - }], - "tools": [{ - "type": "image_generation", - "model": API_MODEL, - "size": size, - "quality": quality, - "output_format": "png", - "background": "opaque", - "partial_images": 1, + "content": input_content, }], + "tools": [tool_options], "tool_choice": { "type": "allowed_tools", "mode": "required", @@ -242,7 +336,14 @@ def flush(): yield payload -def _collect_image_b64(token: str, *, prompt: str, size: str, quality: str) -> Optional[str]: +def _collect_image_b64( + token: str, + *, + prompt: str, + size: str, + quality: str, + reference_images: Optional[Iterable[str]] = None, +) -> Optional[str]: """Stream a Codex Responses image_generation call and return the b64 image.""" import httpx from agent.auxiliary_client import _codex_cloudflare_headers @@ -253,7 +354,12 @@ def _collect_image_b64(token: str, *, prompt: str, size: str, quality: str) -> O "Authorization": f"Bearer {token}", "Content-Type": "application/json", }) - payload = _build_responses_payload(prompt=prompt, size=size, quality=quality) + payload = _build_responses_payload( + prompt=prompt, + size=size, + quality=quality, + reference_images=reference_images, + ) timeout = httpx.Timeout(300.0, connect=30.0, read=300.0, write=30.0, pool=30.0) image_b64: Optional[str] = None @@ -382,12 +488,28 @@ def generate( aspect_ratio=aspect, ) + reference_images = kwargs.get("reference_images") + if isinstance(reference_images, str): + reference_images = [reference_images] + try: + normalized_references = _normalize_reference_images(reference_images) + except (TypeError, ValueError) as exc: + return error_response( + error=f"Invalid reference_images: {exc}", + error_type="invalid_argument", + provider="openai-codex", + model=tier_id, + prompt=prompt, + aspect_ratio=aspect, + ) + try: b64 = _collect_image_b64( token, prompt=prompt, size=size, quality=meta["quality"], + reference_images=normalized_references, ) except Exception as exc: logger.debug("Codex image generation failed", exc_info=True) diff --git a/tests/plugins/image_gen/test_openai_codex_provider.py b/tests/plugins/image_gen/test_openai_codex_provider.py index 2940b300b361..12db088a8438 100644 --- a/tests/plugins/image_gen/test_openai_codex_provider.py +++ b/tests/plugins/image_gen/test_openai_codex_provider.py @@ -8,6 +8,7 @@ from __future__ import annotations +import base64 import importlib from pathlib import Path @@ -129,11 +130,12 @@ def test_codex_stream_request_shape(self, provider, monkeypatch): captured = {} - def _collect(token, *, prompt, size, quality): + def _collect(token, *, prompt, size, quality, reference_images=None): captured.update(codex_plugin._build_responses_payload( prompt=prompt, size=size, quality=quality, + reference_images=reference_images, )) return _b64_png() @@ -159,6 +161,87 @@ def _collect(token, *, prompt, size, quality): assert tool["output_format"] == "png" assert tool["background"] == "opaque" assert tool["partial_images"] == 1 + assert "action" not in tool + + def test_reference_images_are_sent_as_edit_inputs(self, provider, monkeypatch, tmp_path): + monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token") + ref_path = tmp_path / "ref.png" + ref_path.write_bytes(bytes.fromhex(_PNG_HEX)) + + captured = {} + + def _collect(token, *, prompt, size, quality, reference_images=None): + captured.update(codex_plugin._build_responses_payload( + prompt=prompt, + size=size, + quality=quality, + reference_images=reference_images, + )) + return _b64_png() + + monkeypatch.setattr(codex_plugin, "_collect_image_b64", _collect) + + result = provider.generate("same person", reference_images=[str(ref_path)]) + assert result["success"] is True + + content = captured["input"][0]["content"] + assert content[0] == {"type": "input_text", "text": "same person"} + assert content[1]["type"] == "input_image" + assert content[1]["image_url"].startswith("data:image/png;base64,") + assert base64.b64decode(content[1]["image_url"].split(",", 1)[1]) == bytes.fromhex(_PNG_HEX) + assert captured["tools"][0]["action"] == "edit" + + def test_reference_image_url_and_data_url_passthrough(self): + refs = [ + "https://example.com/ref.png", + "data:image/png;base64,YWJj", + ] + payload = codex_plugin._build_responses_payload( + prompt="same object", + size="1024x1024", + quality="medium", + reference_images=refs, + ) + + content = payload["input"][0]["content"] + assert [item["image_url"] for item in content[1:]] == refs + assert payload["tools"][0]["action"] == "edit" + + def test_reference_images_are_limited(self): + refs = ["https://example.com/ref.png"] * (codex_plugin._MAX_REFERENCE_IMAGES + 1) + with pytest.raises(ValueError, match="at most"): + codex_plugin._build_responses_payload( + prompt="same object", + size="1024x1024", + quality="medium", + reference_images=refs, + ) + + def test_invalid_data_url_rejected(self): + with pytest.raises(ValueError, match="invalid base64"): + codex_plugin._build_responses_payload( + prompt="same object", + size="1024x1024", + quality="medium", + reference_images=["data:image/png;base64,not valid"], + ) + + def test_invalid_local_reference_returns_invalid_argument( + self, provider, monkeypatch, tmp_path + ): + monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token") + + def _collect(*args, **kwargs): # pragma: no cover - should not be called + raise AssertionError("invalid refs should fail before API request") + + monkeypatch.setattr(codex_plugin, "_collect_image_b64", _collect) + bad_ref = tmp_path / "not-image.txt" + bad_ref.write_text("not an image") + + result = provider.generate("same person", reference_images=[str(bad_ref)]) + assert result["success"] is False + assert result["error_type"] == "invalid_argument" + assert "Invalid reference_images" in result["error"] def test_partial_image_event_used_when_done_missing(self): """If output_item.done is missing, partial_image_b64 is accepted.""" diff --git a/tests/tools/test_image_generation.py b/tests/tools/test_image_generation.py index b24e6bc1fcc2..f44601515187 100644 --- a/tests/tools/test_image_generation.py +++ b/tests/tools/test_image_generation.py @@ -363,11 +363,11 @@ def test_empty_aspect_defaults_to_landscape(self, image_tool): class TestRegistryIntegration: - def test_schema_exposes_only_prompt_and_aspect_ratio_to_agent(self, image_tool): + def test_schema_exposes_prompt_aspect_ratio_and_reference_images_to_agent(self, image_tool): """The agent-facing schema must stay tight — model selection is a - user-level config choice, not an agent-level arg.""" + user-level config choice, while references are request inputs.""" props = image_tool.IMAGE_GENERATE_SCHEMA["parameters"]["properties"] - assert set(props.keys()) == {"prompt", "aspect_ratio"} + assert set(props.keys()) == {"prompt", "aspect_ratio", "reference_images"} def test_aspect_ratio_enum_is_three_values(self, image_tool): enum = image_tool.IMAGE_GENERATE_SCHEMA["parameters"]["properties"]["aspect_ratio"]["enum"] diff --git a/tests/tools/test_image_generation_plugin_dispatch.py b/tests/tools/test_image_generation_plugin_dispatch.py index fa8ca9d959c9..5bb1a4fa3368 100644 --- a/tests/tools/test_image_generation_plugin_dispatch.py +++ b/tests/tools/test_image_generation_plugin_dispatch.py @@ -15,11 +15,14 @@ def _reset_registry(): class _FakeCodexProvider(ImageGenProvider): + last_kwargs = None + @property def name(self) -> str: return "codex" def generate(self, prompt, aspect_ratio="landscape", **kwargs): + type(self).last_kwargs = kwargs return { "success": True, "image": "/tmp/codex-test.png", @@ -52,6 +55,29 @@ def test_dispatch_routes_to_codex_provider(self, monkeypatch, tmp_path): assert payload["image"] == "/tmp/codex-test.png" assert payload["aspect_ratio"] == "square" + def test_dispatch_passes_reference_images_to_plugin_provider(self, monkeypatch, tmp_path): + from tools import image_generation_tool + from agent import image_gen_registry as registry_module + from hermes_cli import plugins as plugins_module + + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + (tmp_path / "config.yaml").write_text("image_gen:\n provider: codex\n") + _FakeCodexProvider.last_kwargs = None + + monkeypatch.setattr(image_generation_tool, "_read_configured_image_provider", lambda: "codex") + monkeypatch.setattr(plugins_module, "_ensure_plugins_discovered", lambda: None) + monkeypatch.setattr(registry_module, "get_provider", lambda name: _FakeCodexProvider() if name == "codex" else None) + + dispatched = image_generation_tool._dispatch_to_plugin_provider( + "draw cat", + "square", + reference_images=["/tmp/ref.png"], + ) + payload = json.loads(dispatched) + + assert payload["success"] is True + assert _FakeCodexProvider.last_kwargs == {"reference_images": ["/tmp/ref.png"]} + def test_dispatch_reports_missing_registered_provider(self, monkeypatch, tmp_path): from tools import image_generation_tool from hermes_cli import plugins as plugins_module diff --git a/tools/image_generation_tool.py b/tools/image_generation_tool.py index d3263eae8ad2..a82b1cd20c4d 100644 --- a/tools/image_generation_tool.py +++ b/tools/image_generation_tool.py @@ -906,6 +906,15 @@ def check_image_generation_requirements() -> bool: "description": "The aspect ratio of the generated image. 'landscape' is 16:9 wide, 'portrait' is 16:9 tall, 'square' is 1:1.", "default": DEFAULT_ASPECT_RATIO, }, + "reference_images": { + "type": "array", + "items": {"type": "string"}, + "description": ( + "Optional reference image paths, URLs, or data URLs for image-generation backends " + "that support editing or identity/object preservation. Currently honored by " + "the openai-codex provider; other backends may reject or ignore it." + ), + }, }, "required": ["prompt"], }, @@ -951,7 +960,7 @@ def _read_configured_image_provider(): return None -def _dispatch_to_plugin_provider(prompt: str, aspect_ratio: str): +def _dispatch_to_plugin_provider(prompt: str, aspect_ratio: str, **provider_kwargs): """Route the call to a plugin-registered provider when one is selected. Returns a JSON string on dispatch, or ``None`` to fall through to the @@ -1006,6 +1015,7 @@ def _dispatch_to_plugin_provider(prompt: str, aspect_ratio: str): try: kwargs = {"prompt": prompt, "aspect_ratio": aspect_ratio} + kwargs.update({k: v for k, v in provider_kwargs.items() if v is not None}) if configured_model: kwargs["model"] = configured_model result = provider.generate(**kwargs) @@ -1038,7 +1048,15 @@ def _handle_image_generate(args, **kw): # Route to a plugin-registered provider if one is active (and it's # not the in-tree FAL path). - dispatched = _dispatch_to_plugin_provider(prompt, aspect_ratio) + reference_images = args.get("reference_images") + if isinstance(reference_images, str): + reference_images = [reference_images] + + dispatch_kwargs = {} + if reference_images: + dispatch_kwargs["reference_images"] = reference_images + + dispatched = _dispatch_to_plugin_provider(prompt, aspect_ratio, **dispatch_kwargs) if dispatched is not None: return dispatched