fix openai image reference edits

nanobot/providers/image_generation.py

@@ -955,6 +955,56 @@ class OpenAIImageGenerationClient(ImageGenerationProvider):
             return model.split("/", 1)[1]
         return model
 
+    async def _parse_images_response(self, payload: dict[str, Any]) -> list[str]:
+        client = self._client
+        owns_client = client is None
+        if owns_client:
+            client = httpx.AsyncClient(timeout=self.timeout)
+        try:
+            return await _openai_images_from_payload(client, payload)
+        finally:
+            if owns_client:
+                await client.aclose()
+
+    async def _post_image_edit(
+        self,
+        *,
+        headers: dict[str, str],
+        body: dict[str, Any],
+        reference_images: list[str],
+    ) -> httpx.Response:
+        files: list[tuple[str, tuple[str, Any, str]]] = []
+        handles: list[Any] = []
+        try:
+            for path in reference_images:
+                p = Path(path)
+                raw = p.read_bytes()
+                mime = detect_image_mime(raw)
+                if mime is None:
+                    raise ImageGenerationError(f"unsupported reference image: {p}")
+                handle = p.open("rb")
+                handles.append(handle)
+                files.append(("image[]", (p.name, handle, mime)))
+
+            client = self._client
+            if client is not None:
+                return await client.post(
+                    f"{self.api_base}/images/edits",
+                    headers=headers,
+                    data=body,
+                    files=files,
+                )
+            async with httpx.AsyncClient(timeout=self.timeout) as c:
+                return await c.post(
+                    f"{self.api_base}/images/edits",
+                    headers=headers,
+                    data=body,
+                    files=files,
+                )
+        finally:
+            for handle in handles:
+                handle.close()
+
     async def generate(
         self,
         *,
@@ -967,21 +1017,18 @@ class OpenAIImageGenerationClient(ImageGenerationProvider):
         if not self.api_key:
             raise ImageGenerationError(self.missing_key_message)
 
-        if reference_images:
+        clean_model = self._strip_model_prefix(model)
-            logger.warning(
-                "DALL-E models do not support reference images; "
-                "ignoring {} reference image(s) for {}",
-                len(reference_images),
-                model,
-            )
 
-        headers = {
+        generation_headers = {
             "Authorization": f"Bearer {self.api_key}",
             "Content-Type": "application/json",
             **self.extra_headers,
         }
+        edit_headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            **self.extra_headers,
+        }
 
-        clean_model = self._strip_model_prefix(model)
         body: dict[str, Any] = {
             "model": clean_model,
             "prompt": prompt,
@@ -999,13 +1046,37 @@ class OpenAIImageGenerationClient(ImageGenerationProvider):
         # Drop null-valued params so extraBody can opt out of defaults like response_format.
         body = {key: value for key, value in body.items() if value is not None}
 
-        logger.info("OpenAI Images API request: POST {}/images/generations body={}", self.api_base, body)
+        refs = list(reference_images or [])
+        if refs:
+            if not _openai_is_gpt_image_model(clean_model):
+                raise ImageGenerationError(
+                    f"OpenAI model '{clean_model}' does not support reference images; "
+                    "use a GPT Image model"
+                )
+            edit_body = _openai_multipart_form_body(body)
+            logger.info(
+                "OpenAI Images API request: POST {}/images/edits body={} reference_images={}",
+                self.api_base,
+                edit_body,
+                len(refs),
+            )
+            response = await self._post_image_edit(
+                headers=edit_headers,
+                body=edit_body,
+                reference_images=refs,
+            )
+        else:
+            logger.info(
+                "OpenAI Images API request: POST {}/images/generations body={}",
+                self.api_base,
+                body,
+            )
 
-        response = await self._http_post(
+            response = await self._http_post(
-            f"{self.api_base}/images/generations",
+                f"{self.api_base}/images/generations",
-            headers=headers,
+                headers=generation_headers,
-            body=body,
+                body=body,
-        )
+            )
 
         try:
             response.raise_for_status()
@@ -1020,16 +1091,7 @@ class OpenAIImageGenerationClient(ImageGenerationProvider):
         logger.info("OpenAI Images API response ({}): {}", response.status_code,
                        {k: v for k, v in payload.items() if k != "data"})
 
-        client = self._client
+        images = await self._parse_images_response(payload)
-        owns_client = client is None
-        if owns_client:
-            client = httpx.AsyncClient(timeout=self.timeout)
-        try:
-            images = await _openai_images_from_payload(client, payload)
-        finally:
-            if owns_client:
-                await client.aclose()
-
         self._require_images(images, payload)
 
         return GeneratedImageResponse(images=images, content="", raw=payload)
@@ -1260,6 +1322,23 @@ def _openai_size(
     return "1024x1024"
 
 
+def _openai_multipart_form_body(body: dict[str, Any]) -> dict[str, str]:
+    form: dict[str, str] = {}
+    for key, value in body.items():
+        if value is None:
+            continue
+        if isinstance(value, bool):
+            form[key] = "true" if value else "false"
+        elif isinstance(value, str | int | float):
+            form[key] = str(value)
+        else:
+            logger.warning(
+                "OpenAI image edit parameter '{}' is not a scalar form field; ignoring it",
+                key,
+            )
+    return form
+
+
 def _openai_is_gpt_image_model(model: str) -> bool:
     normalized = model.lower()
     return normalized.startswith(("gpt-image", "chatgpt-image"))

tests/providers/test_image_generation.py

@@ -786,6 +786,117 @@ async def test_openai_gpt_image_uses_supported_orientation_sizes() -> None:
     assert fake.calls[1]["json"]["size"] == "1536x1024"
 
 
+@pytest.mark.asyncio
+async def test_openai_reference_images_use_edits_endpoint(tmp_path: Path) -> None:
+    ref = tmp_path / "ref.png"
+    ref.write_bytes(PNG_BYTES)
+    fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
+    client = OpenAIImageGenerationClient(
+        api_key="sk-openai-test",
+        api_base="https://api.openai.com/v1",
+        extra_headers={"X-Test": "1"},
+        client=fake,  # type: ignore[arg-type]
+    )
+
+    response = await client.generate(
+        prompt="make a warmer version",
+        model="gpt-image-1",
+        reference_images=[str(ref)],
+        aspect_ratio="16:9",
+    )
+
+    assert response.images == [PNG_DATA_URL]
+    call = fake.calls[0]
+    assert call["url"] == "https://api.openai.com/v1/images/edits"
+    assert call["headers"]["Authorization"] == "Bearer sk-openai-test"
+    assert call["headers"]["X-Test"] == "1"
+    assert "Content-Type" not in call["headers"]
+    assert "json" not in call
+    assert call["data"]["model"] == "gpt-image-1"
+    assert call["data"]["prompt"] == "make a warmer version"
+    assert call["data"]["size"] == "1536x1024"
+    assert len(call["files"]) == 1
+    assert call["files"][0][0] == "image[]"
+    assert call["files"][0][1][0] == "ref.png"
+    assert call["files"][0][1][2] == "image/png"
+    assert call["files"][0][1][1].closed is True
+
+
+@pytest.mark.asyncio
+async def test_openai_reference_images_send_multiple_multipart_files(
+    tmp_path: Path,
+) -> None:
+    first = tmp_path / "first.png"
+    second = tmp_path / "second.png"
+    first.write_bytes(PNG_BYTES)
+    second.write_bytes(PNG_BYTES)
+    fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
+    client = OpenAIImageGenerationClient(
+        api_key="sk-openai-test",
+        extra_body={
+            "quality": "high",
+            "seed": 0,
+            "safety_checker": False,
+            "metadata": {"ignored": True},
+            "background": None,
+        },
+        client=fake,  # type: ignore[arg-type]
+    )
+
+    await client.generate(
+        prompt="combine these references",
+        model="openai/gpt-image-1",
+        reference_images=[str(first), str(second)],
+    )
+
+    call = fake.calls[0]
+    assert call["url"] == "https://api.openai.com/v1/images/edits"
+    assert call["data"]["model"] == "gpt-image-1"
+    assert call["data"]["prompt"] == "combine these references"
+    assert call["data"]["quality"] == "high"
+    assert call["data"]["seed"] == "0"
+    assert call["data"]["safety_checker"] == "false"
+    assert "metadata" not in call["data"]
+    assert "background" not in call["data"]
+    assert [item[0] for item in call["files"]] == ["image[]", "image[]"]
+    assert [item[1][0] for item in call["files"]] == ["first.png", "second.png"]
+    assert all(item[1][1].closed for item in call["files"])
+
+
+@pytest.mark.asyncio
+async def test_openai_gpt_image_without_reference_images_uses_generations_json() -> None:
+    fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
+    client = OpenAIImageGenerationClient(
+        api_key="sk-openai-test",
+        client=fake,  # type: ignore[arg-type]
+    )
+
+    await client.generate(prompt="draw", model="gpt-image-1", aspect_ratio="16:9")
+
+    call = fake.calls[0]
+    assert call["url"] == "https://api.openai.com/v1/images/generations"
+    assert call["headers"]["Content-Type"] == "application/json"
+    assert call["json"]["model"] == "gpt-image-1"
+    assert call["json"]["prompt"] == "draw"
+    assert call["json"]["size"] == "1536x1024"
+    assert "data" not in call
+    assert "files" not in call
+
+
+@pytest.mark.asyncio
+async def test_openai_dalle_reference_images_raise_clear_error(tmp_path: Path) -> None:
+    ref = tmp_path / "ref.png"
+    ref.write_bytes(PNG_BYTES)
+    client = OpenAIImageGenerationClient(api_key="sk-openai-test")
+
+    with pytest.raises(ImageGenerationError, match="does not support reference images"):
+        await client.generate(
+            prompt="edit this",
+            model="dall-e-3",
+            reference_images=[str(ref)],
+        )
+
+
 @pytest.mark.asyncio
 async def test_openai_default_size_when_no_aspect_ratio() -> None:
     fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))