diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index 4173a3c7..c277896d 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -434,7 +434,7 @@ allowance). OpenMontage estimates cost from the transcribed audio duration. See ### Google — TTS + Imagen + Music + Video (Shared Key) -> **One key, five tools.** Google Cloud TTS has 700+ voices in 50+ languages — the strongest localization option. Imagen 4 generates high-quality images. Google Lyria generates high-quality background music. Gemini Omni Flash supports conversational video editing, and direct Veo generation covers premium short video clips. +> **One key, five tools.** Google Cloud TTS has 700+ voices in 50+ languages — the strongest localization option. `google_imagen` supports both Imagen 4 and Gemini 2.5 Flash Image, including projects without Imagen catalog access. Google Lyria generates high-quality background music. Gemini Omni Flash supports conversational video editing, and direct Veo generation covers premium short video clips. **Tools unlocked:** `google_tts`, `google_imagen`, `google_music`, `gemini_omni_video`, `veo_video` **Env var:** `GOOGLE_API_KEY` (or `GEMINI_API_KEY` — either works; `GEMINI_API_KEY` takes precedence) @@ -475,9 +475,15 @@ The free tiers apply *independently* — you get 1M Standard AND 1M WaveNet AND | Imagen 4 Fast | $0.02 | | Imagen 4 Standard | $0.04 | | Imagen 4 Ultra | $0.06 | +| Gemini 2.5 Flash Image (`gemini-2.5-flash-image`) | $0.039 | **Free tier for Imagen:** None. Paid tier only. +To select the Gemini backend through the governed `image_selector`, pass +`preferred_provider: "google_imagen"` and +`model_name: "gemini-2.5-flash-image"`. The selector maps its neutral +`model_name` field to the provider's `model` input. + #### Gemini Omni Video Pricing | Model | Price | Notes | 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..fc7e960b --- /dev/null +++ b/tests/tools/test_google_imagen_gemini_backend.py @@ -0,0 +1,156 @@ +"""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_image_selector_maps_model_name_to_google_model( + imagen_tool, monkeypatch, tmp_path +): + """The governed selector must be able to reach the Gemini backend.""" + from tools.graphics.image_selector import ImageSelector + + tool, calls = imagen_tool + selector = ImageSelector() + monkeypatch.setattr(selector, "_providers", lambda: [tool]) + + result = selector.execute( + { + "prompt": "a flower", + "preferred_provider": "google_imagen", + "model_name": "gemini-2.5-flash-image", + "output_path": str(tmp_path / "selected.png"), + } + ) + + assert result.success, result.error + assert calls[0]["model"] == "gemini-2.5-flash-image" + assert result.data["selected_tool"] == "google_imagen" + assert result.data["model"] == "gemini-2.5-flash-image" + + +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 2850a8e3..71c2d4dc 100644 --- a/tools/graphics/google_imagen.py +++ b/tools/graphics/google_imagen.py @@ -119,9 +119,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", @@ -177,13 +180,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. @@ -213,27 +314,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] = { diff --git a/tools/graphics/image_selector.py b/tools/graphics/image_selector.py index c3a3c716..3bab0bd8 100644 --- a/tools/graphics/image_selector.py +++ b/tools/graphics/image_selector.py @@ -242,6 +242,14 @@ class ImageSelector(BaseTool): props = tool.input_schema.get("properties", {}) if "query" in props and "query" not in adapted: adapted["query"] = adapted.get("prompt", "") + # The selector exposes a provider-neutral ``model_name`` field, + # while several providers call the same input ``model``. + if ( + "model_name" in adapted + and "model" in props + and "model" not in adapted + ): + adapted["model"] = adapted["model_name"] if "n" in adapted and "num_images" in props and "num_images" not in adapted: adapted["num_images"] = adapted["n"]