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()