From adb737b6142cba8d53a638f804b335c35120fb42 Mon Sep 17 00:00:00 2001 From: sbyinin <2064038+sbyinin@users.noreply.github.com> Date: Thu, 18 Jun 2026 09:26:22 +0800 Subject: [PATCH] fix openai image reference edits --- nanobot/providers/image_generation.py | 129 ++++++++++++++++++----- tests/providers/test_image_generation.py | 111 +++++++++++++++++++ 2 files changed, 215 insertions(+), 25 deletions(-) diff --git a/nanobot/providers/image_generation.py b/nanobot/providers/image_generation.py index 47ed22ae..9f38e763 100644 --- a/nanobot/providers/image_generation.py +++ b/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: - logger.warning( - "DALL-E models do not support reference images; " - "ignoring {} reference image(s) for {}", - len(reference_images), - model, - ) + clean_model = self._strip_model_prefix(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( - f"{self.api_base}/images/generations", - headers=headers, - body=body, - ) + response = await self._http_post( + f"{self.api_base}/images/generations", + headers=generation_headers, + 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 - 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() - + images = await self._parse_images_response(payload) 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")) diff --git a/tests/providers/test_image_generation.py b/tests/providers/test_image_generation.py index 8ca597fa..03dd171a 100644 --- a/tests/providers/test_image_generation.py +++ b/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}]}))