diff --git a/plugins/image_gen/openai-codex/__init__.py b/plugins/image_gen/openai-codex/__init__.py index 6fde2d60bbbe3..24f04b7e88a5d 100644 --- a/plugins/image_gen/openai-codex/__init__.py +++ b/plugins/image_gen/openai-codex/__init__.py @@ -15,12 +15,19 @@ 4. :data:`DEFAULT_MODEL` — ``gpt-image-2-medium`` Output is saved as PNG under ``$HERMES_HOME/cache/images/``. + +Supports optional ``reference_images`` (http(s) URLs, ``data:image/...`` URLs, +or absolute file paths) that are forwarded to the Responses ``image_generation`` +tool as ``input_image`` parts — letting a caller preserve the identity of, say, +a brand mascot in the generated image. """ from __future__ import annotations +import base64 import json import logging +import os from typing import Any, Dict, List, Optional, Tuple from agent.image_gen_provider import ( @@ -70,6 +77,19 @@ "portrait": "1024x1536", } +# Reference-image (input_image) limits. gpt-image-2 processes image inputs at +# high fidelity, so a few references are enough to lock identity (e.g. a brand +# mascot); the caps bound request size and memory for file-path inputs. +MAX_REFERENCE_IMAGES = 4 +MAX_REFERENCE_IMAGE_BYTES = 10 * 1024 * 1024 # 10 MiB per image +_REFERENCE_IMAGE_MIME_BY_EXT = { + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".webp": "image/webp", + ".gif": "image/gif", +} + # Codex Responses surface used for the request. The chat model itself is only # the host that calls the ``image_generation`` tool; the actual image work is # done by ``API_MODEL``. @@ -101,8 +121,6 @@ def _load_image_gen_config() -> Dict[str, Any]: def _resolve_model() -> Tuple[str, Dict[str, Any]]: """Decide which tier to use and return ``(model_id, meta)``.""" - import os - env_override = os.environ.get("OPENAI_IMAGE_MODEL") if env_override and env_override in _MODELS: return env_override, _MODELS[env_override] @@ -143,8 +161,66 @@ def _read_codex_access_token() -> Optional[str]: return None -def _build_responses_payload(*, prompt: str, size: str, quality: str) -> Dict[str, Any]: - """Build the Codex Responses request body for an image_generation call.""" +def _to_input_image_part(ref: str) -> Dict[str, str]: + """Normalize one reference image into a Responses ``input_image`` part. + + Accepts an ``http(s)`` URL or a ``data:image/...`` URL (both passed through + verbatim), or an absolute file path (read as bytes and inlined as a base64 + data URL). Raises :class:`ValueError` on anything else so the caller can + surface a clean ``invalid_reference_image`` error before any network call. + """ + if not isinstance(ref, str) or not ref.strip(): + raise ValueError("reference image must be a non-empty string") + value = ref.strip() + lowered = value.lower() + + if lowered.startswith(("http://", "https://")): + return {"type": "input_image", "image_url": value} + if lowered.startswith("data:image/"): + if len(value) > MAX_REFERENCE_IMAGE_BYTES * 2: + raise ValueError("reference image data URL exceeds the size cap") + return {"type": "input_image", "image_url": value} + if lowered.startswith("data:"): + raise ValueError("only data:image/... data URLs are supported") + + # Otherwise treat it as a filesystem path. Require absolute (the agent's CWD + # is not meaningful to this provider) and resolve symlinks before access. + if not os.path.isabs(value): + raise ValueError(f"reference image path must be absolute: {value}") + real = os.path.realpath(value) + if not os.path.isfile(real): + raise ValueError(f"reference image not found: {value}") + size_bytes = os.path.getsize(real) + if size_bytes > MAX_REFERENCE_IMAGE_BYTES: + raise ValueError( + f"reference image exceeds {MAX_REFERENCE_IMAGE_BYTES // (1024 * 1024)}MB: {value}" + ) + ext = os.path.splitext(real)[1].lower() + mime = _REFERENCE_IMAGE_MIME_BY_EXT.get(ext) + if mime is None: + raise ValueError( + f"unsupported reference image type '{ext or ''}': {value}" + ) + with open(real, "rb") as handle: + encoded = base64.b64encode(handle.read()).decode("ascii") + return {"type": "input_image", "image_url": f"data:{mime};base64,{encoded}"} + + +def _build_responses_payload( + *, + prompt: str, + size: str, + quality: str, + image_parts: Optional[List[Dict[str, str]]] = None, +) -> Dict[str, Any]: + """Build the Codex Responses request body for an image_generation call. + + ``image_parts`` are pre-normalized ``input_image`` content parts (reference + images). When omitted the payload is identical to a text-only request. + """ + content: List[Dict[str, Any]] = [{"type": "input_text", "text": prompt}] + if image_parts: + content.extend(image_parts) return { "model": _CODEX_CHAT_MODEL, "store": False, @@ -152,7 +228,7 @@ def _build_responses_payload(*, prompt: str, size: str, quality: str) -> Dict[st "input": [{ "type": "message", "role": "user", - "content": [{"type": "input_text", "text": prompt}], + "content": content, }], "tools": [{ "type": "image_generation", @@ -161,7 +237,7 @@ def _build_responses_payload(*, prompt: str, size: str, quality: str) -> Dict[st "quality": quality, "output_format": "png", "background": "opaque", - "partial_images": 1, + "partial_images": 3, }], "tool_choice": { "type": "allowed_tools", @@ -172,27 +248,39 @@ def _build_responses_payload(*, prompt: str, size: str, quality: str) -> Dict[st } -def _extract_image_b64(value: Any) -> Optional[str]: - """Return the newest image b64 embedded in a Responses event payload.""" - found: Optional[str] = None +def _extract_image_b64(value: Any) -> Tuple[Optional[str], Optional[str]]: + """Return ``(final_b64, partial_b64)`` embedded in a Responses event payload. + + The final image lives on ``image_generation_call.result``; ``partial_image_b64`` + is a half-rendered streaming preview. They are returned separately so the + caller can *always* prefer the final result and only fall back to a partial + if no final ever arrives — saving a partial as the finished image produces a + soft, half-diffused PNG. + """ + final: Optional[str] = None + partial: Optional[str] = None if isinstance(value, dict): if value.get("type") == "image_generation_call": result = value.get("result") if isinstance(result, str) and result: - found = result - partial = value.get("partial_image_b64") - if isinstance(partial, str) and partial: - found = partial + final = result + p = value.get("partial_image_b64") + if isinstance(p, str) and p: + partial = p for child in value.values(): - nested = _extract_image_b64(child) - if nested: - found = nested + f, pa = _extract_image_b64(child) + if f: + final = f + if pa: + partial = pa elif isinstance(value, list): for child in value: - nested = _extract_image_b64(child) - if nested: - found = nested - return found + f, pa = _extract_image_b64(child) + if f: + final = f + if pa: + partial = pa + return final, partial def _iter_sse_json(response: Any): @@ -242,7 +330,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, + image_parts: Optional[List[Dict[str, 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,10 +348,13 @@ 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, image_parts=image_parts + ) timeout = httpx.Timeout(300.0, connect=30.0, read=300.0, write=30.0, pool=30.0) - image_b64: Optional[str] = None + final_b64: Optional[str] = None + partial_b64: Optional[str] = None with httpx.Client(timeout=timeout, headers=headers) as http: with http.stream("POST", f"{_CODEX_BASE_URL}/responses", json=payload) as response: try: @@ -268,11 +366,15 @@ def _collect_image_b64(token: str, *, prompt: str, size: str, quality: str) -> O f"Codex Responses API returned HTTP {exc.response.status_code}: {body}" ) from exc for event in _iter_sse_json(response): - found = _extract_image_b64(event) - if found: - image_b64 = found + f, p = _extract_image_b64(event) + if f: + final_b64 = f + if p: + partial_b64 = p - return image_b64 + # Always prefer the finished image; a partial is a last-resort fallback so a + # half-rendered preview frame is never saved as the final PNG. + return final_b64 or partial_b64 # --------------------------------------------------------------------------- @@ -368,6 +470,45 @@ def generate( tier_id, meta = _resolve_model() size = _SIZES.get(aspect, _SIZES["square"]) + # Optional reference images (e.g. a brand mascot) to guide/edit the + # generation. Validated up front so a bad input fails before any + # network call. + image_parts: Optional[List[Dict[str, str]]] = None + reference_images = kwargs.get("reference_images") or [] + if reference_images: + if not isinstance(reference_images, (list, tuple)): + return error_response( + error="reference_images must be a list of image URLs or absolute file paths", + error_type="invalid_reference_image", + provider="openai-codex", + model=tier_id, + prompt=prompt, + aspect_ratio=aspect, + ) + if len(reference_images) > MAX_REFERENCE_IMAGES: + return error_response( + error=( + f"At most {MAX_REFERENCE_IMAGES} reference images are supported " + f"(got {len(reference_images)})" + ), + error_type="invalid_reference_image", + provider="openai-codex", + model=tier_id, + prompt=prompt, + aspect_ratio=aspect, + ) + try: + image_parts = [_to_input_image_part(ref) for ref in reference_images] + except ValueError as exc: + return error_response( + error=f"Invalid reference image: {exc}", + error_type="invalid_reference_image", + provider="openai-codex", + model=tier_id, + prompt=prompt, + aspect_ratio=aspect, + ) + token = _read_codex_access_token() if not token: return error_response( @@ -382,13 +523,15 @@ def generate( aspect_ratio=aspect, ) + collect_kwargs: Dict[str, Any] = { + "prompt": prompt, + "size": size, + "quality": meta["quality"], + } + if image_parts: + collect_kwargs["image_parts"] = image_parts try: - b64 = _collect_image_b64( - token, - prompt=prompt, - size=size, - quality=meta["quality"], - ) + b64 = _collect_image_b64(token, **collect_kwargs) except Exception as exc: logger.debug("Codex image generation failed", exc_info=True) return error_response( diff --git a/tests/plugins/image_gen/test_openai_codex_provider.py b/tests/plugins/image_gen/test_openai_codex_provider.py index a3eb01f237463..ef2441aa7d3dd 100644 --- a/tests/plugins/image_gen/test_openai_codex_provider.py +++ b/tests/plugins/image_gen/test_openai_codex_provider.py @@ -234,3 +234,132 @@ def register_image_gen_provider(self, prov): codex_plugin.register(_Ctx()) assert len(registered) == 1 assert registered[0].name == "openai-codex" + + +# ── Reference images (input_image) ─────────────────────────────────────────── + + +class TestToInputImagePart: + def test_http_url_passthrough(self): + part = codex_plugin._to_input_image_part("https://example.com/m.png") + assert part == {"type": "input_image", "image_url": "https://example.com/m.png"} + + def test_data_image_url_passthrough(self): + ref = "data:image/png;base64," + _b64_png() + part = codex_plugin._to_input_image_part(ref) + assert part == {"type": "input_image", "image_url": ref} + + def test_absolute_path_becomes_data_url(self, tmp_path): + img = tmp_path / "mascot.png" + img.write_bytes(bytes.fromhex(_PNG_HEX)) + part = codex_plugin._to_input_image_part(str(img)) + assert part["type"] == "input_image" + assert part["image_url"].startswith("data:image/png;base64,") + + def test_jpg_extension_maps_to_jpeg_mime(self, tmp_path): + img = tmp_path / "mascot.jpg" + img.write_bytes(bytes.fromhex(_PNG_HEX)) # content bytes irrelevant for mime mapping + part = codex_plugin._to_input_image_part(str(img)) + assert part["image_url"].startswith("data:image/jpeg;base64,") + + def test_relative_path_rejected(self): + with pytest.raises(ValueError, match="absolute"): + codex_plugin._to_input_image_part("mascot.png") + + def test_missing_file_rejected(self, tmp_path): + with pytest.raises(ValueError, match="not found"): + codex_plugin._to_input_image_part(str(tmp_path / "nope.png")) + + def test_unsupported_extension_rejected(self, tmp_path): + bad = tmp_path / "mascot.bmp" + bad.write_bytes(bytes.fromhex(_PNG_HEX)) + with pytest.raises(ValueError, match="unsupported"): + codex_plugin._to_input_image_part(str(bad)) + + def test_oversize_file_rejected(self, tmp_path, monkeypatch): + monkeypatch.setattr(codex_plugin, "MAX_REFERENCE_IMAGE_BYTES", 4) + img = tmp_path / "big.png" + img.write_bytes(bytes.fromhex(_PNG_HEX)) # > 4 bytes + with pytest.raises(ValueError, match="exceeds"): + codex_plugin._to_input_image_part(str(img)) + + def test_non_image_data_url_rejected(self): + with pytest.raises(ValueError, match="data:image"): + codex_plugin._to_input_image_part("data:text/plain;base64,QQ==") + + def test_empty_rejected(self): + with pytest.raises(ValueError, match="non-empty"): + codex_plugin._to_input_image_part(" ") + + +class TestPayloadReferenceImages: + def test_text_only_when_no_parts(self): + payload = codex_plugin._build_responses_payload( + prompt="a cat", size="1024x1024", quality="medium" + ) + content = payload["input"][0]["content"] + assert content == [{"type": "input_text", "text": "a cat"}] + + def test_parts_appended_after_text(self): + parts = [{"type": "input_image", "image_url": "https://x/y.png"}] + payload = codex_plugin._build_responses_payload( + prompt="a cat", size="1024x1024", quality="medium", image_parts=parts + ) + content = payload["input"][0]["content"] + assert content[0] == {"type": "input_text", "text": "a cat"} + assert content[1:] == parts + + +class TestGenerateReferenceImages: + def test_generate_forwards_reference_images_into_payload(self, provider, monkeypatch): + monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token") + captured = {} + + def _collect(token, *, prompt, size, quality, image_parts=None): + captured["payload"] = codex_plugin._build_responses_payload( + prompt=prompt, size=size, quality=quality, image_parts=image_parts + ) + return _b64_png() + + monkeypatch.setattr(codex_plugin, "_collect_image_b64", _collect) + + result = provider.generate( + "hamster mascot on a bike", + aspect_ratio="square", + reference_images=["https://mojapteczka.pl/brand/mascot.jpeg"], + ) + + assert result["success"] is True + content = captured["payload"]["input"][0]["content"] + assert content[0]["type"] == "input_text" + assert { + "type": "input_image", + "image_url": "https://mojapteczka.pl/brand/mascot.jpeg", + } in content[1:] + + def test_generate_rejects_too_many_references(self, provider, monkeypatch): + monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token") + + def _boom(*a, **kw): + raise AssertionError("network call must not happen for invalid refs") + + monkeypatch.setattr(codex_plugin, "_collect_image_b64", _boom) + + refs = [f"https://x/{i}.png" for i in range(codex_plugin.MAX_REFERENCE_IMAGES + 1)] + result = provider.generate("a cat", reference_images=refs) + + assert result["success"] is False + assert result["error_type"] == "invalid_reference_image" + + def test_generate_rejects_bad_reference(self, provider, monkeypatch): + monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token") + + def _boom(*a, **kw): + raise AssertionError("network call must not happen for invalid refs") + + monkeypatch.setattr(codex_plugin, "_collect_image_b64", _boom) + + result = provider.generate("a cat", reference_images=["relative/mascot.png"]) + + assert result["success"] is False + assert result["error_type"] == "invalid_reference_image" diff --git a/tests/tools/test_image_generation.py b/tests/tools/test_image_generation.py index b24e6bc1fcc22..a3ac3bbc1b982 100644 --- a/tests/tools/test_image_generation.py +++ b/tests/tools/test_image_generation.py @@ -363,11 +363,15 @@ 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): - """The agent-facing schema must stay tight — model selection is a + def test_schema_exposes_prompt_aspect_ratio_and_reference_images(self, image_tool): + """The agent-facing schema stays tight — prompt, aspect_ratio, and the + optional reference_images content arg. Model/backend selection remains a user-level config choice, not an agent-level arg.""" props = image_tool.IMAGE_GENERATE_SCHEMA["parameters"]["properties"] - assert set(props.keys()) == {"prompt", "aspect_ratio"} + assert set(props.keys()) == {"prompt", "aspect_ratio", "reference_images"} + # Backend/model selection must NOT leak into the agent schema. + assert "model" not in props + assert "provider" not in props 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_artifacts.py b/tests/tools/test_image_generation_artifacts.py index 2a1ce11135361..b9ebbbf2caf5b 100644 --- a/tests/tools/test_image_generation_artifacts.py +++ b/tests/tools/test_image_generation_artifacts.py @@ -110,7 +110,9 @@ def fake_active_env(task_id): monkeypatch.setattr( image_generation_tool, "_dispatch_to_plugin_provider", - lambda prompt, aspect_ratio: json.dumps({"success": True, "image": str(image_path)}), + lambda prompt, aspect_ratio, reference_images=None: json.dumps( + {"success": True, "image": str(image_path)} + ), ) result = json.loads( diff --git a/tests/tools/test_image_generation_plugin_dispatch.py b/tests/tools/test_image_generation_plugin_dispatch.py index fa8ca9d959c92..b4b7502d4edb7 100644 --- a/tests/tools/test_image_generation_plugin_dispatch.py +++ b/tests/tools/test_image_generation_plugin_dispatch.py @@ -97,3 +97,76 @@ def fake_ensure_plugins_discovered(force=False): assert payload["success"] is True assert payload["provider"] == "codex" assert payload["aspect_ratio"] == "portrait" + + +class _CapturingProvider(ImageGenProvider): + def __init__(self): + self.seen = {} + + @property + def name(self) -> str: + return "codex" + + def generate(self, prompt, aspect_ratio="landscape", **kwargs): + self.seen = {"prompt": prompt, "aspect_ratio": aspect_ratio, **kwargs} + return { + "success": True, + "image": "/tmp/codex-test.png", + "model": "gpt-image-2-medium", + "prompt": prompt, + "aspect_ratio": aspect_ratio, + "provider": "codex", + } + + +class TestReferenceImageForwarding: + def test_reference_images_forwarded_to_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") + + capturing = _CapturingProvider() + monkeypatch.setattr(image_generation_tool, "_read_configured_image_provider", lambda: "codex") + monkeypatch.setattr(plugins_module, "_ensure_plugins_discovered", lambda *a, **k: None) + monkeypatch.setattr(registry_module, "get_provider", lambda name: capturing) + + refs = ["https://mojapteczka.pl/brand/mascot.jpeg"] + dispatched = image_generation_tool._dispatch_to_plugin_provider( + "draw mascot", "square", reference_images=refs + ) + payload = json.loads(dispatched) + + assert payload["success"] is True + assert capturing.seen.get("reference_images") == refs + + def test_no_reference_images_omits_kwarg(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") + + capturing = _CapturingProvider() + monkeypatch.setattr(image_generation_tool, "_read_configured_image_provider", lambda: "codex") + monkeypatch.setattr(plugins_module, "_ensure_plugins_discovered", lambda *a, **k: None) + monkeypatch.setattr(registry_module, "get_provider", lambda name: capturing) + + image_generation_tool._dispatch_to_plugin_provider("draw mascot", "square") + + assert "reference_images" not in capturing.seen + + +class TestSchema: + def test_schema_exposes_optional_reference_images(self): + from tools.image_generation_tool import IMAGE_GENERATE_SCHEMA + + props = IMAGE_GENERATE_SCHEMA["parameters"]["properties"] + assert "reference_images" in props + assert props["reference_images"]["type"] == "array" + assert props["reference_images"]["items"]["type"] == "string" + # Optional — must not be in required. + assert "reference_images" not in IMAGE_GENERATE_SCHEMA["parameters"]["required"] diff --git a/tools/image_generation_tool.py b/tools/image_generation_tool.py index d7eeb30d1750e..c4cc6cff49907 100644 --- a/tools/image_generation_tool.py +++ b/tools/image_generation_tool.py @@ -1024,6 +1024,17 @@ 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/input images to guide or edit the generation " + "(e.g. a brand mascot to preserve its identity in the result). Each " + "item is an http(s) URL, a data:image/...;base64 URL, or an absolute " + "file path. Honored by backends that support image input (currently " + "openai-codex); ignored by backends that do not." + ), + }, }, "required": ["prompt"], }, @@ -1069,12 +1080,16 @@ 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, reference_images=None): """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 in-tree FAL fallback in ``image_generate_tool``. + ``reference_images`` (optional list of URLs / data URLs / absolute paths) is + forwarded to the provider when present; backends that do not support image + input ignore it via their ``**kwargs`` signature. + Dispatch fires when ``image_gen.provider`` is explicitly set — including ``"fal"`` itself, which now resolves to the ``plugins/image_gen/fal/`` plugin (the plugin re-enters this module's @@ -1126,6 +1141,8 @@ def _dispatch_to_plugin_provider(prompt: str, aspect_ratio: str): kwargs = {"prompt": prompt, "aspect_ratio": aspect_ratio} if configured_model: kwargs["model"] = configured_model + if reference_images: + kwargs["reference_images"] = reference_images result = provider.generate(**kwargs) except Exception as exc: logger.warning( @@ -1153,11 +1170,14 @@ def _handle_image_generate(args, **kw): if not prompt: return tool_error("prompt is required for image generation") aspect_ratio = args.get("aspect_ratio", DEFAULT_ASPECT_RATIO) + reference_images = args.get("reference_images") task_id = kw.get("task_id") # 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) + dispatched = _dispatch_to_plugin_provider( + prompt, aspect_ratio, reference_images=reference_images + ) if dispatched is not None: return _postprocess_image_generate_result(dispatched, task_id=task_id)