feat(providers): add Google GenerateContent and Imagen API packages

packages/capabilities/image-generation/pyproject.toml

@@ -27,6 +27,7 @@ Issues = "https://github.com/withceleste/celeste-python/issues"
 
 [tool.uv.sources]
 celeste-ai = { workspace = true }
+celeste-google = { workspace = true }
 
 [project.entry-points."celeste.packages"]
 image-generation = "celeste_image_generation:register_package"

packages/capabilities/image-generation/src/celeste_image_generation/client.py

@@ -35,24 +35,22 @@ def _parse_content(
         self,
         response_data: dict[str, Any],
         **parameters: Unpack[ImageGenerationParameters],
-    ) -> ImageArtifact:
+    ) -> ImageArtifact | list[ImageArtifact]:
         """Parse content from provider response."""
 
     @abstractmethod
     def _parse_finish_reason(
         self, response_data: dict[str, Any]
-    ) -> ImageGenerationFinishReason | None:
+    ) -> ImageGenerationFinishReason:
         """Parse finish reason from provider response."""
 
     def _create_inputs(
-        self,
-        *args: str,
-        prompt: str | None = None,
-        **parameters: Unpack[ImageGenerationParameters],
+        self, *args: str, **parameters: Unpack[ImageGenerationParameters]
     ) -> ImageGenerationInput:
         """Map positional arguments to Input type."""
         if args:
             return ImageGenerationInput(prompt=args[0])
+        prompt: str | None = parameters.get("prompt")
         if prompt is None:
             msg = (
                 "prompt is required (either as positional argument or keyword argument)"

packages/capabilities/image-generation/src/celeste_image_generation/io.py

@@ -16,7 +16,7 @@ class ImageGenerationFinishReason(FinishReason):
     Stores raw provider reason. Providers map their values in implementation.
     """
 
-    reason: str
+    reason: str | None = None
     message: str | None = None
 
 
@@ -30,12 +30,12 @@ class ImageGenerationUsage(Usage):
     input_tokens: int | None = None
     output_tokens: int | None = None
     reasoning_tokens: int | None = None
-    generated_images: int | None = None
+    num_images: int | None = None
     billed_units: float | None = None
 
 
-class ImageGenerationOutput(Output[ImageArtifact]):
-    """Output with ImageArtifact content."""
+class ImageGenerationOutput(Output[ImageArtifact | list[ImageArtifact]]):
+    """Output with ImageArtifact content (single or multiple)."""
 
     pass
 

packages/capabilities/image-generation/src/celeste_image_generation/parameters.py

@@ -2,34 +2,39 @@
 
 from enum import StrEnum
 
+from celeste.artifacts import ImageArtifact
 from celeste.parameters import Parameters
 
 
 class ImageGenerationParameter(StrEnum):
     """Unified parameter names for image generation capability."""
 
     ASPECT_RATIO = "aspect_ratio"
-    GUIDANCE = "guidance"
-    OUTPUT_FORMAT = "output_format"
+    NUM_IMAGES = "num_images"
     PARTIAL_IMAGES = "partial_images"
-    PROMPT_UPSAMPLING = "prompt_upsampling"
     QUALITY = "quality"
-    SAFETY_TOLERANCE = "safety_tolerance"
+    WATERMARK = "watermark"
+    REFERENCE_IMAGES = "reference_images"
+    PROMPT_UPSAMPLING = "prompt_upsampling"
     SEED = "seed"
+    SAFETY_TOLERANCE = "safety_tolerance"
+    OUTPUT_FORMAT = "output_format"
     STEPS = "steps"
-    WATERMARK = "watermark"
+    GUIDANCE = "guidance"
 
 
 class ImageGenerationParameters(Parameters):
     """Parameters for image generation."""
 
     aspect_ratio: str | None
-    guidance: float | None
-    output_format: str | None
+    num_images: int | None
     partial_images: int | None
-    prompt_upsampling: bool | None
     quality: str | None
-    safety_tolerance: int | None
+    watermark: bool | None
+    reference_images: list[ImageArtifact] | None
+    prompt_upsampling: bool | None
     seed: int | None
+    safety_tolerance: int | None
+    output_format: str | None
     steps: int | None
-    watermark: bool | None
+    guidance: float | None

packages/capabilities/image-generation/src/celeste_image_generation/providers/google/__init__.py

@@ -1,6 +1,11 @@
 """Google provider for image generation."""
 
 from .client import GoogleImageGenerationClient
-from .models import MODELS
+from .models import GEMINI_MODELS, IMAGEN_MODELS, MODELS
 
-__all__ = ["MODELS", "GoogleImageGenerationClient"]
+__all__ = [
+    "GEMINI_MODELS",
+    "IMAGEN_MODELS",
+    "MODELS",
+    "GoogleImageGenerationClient",
+]

packages/capabilities/image-generation/src/celeste_image_generation/providers/google/client.py

@@ -1,15 +1,10 @@
 """Google client implementation for image generation."""
 
-import base64
 from typing import Any, Unpack
 
 import httpx
-from pydantic import ConfigDict
 
 from celeste.artifacts import ImageArtifact
-from celeste.core import Provider
-from celeste.exceptions import ModelNotFoundError
-from celeste.mime_types import ApplicationMimeType, ImageMimeType
 from celeste.parameters import ParameterMapper
 from celeste_image_generation.client import ImageGenerationClient
 from celeste_image_generation.io import (
@@ -19,123 +14,69 @@
 )
 from celeste_image_generation.parameters import ImageGenerationParameters
 
-from . import config
+from .gemini import GeminiImageGenerationClient
+from .imagen import ImagenImageGenerationClient
+from .models import GEMINI_MODELS, IMAGEN_MODELS
 from .parameters import GOOGLE_PARAMETER_MAPPERS
 
+# Model ID → Client class mapping (extensible - add new model types here)
+GOOGLE_MODEL_MAP = {
+    **{m.id: ImagenImageGenerationClient for m in IMAGEN_MODELS},
+    **{m.id: GeminiImageGenerationClient for m in GEMINI_MODELS},
+}
+
 
 class GoogleImageGenerationClient(ImageGenerationClient):
     """Google client for image generation."""
 
-    model_config = ConfigDict(extra="allow")
+    _strategy: GeminiImageGenerationClient | ImagenImageGenerationClient | None = None
 
-    def model_post_init(self, __context: Any) -> None:
-        """Initialize API adapter based on model type."""
+    def model_post_init(self, __context: object) -> None:
+        """Initialize strategy based on model."""
         super().model_post_init(__context)
 
-        adapter_class, _ = _get_adapter_for_model(self.model.id)
-        self.api = adapter_class()
-        self.endpoint = self.api.endpoint(self.model.id)
+        StrategyClass = GOOGLE_MODEL_MAP[self.model.id]
+        strategy = StrategyClass(
+            model=self.model,
+            provider=self.provider,
+            capability=self.capability,
+            auth=self.auth,
+        )
+        object.__setattr__(self, "_strategy", strategy)
 
     @classmethod
     def parameter_mappers(cls) -> list[ParameterMapper]:
         return GOOGLE_PARAMETER_MAPPERS
 
     def _init_request(self, inputs: ImageGenerationInput) -> dict[str, Any]:
-        """Initialize request from Google API format."""
-        return self.api.build_request(inputs.prompt, {})
+        """Delegate to strategy."""
+        return self._strategy._init_request(inputs)  # type: ignore[union-attr]
 
     def _parse_usage(self, response_data: dict[str, Any]) -> ImageGenerationUsage:
-        """Parse usage from response."""
-        return self.api.parse_usage(response_data)
+        """Delegate to strategy."""
+        return self._strategy._parse_usage(response_data)  # type: ignore[union-attr]
 
     def _parse_content(
         self,
         response_data: dict[str, Any],
         **parameters: Unpack[ImageGenerationParameters],
-    ) -> ImageArtifact:
-        """Parse content from response."""
-        prediction = self.api.parse_response(response_data)
-
-        if prediction is None:
-            return ImageArtifact()
-
-        base64_data = prediction.get("bytesBase64Encoded") or prediction["data"]
-        mime_type = ImageMimeType(prediction.get("mimeType", "image/png"))
-        image_bytes = base64.b64decode(base64_data)
-
-        return ImageArtifact(data=image_bytes, mime_type=mime_type)
+    ) -> ImageArtifact | list[ImageArtifact]:
+        """Delegate to strategy."""
+        return self._strategy._parse_content(response_data, **parameters)  # type: ignore[union-attr]
 
     def _parse_finish_reason(
         self, response_data: dict[str, Any]
-    ) -> ImageGenerationFinishReason | None:
-        """Parse finish reason from response.
-
-        For Gemini models, extracts finishReason from candidates[0].
-        For Imagen models, returns None (not provided).
-        """
-        # Check if this is a Gemini response (has "candidates")
-        candidates = response_data.get("candidates")
-        if candidates:
-            candidate = candidates[0]
-            finish_reason_str = candidate.get("finishReason")
-            if finish_reason_str:
-                finish_message = candidate.get("finishMessage")
-                return ImageGenerationFinishReason(
-                    reason=finish_reason_str,
-                    message=finish_message,
-                )
-        # Imagen models don't provide finish reasons
-        return None
-
-    def _build_metadata(self, response_data: dict[str, Any]) -> dict[str, Any]:
-        """Build metadata dictionary from response data."""
-        # Filter content field before calling super
-        content_fields = {"candidates"}
-        filtered_data = {
-            k: v for k, v in response_data.items() if k not in content_fields
-        }
-        return super()._build_metadata(filtered_data)
+    ) -> ImageGenerationFinishReason:
+        """Delegate to strategy."""
+        return self._strategy._parse_finish_reason(response_data)  # type: ignore[union-attr]
 
     async def _make_request(
         self,
         request_body: dict[str, Any],
         **parameters: Unpack[ImageGenerationParameters],
     ) -> httpx.Response:
-        """Make HTTP request(s) and return response object."""
-        headers = {
-            **self.auth.get_headers(),
-            "Content-Type": ApplicationMimeType.JSON,
-        }
-
-        return await self.http_client.post(
-            f"{config.BASE_URL}{self.endpoint}",
-            headers=headers,
-            json_body=request_body,
-        )
-
-
-def _get_adapter_for_model(model_id: str) -> tuple[type, str]:
-    """Get adapter class and endpoint for model ID.
-
-    Returns:
-        Tuple of (adapter_class, endpoint_template).
-    """
-    from .models import GEMINI_MODELS, IMAGEN_MODELS
-
-    # Create sets for O(1) lookup (computed once per import)
-    imagen_model_ids = {model.id for model in IMAGEN_MODELS}
-    gemini_model_ids = {model.id for model in GEMINI_MODELS}
-
-    if model_id in imagen_model_ids:
-        from .imagen_api import ImagenAPIAdapter
-
-        return ImagenAPIAdapter, config.IMAGEN_ENDPOINT
-    if model_id in gemini_model_ids:
-        from .gemini_api import GeminiImageAPIAdapter
-
-        return GeminiImageAPIAdapter, config.GEMINI_ENDPOINT
-
-    raise ModelNotFoundError(model_id=model_id, provider=Provider.GOOGLE)
+        """Delegate to strategy."""
+        return await self._strategy._make_request(request_body, **parameters)  # type: ignore[union-attr]
 
 
 __all__ = ["GoogleImageGenerationClient"]

packages/capabilities/image-generation/src/celeste_image_generation/providers/google/gemini.py

@@ -0,0 +1,90 @@
+"""Gemini client for Google image generation."""
+
+import base64
+from typing import Any, Unpack
+
+from celeste_google.generate_content.client import GoogleGenerateContentClient
+
+from celeste.artifacts import ImageArtifact
+from celeste.mime_types import ImageMimeType
+from celeste.parameters import ParameterMapper
+from celeste_image_generation.client import ImageGenerationClient
+from celeste_image_generation.io import (
+    ImageGenerationFinishReason,
+    ImageGenerationInput,
+    ImageGenerationUsage,
+)
+from celeste_image_generation.parameters import ImageGenerationParameters
+
+
+class GeminiImageGenerationClient(GoogleGenerateContentClient, ImageGenerationClient):
+    """Google Gemini client for image generation."""
+
+    @classmethod
+    def parameter_mappers(cls) -> list[ParameterMapper]:
+        """Parameter mappers for Gemini image generation."""
+        return []  # Parameter mapping handled by GoogleImageGenerationClient wrapper
+
+    def _init_request(self, inputs: ImageGenerationInput) -> dict[str, Any]:
+        """Initialize request for Gemini image generation."""
+        return {
+            "contents": [{"parts": [{"text": inputs.prompt}]}],
+            "generationConfig": {
+                "responseModalities": ["Image"],
+                "imageConfig": {},
+            },
+        }
+
+    def _parse_usage(self, response_data: dict[str, Any]) -> ImageGenerationUsage:
+        """Parse usage from response."""
+        usage = super()._parse_usage(response_data)
+        candidates = response_data.get("candidates", [])
+        return ImageGenerationUsage(**usage, num_images=len(candidates))
+
+    def _parse_content(
+        self,
+        response_data: dict[str, Any],
+        **parameters: Unpack[ImageGenerationParameters],
+    ) -> ImageArtifact | list[ImageArtifact]:
+        """Parse content from response."""
+        candidates = super()._parse_content(response_data)
+        artifacts = []
+
+        for candidate in candidates:
+            content = candidate.get("content", {})
+            parts = content.get("parts", [])
+            for part in parts:
+                inline_data = part.get("inlineData", {})
+                base64_data = inline_data.get("data")
+
+                if base64_data:
+                    mime_type = ImageMimeType(inline_data.get("mimeType", "image/png"))
+                    image_bytes = base64.b64decode(base64_data)
+                    artifacts.append(
+                        ImageArtifact(data=image_bytes, mime_type=mime_type)
+                    )
+
+        if not artifacts:
+            return ImageArtifact()
+
+        if len(artifacts) == 1:
+            return artifacts[0]
+
+        return artifacts
+
+    def _parse_finish_reason(
+        self, response_data: dict[str, Any]
+    ) -> ImageGenerationFinishReason:
+        """Parse finish reason from response."""
+        finish_reason = super()._parse_finish_reason(response_data)
+        candidates = response_data.get("candidates", [])
+        finish_message = None
+        if candidates:
+            finish_message = candidates[0].get("finishMessage")
+        return ImageGenerationFinishReason(
+            reason=finish_reason.reason,
+            message=finish_message,
+        )
+
+
+__all__ = ["GeminiImageGenerationClient"]