From 1d40f3a7dd9c8f4199cb9cfbc1445f4cfd849aee Mon Sep 17 00:00:00 2001 From: Casey West Date: Sun, 12 Jul 2026 12:15:58 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(image=5Fgen):=20route=20nano?= =?UTF-8?q?-banana=20text-to-image=20through=20/v1/images/generations?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit gemini-3-pro-image silently returned ~1K (1408x768) regardless of the configured resolution because the chat/completions path has NO imageConfig->generationConfig mapping (verified against LiteLLM 1.91.2 and 1.92.0 source). LiteLLM maps resolution ONLY on /v1/images/generations. Route text-to-image to POST {base_url}/v1/images/generations with a NESTED imageConfig ({"imageSize": resolution, "aspectRatio": ratio}), which maps to Vertex generationConfig.imageConfig for pro/flash/lite (model-agnostic, no flag), so Pro (and all image models) honor 4K. Parse the images-API response (b64_json preferred, url fallback) and save it. Keep the EDIT / reference path on chat/completions: the images endpoint has no clean input-image contract at 1.92.0 (image input lives on the multipart /v1/images/edits route, unverified against this proxy), the chat edit path works today, and resolution matters less for edits since the model preserves source dimensions. The split is documented in the module docstring + comments. Preserve resolution-fallback resilience: if an older proxy rejects the nested imageConfig, retry once without it so generation still lands. Reuse _resolve_model/_resolve_resolution/_clamp_resolution and the aspect map unchanged; the Lite 1K cap still clamps on the images path. Tests mock the HTTP layer (no running proxy): assert text-to-image hits /v1/images/generations with the nested imageConfig carrying imageSize + aspectRatio, response parse (b64_json/url), imageConfig-rejection fallback, and that the edit/reference path still routes to chat/completions unchanged. --- plugins/image_gen/nano-banana/__init__.py | 300 ++++++++++++++--- .../image_gen/test_nano_banana_provider.py | 310 +++++++++++++----- 2 files changed, 486 insertions(+), 124 deletions(-) diff --git a/plugins/image_gen/nano-banana/__init__.py b/plugins/image_gen/nano-banana/__init__.py index 79f1bdad06d8..f3d7e7d18010 100644 --- a/plugins/image_gen/nano-banana/__init__.py +++ b/plugins/image_gen/nano-banana/__init__.py @@ -6,22 +6,44 @@ ``Authorization: Bearer `` — no Google credential is ever handled client-side. -Protocol (OpenAI ``/chat/completions`` image output — the same shape the -``openrouter`` backend speaks): - -- ``POST {base_url}/chat/completions`` with ``modalities: ["image", "text"]``, - the prompt (and any source/reference images) as ``messages[0].content`` parts, - and an optional ``image_config.aspect_ratio``. -- The generated image comes back at - ``choices[0].message.images[0].image_url.url`` as a ``data:image/...;base64`` - URI (``message.content`` is ``null`` on this protocol). It is decoded and - saved under ``$HERMES_HOME/cache/images/`` via the framework's - ``save_b64_image``; the tool returns the path, delivered as ``MEDIA:/path``. - -Unified generate + edit + reference: when ``image_url`` (primary source to edit) -or ``reference_image_urls`` (style/subject references, e.g. a likeness for -character consistency) are present, they are inlined as ``image_url`` content -parts and the call routes to image-to-image / editing; otherwise text-to-image. +Two request protocols, split by whether the call has an input image — because +the proxy honors output resolution on only ONE of them: + +TEXT-TO-IMAGE → ``POST {base_url}/v1/images/generations`` (OpenAI images API): + +- Body: ``{"model": id, "prompt": ..., "imageConfig": {"imageSize": res, + "aspectRatio": ratio}}`` — a NESTED ``imageConfig``. This is the path where + LiteLLM (>=1.92.0) maps ``imageConfig`` → Vertex + ``generationConfig.imageConfig`` for pro/flash/lite (model-agnostic, no + capability flag), so ``gemini-3-pro-image`` actually honors 4K here. +- The chat/completions path has NO ``imageConfig``→``generationConfig`` mapping + (verified against LiteLLM 1.91.2 AND 1.92.0 source), so ``image_size`` on chat + is INERT for Pro — it silently returns ~1K regardless. That is the bug this + split fixes. FLAT ``imageSize``/``image_size`` at the top level are DROPPED by + the proxy; only the nested ``imageConfig`` is honored. +- Response is the images-API shape: ``data[0].b64_json`` (base64 PNG) and/or + ``data[0].url``. ``b64_json`` is decoded via ``save_b64_image``; a bare ``url`` + is fetched via ``save_url_image``. The tool returns the saved path + (``MEDIA:/path``). +- Requires the proxy on LiteLLM >=1.92.0. If an older proxy rejects the nested + ``imageConfig``, the request is retried once WITHOUT it so generation still + lands (falls back to default geometry rather than hard-failing). + +EDIT / reference → ``POST {base_url}/chat/completions`` (unchanged): + +- When ``image_url`` (primary source to edit) or ``reference_image_urls`` + (style/subject references, e.g. a likeness for character consistency) are + present, they are inlined as ``image_url`` content parts and the call routes + to image-to-image / editing on the chat path with ``modalities: ["image", + "text"]``. The image comes back at + ``choices[0].message.images[0].image_url.url`` (``message.content`` is + ``null`` on this protocol). +- Why keep edits on chat: the images-API endpoint has no clean input-image + contract at 1.92.0 (image input lives on the multipart ``/v1/images/edits`` + route, unverified against this proxy). The chat edit path works today, and + resolution matters less for edits because the model preserves the source + image's dimensions regardless. So text-to-image — where 4K actually matters — + moves to the images path, and edit/reference stays on chat. Model routing is Pro-default and fully config-driven (precedence, first hit wins): @@ -38,15 +60,14 @@ Resolution is config-only (no per-call parameter on the ``image_generate`` tool schema): ``image_gen.nano-banana.resolution`` in config.yaml selects the output -size sent as ``image_config.image_size`` (``1K``/``2K``/``4K``, uppercase K), -defaulting to ``4K``. Verified against the live proxy — ``image_size`` changes -the decoded PNG dimensions on the text-to-image path and composes with -``aspect_ratio``. Caveat: on the EDIT path the model tends to preserve the -source image's dimensions regardless, so resolution primarily affects -text-to-image. A per-model cap degrades gracefully (a capped model like Lite = -1K clamps down rather than erroring), and if the proxy ever rejects the -``image_size`` field the request is retried once WITHOUT it (current -no-resolution behavior) so the generation still lands. +size (``1K``/``2K``/``4K``, uppercase K), defaulting to ``4K``. On the +text-to-image path it is sent as ``imageConfig.imageSize`` on +``/v1/images/generations`` — the only place the proxy actually honors it (see +above). A per-model cap degrades gracefully (a capped model like Lite = 1K +clamps down rather than erroring), and if the proxy rejects the nested +``imageConfig`` (older LiteLLM) the request is retried once WITHOUT it so the +generation still lands. On the EDIT path resolution is not sent — the model +preserves the source image's dimensions regardless. Prompt craft lives in the ``nano-banana-prompting`` skill, not in this backend. """ @@ -67,6 +88,7 @@ normalize_reference_images, resolve_aspect_ratio, save_b64_image, + save_url_image, success_response, ) from hermes_cli.runtime_provider import resolve_runtime_provider @@ -213,6 +235,23 @@ def _extract_images(payload: Dict[str, Any]) -> List[str]: return out +def _extract_images_api(payload: Dict[str, Any]) -> List[Dict[str, Any]]: + """Pull image items from an OpenAI images-API (/v1/images/generations) response. + + Shape: ``{"data": [{"b64_json": "..."} | {"url": "https://..."}]}``. Returns + the raw ``data`` item dicts (each carrying ``b64_json`` and/or ``url``) so the + caller can prefer inline base64 over a bare URL. + """ + out: List[Dict[str, Any]] = [] + data = payload.get("data") if isinstance(payload, dict) else None + if not isinstance(data, list): + return out + for item in data: + if isinstance(item, dict) and (item.get("b64_json") or item.get("url")): + out.append(item) + return out + + class NanoBananaImageGenProvider(ImageGenProvider): """Google Gemini image models served through the local proxy.""" @@ -372,8 +411,6 @@ def generate( reference_image_urls: Optional[List[str]] = None, **kwargs: Any, ) -> Dict[str, Any]: - import requests - prompt = (prompt or "").strip() aspect = resolve_aspect_ratio(aspect_ratio) proxy_aspect = _ASPECT_RATIOS.get(aspect, "1:1") @@ -391,7 +428,6 @@ def generate( for ref in normalize_reference_images(reference_image_urls) or []: references.append(str(ref)) has_source = bool(references) - modality = "image" if has_source else "text" if not prompt and not has_source: return error_response( @@ -429,6 +465,140 @@ def generate( aspect_ratio=aspect, ) + headers = { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + } + + # Split by whether we have an input image. Text-to-image goes to the + # images API (/v1/images/generations) where the proxy honors 4K via a + # nested imageConfig; edit/reference stays on chat/completions (the + # images path has no clean input-image contract at LiteLLM 1.92.0). See + # the module docstring for the full rationale. + if has_source: + return self._generate_edit( + base_url=base_url, + headers=headers, + prompt=prompt, + references=references, + model_id=model_id, + proxy_aspect=proxy_aspect, + resolution=resolution, + aspect=aspect, + ) + return self._generate_text_to_image( + base_url=base_url, + headers=headers, + prompt=prompt, + model_id=model_id, + proxy_aspect=proxy_aspect, + resolution=resolution, + aspect=aspect, + ) + + def _generate_text_to_image( + self, + *, + base_url: str, + headers: Dict[str, str], + prompt: str, + model_id: str, + proxy_aspect: str, + resolution: str, + aspect: str, + ) -> Dict[str, Any]: + """Text-to-image via the OpenAI images API (/v1/images/generations). + + This is the ONLY path where LiteLLM maps the nested ``imageConfig`` to + Vertex ``generationConfig.imageConfig``, so Pro honors the configured + resolution (4K) here. The chat path silently returns ~1K for Pro. + """ + payload: Dict[str, Any] = { + "model": model_id, + "prompt": prompt, + # NESTED imageConfig — flat imageSize/image_size at the top level are + # dropped by the proxy and would leave Pro stuck at ~1K. + "imageConfig": {"imageSize": resolution, "aspectRatio": proxy_aspect}, + } + # base_url already ends in /v1 (the proxy's OpenAI-compatible root), so + # this resolves to {host}/v1/images/generations — the images-API path. + url = f"{base_url}/images/generations" + + result = self._post_with_image_config_fallback( + url=url, + headers=headers, + payload=payload, + model_id=model_id, + prompt=prompt, + aspect=aspect, + resolution=resolution, + ) + if isinstance(result, dict) and result.get("success") is False: + return result + + items = _extract_images_api(result) + if not items: + return error_response( + error=( + f"nano-banana returned no image. Ensure the model '{model_id}' " + "supports image output on the proxy." + ), + error_type="empty_response", + provider="nano-banana", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + + first = items[0] + b64 = first.get("b64_json") + item_url = first.get("url") + try: + # Prefer inline base64 (the proxy's default) over a bare URL. + if isinstance(b64, str) and b64.strip(): + saved_path = save_b64_image(b64, prefix="nano_banana") + else: + saved_path = save_url_image(str(item_url), prefix="nano_banana") + except Exception as exc: # noqa: BLE001 + return error_response( + error=f"Could not save generated image: {exc}", + error_type="io_error", + provider="nano-banana", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + + return success_response( + image=str(saved_path), + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + provider="nano-banana", + modality="text", + extra={"resolution": resolution}, + ) + + def _generate_edit( + self, + *, + base_url: str, + headers: Dict[str, str], + prompt: str, + references: List[str], + model_id: str, + proxy_aspect: str, + resolution: str, + aspect: str, + ) -> Dict[str, Any]: + """Edit / reference via chat/completions (unchanged behavior). + + Kept on the chat path because the images API has no clean input-image + contract at LiteLLM 1.92.0. Resolution matters less here — the model + preserves the source image's dimensions regardless — so the existing + chat ``image_config`` (with its resolution-field fallback) is retained + to avoid any regression in the currently-working edit flow. + """ content: List[Dict[str, Any]] = [{"type": "text", "text": prompt}] for ref in references[:_MAX_REFERENCE_IMAGES]: part = _to_image_url_part(ref) @@ -441,10 +611,6 @@ def generate( "messages": [{"role": "user", "content": content}], "image_config": {"aspect_ratio": proxy_aspect, "image_size": resolution}, } - headers = { - "Authorization": f"Bearer {api_key}", - "Content-Type": "application/json", - } url = f"{base_url}/chat/completions" result = self._post_with_resolution_fallback( @@ -482,8 +648,6 @@ def generate( else: # The proxy returns inline base64 data URIs; a bare URL is # unexpected but handled for robustness. - from agent.image_gen_provider import save_url_image - saved_path = save_url_image(first, prefix="nano_banana") except Exception as exc: # noqa: BLE001 return error_response( @@ -501,10 +665,61 @@ def generate( prompt=prompt, aspect_ratio=aspect, provider="nano-banana", - modality=modality, + modality="image", extra={"resolution": resolution}, ) + def _post_with_image_config_fallback( + self, + *, + url: str, + headers: Dict[str, str], + payload: Dict[str, Any], + model_id: str, + prompt: str, + aspect: str, + resolution: str, + ) -> Dict[str, Any]: + """POST the text-to-image request to the images API; on an + ``imageConfig`` rejection (an older proxy that predates the nested + imageConfig mapping), retry once WITHOUT it so the generation still + lands rather than hard-failing on an un-upgraded proxy. + + Returns the parsed JSON result on success, or an ``error_response`` dict + (``success`` is ``False``) on a non-recoverable failure. + """ + result = self._post_once( + url=url, + headers=headers, + payload=payload, + model_id=model_id, + prompt=prompt, + aspect=aspect, + ) + if ( + isinstance(result, dict) + and result.get("success") is False + and result.get("error_type") == "api_error" + and self._is_resolution_field_error(result.get("error", "")) + and "imageConfig" in payload + ): + logger.warning( + "nano-banana: proxy rejected imageConfig (imageSize=%s); retrying " + "without it (falling back to default geometry on an un-upgraded " + "proxy)", + resolution, + ) + fallback = {k: v for k, v in payload.items() if k != "imageConfig"} + result = self._post_once( + url=url, + headers=headers, + payload=fallback, + model_id=model_id, + prompt=prompt, + aspect=aspect, + ) + return result + def _post_with_resolution_fallback( self, *, @@ -565,14 +780,17 @@ def _post_with_resolution_fallback( @staticmethod def _is_resolution_field_error(message: str) -> bool: - """True when an API error is about the ``image_size`` field itself. - - Scopes the fallback to genuine field-contract rejections so an unrelated - 400 (e.g. a safety block) still surfaces as an error instead of being - masked by a retry. + """True when an API error is about the resolution field itself. + + Matches the field names used across both request protocols and proxy + versions — the chat path's ``image_size`` and the images path's nested + ``imageConfig`` / ``imageSize`` (LiteLLM error text is not always + snake_case). Scopes the fallback to genuine field-contract rejections so + an unrelated 400 (e.g. a safety block) still surfaces as an error + instead of being masked by a retry. """ low = str(message or "").lower() - return "image_size" in low + return "image_size" in low or "imagesize" in low or "imageconfig" in low def _post_once( self, diff --git a/tests/plugins/image_gen/test_nano_banana_provider.py b/tests/plugins/image_gen/test_nano_banana_provider.py index 8eb0a386d65b..bb439673f1d1 100644 --- a/tests/plugins/image_gen/test_nano_banana_provider.py +++ b/tests/plugins/image_gen/test_nano_banana_provider.py @@ -15,6 +15,8 @@ nb = importlib.import_module("plugins.image_gen.nano-banana") _PNG_DATA_URI = "data:image/png;base64,dGVzdC1pbWFnZS1kYXRh" # "test-image-data" +# Raw base64 (no data: prefix) — the images-API b64_json field carries this. +_PNG_B64 = "dGVzdC1pbWFnZS1kYXRh" # "test-image-data" def _runtime_ok(**over): @@ -30,6 +32,7 @@ def _runtime_ok(**over): def _mock_chat_response(images, *, content=None): + """Mock a /chat/completions image response (the EDIT / reference path).""" resp = MagicMock() resp.status_code = 200 resp.raise_for_status = MagicMock() @@ -49,6 +52,21 @@ def _mock_chat_response(images, *, content=None): return resp +def _mock_images_response(items): + """Mock a /v1/images/generations response (the TEXT-TO-IMAGE path). + + Shape: ``{"data": [{"b64_json": ...} | {"url": ...}]}`` — the OpenAI + images-API shape LiteLLM returns. This is the path where the nested + ``imageConfig`` -> Vertex ``generationConfig.imageConfig`` mapping happens, + so resolution (4K) is actually honored here (unlike the chat path). + """ + resp = MagicMock() + resp.status_code = 200 + resp.raise_for_status = MagicMock() + resp.json.return_value = {"data": list(items)} + return resp + + def _provider(): return nb.NanoBananaImageGenProvider() @@ -160,7 +178,7 @@ def test_unavailable_on_resolution_error(self): # --------------------------------------------------------------------------- -# generate() +# generate() — TEXT-TO-IMAGE (routes to /v1/images/generations) # --------------------------------------------------------------------------- @@ -171,9 +189,9 @@ def test_missing_credentials(self): assert result["success"] is False assert result["error_type"] == "missing_api_key" - def test_success_data_uri(self): + def test_success_b64_json(self): with _patch_runtime(), \ - patch("requests.post", return_value=_mock_chat_response([_PNG_DATA_URI])), \ + patch("requests.post", return_value=_mock_images_response([{"b64_json": _PNG_B64}])), \ _patch_save() as mock_save: result = _provider().generate(prompt="a banana") assert result["success"] is True @@ -182,79 +200,60 @@ def test_success_data_uri(self): assert result["model"] == "gemini-3-pro-image" mock_save.assert_called_once() - def test_reads_image_from_choices_message_images(self): - """The image MUST be read from choices[0].message.images[0].image_url.url - (content is null on this protocol).""" - resp = _mock_chat_response([_PNG_DATA_URI], content=None) - with _patch_runtime(), \ - patch("requests.post", return_value=resp), \ - _patch_save("/tmp/x.png"): - result = _provider().generate(prompt="a banana") - assert result["success"] is True - - def test_empty_images_is_empty_response(self): + def test_empty_data_is_empty_response(self): with _patch_runtime(), \ - patch("requests.post", return_value=_mock_chat_response([])): + patch("requests.post", return_value=_mock_images_response([])): result = _provider().generate(prompt="a banana") assert result["success"] is False assert result["error_type"] == "empty_response" - def test_payload_shape_text_to_image(self): - with _patch_runtime(), \ - patch("requests.post", return_value=_mock_chat_response([_PNG_DATA_URI])) as mock_post, \ - _patch_save("/tmp/x.png"): - _provider().generate(prompt="a banana", aspect_ratio="portrait") - payload = mock_post.call_args.kwargs["json"] - assert payload["model"] == "gemini-3-pro-image" - assert payload["modalities"] == ["image", "text"] - assert payload["image_config"]["aspect_ratio"] == "9:16" - content = payload["messages"][0]["content"] - assert content[0] == {"type": "text", "text": "a banana"} - assert all(c["type"] != "image_url" for c in content) - - def test_posts_to_resolved_base_url(self): + def test_posts_to_images_generations_endpoint(self): + """TEXT-TO-IMAGE MUST hit /v1/images/generations (NOT chat/completions): + that is the only path where LiteLLM maps imageConfig -> Vertex, so Pro + honors 4K. This contract would have caught the original ~1K bug.""" with _patch_runtime(), \ - patch("requests.post", return_value=_mock_chat_response([_PNG_DATA_URI])) as mock_post, \ + patch("requests.post", return_value=_mock_images_response([{"b64_json": _PNG_B64}])) as mock_post, \ _patch_save("/tmp/x.png"): _provider().generate(prompt="a banana") url = mock_post.call_args[0][0] - assert url == "http://127.0.0.1:4000/v1/chat/completions" + assert url == "http://127.0.0.1:4000/v1/images/generations" + + def test_payload_shape_text_to_image_nested_image_config(self): + """The payload MUST carry model + prompt + a NESTED imageConfig with + imageSize (= configured resolution) and aspectRatio (proxy string). + FLAT imageSize/image_size at the top level are DROPPED by the proxy, so + they must NOT be sent.""" + with _patch_cfg({"nano-banana": {"resolution": "4K"}}), _patch_runtime(), \ + patch("requests.post", return_value=_mock_images_response([{"b64_json": _PNG_B64}])) as mock_post, \ + _patch_save("/tmp/x.png"): + _provider().generate(prompt="a banana", aspect_ratio="landscape") + payload = mock_post.call_args.kwargs["json"] + assert payload["model"] == "gemini-3-pro-image" + assert payload["prompt"] == "a banana" + # NESTED imageConfig — the load-bearing contract. + assert payload["imageConfig"] == {"imageSize": "4K", "aspectRatio": "16:9"} + # No FLAT resolution keys leak to the top level (proxy drops them). + assert "imageSize" not in payload + assert "image_size" not in payload + # This is the images API, not chat: no chat-only keys. + assert "messages" not in payload + assert "modalities" not in payload def test_auth_header_bearer_token(self): with _patch_runtime(), \ - patch("requests.post", return_value=_mock_chat_response([_PNG_DATA_URI])) as mock_post, \ + patch("requests.post", return_value=_mock_images_response([{"b64_json": _PNG_B64}])) as mock_post, \ _patch_save("/tmp/x.png"): _provider().generate(prompt="a banana") headers = mock_post.call_args.kwargs["headers"] assert headers["Authorization"] == "Bearer sk-local" - def test_edit_routing_attaches_image_url(self, tmp_path): - src = tmp_path / "src.png" - src.write_bytes(b"\x89PNG\r\n") - with _patch_runtime(), \ - patch("requests.post", return_value=_mock_chat_response([_PNG_DATA_URI])) as mock_post, \ - _patch_save("/tmp/x.png"): - result = _provider().generate(prompt="make it red", image_url=str(src)) - assert result["success"] is True - assert result["modality"] == "image" - content = mock_post.call_args.kwargs["json"]["messages"][0]["content"] - image_parts = [c for c in content if c["type"] == "image_url"] - assert len(image_parts) == 1 - assert image_parts[0]["image_url"]["url"].startswith("data:image/png;base64,") - - def test_multiple_references_clamped(self, tmp_path): - refs = [] - for i in range(nb._MAX_REFERENCE_IMAGES + 2): - f = tmp_path / f"r{i}.png" - f.write_bytes(b"\x89PNG\r\n") - refs.append(str(f)) + def test_model_kwarg_flows_into_payload(self): with _patch_runtime(), \ - patch("requests.post", return_value=_mock_chat_response([_PNG_DATA_URI])) as mock_post, \ + patch("requests.post", return_value=_mock_images_response([{"b64_json": _PNG_B64}])) as mock_post, \ _patch_save("/tmp/x.png"): - _provider().generate(prompt="keep the character", reference_image_urls=refs) - content = mock_post.call_args.kwargs["json"]["messages"][0]["content"] - image_parts = [c for c in content if c["type"] == "image_url"] - assert len(image_parts) == nb._MAX_REFERENCE_IMAGES + result = _provider().generate(prompt="a banana", model="gemini-3.1-flash-image") + assert result["model"] == "gemini-3.1-flash-image" + assert mock_post.call_args.kwargs["json"]["model"] == "gemini-3.1-flash-image" def test_http_error_is_api_error(self): import requests as req_lib @@ -291,13 +290,124 @@ def test_empty_prompt_with_no_image_rejected(self): assert result["success"] is False assert result["error_type"] == "invalid_argument" - def test_model_kwarg_flows_into_payload(self): + +# --------------------------------------------------------------------------- +# generate() — TEXT-TO-IMAGE response parsing (images-API shape) +# --------------------------------------------------------------------------- + + +class TestTextToImageResponseParse: + def test_b64_json_is_saved(self): + with _patch_runtime(), \ + patch("requests.post", return_value=_mock_images_response([{"b64_json": _PNG_B64}])), \ + _patch_save("/tmp/b64.png") as mock_save: + result = _provider().generate(prompt="a banana") + assert result["success"] is True + assert result["image"] == "/tmp/b64.png" + # b64_json is decoded and saved directly (no data: prefix stripping). + mock_save.assert_called_once() + assert mock_save.call_args[0][0] == _PNG_B64 + + def test_url_is_fetched_and_saved(self): + with _patch_runtime(), \ + patch("requests.post", return_value=_mock_images_response([{"url": "https://cdn.example/img.png"}])), \ + patch.object(nb, "save_url_image", return_value=Path("/tmp/url.png")) as mock_url_save: + result = _provider().generate(prompt="a banana") + assert result["success"] is True + assert result["image"] == "/tmp/url.png" + mock_url_save.assert_called_once() + assert mock_url_save.call_args[0][0] == "https://cdn.example/img.png" + + def test_b64_json_preferred_when_both_present(self): + with _patch_runtime(), \ + patch("requests.post", return_value=_mock_images_response( + [{"b64_json": _PNG_B64, "url": "https://cdn.example/img.png"}])), \ + _patch_save("/tmp/b64.png") as mock_save: + result = _provider().generate(prompt="a banana") + assert result["success"] is True + assert result["image"] == "/tmp/b64.png" + mock_save.assert_called_once() + + +# --------------------------------------------------------------------------- +# generate() — EDIT / reference path (KEPT on /chat/completions) +# --------------------------------------------------------------------------- +# +# Design decision (spec change #2): /v1/images/generations has no clean input- +# image contract at LiteLLM 1.92.0 (image input lives on the multipart +# /v1/images/edits route, unverified against this proxy). The edit/reference +# case already works on chat/completions today, and resolution matters less +# there because the model preserves source dimensions. So text-to-image moves +# to the images path (where 4K is honored) and edit/reference STAYS on chat. + + +class TestEditPathStaysOnChat: + def test_edit_routing_hits_chat_completions(self, tmp_path): + src = tmp_path / "src.png" + src.write_bytes(b"\x89PNG\r\n") with _patch_runtime(), \ patch("requests.post", return_value=_mock_chat_response([_PNG_DATA_URI])) as mock_post, \ _patch_save("/tmp/x.png"): - result = _provider().generate(prompt="a banana", model="gemini-3.1-flash-image") - assert result["model"] == "gemini-3.1-flash-image" - assert mock_post.call_args.kwargs["json"]["model"] == "gemini-3.1-flash-image" + result = _provider().generate(prompt="make it red", image_url=str(src)) + assert result["success"] is True + assert result["modality"] == "image" + # Edit stays on the chat path (input-image contract works there). + url = mock_post.call_args[0][0] + assert url == "http://127.0.0.1:4000/v1/chat/completions" + + def test_edit_routing_attaches_image_url(self, tmp_path): + src = tmp_path / "src.png" + src.write_bytes(b"\x89PNG\r\n") + with _patch_runtime(), \ + patch("requests.post", return_value=_mock_chat_response([_PNG_DATA_URI])) as mock_post, \ + _patch_save("/tmp/x.png"): + result = _provider().generate(prompt="make it red", image_url=str(src)) + assert result["success"] is True + content = mock_post.call_args.kwargs["json"]["messages"][0]["content"] + image_parts = [c for c in content if c["type"] == "image_url"] + assert len(image_parts) == 1 + assert image_parts[0]["image_url"]["url"].startswith("data:image/png;base64,") + + def test_edit_payload_is_chat_shape(self, tmp_path): + src = tmp_path / "src.png" + src.write_bytes(b"\x89PNG\r\n") + with _patch_runtime(), \ + patch("requests.post", return_value=_mock_chat_response([_PNG_DATA_URI])) as mock_post, \ + _patch_save("/tmp/x.png"): + _provider().generate(prompt="make it red", image_url=str(src)) + payload = mock_post.call_args.kwargs["json"] + assert payload["modalities"] == ["image", "text"] + assert "messages" in payload + # The chat path keeps the flat image_config it already used. + assert "image_config" in payload + # It does NOT carry the images-API nested imageConfig. + assert "imageConfig" not in payload + + def test_multiple_references_clamped(self, tmp_path): + refs = [] + for i in range(nb._MAX_REFERENCE_IMAGES + 2): + f = tmp_path / f"r{i}.png" + f.write_bytes(b"\x89PNG\r\n") + refs.append(str(f)) + with _patch_runtime(), \ + patch("requests.post", return_value=_mock_chat_response([_PNG_DATA_URI])) as mock_post, \ + _patch_save("/tmp/x.png"): + _provider().generate(prompt="keep the character", reference_image_urls=refs) + content = mock_post.call_args.kwargs["json"]["messages"][0]["content"] + image_parts = [c for c in content if c["type"] == "image_url"] + assert len(image_parts) == nb._MAX_REFERENCE_IMAGES + + def test_edit_reads_image_from_choices_message_images(self, tmp_path): + """The edit path still reads the image from + choices[0].message.images[0].image_url.url (content is null).""" + src = tmp_path / "src.png" + src.write_bytes(b"\x89PNG\r\n") + resp = _mock_chat_response([_PNG_DATA_URI], content=None) + with _patch_runtime(), \ + patch("requests.post", return_value=resp), \ + _patch_save("/tmp/x.png"): + result = _provider().generate(prompt="make it red", image_url=str(src)) + assert result["success"] is True # --------------------------------------------------------------------------- @@ -359,49 +469,87 @@ def test_capped_model_below_cap_unchanged(self): class TestResolutionPayload: + """image_size now travels as a NESTED imageConfig.imageSize on the images + path (text-to-image), not the flat chat image_config.image_size.""" + def test_image_size_sent_on_text_to_image(self): with _patch_cfg({"nano-banana": {"resolution": "4K"}}), _patch_runtime(), \ - patch("requests.post", return_value=_mock_chat_response([_PNG_DATA_URI])) as mock_post, \ + patch("requests.post", return_value=_mock_images_response([{"b64_json": _PNG_B64}])) as mock_post, \ _patch_save("/tmp/x.png"): _provider().generate(prompt="a banana", aspect_ratio="landscape") - ic = mock_post.call_args.kwargs["json"]["image_config"] - assert ic["image_size"] == "4K" - assert ic["aspect_ratio"] == "16:9" + ic = mock_post.call_args.kwargs["json"]["imageConfig"] + assert ic["imageSize"] == "4K" + assert ic["aspectRatio"] == "16:9" def test_default_4k_sent_when_config_absent(self): with _patch_cfg({}), _patch_runtime(), \ - patch("requests.post", return_value=_mock_chat_response([_PNG_DATA_URI])) as mock_post, \ + patch("requests.post", return_value=_mock_images_response([{"b64_json": _PNG_B64}])) as mock_post, \ _patch_save("/tmp/x.png"): _provider().generate(prompt="a banana") - assert mock_post.call_args.kwargs["json"]["image_config"]["image_size"] == "4K" + assert mock_post.call_args.kwargs["json"]["imageConfig"]["imageSize"] == "4K" def test_config_override_2k_sent(self): with _patch_cfg({"nano-banana": {"resolution": "2K"}}), _patch_runtime(), \ - patch("requests.post", return_value=_mock_chat_response([_PNG_DATA_URI])) as mock_post, \ + patch("requests.post", return_value=_mock_images_response([{"b64_json": _PNG_B64}])) as mock_post, \ _patch_save("/tmp/x.png"): _provider().generate(prompt="a banana") - assert mock_post.call_args.kwargs["json"]["image_config"]["image_size"] == "2K" + assert mock_post.call_args.kwargs["json"]["imageConfig"]["imageSize"] == "2K" def test_success_reports_resolution(self): with _patch_cfg({"nano-banana": {"resolution": "2K"}}), _patch_runtime(), \ - patch("requests.post", return_value=_mock_chat_response([_PNG_DATA_URI])), \ + patch("requests.post", return_value=_mock_images_response([{"b64_json": _PNG_B64}])), \ _patch_save("/tmp/x.png"): result = _provider().generate(prompt="a banana") assert result["resolution"] == "2K" + def test_lite_cap_still_applies_on_images_path(self): + """The per-model 1K cap for Lite still clamps on the images path.""" + with _patch_cfg({"nano-banana": {"resolution": "4K", "model": "gemini-3.1-flash-lite-image"}}), \ + _patch_runtime(), \ + patch("requests.post", return_value=_mock_images_response([{"b64_json": _PNG_B64}])) as mock_post, \ + _patch_save("/tmp/x.png"): + _provider().generate(prompt="a banana") + assert mock_post.call_args.kwargs["json"]["imageConfig"]["imageSize"] == "1K" + class TestResolutionGracefulDegradation: - def test_proxy_rejects_image_size_falls_back_without_it(self): - """If the proxy 400s on the image_size field, retry once WITHOUT it - (current no-resolution behavior) rather than failing the generation.""" + """If the proxy is on an older LiteLLM that rejects the nested imageConfig, + retry once WITHOUT it so generation still lands (don't hard-fail).""" + + def test_proxy_rejects_image_config_falls_back_without_it(self): + import requests as req_lib + + bad = MagicMock() + bad.status_code = 400 + bad.text = "Unknown name \"imageConfig\"" + bad.json.return_value = {"error": {"message": "Unknown name \"imageConfig\""}} + bad.raise_for_status.side_effect = req_lib.HTTPError(response=bad) + good = _mock_images_response([{"b64_json": _PNG_B64}]) + + with _patch_cfg({"nano-banana": {"resolution": "4K"}}), _patch_runtime(), \ + patch("requests.post", side_effect=[bad, good]) as mock_post, \ + _patch_save("/tmp/x.png"): + result = _provider().generate(prompt="a banana") + + assert result["success"] is True + # Two calls: first with imageConfig, retry without it. + assert mock_post.call_count == 2 + first_payload = mock_post.call_args_list[0].kwargs["json"] + second_payload = mock_post.call_args_list[1].kwargs["json"] + assert "imageConfig" in first_payload + assert "imageConfig" not in second_payload + + def test_rejects_image_size_error_also_triggers_fallback(self): + """An older proxy may name the field imageSize/image_size in its 400; + the fallback still fires.""" import requests as req_lib bad = MagicMock() bad.status_code = 400 - bad.text = "Unknown name \"image_size\"" - bad.json.return_value = {"error": {"message": "Unknown name \"image_size\""}} + bad.text = "Unknown name \"imageSize\"" + bad.json.return_value = {"error": {"message": "Unknown name \"imageSize\""}} bad.raise_for_status.side_effect = req_lib.HTTPError(response=bad) - good = _mock_chat_response([_PNG_DATA_URI]) + good = _mock_images_response([{"b64_json": _PNG_B64}]) with _patch_cfg({"nano-banana": {"resolution": "4K"}}), _patch_runtime(), \ patch("requests.post", side_effect=[bad, good]) as mock_post, \ @@ -409,16 +557,12 @@ def test_proxy_rejects_image_size_falls_back_without_it(self): result = _provider().generate(prompt="a banana") assert result["success"] is True - # Two calls: first with image_size, retry without it. assert mock_post.call_count == 2 - first_ic = mock_post.call_args_list[0].kwargs["json"]["image_config"] - second_ic = mock_post.call_args_list[1].kwargs["json"]["image_config"] - assert "image_size" in first_ic - assert "image_size" not in second_ic + assert "imageConfig" not in mock_post.call_args_list[1].kwargs["json"] def test_unrelated_400_is_not_masked_by_fallback(self): - """A 400 that is NOT about the resolution field must surface as an error, - not trigger an infinite/masking retry.""" + """A 400 that is NOT about the imageConfig/resolution field must surface + as an error, not trigger a masking retry.""" import requests as req_lib bad = MagicMock()