From 5d152e4699b770c889b595e4b8a13c0ca5f4ad5d Mon Sep 17 00:00:00 2001 From: shewulong Date: Wed, 29 Jul 2026 20:41:31 -0500 Subject: [PATCH] feat: add gemini-2.5-flash-image backend to google_imagen --- .../test_google_imagen_gemini_backend.py | 131 ++++++++++++++++++ tools/graphics/google_imagen.py | 123 +++++++++++++--- 2 files changed, 234 insertions(+), 20 deletions(-) create mode 100644 tests/tools/test_google_imagen_gemini_backend.py diff --git a/tests/tools/test_google_imagen_gemini_backend.py b/tests/tools/test_google_imagen_gemini_backend.py new file mode 100644 index 00000000..9b899b1d --- /dev/null +++ b/tests/tools/test_google_imagen_gemini_backend.py @@ -0,0 +1,131 @@ +"""Tests for the Gemini image backend in google_imagen. + +Models named `gemini-*` (e.g. gemini-2.5-flash-image) are not served by the +Imagen `:predict` endpoint — they generate images through generate_content +with an image_config. This backend matters on Vertex projects that have no +Imagen catalog access, where it is the only working Google image path. +""" + +import sys +import types as pytypes +from pathlib import Path + +import pytest + +PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent +sys.path.insert(0, str(PROJECT_ROOT)) + + +class _FakeInline: + def __init__(self, data: bytes): + self.data = data + + +class _FakePart: + def __init__(self, data: bytes): + self.inline_data = _FakeInline(data) + + +class _FakeContent: + def __init__(self, parts): + self.parts = parts + + +class _FakeCandidate: + def __init__(self, parts): + self.content = _FakeContent(parts) + + +class _FakeResponse: + def __init__(self, parts): + self.candidates = [_FakeCandidate(parts)] + + +@pytest.fixture +def imagen_tool(monkeypatch): + monkeypatch.setenv("GOOGLE_API_KEY", "test-key") + + calls: list[dict] = [] + + class _FakeModels: + def generate_content(self, model=None, contents=None, config=None): + calls.append({"model": model, "contents": contents, "config": config}) + return _FakeResponse([_FakePart(b"GEMINI_IMG")]) + + class _FakeClient: + models = _FakeModels() + + import tools.google_credentials as gc + + monkeypatch.setattr( + gc, "get_genai_client", lambda http_options=None, location=None: _FakeClient() + ) + + from tools.graphics.google_imagen import GoogleImagen + + return GoogleImagen(), calls + + +def test_gemini_model_routes_to_generate_content(imagen_tool, tmp_path): + tool, calls = imagen_tool + out = tmp_path / "img.png" + + result = tool.execute( + { + "prompt": "a flower", + "model": "gemini-2.5-flash-image", + "aspect_ratio": "16:9", + "output_path": str(out), + } + ) + + assert result.success + assert result.data["model"] == "gemini-2.5-flash-image" + assert out.read_bytes() == b"GEMINI_IMG" + + assert len(calls) == 1 + assert calls[0]["model"] == "gemini-2.5-flash-image" + # Aspect ratio must reach the API through image_config, not be dropped. + assert calls[0]["config"].image_config.aspect_ratio == "16:9" + + +def test_gemini_cost_estimate_is_per_image(): + from tools.graphics.google_imagen import GoogleImagen + + tool = GoogleImagen() + assert tool.estimate_cost( + {"model": "gemini-2.5-flash-image", "number_of_images": 2} + ) == pytest.approx(0.039 * 2) + + +def test_text_only_response_is_a_clear_error(monkeypatch, tmp_path): + monkeypatch.setenv("GOOGLE_API_KEY", "test-key") + + class _TextPart: + inline_data = None + + class _FakeModels: + def generate_content(self, model=None, contents=None, config=None): + return _FakeResponse([_TextPart()]) + + class _FakeClient: + models = _FakeModels() + + import tools.google_credentials as gc + + monkeypatch.setattr( + gc, "get_genai_client", lambda http_options=None, location=None: _FakeClient() + ) + + from tools.graphics.google_imagen import GoogleImagen + + result = GoogleImagen().execute( + { + "prompt": "a flower", + "model": "gemini-2.5-flash-image", + "output_path": str(tmp_path / "img.png"), + } + ) + + assert not result.success + assert "No image data" in result.error diff --git a/tools/graphics/google_imagen.py b/tools/graphics/google_imagen.py index ac2717ba..4994b030 100644 --- a/tools/graphics/google_imagen.py +++ b/tools/graphics/google_imagen.py @@ -118,9 +118,12 @@ class GoogleImagen(BaseTool): "imagen-4.0-generate-001", "imagen-4.0-fast-generate-001", "imagen-4.0-ultra-generate-001", + "gemini-2.5-flash-image", ], "default": "imagen-4.0-generate-001", - "description": "Imagen model variant", + "description": "Imagen model variant, or a Gemini image model " + "(gemini-*) routed through generate_content. Use " + "gemini-2.5-flash-image when the project has no Imagen access.", }, "number_of_images": { "type": "integer", @@ -176,13 +179,111 @@ class GoogleImagen(BaseTool): def estimate_cost(self, inputs: dict[str, Any]) -> float: model = inputs.get("model", "imagen-4.0-generate-001") n = inputs.get("number_of_images", 1) + if model.startswith("gemini-"): + # ~1290 output tokens per image at $30/1M tokens + return 0.039 * n if "ultra" in model: return 0.06 * n if "fast" in model: return 0.02 * n return 0.04 * n + def _resolve_aspect_ratio(self, inputs: dict[str, Any]) -> str: + """Explicit aspect_ratio > derived from width/height > default 1:1.""" + if "aspect_ratio" in inputs: + return inputs["aspect_ratio"] + if "width" in inputs and "height" in inputs: + import logging + + aspect_ratio = _dims_to_aspect_ratio(inputs["width"], inputs["height"]) + logging.getLogger(__name__).info( + "google_imagen: remapped %sx%s to nearest supported aspect ratio %s", + inputs["width"], + inputs["height"], + aspect_ratio, + ) + return aspect_ratio + return "1:1" + + def _execute_gemini(self, inputs: dict[str, Any], model: str) -> ToolResult: + """Generate via a Gemini image model (e.g. gemini-2.5-flash-image). + + These models use generate_content with an image_config instead of the + Imagen :predict endpoint, and work on both auth paths (API key and + Vertex service account) through the shared genai client. + """ + start = time.time() + try: + from google.genai import types + from tools.google_credentials import get_genai_client + + client = get_genai_client() + except Exception as e: + return ToolResult( + success=False, + error=f"Failed to initialize Google GenAI client: {e}", + ) + + prompt = inputs["prompt"] + aspect_ratio = self._resolve_aspect_ratio(inputs) + number_of_images = inputs.get("number_of_images", 1) + config = types.GenerateContentConfig( + image_config=types.ImageConfig(aspect_ratio=aspect_ratio), + ) + + image_bytes: list[bytes] = [] + try: + for _ in range(number_of_images): + response = client.models.generate_content( + model=model, contents=prompt, config=config + ) + for part in response.candidates[0].content.parts or []: + inline = getattr(part, "inline_data", None) + if inline and inline.data: + image_bytes.append(inline.data) + break + except Exception as e: + return ToolResult( + success=False, error=f"Gemini image generation failed: {e}" + ) + + if not image_bytes: + return ToolResult( + success=False, + error=f"No image data returned by {model} (text-only response).", + ) + + output_paths = self._output_paths(inputs.get("output_path"), len(image_bytes)) + outputs: list[str] = [] + for data, out_path in zip(image_bytes, output_paths): + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_bytes(data) + outputs.append(str(out_path)) + + return ToolResult( + success=True, + data={ + "provider": "google_imagen", + "model": model, + "prompt": prompt, + "aspect_ratio": aspect_ratio, + "output": outputs[0], + "outputs": outputs, + "images_generated": len(outputs), + }, + artifacts=outputs, + cost_usd=self.estimate_cost(inputs), + duration_seconds=round(time.time() - start, 2), + model=model, + ) + def execute(self, inputs: dict[str, Any]) -> ToolResult: + # Gemini image models go through generate_content via the shared genai + # client, which resolves auth (API key or Vertex service account) itself. + model = inputs.get("model", "imagen-4.0-generate-001") + if model.startswith("gemini-"): + return self._execute_gemini(inputs, model) + # Two auth paths: an AI Studio API key, or a service-account JSON that # routes to Vertex AI (the AI Studio endpoint does not accept service # accounts). API key wins when both are present. @@ -212,27 +313,9 @@ class GoogleImagen(BaseTool): import requests start = time.time() - model = inputs.get("model", "imagen-4.0-generate-001") prompt = inputs["prompt"] - import logging - - logger = logging.getLogger(__name__) - - # Resolve aspect ratio: explicit > derived from width/height > default - if "aspect_ratio" in inputs: - aspect_ratio = inputs["aspect_ratio"] - elif "width" in inputs and "height" in inputs: - requested_ratio = f"{inputs['width']}x{inputs['height']}" - aspect_ratio = _dims_to_aspect_ratio(inputs["width"], inputs["height"]) - logger.info( - "google_imagen: remapped %s to nearest supported aspect ratio %s", - requested_ratio, - aspect_ratio, - ) - else: - aspect_ratio = "1:1" - + aspect_ratio = self._resolve_aspect_ratio(inputs) number_of_images = inputs.get("number_of_images", 1) parameters: dict[str, Any] = {