diff --git a/nanobot/providers/image_generation.py b/nanobot/providers/image_generation.py index 9e9332fd..b3843b11 100644 --- a/nanobot/providers/image_generation.py +++ b/nanobot/providers/image_generation.py @@ -33,13 +33,18 @@ _AIHUBMIX_ASPECT_RATIO_SIZES = { } _GEMINI_DEFAULT_TIMEOUT_S = 120.0 _GEMINI_IMAGEN_ASPECT_RATIOS = {"1:1", "9:16", "16:9", "3:4", "4:3"} -# Aspect ratios documented for every Gemini Flash image (generateContent) model. -# The extreme ratios (1:4, 4:1, 1:8, 8:1) are only listed for the 3.1 Flash / -# Flash Lite tables, so they are left out to avoid sending an unsupported value -# to 2.5 Flash Image or 3.1 Pro Image. -_GEMINI_FLASH_ASPECT_RATIOS = { +# Aspect ratios documented for every Gemini image model using generateContent. +_GEMINI_FLASH_COMMON_ASPECT_RATIOS = { "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9", } +# Gemini 3.1 Flash and Flash Lite additionally accept extreme aspect ratios. +_GEMINI_31_FLASH_ASPECT_RATIOS = { + *_GEMINI_FLASH_COMMON_ASPECT_RATIOS, + "1:4", + "4:1", + "1:8", + "8:1", +} # Gemini 3 Pro image models accept these sizes. Gemini 3.1 Flash adds 512, # while Gemini 3.1 Flash Lite supports only 1K. _GEMINI_3_IMAGE_SIZES = {"1K", "2K", "4K"} @@ -778,11 +783,12 @@ def _gemini_flash_image_config( ) -> dict[str, str]: """Build the ``responseFormat.image`` config for Gemini Flash image models. - Aspect ratio applies to all Flash image models; image size is only honored - by Gemini 3+ image models (``gemini-2.5-flash-image`` ignores it). + Capabilities are model-specific: Gemini 3.1 Flash variants support four + additional extreme ratios, while configurable image sizes are limited to + the documented Gemini 3 image model families. """ config: dict[str, str] = {} - if aspect_ratio and aspect_ratio in _GEMINI_FLASH_ASPECT_RATIOS: + if aspect_ratio and aspect_ratio in _gemini_flash_supported_aspect_ratios(model): config["aspectRatio"] = aspect_ratio if image_size: normalized = image_size.strip().upper() @@ -791,6 +797,19 @@ def _gemini_flash_image_config( return config +def _gemini_flash_supported_aspect_ratios(model: str) -> set[str]: + """Return the documented aspect ratios for a generateContent image model.""" + normalized = model.lower() + if ( + "gemini-3.1-flash-lite-image" in normalized + or "gemini-3.1-flash-image" in normalized + ): + return _GEMINI_31_FLASH_ASPECT_RATIOS + if "gemini-" in normalized and "image" in normalized: + return _GEMINI_FLASH_COMMON_ASPECT_RATIOS + return set() + + def _gemini_flash_supported_image_sizes(model: str) -> set[str]: """Return the ``imageSize`` values documented for a Flash-path model. @@ -802,7 +821,7 @@ def _gemini_flash_supported_image_sizes(model: str) -> set[str]: return _GEMINI_31_FLASH_LITE_IMAGE_SIZES if "gemini-3.1-flash-image" in normalized: return _GEMINI_31_FLASH_IMAGE_SIZES - if "gemini-3" in normalized: + if "gemini-3-pro-image" in normalized: return _GEMINI_3_IMAGE_SIZES return set() diff --git a/tests/providers/test_image_generation.py b/tests/providers/test_image_generation.py index 7fe63e43..0ec35c0e 100644 --- a/tests/providers/test_image_generation.py +++ b/tests/providers/test_image_generation.py @@ -480,10 +480,39 @@ async def test_gemini_flash_2_0_drops_image_size() -> None: assert image_config == {"aspectRatio": "16:9"} +@pytest.mark.parametrize( + ("model", "aspect_ratio", "expected"), + [ + ("gemini-3.1-flash-image", "1:8", {"aspectRatio": "1:8"}), + ("gemini-3.1-flash-lite-image", "4:1", {"aspectRatio": "4:1"}), + ("gemini-3-pro-image", "1:8", None), + ("gemini-2.5-flash-image", "4:1", None), + ], +) +@pytest.mark.asyncio +async def test_gemini_flash_scopes_extreme_aspect_ratios_by_model( + model: str, + aspect_ratio: str, + expected: dict[str, str] | None, +) -> None: + fake = FakeClient(_gemini_flash_image_response()) + client = GeminiImageGenerationClient(api_key="AIza-test", client=fake) # type: ignore[arg-type] + + await client.generate( + prompt="draw a cat", + model=model, + aspect_ratio=aspect_ratio, + ) + + response_format = fake.calls[0]["json"]["generationConfig"].get("responseFormat") + assert response_format == ({"image": expected} if expected else None) + + @pytest.mark.parametrize( ("model", "image_size", "expected"), [ ("gemini-3-pro-image", "512", None), + ("gemini-3-pro", "2K", None), ("gemini-3.1-flash-lite-image", "2K", None), ("gemini-3.1-flash-lite-image", "1K", {"imageSize": "1K"}), ("gemini-3.1-flash-image", "512", {"imageSize": "512"}),