codex-ai 0.2.0__tar.gz → 0.2.2__tar.gz

{codex_ai-0.2.0 → codex_ai-0.2.2}/CHANGELOG.md

@@ -4,6 +4,21 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
 
+## [0.2.2] - 2026-05-15
+
+### Fixed
+- Changed the Gemini image default model to the official API model id `gemini-2.5-flash-image`.
+
+## [0.2.1] - 2026-05-15
+
+### Added
+- Added explicit `GeminiProvider.generate_imagen_bytes()` for Imagen models through `generate_images`.
+- Added `ImagenGenerationProvider` and dispatcher delegation for explicit Imagen generation.
+
+### Fixed
+- Stopped passing image MIME values such as `image/webp` into Gemini `GenerateContentConfig.response_mime_type` on the Gemini image path.
+- Pinned `google-genai` to `1.68.0` to avoid accidental SDK API drift in the alpha image-generation contract.
+
 ## [0.2.0] - 2026-05-15
 
 ### Added

{codex_ai-0.2.0 → codex_ai-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: codex-ai
-Version: 0.2.0
+Version: 0.2.2
 Summary: Gemini-first and OpenAI provider helpers for Codex
 Project-URL: Homepage, https://github.com/codexdlc/codex-ai
 Project-URL: Documentation, https://codexdlc.github.io/codex-ai/
@@ -19,12 +19,12 @@ Requires-Python: >=3.12
 Requires-Dist: codex-core<0.4.0,>=0.2.2
 Requires-Dist: pydantic<3.0,>=2.0
 Provides-Extra: all
-Requires-Dist: google-genai>=1.0; extra == 'all'
+Requires-Dist: google-genai==1.68.0; extra == 'all'
 Requires-Dist: openai<2.0,>=1.0; extra == 'all'
 Provides-Extra: dev
 Requires-Dist: bandit>=1.7; extra == 'dev'
 Requires-Dist: detect-secrets>=1.5; extra == 'dev'
-Requires-Dist: google-genai>=1.0; extra == 'dev'
+Requires-Dist: google-genai==1.68.0; extra == 'dev'
 Requires-Dist: mypy>=1.10; extra == 'dev'
 Requires-Dist: openai<2.0,>=1.0; extra == 'dev'
 Requires-Dist: pip-audit>=2.7; extra == 'dev'
@@ -40,7 +40,7 @@ Requires-Dist: mkdocs-material>=9.0; extra == 'docs'
 Requires-Dist: mkdocs>=1.5; extra == 'docs'
 Requires-Dist: mkdocstrings[python]>=0.24; extra == 'docs'
 Provides-Extra: gemini
-Requires-Dist: google-genai>=1.0; extra == 'gemini'
+Requires-Dist: google-genai==1.68.0; extra == 'gemini'
 Provides-Extra: openai
 Requires-Dist: openai<2.0,>=1.0; extra == 'openai'
 Description-Content-Type: text/markdown
@@ -84,12 +84,24 @@ text = await gemini.generate_text("Write one short tavern rumor.")
 loot = await gemini.generate_json("Create one loot item.", schema=LootItem)
 image_bytes, content_type = await gemini.generate_image_bytes(
     "A fantasy clan banner, game icon style.",
+    model="gemini-2.5-flash-image",
     response_mime_type="image/webp",
 )
+
+imagen_bytes, imagen_content_type = await gemini.generate_imagen_bytes(
+    "A fantasy clan banner, game icon style.",
+    response_mime_type="image/jpeg",
+)
 ```
 
 `answer(prompt)` remains available as a compatibility wrapper for text generation.
 
+`generate_image_bytes()` targets Gemini image models through `generate_content` and treats
+`response_mime_type` as a preferred/fallback MIME type. It does not pass image MIME values
+to Gemini's text `response_mime_type` config field. Use `generate_imagen_bytes()` for
+Imagen models; that path uses `generate_images` and passes the requested MIME as
+`output_mime_type`.
+
 ## Router Pipeline
 
 ```python
@@ -117,7 +129,7 @@ response = await dispatcher.process("chat", text="Hello!")
 | Module | Extra | Description |
 | :--- | :--- | :--- |
 | `codex_ai.core` | - | Dispatcher, router, protocol types, sync wrapper, and shared exception contract |
-| `codex_ai.providers.gemini` | `[gemini]` | Google Gemini text, JSON, and image generation via `google-genai` |
+| `codex_ai.providers.gemini` | `[gemini]` | Google Gemini text, JSON, Gemini image, and Imagen generation via pinned `google-genai` |
 | `codex_ai.providers.openai` | `[openai]` | OpenAI Chat Completions text provider |
 
 ## Development

{codex_ai-0.2.0 → codex_ai-0.2.2}/README.md

@@ -37,12 +37,24 @@ text = await gemini.generate_text("Write one short tavern rumor.")
 loot = await gemini.generate_json("Create one loot item.", schema=LootItem)
 image_bytes, content_type = await gemini.generate_image_bytes(
     "A fantasy clan banner, game icon style.",
+    model="gemini-2.5-flash-image",
     response_mime_type="image/webp",
 )
+
+imagen_bytes, imagen_content_type = await gemini.generate_imagen_bytes(
+    "A fantasy clan banner, game icon style.",
+    response_mime_type="image/jpeg",
+)
 ```
 
 `answer(prompt)` remains available as a compatibility wrapper for text generation.
 
+`generate_image_bytes()` targets Gemini image models through `generate_content` and treats
+`response_mime_type` as a preferred/fallback MIME type. It does not pass image MIME values
+to Gemini's text `response_mime_type` config field. Use `generate_imagen_bytes()` for
+Imagen models; that path uses `generate_images` and passes the requested MIME as
+`output_mime_type`.
+
 ## Router Pipeline
 
 ```python
@@ -70,7 +82,7 @@ response = await dispatcher.process("chat", text="Hello!")
 | Module | Extra | Description |
 | :--- | :--- | :--- |
 | `codex_ai.core` | - | Dispatcher, router, protocol types, sync wrapper, and shared exception contract |
-| `codex_ai.providers.gemini` | `[gemini]` | Google Gemini text, JSON, and image generation via `google-genai` |
+| `codex_ai.providers.gemini` | `[gemini]` | Google Gemini text, JSON, Gemini image, and Imagen generation via pinned `google-genai` |
 | `codex_ai.providers.openai` | `[openai]` | OpenAI Chat Completions text provider |
 
 ## Development

{codex_ai-0.2.0 → codex_ai-0.2.2}/docs/en/architecture/providers/README.md

@@ -10,6 +10,7 @@ Gemini is the primary target and exposes direct methods for text, JSON, and imag
 await gemini.generate_text(...)
 await gemini.generate_json(...)
 await gemini.generate_image_bytes(...)
+await gemini.generate_imagen_bytes(...)
 ```
 
 OpenAI is kept as a text provider with the same `generate_text(...)` convenience.
@@ -22,6 +23,7 @@ PromptResult/String
       ├── GeminiProvider.generate_text(...)        -> str
       ├── GeminiProvider.generate_json(...)        -> dict | BaseModel
       ├── GeminiProvider.generate_image_bytes(...) -> tuple[bytes, str]
+      ├── GeminiProvider.generate_imagen_bytes(...) -> tuple[bytes, str]
       └── OpenAIProvider.generate_text(...)        -> str
 ```
 
@@ -44,5 +46,7 @@ LLMRouter builder -> PromptResult -> LLMDispatcher.process() -> provider.answer(
 
 - Gemini-specific capabilities are represented directly instead of being hidden behind a broad universal abstraction.
 - JSON generation uses provider-native JSON configuration and still validates locally with `json.loads` and optional Pydantic models.
-- Image generation returns raw bytes plus the actual MIME type returned by Gemini.
+- Gemini image generation and Imagen generation are separate explicit methods because they use different SDK calls.
+- `generate_image_bytes()` uses Gemini `generate_content` with image modality. Its `response_mime_type` is only a preferred/fallback MIME type.
+- `generate_imagen_bytes()` uses Imagen `generate_images` and passes the requested MIME as `output_mime_type`.
 - Anthropic, OpenRouter, and multi-provider failover are not active APIs in this alpha line.

{codex_ai-0.2.0 → codex_ai-0.2.2}/docs/en/architecture/providers/data_flow.md

@@ -33,6 +33,18 @@ prompt: str
     -> (bytes, actual_mime_type)
 ```
 
+`response_mime_type` is not passed to `GenerateContentConfig.response_mime_type` on this path; it is only a fallback content type when Gemini omits `inline_data.mime_type`.
+
+## Imagen Images
+
+```
+prompt: str
+    -> GeminiProvider.generate_imagen_bytes()
+    -> GenerateImagesConfig(output_mime_type=requested_mime)
+    -> first generated_images image
+    -> (bytes, actual_mime_type)
+```
+
 ## OpenAI Text
 
 ```

{codex_ai-0.2.0 → codex_ai-0.2.2}/docs/ru/architecture/providers/README.md

@@ -10,6 +10,7 @@ Gemini является основным направлением и дает п
 await gemini.generate_text(...)
 await gemini.generate_json(...)
 await gemini.generate_image_bytes(...)
+await gemini.generate_imagen_bytes(...)
 ```
 
 OpenAI оставлен как текстовый провайдер с `generate_text(...)`.
@@ -22,6 +23,7 @@ PromptResult/String
       ├── GeminiProvider.generate_text(...)        -> str
       ├── GeminiProvider.generate_json(...)        -> dict | BaseModel
       ├── GeminiProvider.generate_image_bytes(...) -> tuple[bytes, str]
+      ├── GeminiProvider.generate_imagen_bytes(...) -> tuple[bytes, str]
       └── OpenAIProvider.generate_text(...)        -> str
 ```
 
@@ -44,5 +46,7 @@ LLMRouter builder -> PromptResult -> LLMDispatcher.process() -> provider.answer(
 
 - Возможности Gemini представлены напрямую, без широкой универсальной абстракции.
 - JSON generation использует native JSON config провайдера и локальную проверку через `json.loads` и optional Pydantic schema.
-- Image generation возвращает raw bytes и фактический MIME type от Gemini.
+- Gemini image generation и Imagen generation разведены в отдельные явные методы, потому что они используют разные SDK calls.
+- `generate_image_bytes()` использует Gemini `generate_content` с image modality. Его `response_mime_type` является только preferred/fallback MIME type.
+- `generate_imagen_bytes()` использует Imagen `generate_images` и передает requested MIME как `output_mime_type`.
 - Anthropic, OpenRouter и multi-provider failover не являются активными API в этой alpha-линейке.

{codex_ai-0.2.0 → codex_ai-0.2.2}/docs/ru/architecture/providers/data_flow.md

@@ -33,6 +33,18 @@ prompt: str
     -> (bytes, actual_mime_type)
 ```
 
+`response_mime_type` в этом пути не передается в `GenerateContentConfig.response_mime_type`; он используется только как fallback content type, если Gemini не вернул `inline_data.mime_type`.
+
+## Imagen Images
+
+```
+prompt: str
+    -> GeminiProvider.generate_imagen_bytes()
+    -> GenerateImagesConfig(output_mime_type=requested_mime)
+    -> first generated_images image
+    -> (bytes, actual_mime_type)
+```
+
 ## OpenAI Text
 
 ```

{codex_ai-0.2.0 → codex_ai-0.2.2}/pyproject.toml

@@ -31,7 +31,7 @@ Issues = "https://github.com/codexdlc/codex-ai/issues"
 
 [project.optional-dependencies]
 openai = ["openai>=1.0,<2.0"]
-gemini = ["google-genai>=1.0"]
+gemini = ["google-genai==1.68.0"]
 all = [
     "codex-ai[openai,gemini]",
 ]

{codex_ai-0.2.0 → codex_ai-0.2.2}/src/codex_ai/__init__.py

@@ -2,6 +2,7 @@
 
 from codex_ai.core import (
     ImageGenerationProvider,
+    ImagenGenerationProvider,
     JsonGenerationProvider,
     LLMDispatcher,
     LLMMessage,
@@ -18,6 +19,7 @@ __all__ = [
     # Core
     "LLMDispatcher",
     "ImageGenerationProvider",
+    "ImagenGenerationProvider",
     "JsonGenerationProvider",
     "LLMMessage",
     "LLMProviderError",

{codex_ai-0.2.0 → codex_ai-0.2.2}/src/codex_ai/core/__init__.py

@@ -8,6 +8,7 @@ from .dispatcher import LLMDispatcher
 from .exceptions import LLMProviderError
 from .protocol import (
     ImageGenerationProvider,
+    ImagenGenerationProvider,
     JsonGenerationProvider,
     LLMMessage,
     LLMProviderProtocol,
@@ -22,6 +23,7 @@ __all__ = [
     "LLMDispatcher",
     "LLMProviderError",
     "ImageGenerationProvider",
+    "ImagenGenerationProvider",
     "JsonGenerationProvider",
     "LLMMessage",
     "LLMProviderProtocol",

{codex_ai-0.2.0 → codex_ai-0.2.2}/src/codex_ai/core/dispatcher.py

@@ -14,6 +14,7 @@ from typing import Any
 
 from .protocol import (
     ImageGenerationProvider,
+    ImagenGenerationProvider,
     JsonGenerationProvider,
     LLMProviderProtocol,
     PromptResult,
@@ -116,6 +117,33 @@ class LLMDispatcher:
             **kwargs,
         )
 
+    async def generate_imagen_bytes(
+        self,
+        prompt: str,
+        *,
+        model: str | None = None,
+        response_mime_type: str = "image/jpeg",
+        **kwargs: Any,
+    ) -> tuple[bytes, str]:
+        """
+        Generate Imagen bytes through a provider that supports Imagen generation.
+
+        This method is intentionally separate from ``generate_image_bytes()``
+        because Imagen uses ``generate_images`` and an image output MIME config.
+        """
+        if not isinstance(self._provider, ImagenGenerationProvider):
+            provider_name = type(self._provider).__name__
+            raise TypeError(
+                f"Provider {provider_name} does not support Imagen generation; expected generate_imagen_bytes(...)"
+            )
+
+        return await self._provider.generate_imagen_bytes(
+            prompt,
+            model=model,
+            response_mime_type=response_mime_type,
+            **kwargs,
+        )
+
     async def generate_text(
         self,
         prompt: PromptResult | str,

{codex_ai-0.2.0 → codex_ai-0.2.2}/src/codex_ai/core/protocol.py

@@ -8,6 +8,7 @@ LLMProviderProtocol — adapter contract for LLM backends (OpenAI, Gemini, etc.)
 TextGenerationProvider — optional adapter contract for direct text generation.
 JsonGenerationProvider — optional adapter contract for direct JSON generation.
 ImageGenerationProvider — optional adapter contract for binary image generation.
+ImagenGenerationProvider — optional adapter contract for Imagen image generation.
 PromptBuilder — type alias for async builder functions registered via LLMRouter.
 """
 
@@ -164,5 +165,37 @@ class ImageGenerationProvider(Protocol):
         ...
 
 
+@runtime_checkable
+class ImagenGenerationProvider(Protocol):
+    """
+    Optional adapter contract for providers that expose Imagen image generation.
+
+    This is separate from ``ImageGenerationProvider`` because Gemini image models
+    and Imagen models use different SDK methods and MIME configuration fields.
+    """
+
+    async def generate_imagen_bytes(
+        self,
+        prompt: str,
+        *,
+        model: str | None = None,
+        response_mime_type: str = "image/jpeg",
+        **kwargs: Any,
+    ) -> tuple[bytes, str]:
+        """
+        Generate an Imagen image and return raw bytes plus the actual content type.
+
+        Args:
+            prompt: Plain image-generation prompt.
+            model: Optional Imagen model override.
+            response_mime_type: Requested image MIME type passed to Imagen when supported.
+            **kwargs: Extra provider-specific kwargs.
+
+        Returns:
+            Tuple of ``(image_bytes, content_type)``.
+        """
+        ...
+
+
 # Callable registered via @LLMRouter.prompt(mode)
 PromptBuilder = Callable[..., Awaitable[PromptResult]]

{codex_ai-0.2.0 → codex_ai-0.2.2}/src/codex_ai/providers/gemini.py

@@ -8,6 +8,7 @@ Requires: ``pip install codex-ai[gemini]``
 
 from __future__ import annotations
 
+import base64
 import json
 from typing import Any, cast
 
@@ -27,7 +28,8 @@ from codex_ai.core.exceptions import LLMProviderError
 from codex_ai.core.protocol import PromptResult
 
 _DEFAULT_MODEL = "gemini-2.5-flash-lite"
-_DEFAULT_IMAGE_MODEL = "gemini-3.1-flash-image-preview"
+_DEFAULT_IMAGE_MODEL = "gemini-2.5-flash-image"
+_DEFAULT_IMAGEN_MODEL = "imagen-3.0-generate-002"
 
 
 class GeminiProvider:
@@ -39,7 +41,7 @@ class GeminiProvider:
     Args:
         api_key: Google AI API key.
         model: Gemini text model name. Defaults to ``"gemini-2.5-flash-lite"``.
-        image_model: Gemini image model name. Defaults to ``"gemini-3.1-flash-image-preview"``.
+        image_model: Gemini image model name. Defaults to ``"gemini-2.5-flash-image"``.
 
     Example:
         ```python
@@ -52,13 +54,20 @@ class GeminiProvider:
         ```
     """
 
-    def __init__(self, api_key: str, model: str = _DEFAULT_MODEL, image_model: str = _DEFAULT_IMAGE_MODEL) -> None:
+    def __init__(
+        self,
+        api_key: str,
+        model: str = _DEFAULT_MODEL,
+        image_model: str = _DEFAULT_IMAGE_MODEL,
+        imagen_model: str = _DEFAULT_IMAGEN_MODEL,
+    ) -> None:
         self._client = genai.Client(
             api_key=api_key,
             http_options=genai_types.HttpOptions(api_version="v1alpha"),
         )
         self._model = model
         self._image_model = image_model
+        self._imagen_model = imagen_model
 
     async def answer(self, prompt: PromptResult, **kw: Any) -> str:
         """
@@ -180,10 +189,16 @@ class GeminiProvider:
         **kwargs: Any,
     ) -> tuple[bytes, str]:
         """
-        Generate an image with Gemini and return raw bytes plus content type.
+        Generate an image with Gemini image models and return bytes plus content type.
+
+        This uses the Gemini ``generate_content`` image path with
+        ``GenerateContentConfig.response_modalities=[IMAGE]``. Use it for Gemini
+        image preview / flash-image / nano-banana style models.
 
-        ``response_mime_type`` is treated as the requested MIME type. The actual
-        MIME type returned by Gemini wins when present.
+        ``response_mime_type`` is a preferred/fallback MIME type only. Gemini's
+        ``GenerateContentConfig.response_mime_type`` accepts text response MIME
+        values, so image MIME values are not sent there. The actual MIME type
+        returned in ``inline_data.mime_type`` wins when present.
         """
         selected_model = model or self._image_model
         requested_mime = response_mime_type
@@ -194,7 +209,6 @@ class GeminiProvider:
 
         config = genai_types.GenerateContentConfig(
             response_modalities=[genai_types.Modality.IMAGE],
-            response_mime_type=requested_mime,
             **runtime_kw,
         )
 
@@ -215,6 +229,52 @@ class GeminiProvider:
         except Exception as exc:
             raise LLMProviderError(f"Gemini image generation error: {exc}") from exc
 
+    async def generate_imagen_bytes(
+        self,
+        prompt: str,
+        *,
+        model: str | None = None,
+        response_mime_type: str = "image/jpeg",
+        **kwargs: Any,
+    ) -> tuple[bytes, str]:
+        """
+        Generate an image with Imagen models and return bytes plus content type.
+
+        This uses the Imagen ``generate_images`` SDK path and passes
+        ``response_mime_type`` as ``GenerateImagesConfig.output_mime_type``.
+        Use it for ``imagen-*`` models, not Gemini flash-image / nano-banana
+        models.
+        """
+        selected_model = model or self._imagen_model
+        requested_mime = response_mime_type
+
+        runtime_kw = kwargs.copy()
+        runtime_kw.pop("model", None)
+        runtime_kw.pop("response_mime_type", None)
+        runtime_kw.pop("output_mime_type", None)
+
+        config = genai_types.GenerateImagesConfig(
+            output_mime_type=requested_mime,
+            **runtime_kw,
+        )
+
+        try:
+            response = await self._client.aio.models.generate_images(
+                model=selected_model,
+                prompt=prompt,
+                config=config,
+            )
+            image = self._extract_first_imagen_image(response, fallback_mime=requested_mime)
+            if image is not None:
+                return image
+
+            detail = self._describe_imagen_non_image_response(response)
+            raise LLMProviderError(f"Gemini Imagen generation did not return image data{detail}")
+        except LLMProviderError:
+            raise
+        except Exception as exc:
+            raise LLMProviderError(f"Gemini Imagen generation error: {exc}") from exc
+
     @staticmethod
     def _extract_first_inline_image(response: Any, *, fallback_mime: str) -> tuple[bytes, str] | None:
         for part in GeminiProvider._iter_response_parts(response):
@@ -228,6 +288,45 @@ class GeminiProvider:
 
         return None
 
+    @staticmethod
+    def _extract_first_imagen_image(response: Any, *, fallback_mime: str) -> tuple[bytes, str] | None:
+        for generated_image in getattr(response, "generated_images", None) or []:
+            image = getattr(generated_image, "image", None)
+            if image is None:
+                continue
+
+            data = getattr(image, "image_bytes", None)
+            if data is None:
+                data = getattr(image, "data", None)
+            if data is None:
+                continue
+
+            image_bytes = base64.b64decode(data) if isinstance(data, str) else bytes(data)
+
+            mime_type = getattr(image, "mime_type", None) or getattr(generated_image, "mime_type", None)
+            mime_type = mime_type or fallback_mime
+            return image_bytes, mime_type
+
+        return None
+
+    @staticmethod
+    def _describe_imagen_non_image_response(response: Any) -> str:
+        details: list[str] = []
+
+        for generated_image in getattr(response, "generated_images", None) or []:
+            rai_reason = getattr(generated_image, "rai_filtered_reason", None)
+            if rai_reason:
+                details.append(f"rai_filtered_reason={rai_reason}")
+
+            safety_attributes = getattr(generated_image, "safety_attributes", None)
+            if safety_attributes:
+                details.append(f"safety_attributes={safety_attributes}")
+
+        if not details:
+            return ""
+
+        return f": {'; '.join(details)}"
+
     @staticmethod
     def _describe_non_image_response(response: Any) -> str:
         details: list[str] = []

{codex_ai-0.2.0 → codex_ai-0.2.2}/tests/unit/core/test_dispatcher.py

@@ -131,6 +131,17 @@ class ImageProvider:
         self.calls.append((prompt, model, response_mime_type, kwargs))
         return b"image-bytes", "image/png"
 
+    async def generate_imagen_bytes(
+        self,
+        prompt: str,
+        *,
+        model: str | None = None,
+        response_mime_type: str = "image/jpeg",
+        **kwargs,
+    ) -> tuple[bytes, str]:
+        self.calls.append((prompt, model, response_mime_type, kwargs))
+        return b"imagen-bytes", "image/jpeg"
+
     async def generate_text(self, prompt: PromptResult | str, *, model: str | None = None, **kwargs) -> str:
         self.calls.append((prompt, model, kwargs))
         return "direct text"
@@ -167,6 +178,33 @@ async def test_dispatcher_generate_image_bytes_raises_for_unsupported_provider(m
     assert mock_provider.calls == []
 
 
+async def test_dispatcher_generate_imagen_bytes_delegates_to_imagen_provider():
+    provider = ImageProvider()
+    dispatcher = LLMDispatcher(provider=provider)
+
+    result = await dispatcher.generate_imagen_bytes(
+        "draw a castle",
+        model="imagen-model",
+        response_mime_type="image/jpeg",
+        seed=123,
+    )
+
+    assert result == (b"imagen-bytes", "image/jpeg")
+    assert provider.calls == [("draw a castle", "imagen-model", "image/jpeg", {"seed": 123})]
+
+
+async def test_dispatcher_generate_imagen_bytes_raises_for_unsupported_provider(mock_provider):
+    dispatcher = LLMDispatcher(provider=mock_provider)
+
+    with pytest.raises(
+        TypeError,
+        match=r"Provider MockProvider does not support Imagen generation; expected generate_imagen_bytes\(\.\.\.\)",
+    ):
+        await dispatcher.generate_imagen_bytes("draw a castle")
+
+    assert mock_provider.calls == []
+
+
 async def test_dispatcher_generate_text_delegates_to_text_provider():
     provider = ImageProvider()
     dispatcher = LLMDispatcher(provider=provider)

{codex_ai-0.2.0 → codex_ai-0.2.2}/tests/unit/core/test_protocol.py

@@ -3,6 +3,7 @@ from pydantic import ValidationError
 
 from codex_ai.core.protocol import (
     ImageGenerationProvider,
+    ImagenGenerationProvider,
     JsonGenerationProvider,
     LLMMessage,
     LLMProviderProtocol,
@@ -116,6 +117,25 @@ def test_object_without_generate_image_bytes_fails_image_provider_check():
     assert not isinstance(object(), ImageGenerationProvider)
 
 
+def test_imagen_generation_provider_structural_check():
+    class MockImagenProvider:
+        async def generate_imagen_bytes(
+            self,
+            prompt: str,
+            *,
+            model: str | None = None,
+            response_mime_type: str = "image/jpeg",
+            **kwargs,
+        ) -> tuple[bytes, str]:
+            return b"image", response_mime_type
+
+    assert isinstance(MockImagenProvider(), ImagenGenerationProvider)
+
+
+def test_object_without_generate_imagen_bytes_fails_imagen_provider_check():
+    assert not isinstance(object(), ImagenGenerationProvider)
+
+
 def test_text_generation_provider_structural_check():
     class MockTextProvider:
         async def generate_text(self, prompt: PromptResult | str, *, model: str | None = None, **kwargs) -> str:

{codex_ai-0.2.0 → codex_ai-0.2.2}/tests/unit/providers/test_gemini_provider.py

@@ -9,6 +9,7 @@ from pydantic import BaseModel
 from codex_ai.core.exceptions import LLMProviderError
 from codex_ai.core.protocol import (
     ImageGenerationProvider,
+    ImagenGenerationProvider,
     JsonGenerationProvider,
     LLMMessage,
     LLMProviderProtocol,
@@ -49,6 +50,12 @@ def test_gemini_provider_satisfies_image_generation_protocol():
     assert isinstance(provider, ImageGenerationProvider)
 
 
+def test_gemini_provider_satisfies_imagen_generation_protocol():
+    with patch("codex_ai.providers.gemini.genai_types"):
+        provider = GeminiProvider(api_key="x")
+    assert isinstance(provider, ImagenGenerationProvider)
+
+
 def test_gemini_provider_satisfies_text_generation_protocol():
     with patch("codex_ai.providers.gemini.genai_types"):
         provider = GeminiProvider(api_key="x")
@@ -355,6 +362,12 @@ def _image_response(data: bytes | bytearray = b"image", mime_type: str | None =
     return SimpleNamespace(parts=[part], text=None)
 
 
+def _imagen_response(data: bytes | str = b"image", mime_type: str | None = "image/jpeg") -> SimpleNamespace:
+    image = SimpleNamespace(image_bytes=data, mime_type=mime_type)
+    generated_image = SimpleNamespace(image=image)
+    return SimpleNamespace(generated_images=[generated_image])
+
+
 async def test_gemini_generate_image_bytes_uses_image_model_not_text_model():
     provider, mock_generate, _ = _make_provider()
     provider._model = "text-model"
@@ -371,6 +384,19 @@ async def test_gemini_generate_image_bytes_uses_image_model_not_text_model():
     assert kwargs["contents"] == "draw a castle"
 
 
+async def test_gemini_generate_image_bytes_default_image_model_matches_api_id():
+    provider, mock_generate, _ = _make_provider()
+    mock_generate.return_value = _image_response()
+
+    with patch("codex_ai.providers.gemini.genai_types") as mock_types:
+        mock_types.Modality.IMAGE = "IMAGE"
+        mock_types.GenerateContentConfig.return_value = MagicMock()
+        await provider.generate_image_bytes("draw a castle")
+
+    _, kwargs = mock_generate.call_args
+    assert kwargs["model"] == "gemini-2.5-flash-image"
+
+
 async def test_gemini_generate_image_bytes_model_override_wins():
     provider, mock_generate, _ = _make_provider()
     provider._image_model = "image-model"
@@ -385,7 +411,7 @@ async def test_gemini_generate_image_bytes_model_override_wins():
     assert kwargs["model"] == "explicit-image-model"
 
 
-async def test_gemini_generate_image_bytes_config_requests_image_modality_and_mime():
+async def test_gemini_generate_image_bytes_config_requests_image_modality_not_text_mime():
     provider, mock_generate, _ = _make_provider()
     mock_generate.return_value = _image_response()
 
@@ -396,7 +422,7 @@ async def test_gemini_generate_image_bytes_config_requests_image_modality_and_mi
 
     config_kwargs = mock_types.GenerateContentConfig.call_args.kwargs
     assert config_kwargs["response_modalities"] == ["IMAGE"]
-    assert config_kwargs["response_mime_type"] == "image/webp"
+    assert "response_mime_type" not in config_kwargs
     assert config_kwargs["seed"] == 7
 
 
@@ -474,3 +500,111 @@ async def test_gemini_generate_image_bytes_wraps_sdk_errors():
             await provider.generate_image_bytes("draw a castle")
 
     assert exc_info.value.__cause__ is original
+
+
+async def test_gemini_generate_imagen_bytes_uses_imagen_model_not_gemini_image_model():
+    provider, _, _ = _make_provider()
+    provider._image_model = "gemini-image-model"
+    provider._imagen_model = "imagen-model"
+    mock_generate_images = AsyncMock(return_value=_imagen_response())
+    provider._client.aio.models.generate_images = mock_generate_images
+
+    with patch("codex_ai.providers.gemini.genai_types") as mock_types:
+        mock_types.GenerateImagesConfig.return_value = MagicMock()
+        await provider.generate_imagen_bytes("draw a castle")
+
+    _, kwargs = mock_generate_images.call_args
+    assert kwargs["model"] == "imagen-model"
+    assert kwargs["prompt"] == "draw a castle"
+
+
+async def test_gemini_generate_imagen_bytes_model_override_wins():
+    provider, _, _ = _make_provider()
+    provider._imagen_model = "imagen-model"
+    mock_generate_images = AsyncMock(return_value=_imagen_response())
+    provider._client.aio.models.generate_images = mock_generate_images
+
+    with patch("codex_ai.providers.gemini.genai_types") as mock_types:
+        mock_types.GenerateImagesConfig.return_value = MagicMock()
+        await provider.generate_imagen_bytes("draw a castle", model="explicit-imagen-model")
+
+    _, kwargs = mock_generate_images.call_args
+    assert kwargs["model"] == "explicit-imagen-model"
+
+
+async def test_gemini_generate_imagen_bytes_config_sets_output_mime_type():
+    provider, _, _ = _make_provider()
+    provider._client.aio.models.generate_images = AsyncMock(return_value=_imagen_response())
+
+    with patch("codex_ai.providers.gemini.genai_types") as mock_types:
+        mock_types.GenerateImagesConfig.return_value = MagicMock()
+        await provider.generate_imagen_bytes("draw a castle", response_mime_type="image/jpeg", seed=7)
+
+    config_kwargs = mock_types.GenerateImagesConfig.call_args.kwargs
+    assert config_kwargs["output_mime_type"] == "image/jpeg"
+    assert config_kwargs["seed"] == 7
+
+
+async def test_gemini_generate_imagen_bytes_returns_image_bytes_and_actual_mime():
+    provider, _, _ = _make_provider()
+    provider._client.aio.models.generate_images = AsyncMock(
+        return_value=_imagen_response(data=b"jpeg-bytes", mime_type="image/jpeg")
+    )
+
+    with patch("codex_ai.providers.gemini.genai_types") as mock_types:
+        mock_types.GenerateImagesConfig.return_value = MagicMock()
+        result = await provider.generate_imagen_bytes("draw a castle", response_mime_type="image/png")
+
+    assert result == (b"jpeg-bytes", "image/jpeg")
+
+
+async def test_gemini_generate_imagen_bytes_decodes_base64_image_bytes():
+    provider, _, _ = _make_provider()
+    provider._client.aio.models.generate_images = AsyncMock(
+        return_value=_imagen_response(data="anBlZy1ieXRlcw==", mime_type="image/jpeg")
+    )
+
+    with patch("codex_ai.providers.gemini.genai_types") as mock_types:
+        mock_types.GenerateImagesConfig.return_value = MagicMock()
+        result = await provider.generate_imagen_bytes("draw a castle")
+
+    assert result == (b"jpeg-bytes", "image/jpeg")
+
+
+async def test_gemini_generate_imagen_bytes_falls_back_to_requested_mime_when_missing():
+    provider, _, _ = _make_provider()
+    provider._client.aio.models.generate_images = AsyncMock(
+        return_value=_imagen_response(data=b"image-bytes", mime_type=None)
+    )
+
+    with patch("codex_ai.providers.gemini.genai_types") as mock_types:
+        mock_types.GenerateImagesConfig.return_value = MagicMock()
+        result = await provider.generate_imagen_bytes("draw a castle", response_mime_type="image/png")
+
+    assert result == (b"image-bytes", "image/png")
+
+
+async def test_gemini_generate_imagen_bytes_raises_when_no_image_data():
+    provider, _, _ = _make_provider()
+    generated_image = SimpleNamespace(image=SimpleNamespace(image_bytes=None, mime_type=None))
+    provider._client.aio.models.generate_images = AsyncMock(
+        return_value=SimpleNamespace(generated_images=[generated_image])
+    )
+
+    with patch("codex_ai.providers.gemini.genai_types") as mock_types:
+        mock_types.GenerateImagesConfig.return_value = MagicMock()
+        with pytest.raises(LLMProviderError, match="Imagen generation did not return image data"):
+            await provider.generate_imagen_bytes("draw a castle")
+
+
+async def test_gemini_generate_imagen_bytes_wraps_sdk_errors():
+    provider, _, _ = _make_provider()
+    original = RuntimeError("quota exceeded")
+    provider._client.aio.models.generate_images = AsyncMock(side_effect=original)
+
+    with patch("codex_ai.providers.gemini.genai_types") as mock_types:
+        mock_types.GenerateImagesConfig.return_value = MagicMock()
+        with pytest.raises(LLMProviderError, match="Gemini Imagen generation error") as exc_info:
+            await provider.generate_imagen_bytes("draw a castle")
+
+    assert exc_info.value.__cause__ is original

{codex_ai-0.2.0 → codex_ai-0.2.2}/tests/unit/test_public_api.py

@@ -20,6 +20,7 @@ def test_top_level_core_exports():
 
     assert hasattr(codex_ai, "LLMDispatcher")
     assert hasattr(codex_ai, "ImageGenerationProvider")
+    assert hasattr(codex_ai, "ImagenGenerationProvider")
     assert hasattr(codex_ai, "JsonGenerationProvider")
     assert hasattr(codex_ai, "LLMRouter")
     assert hasattr(codex_ai, "LLMMessage")

{codex_ai-0.2.0 → codex_ai-0.2.2}/uv.lock

@@ -344,9 +344,9 @@ requires-dist = [
     { name = "bandit", marker = "extra == 'dev'", specifier = ">=1.7" },
     { name = "codex-core", specifier = ">=0.2.2,<0.4.0" },
     { name = "detect-secrets", marker = "extra == 'dev'", specifier = ">=1.5" },
-    { name = "google-genai", marker = "extra == 'all'", specifier = ">=1.0" },
-    { name = "google-genai", marker = "extra == 'dev'", specifier = ">=1.0" },
-    { name = "google-genai", marker = "extra == 'gemini'", specifier = ">=1.0" },
+    { name = "google-genai", marker = "extra == 'all'", specifier = "==1.68.0" },
+    { name = "google-genai", marker = "extra == 'dev'", specifier = "==1.68.0" },
+    { name = "google-genai", marker = "extra == 'gemini'", specifier = "==1.68.0" },
     { name = "mike", marker = "extra == 'docs'", specifier = ">=2.0" },
     { name = "mkdocs", marker = "extra == 'docs'", specifier = ">=1.5" },
     { name = "mkdocs-include-markdown-plugin", marker = "extra == 'docs'" },
@@ -622,7 +622,7 @@ requests = [
 
 [[package]]
 name = "google-genai"
-version = "2.3.0"
+version = "1.68.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "anyio" },
@@ -636,9 +636,9 @@ dependencies = [
     { name = "typing-extensions" },
     { name = "websockets" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/02/8e/dfa4b34dd4c0baffccf6466fc68d6d35011662d43e7d79accb902320db74/google_genai-2.3.0.tar.gz", hash = "sha256:e877c750a4ccacdd9928fc3aa8ca8820ce85cade0ca51bd83feceacf5959b579", size = 546930, upload-time = "2026-05-15T06:22:36.264Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/9c/2c/f059982dbcb658cc535c81bbcbe7e2c040d675f4b563b03cdb01018a4bc3/google_genai-1.68.0.tar.gz", hash = "sha256:ac30c0b8bc630f9372993a97e4a11dae0e36f2e10d7c55eacdca95a9fa14ca96", size = 511285, upload-time = "2026-03-18T01:03:18.243Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/b4/6e/aa6b30b09f58b946750fc4089c5248fbd3576f746e0e818d88633559dc84/google_genai-2.3.0-py3-none-any.whl", hash = "sha256:89d3c71c9f5f5b931b405b88a5837aea2bd4d27ed90323b9599f5760bbb91d92", size = 805484, upload-time = "2026-05-15T06:22:34.247Z" },
+    { url = "https://files.pythonhosted.org/packages/84/de/7d3ee9c94b74c3578ea4f88d45e8de9405902f857932334d81e89bce3dfa/google_genai-1.68.0-py3-none-any.whl", hash = "sha256:a1bc9919c0e2ea2907d1e319b65471d3d6d58c54822039a249fe1323e4178d15", size = 750912, upload-time = "2026-03-18T01:03:15.983Z" },
 ]
 
 [[package]]