feat: add gemini-2.5-flash-image backend to google_imagen

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

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] = {