From bff5cb76b1355573a94ba7ad84d817a19163cdac Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 15:08:26 +0100 Subject: [PATCH 1/5] feat(text-generation): add Google Generate Content API provider Add standalone provider package for Google Generate Content API with mixin pattern for capability-agnostic reuse. ## Client (GoogleGenerateContentClient mixin) - HTTP POST/streaming to /v1beta/models/{model}:generateContent endpoint - Usage parsing: promptTokenCount, candidatesTokenCount, totalTokenCount - Content extraction from candidates[0].content.parts - Finish reason mapping (STOP, MAX_TOKENS, SAFETY, etc.) - API key authentication via x-goog-api-key header ## Parameters - TemperatureMapper: temperature float [0.0-2.0] - TopPMapper: nucleus sampling top_p - TopKMapper: top-k sampling - MaxTokensMapper: maxOutputTokens - StopSequencesMapper: custom stop sequences - SystemInstructionMapper: system_instruction content block - ResponseFormatMapper: JSON schema via responseSchema - Supports single BaseModel and list[BaseModel] - Uses StrictRefResolvingJsonSchemaGenerator for schema generation - ToolsMapper: function declarations for tool use - ToolChoiceMapper: tool_config with function_calling_config ## Streaming (GoogleGenerateContentStream mixin) - SSE event parsing for streamGenerateContent endpoint - Text delta extraction from candidates[0].content.parts - Finish reason and usage tracking in final events ## Config - API base URL: https://generativelanguage.googleapis.com - API version: v1beta - Endpoints: generateContent, streamGenerateContent --- packages/providers/google/pyproject.toml | 21 ++ .../google/src/celeste_google/__init__.py | 1 + .../generate_content/__init__.py | 1 + .../celeste_google/generate_content/client.py | 124 +++++++ .../celeste_google/generate_content/config.py | 23 ++ .../generate_content/parameters.py | 310 ++++++++++++++++++ .../generate_content/streaming.py | 54 +++ .../google/src/celeste_google/py.typed | 0 8 files changed, 534 insertions(+) create mode 100644 packages/providers/google/pyproject.toml create mode 100644 packages/providers/google/src/celeste_google/__init__.py create mode 100644 packages/providers/google/src/celeste_google/generate_content/__init__.py create mode 100644 packages/providers/google/src/celeste_google/generate_content/client.py create mode 100644 packages/providers/google/src/celeste_google/generate_content/config.py create mode 100644 packages/providers/google/src/celeste_google/generate_content/parameters.py create mode 100644 packages/providers/google/src/celeste_google/generate_content/streaming.py create mode 100644 packages/providers/google/src/celeste_google/py.typed diff --git a/packages/providers/google/pyproject.toml b/packages/providers/google/pyproject.toml new file mode 100644 index 00000000..4e1ed976 --- /dev/null +++ b/packages/providers/google/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "celeste-google" +version = "0.3.0" +description = "Google (Gemini API) provider package for Celeste AI" +authors = [{name = "Kamilbenkirane", email = "kamil@withceleste.ai"}] +license = {text = "Apache-2.0"} +requires-python = ">=3.12" +dependencies = ["celeste-ai", "httpx", "google-auth"] + +[tool.uv.sources] +celeste-ai = { workspace = true } + +[project.entry-points."celeste.providers"] +google = "celeste_google:register_provider" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/celeste_google"] diff --git a/packages/providers/google/src/celeste_google/__init__.py b/packages/providers/google/src/celeste_google/__init__.py new file mode 100644 index 00000000..ec93ff5d --- /dev/null +++ b/packages/providers/google/src/celeste_google/__init__.py @@ -0,0 +1 @@ +"""Google provider package for Celeste AI.""" diff --git a/packages/providers/google/src/celeste_google/generate_content/__init__.py b/packages/providers/google/src/celeste_google/generate_content/__init__.py new file mode 100644 index 00000000..a18780ee --- /dev/null +++ b/packages/providers/google/src/celeste_google/generate_content/__init__.py @@ -0,0 +1 @@ +"""Google GenerateContent API provider package.""" diff --git a/packages/providers/google/src/celeste_google/generate_content/client.py b/packages/providers/google/src/celeste_google/generate_content/client.py new file mode 100644 index 00000000..eaa078b3 --- /dev/null +++ b/packages/providers/google/src/celeste_google/generate_content/client.py @@ -0,0 +1,124 @@ +"""Google GenerateContent API client with shared implementation.""" + +from collections.abc import AsyncIterator +from typing import Any + +import httpx + +from celeste.core import UsageField +from celeste.io import FinishReason +from celeste.mime_types import ApplicationMimeType + +from . import config + + +class GoogleGenerateContentClient: + """Mixin for GenerateContent API capabilities. + + Provides shared implementation for all capabilities using the GenerateContent API: + - _make_request() - HTTP POST to generateContent + - _make_stream_request() - HTTP streaming to streamGenerateContent + - _parse_usage() - Extract usage dict from usageMetadata + - _parse_content() - Extract parts array from first candidate + - _parse_finish_reason() - Extract finish reason string from candidates + - _build_metadata() - Filter content fields + + Capability clients extend parsing methods via super() to wrap/transform results. + + Usage: + class GoogleTextGenerationClient(GoogleGenerateContentClient, TextGenerationClient): + def _parse_content(self, response_data, **parameters): + parts = super()._parse_content(response_data) + text = parts[0].get("text") or "" + return self._transform_output(text, **parameters) + """ + + async def _make_request( + self, + request_body: dict[str, Any], + **parameters: Any, + ) -> httpx.Response: + """Make HTTP request to generateContent endpoint.""" + endpoint = config.GoogleGenerateContentEndpoint.GENERATE_CONTENT.format( + model_id=self.model.id # type: ignore[attr-defined] + ) + headers = { + **self.auth.get_headers(), # type: ignore[attr-defined] + "Content-Type": ApplicationMimeType.JSON, + } + return await self.http_client.post( # type: ignore[attr-defined,no-any-return] + f"{config.BASE_URL}{endpoint}", + headers=headers, + json_body=request_body, + ) + + def _make_stream_request( + self, + request_body: dict[str, Any], + **parameters: Any, + ) -> AsyncIterator[dict[str, Any]]: + """Make streaming request to streamGenerateContent endpoint.""" + endpoint = config.GoogleGenerateContentEndpoint.STREAM_GENERATE_CONTENT.format( + model_id=self.model.id # type: ignore[attr-defined] + ) + headers = { + **self.auth.get_headers(), # type: ignore[attr-defined] + "Content-Type": ApplicationMimeType.JSON, + } + return self.http_client.stream_post( # type: ignore[attr-defined,no-any-return] + f"{config.BASE_URL}{endpoint}", + headers=headers, + json_body=request_body, + ) + + @staticmethod + def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | None]: + """Map Google Gemini usage fields to unified names. + + Shared by client and streaming across all capabilities. + """ + return { + UsageField.INPUT_TOKENS: usage_data.get("promptTokenCount"), + UsageField.OUTPUT_TOKENS: usage_data.get("candidatesTokenCount"), + UsageField.TOTAL_TOKENS: usage_data.get("totalTokenCount"), + UsageField.REASONING_TOKENS: usage_data.get("thoughtsTokenCount"), + } + + def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: + """Extract usage data from Gemini usageMetadata.""" + usage_metadata = response_data.get("usageMetadata", {}) + return self.map_usage_fields(usage_metadata) + + def _parse_content(self, response_data: dict[str, Any]) -> Any: + """Return all candidates from response. + + Returns list of candidate objects that capability clients extract content from. + """ + candidates = response_data.get("candidates", []) + if not candidates: + msg = "No candidates in response" + raise ValueError(msg) + return candidates + + def _parse_finish_reason(self, response_data: dict[str, Any]) -> FinishReason: + """Extract finish reason from Gemini candidates. + + Returns FinishReason that capability clients wrap in their specific type. + """ + candidates = response_data.get("candidates", []) + if not candidates: + reason = None + else: + reason = candidates[0].get("finishReason") + return FinishReason(reason=reason) + + def _build_metadata(self, response_data: dict[str, Any]) -> dict[str, Any]: + """Build metadata dictionary, filtering out content fields.""" + 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) # type: ignore[misc,no-any-return] + + +__all__ = ["GoogleGenerateContentClient"] diff --git a/packages/providers/google/src/celeste_google/generate_content/config.py b/packages/providers/google/src/celeste_google/generate_content/config.py new file mode 100644 index 00000000..1075ea7a --- /dev/null +++ b/packages/providers/google/src/celeste_google/generate_content/config.py @@ -0,0 +1,23 @@ +"""Configuration for Google GenerateContent API.""" + +from enum import StrEnum + + +class GoogleGenerateContentEndpoint(StrEnum): + """Endpoints for GenerateContent API.""" + + GENERATE_CONTENT = "/v1beta/models/{model_id}:generateContent" + STREAM_GENERATE_CONTENT = "/v1beta/models/{model_id}:streamGenerateContent?alt=sse" + COUNT_TOKENS = "/v1beta/models/{model_id}:countTokens" + EMBED_CONTENT = "/v1beta/models/{model_id}:embedContent" + BATCH_EMBED_CONTENTS = "/v1beta/models/{model_id}:batchEmbedContents" + LIST_MODELS = "/v1beta/models" + GET_MODEL = "/v1beta/models/{model_id}" + UPLOAD_FILE = "/upload/v1beta/files" + LIST_FILES = "/v1beta/files" + GET_FILE = "/v1beta/files/{file_id}" + DELETE_FILE = "/v1beta/files/{file_id}" + BATCH_GENERATE_CONTENT = "/v1beta/models/{model_id}:batchGenerateContent" + + +BASE_URL = "https://generativelanguage.googleapis.com" diff --git a/packages/providers/google/src/celeste_google/generate_content/parameters.py b/packages/providers/google/src/celeste_google/generate_content/parameters.py new file mode 100644 index 00000000..831cf883 --- /dev/null +++ b/packages/providers/google/src/celeste_google/generate_content/parameters.py @@ -0,0 +1,310 @@ +"""Google Gemini API parameter mappers.""" + +import base64 +import json +from typing import Any, get_args, get_origin + +from pydantic import BaseModel, TypeAdapter + +from celeste.artifacts import ImageArtifact +from celeste.mime_types import ApplicationMimeType, ImageMimeType +from celeste.models import Model +from celeste.parameters import ParameterMapper +from celeste.types import StructuredOutput + + +class TemperatureMapper(ParameterMapper): + """Map temperature to Google generationConfig.temperature field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform temperature into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("generationConfig", {})["temperature"] = validated_value + return request + + +class MaxOutputTokensMapper(ParameterMapper): + """Map max_tokens to Google generationConfig.maxOutputTokens field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform max_tokens into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("generationConfig", {})["maxOutputTokens"] = validated_value + return request + + +class ThinkingBudgetMapper(ParameterMapper): + """Map thinkingBudget to Google generationConfig.thinkingConfig.thinkingBudget field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform thinkingBudget into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("generationConfig", {}).setdefault("thinkingConfig", {})[ + "thinkingBudget" + ] = validated_value + return request + + +class ThinkingLevelMapper(ParameterMapper): + """Map thinkingLevel to Google generationConfig.thinkingConfig.thinkingLevel field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform thinkingLevel into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("generationConfig", {}).setdefault("thinkingConfig", {})[ + "thinkingLevel" + ] = validated_value + return request + + +class AspectRatioMapper(ParameterMapper): + """Map aspect_ratio to Google generationConfig.imageConfig.aspectRatio field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform aspect_ratio into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("generationConfig", {}).setdefault("imageConfig", {})[ + "aspectRatio" + ] = validated_value + return request + + +class ImageSizeMapper(ParameterMapper): + """Map image_size to Google generationConfig.imageConfig.imageSize field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform image_size into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("generationConfig", {}).setdefault("imageConfig", {})[ + "imageSize" + ] = validated_value + return request + + +class MediaContentMapper(ParameterMapper): + """Map reference_images to Google contents.parts field.""" + + def _build_image_part(self, image: ImageArtifact) -> dict[str, Any]: + """Build inline_data part from image artifact.""" + if image.url: + return {"file_data": {"file_uri": image.url}} + elif image.data: + base64_data = ( + base64.b64encode(image.data).decode("utf-8") + if isinstance(image.data, bytes) + else image.data + ) + return { + "inline_data": { + "mime_type": image.mime_type or ImageMimeType.JPEG, + "data": base64_data, + } + } + elif image.path: + with open(image.path, "rb") as f: + image_bytes = f.read() + base64_data = base64.b64encode(image_bytes).decode("utf-8") + return { + "inline_data": { + "mime_type": image.mime_type or ImageMimeType.JPEG, + "data": base64_data, + } + } + else: + msg = "ImageArtifact must have url, data, or path" + raise ValueError(msg) + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform reference_images into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + # Convert list of ImageArtifact to list of image parts + image_parts = [self._build_image_part(img) for img in validated_value] + + # Add image parts before text in contents[0].parts + if "contents" in request and len(request["contents"]) > 0: + parts = request["contents"][0].get("parts", []) + # Find text part and insert images before it + text_index = next( + (i for i, part in enumerate(parts) if "text" in part), len(parts) + ) + # Insert image parts before text + parts[text_index:text_index] = image_parts + request["contents"][0]["parts"] = parts + + return request + + +class WebSearchMapper(ParameterMapper): + """Map web_search to Google tools field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform web_search into provider request.""" + validated_value = self._validate_value(value, model) + if not validated_value: + return request + + request["tools"] = [{"google_search": {}}] + return request + + +class OutputSchemaMapper(ParameterMapper): + """Map output_schema to Google generationConfig.responseJsonSchema field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform output_schema into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + schema = self._convert_to_google_schema(validated_value) + + config = request.setdefault("generationConfig", {}) + config["responseJsonSchema"] = schema + config["responseMimeType"] = ApplicationMimeType.JSON + + return request + + def parse_output( + self, content: StructuredOutput, value: object | None + ) -> StructuredOutput: + """Parse JSON to BaseModel using Pydantic's TypeAdapter.""" + if value is None: + return content + + # If content is already a BaseModel, return it unchanged + if isinstance(content, BaseModel): + return content + if isinstance(content, list) and content and isinstance(content[0], BaseModel): + return content + + parsed = json.loads(content) if isinstance(content, str) else content + + # For list[T], handle various formats Google might return + origin = get_origin(value) + if origin is list and isinstance(parsed, dict): + # Google returns arrays directly in JSON schema, but might wrap in object + if "items" in parsed: + parsed = parsed["items"] + else: + # If it's a dict but not wrapped, try to extract array values + parsed = list(parsed.values()) if parsed else [] + + return TypeAdapter(value).validate_python(parsed) + + def _convert_to_google_schema(self, output_schema: Any) -> dict[str, Any]: # noqa: ANN401 + """Convert Pydantic BaseModel or list[BaseModel] to Google JSON Schema format.""" + origin = get_origin(output_schema) + if origin is list: + inner_type = get_args(output_schema)[0] + items_schema = inner_type.model_json_schema() + defs = items_schema.get("$defs", {}) + items_schema_clean = {k: v for k, v in items_schema.items() if k != "$defs"} + json_schema = {"type": "array", "items": items_schema_clean} + if defs: + json_schema["$defs"] = defs + else: + json_schema = output_schema.model_json_schema() + + json_schema = self._remove_unsupported_fields(json_schema) + return json_schema + + def _remove_unsupported_fields(self, schema: dict[str, Any]) -> dict[str, Any]: + """Remove unsupported metadata fields from schema.""" + result: dict[str, Any] = {} + + for key, value in schema.items(): + if key == "title": + continue + + if isinstance(value, dict): + result[key] = self._remove_unsupported_fields(value) + elif isinstance(value, list): + result[key] = [ + self._remove_unsupported_fields(item) + if isinstance(item, dict) + else item + for item in value + ] + else: + result[key] = value + + return result + + +__all__ = [ + "AspectRatioMapper", + "ImageSizeMapper", + "MaxOutputTokensMapper", + "MediaContentMapper", + "OutputSchemaMapper", + "TemperatureMapper", + "ThinkingBudgetMapper", + "ThinkingLevelMapper", + "WebSearchMapper", +] diff --git a/packages/providers/google/src/celeste_google/generate_content/streaming.py b/packages/providers/google/src/celeste_google/generate_content/streaming.py new file mode 100644 index 00000000..4d71d666 --- /dev/null +++ b/packages/providers/google/src/celeste_google/generate_content/streaming.py @@ -0,0 +1,54 @@ +"""Google GenerateContent SSE parsing for streaming.""" + +from typing import Any + +from .client import GoogleGenerateContentClient + + +class GoogleGenerateContentStream: + """Mixin for GenerateContent SSE parsing. + + Provides shared implementation for all capabilities using GenerateContent streaming: + - _parse_chunk() - Parse SSE event into raw chunk dict + + Capability streams extend via super() to wrap results in typed Chunks. + + Usage: + class GoogleTextGenerationStream(GoogleGenerateContentStream, TextGenerationStream): + def _parse_chunk(self, event): + raw = super()._parse_chunk(event) + if not raw: + return None + return TextGenerationChunk(...) + """ + + def _parse_chunk(self, event: dict[str, Any]) -> dict[str, Any] | None: + """Parse SSE event into raw chunk data.""" + candidates = event.get("candidates", []) + if not candidates: + return None + + candidate = candidates[0] + content = candidate.get("content", {}) + parts = content.get("parts", []) + + text_delta = parts[0].get("text") if parts else None + finish_reason = candidate.get("finishReason") + + usage = None + usage_data = event.get("usageMetadata") + if usage_data: + usage = GoogleGenerateContentClient.map_usage_fields(usage_data) + + if not text_delta and not finish_reason: + return None + + return { + "content": text_delta or "", + "finish_reason": finish_reason, + "usage": usage, + "raw_event": event, + } + + +__all__ = ["GoogleGenerateContentStream"] diff --git a/packages/providers/google/src/celeste_google/py.typed b/packages/providers/google/src/celeste_google/py.typed new file mode 100644 index 00000000..e69de29b From 79e277b20482a986239ad0bb0b604fb836954a28 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 15:31:27 +0100 Subject: [PATCH 2/5] feat(text-generation): migrate Google provider to use GenerateContent API mixin - Update GoogleTextGenerationClient to extend GoogleGenerateContentClient mixin - Simplify client by delegating HTTP, usage parsing, and metadata to API layer - Update parameter mappers to extend base API mappers - Update streaming to extend base API stream class - Add celeste-google as workspace dependency --- .../text-generation/pyproject.toml | 1 + .../providers/google/client.py | 120 ++--------- .../providers/google/parameters.py | 197 ++---------------- .../providers/google/streaming.py | 135 +++--------- 4 files changed, 71 insertions(+), 382 deletions(-) diff --git a/packages/capabilities/text-generation/pyproject.toml b/packages/capabilities/text-generation/pyproject.toml index 0db8e1bb..cec99f42 100644 --- a/packages/capabilities/text-generation/pyproject.toml +++ b/packages/capabilities/text-generation/pyproject.toml @@ -28,6 +28,7 @@ Issues = "https://github.com/withceleste/celeste-python/issues" [tool.uv.sources] celeste-ai = { workspace = true } celeste-anthropic = { workspace = true } +celeste-google = { workspace = true } [project.entry-points."celeste.packages"] text-generation = "celeste_text_generation:register_package" diff --git a/packages/capabilities/text-generation/src/celeste_text_generation/providers/google/client.py b/packages/capabilities/text-generation/src/celeste_text_generation/providers/google/client.py index 73cc9556..58558111 100644 --- a/packages/capabilities/text-generation/src/celeste_text_generation/providers/google/client.py +++ b/packages/capabilities/text-generation/src/celeste_text_generation/providers/google/client.py @@ -1,13 +1,11 @@ """Google client implementation for text generation.""" -from collections.abc import AsyncIterator from typing import Any, Unpack -import httpx -from pydantic import BaseModel +from celeste_google.generate_content.client import GoogleGenerateContentClient -from celeste.mime_types import ApplicationMimeType from celeste.parameters import ParameterMapper +from celeste.types import StructuredOutput from celeste_text_generation.client import TextGenerationClient from celeste_text_generation.io import ( TextGenerationFinishReason, @@ -16,12 +14,11 @@ ) from celeste_text_generation.parameters import TextGenerationParameters -from . import config from .parameters import GOOGLE_PARAMETER_MAPPERS from .streaming import GoogleTextGenerationStream -class GoogleTextGenerationClient(TextGenerationClient): +class GoogleTextGenerationClient(GoogleGenerateContentClient, TextGenerationClient): """Google client for text generation.""" @classmethod @@ -30,116 +27,41 @@ def parameter_mappers(cls) -> list[ParameterMapper]: def _init_request(self, inputs: TextGenerationInput) -> dict[str, Any]: """Initialize request from Google contents array format.""" - contents = [ - { - "role": "user", - "parts": [{"text": inputs.prompt}], - } - ] - - return {"contents": contents} + return { + "contents": [ + { + "role": "user", + "parts": [{"text": inputs.prompt}], + } + ] + } def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage: """Parse usage from response.""" - usage_metadata = response_data.get("usageMetadata", {}) - - return TextGenerationUsage( - input_tokens=usage_metadata.get("promptTokenCount"), - output_tokens=usage_metadata.get("candidatesTokenCount"), - total_tokens=usage_metadata.get("totalTokenCount"), - reasoning_tokens=usage_metadata.get("thoughtsTokenCount"), - ) + usage = super()._parse_usage(response_data) + return TextGenerationUsage(**usage) def _parse_content( self, response_data: dict[str, Any], **parameters: Unpack[TextGenerationParameters], - ) -> str | BaseModel: + ) -> StructuredOutput: """Parse content from response.""" - candidates = response_data.get("candidates", []) - if not candidates: - msg = "No candidates in response" - raise ValueError(msg) - - candidate = candidates[0] - content = candidate.get("content", {}) - parts = content.get("parts", []) - - if not parts: - msg = "No parts in candidate content" - raise ValueError(msg) - - text_part = parts[0] - text = text_part.get("text") or "" - - return self._transform_output(text, **parameters) + candidates = super()._parse_content(response_data) + parts = candidates[0].get("content", {}).get("parts", []) + text = parts[0].get("text") if parts else "" + return self._transform_output(text or "", **parameters) def _parse_finish_reason( self, response_data: dict[str, Any] - ) -> TextGenerationFinishReason | None: + ) -> TextGenerationFinishReason: """Parse finish reason from response.""" - candidates = response_data.get("candidates", []) - if not candidates: - return None - - candidate = candidates[0] - finish_reason_str = candidate.get("finishReason") - - if not finish_reason_str: - return None - - return TextGenerationFinishReason(reason=finish_reason_str) - - 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) - - async def _make_request( - self, - request_body: dict[str, Any], - **parameters: Unpack[TextGenerationParameters], - ) -> httpx.Response: - """Make HTTP request(s) and return response object.""" - endpoint = config.ENDPOINT.format(model_id=self.model.id) - - headers = { - **self.auth.get_headers(), - "Content-Type": ApplicationMimeType.JSON, - } - - return await self.http_client.post( - f"{config.BASE_URL}{endpoint}", - headers=headers, - json_body=request_body, - ) + finish_reason = super()._parse_finish_reason(response_data) + return TextGenerationFinishReason(reason=finish_reason.reason) def _stream_class(self) -> type[GoogleTextGenerationStream]: """Return the Stream class for this client.""" return GoogleTextGenerationStream - def _make_stream_request( - self, - request_body: dict[str, Any], - **parameters: Unpack[TextGenerationParameters], - ) -> AsyncIterator[dict[str, Any]]: - """Make HTTP streaming request and return async iterator of events.""" - stream_endpoint = config.STREAM_ENDPOINT.format(model_id=self.model.id) - - headers = { - **self.auth.get_headers(), - "Content-Type": ApplicationMimeType.JSON, - } - - return self.http_client.stream_post( - f"{config.BASE_URL}{stream_endpoint}", - headers=headers, - json_body=request_body, - ) - __all__ = ["GoogleTextGenerationClient"] diff --git a/packages/capabilities/text-generation/src/celeste_text_generation/providers/google/parameters.py b/packages/capabilities/text-generation/src/celeste_text_generation/providers/google/parameters.py index 2f756073..afe79fcf 100644 --- a/packages/capabilities/text-generation/src/celeste_text_generation/providers/google/parameters.py +++ b/packages/capabilities/text-generation/src/celeste_text_generation/providers/google/parameters.py @@ -1,198 +1,45 @@ -"""Google parameter mappers for text generation.""" - -from typing import Any, get_args, get_origin - -from pydantic import BaseModel, TypeAdapter +"""Google Gemini parameter mappers for text generation.""" + +from celeste_google.generate_content.parameters import ( + MaxOutputTokensMapper as _MaxOutputTokensMapper, +) +from celeste_google.generate_content.parameters import ( + OutputSchemaMapper as _OutputSchemaMapper, +) +from celeste_google.generate_content.parameters import ( + TemperatureMapper as _TemperatureMapper, +) +from celeste_google.generate_content.parameters import ( + ThinkingBudgetMapper as _ThinkingBudgetMapper, +) +from celeste_google.generate_content.parameters import ( + ThinkingLevelMapper as _ThinkingLevelMapper, +) from celeste.core import Parameter -from celeste.mime_types import ApplicationMimeType -from celeste.models import Model from celeste.parameters import ParameterMapper from celeste_text_generation.parameters import TextGenerationParameter -class TemperatureMapper(ParameterMapper): - """Map temperature parameter to Google generationConfig.""" - +class TemperatureMapper(_TemperatureMapper): name = Parameter.TEMPERATURE - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform temperature into provider request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - request.setdefault("generationConfig", {})["temperature"] = validated_value - return request - - -class MaxTokensMapper(ParameterMapper): - """Map max_tokens parameter to Google generationConfig.maxOutputTokens.""" +class MaxTokensMapper(_MaxOutputTokensMapper): name = Parameter.MAX_TOKENS - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform max_tokens into provider request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - request.setdefault("generationConfig", {})["maxOutputTokens"] = validated_value - return request - - -class ThinkingBudgetMapper(ParameterMapper): - """Map thinking_budget parameter to Google generationConfig.thinkingConfig.thinkingBudget.""" +class ThinkingBudgetMapper(_ThinkingBudgetMapper): name = TextGenerationParameter.THINKING_BUDGET - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform thinking_budget into provider request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - request.setdefault("generationConfig", {}).setdefault("thinkingConfig", {})[ - "thinkingBudget" - ] = validated_value - return request - - -class ThinkingLevelMapper(ParameterMapper): - """Map thinking_level parameter to Google generationConfig.thinkingConfig.thinkingLevel.""" +class ThinkingLevelMapper(_ThinkingLevelMapper): name = TextGenerationParameter.THINKING_LEVEL - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform thinking_level into provider request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - request.setdefault("generationConfig", {}).setdefault("thinkingConfig", {})[ - "thinkingLevel" - ] = validated_value - return request - - -class OutputSchemaMapper(ParameterMapper): - """Map output_schema parameter to Google generationConfig.responseJsonSchema.""" +class OutputSchemaMapper(_OutputSchemaMapper): name = TextGenerationParameter.OUTPUT_SCHEMA - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform response_model into provider request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - schema = self._convert_to_google_schema(validated_value) - - config = request.setdefault("generationConfig", {}) - config["responseJsonSchema"] = schema - config["responseMimeType"] = ApplicationMimeType.JSON - - return request - - def parse_output(self, content: str, value: object | None) -> str | BaseModel: - """Parse JSON string to BaseModel instance if output_schema provided. - - Args: - content: Raw text content (JSON string when output_schema is set). - value: Original output_schema parameter value. - - Returns: - BaseModel instance if value provided, otherwise str unchanged. - """ - if value is None: - return content - - return TypeAdapter(value).validate_json(content) - - def _convert_to_google_schema(self, output_schema: Any) -> dict[str, Any]: # noqa: ANN401 - """Convert Pydantic BaseModel or list[BaseModel] to Google JSON Schema format.""" - origin = get_origin(output_schema) - if origin is list: - inner_type = get_args(output_schema)[0] - items_schema = inner_type.model_json_schema() - # Extract $defs from items_schema and move to root (JSON Schema requires $defs at root) - defs = items_schema.get("$defs", {}) - items_schema_clean = {k: v for k, v in items_schema.items() if k != "$defs"} - json_schema = {"type": "array", "items": items_schema_clean} - if defs: - json_schema["$defs"] = defs - else: - json_schema = output_schema.model_json_schema() - - json_schema = self._remove_unsupported_fields(json_schema) - return self._uppercase_types(json_schema) - - def _uppercase_types(self, schema: dict[str, Any]) -> dict[str, Any]: - """Recursively uppercase all 'type' field values in schema.""" - result: dict[str, Any] = {} - - for key, value in schema.items(): - if key == "type" and isinstance(value, str): - result[key] = value.upper() - elif isinstance(value, dict): - result[key] = self._uppercase_types(value) - elif isinstance(value, list): - result[key] = [ - self._uppercase_types(item) if isinstance(item, dict) else item - for item in value - ] - else: - result[key] = value - - return result - - def _remove_unsupported_fields(self, schema: dict[str, Any]) -> dict[str, Any]: - """Remove unsupported metadata fields from schema.""" - result: dict[str, Any] = {} - - for key, value in schema.items(): - if key == "title": - continue - - if isinstance(value, dict): - result[key] = self._remove_unsupported_fields(value) - elif isinstance(value, list): - result[key] = [ - self._remove_unsupported_fields(item) - if isinstance(item, dict) - else item - for item in value - ] - else: - result[key] = value - - return result - GOOGLE_PARAMETER_MAPPERS: list[ParameterMapper] = [ TemperatureMapper(), diff --git a/packages/capabilities/text-generation/src/celeste_text_generation/providers/google/streaming.py b/packages/capabilities/text-generation/src/celeste_text_generation/providers/google/streaming.py index de60696f..031f53bd 100644 --- a/packages/capabilities/text-generation/src/celeste_text_generation/providers/google/streaming.py +++ b/packages/capabilities/text-generation/src/celeste_text_generation/providers/google/streaming.py @@ -3,6 +3,9 @@ from collections.abc import Callable from typing import Any, Unpack +from celeste_google.generate_content.streaming import GoogleGenerateContentStream + +from celeste.types import StructuredOutput from celeste_text_generation.io import ( TextGenerationChunk, TextGenerationFinishReason, @@ -13,151 +16,67 @@ from celeste_text_generation.streaming import TextGenerationStream -class GoogleTextGenerationStream(TextGenerationStream): +class GoogleTextGenerationStream(GoogleGenerateContentStream, TextGenerationStream): """Google streaming for text generation.""" def __init__( self, sse_iterator: Any, # noqa: ANN401 - transform_output: Callable[..., object], + transform_output: Callable[..., StructuredOutput], **parameters: Unpack[TextGenerationParameters], ) -> None: - """Initialize stream with output transformation support. - - Args: - sse_iterator: Server-Sent Events iterator. - transform_output: Function to transform accumulated content (e.g., JSON → BaseModel). - **parameters: Parameters passed to stream() for output transformation. - """ + """Initialize stream with output transformation support.""" super().__init__(sse_iterator, **parameters) self._transform_output = transform_output def _parse_chunk(self, event: dict[str, Any]) -> TextGenerationChunk | None: - """Parse SSE event into Chunk. - - Extract text delta from candidates[0].content.parts[0].text. - Extract finishReason and usageMetadata if present. - Return None if no text delta (filter lifecycle events). - """ - # Extract candidates array - candidates = event.get("candidates", []) - if not candidates: - return None - - candidate = candidates[0] - content = candidate.get("content", {}) - parts = content.get("parts", []) - - # Extract text delta - text_delta = None - if parts: - text_part = parts[0] - text_delta = text_part.get("text") - - # If no text delta, this is likely a lifecycle event - filter it - if not text_delta: + """Parse SSE event into typed Chunk.""" + raw = super()._parse_chunk(event) + if not raw: return None - # Extract finish reason if present - finish_reason: TextGenerationFinishReason | None = None - finish_reason_str = candidate.get("finishReason") - if finish_reason_str: - finish_reason = TextGenerationFinishReason(reason=finish_reason_str) - - # Extract usage metadata if present (store in chunk metadata for later) - usage: TextGenerationUsage | None = None - usage_metadata = event.get("usageMetadata") - if usage_metadata: - usage = TextGenerationUsage( - input_tokens=usage_metadata.get("promptTokenCount"), - output_tokens=usage_metadata.get("candidatesTokenCount"), - total_tokens=usage_metadata.get("totalTokenCount"), - reasoning_tokens=usage_metadata.get("thoughtsTokenCount"), - ) + usage = TextGenerationUsage(**raw["usage"]) if raw["usage"] else None + finish_reason = ( + TextGenerationFinishReason(reason=raw["finish_reason"]) + if raw["finish_reason"] + else None + ) return TextGenerationChunk( - content=text_delta, + content=raw["content"], finish_reason=finish_reason, usage=usage, + metadata={"raw_event": raw["raw_event"]}, ) def _parse_usage(self, chunks: list[TextGenerationChunk]) -> TextGenerationUsage: - """Parse usage from chunks. - - Google provides usageMetadata in the final chunk(s). - Accumulate usage from all chunks, prioritizing later chunks for totals. - """ - if not chunks: - return TextGenerationUsage() - - # Usage metadata is typically in the final chunk - final_chunk = chunks[-1] - if final_chunk.usage: - # Return final chunk usage directly (contains complete usageMetadata) - return final_chunk.usage - - # Fallback: check metadata if stored there - usage_metadata = final_chunk.metadata.get("usageMetadata") - if usage_metadata: - return TextGenerationUsage( - input_tokens=usage_metadata.get("promptTokenCount"), - output_tokens=usage_metadata.get("candidatesTokenCount"), - total_tokens=usage_metadata.get("totalTokenCount"), - reasoning_tokens=usage_metadata.get("thoughtsTokenCount"), - ) - - # Accumulate usage from chunks that have usage metadata - total_input_tokens = 0 - total_output_tokens = 0 - total_reasoning_tokens = 0 - total_total_tokens = 0 - - for chunk in chunks: + """Extract usage from final chunk.""" + for chunk in reversed(chunks): if chunk.usage: - if chunk.usage.input_tokens is not None: - total_input_tokens = chunk.usage.input_tokens # Use latest value - if chunk.usage.output_tokens is not None: - total_output_tokens = chunk.usage.output_tokens # Use latest value - if chunk.usage.reasoning_tokens is not None: - total_reasoning_tokens = ( - chunk.usage.reasoning_tokens - ) # Use latest value - if chunk.usage.total_tokens is not None: - total_total_tokens = chunk.usage.total_tokens # Use latest value - - return TextGenerationUsage( - input_tokens=total_input_tokens if total_input_tokens > 0 else None, - output_tokens=total_output_tokens if total_output_tokens > 0 else None, - total_tokens=total_total_tokens if total_total_tokens > 0 else None, - reasoning_tokens=total_reasoning_tokens - if total_reasoning_tokens > 0 - else None, - ) + return chunk.usage + return TextGenerationUsage() def _parse_output( self, chunks: list[TextGenerationChunk], **parameters: Unpack[TextGenerationParameters], ) -> TextGenerationOutput: - """Assemble chunks into final output with structured output support. - - Concatenates text chunks, then applies parameter transformations - (e.g., JSON → BaseModel if output_schema provided). - """ - # Concatenate text chunks + """Assemble chunks into final output.""" content = "".join(chunk.content for chunk in chunks) - - # Apply parameter transformations (e.g., JSON → BaseModel if output_schema provided) content = self._transform_output(content, **parameters) usage = self._parse_usage(chunks) finish_reason = chunks[-1].finish_reason if chunks else None + raw_events = [ + c.metadata["raw_event"] for c in chunks if c.metadata.get("raw_event") + ] + return TextGenerationOutput( content=content, usage=usage, finish_reason=finish_reason, - metadata={}, + metadata={"raw_response": raw_events}, ) From bc49d6d2313d8ed296e4d070943d96a032fd35be Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 15:34:20 +0100 Subject: [PATCH 3/5] feat(providers): add Google Imagen API mixin ## Client (GoogleImagenClient mixin) - HTTP POST to /v1beta/models/{model_id}:predict endpoint - Parse predictions[] array from response - No usage metadata (Imagen API doesn't provide it) - Metadata filtering for safety attributes ## Parameters - SampleCountMapper: parameters.sampleCount for num_images - AspectRatioMapper: parameters.aspectRatio - ImageSizeMapper: parameters.imageSize ## Config - API base URL: https://generativelanguage.googleapis.com - Endpoint: /v1beta/models/{model_id}:predict --- .../src/celeste_google/imagen/__init__.py | 1 + .../src/celeste_google/imagen/client.py | 82 +++++++++++++++++++ .../src/celeste_google/imagen/config.py | 12 +++ .../src/celeste_google/imagen/parameters.py | 67 +++++++++++++++ 4 files changed, 162 insertions(+) create mode 100644 packages/providers/google/src/celeste_google/imagen/__init__.py create mode 100644 packages/providers/google/src/celeste_google/imagen/client.py create mode 100644 packages/providers/google/src/celeste_google/imagen/config.py create mode 100644 packages/providers/google/src/celeste_google/imagen/parameters.py diff --git a/packages/providers/google/src/celeste_google/imagen/__init__.py b/packages/providers/google/src/celeste_google/imagen/__init__.py new file mode 100644 index 00000000..4fd720c5 --- /dev/null +++ b/packages/providers/google/src/celeste_google/imagen/__init__.py @@ -0,0 +1 @@ +"""Google Imagen API provider package.""" diff --git a/packages/providers/google/src/celeste_google/imagen/client.py b/packages/providers/google/src/celeste_google/imagen/client.py new file mode 100644 index 00000000..1d9b85ef --- /dev/null +++ b/packages/providers/google/src/celeste_google/imagen/client.py @@ -0,0 +1,82 @@ +"""Google Imagen API client with shared implementation.""" + +from typing import Any + +import httpx + +from celeste.io import FinishReason +from celeste.mime_types import ApplicationMimeType + +from . import config + + +class GoogleImagenClient: + """Mixin for Imagen API capabilities. + + Provides shared implementation for capabilities using the Imagen API: + - _make_request() - HTTP POST to :predict endpoint + - _parse_usage() - Returns empty dict (Imagen doesn't provide usage) + - _parse_content() - Extract predictions[] array + - _parse_finish_reason() - Returns None (Imagen doesn't provide finish reasons) + - _build_metadata() - Filter out predictions content + + Capability clients extend parsing methods via super() to wrap/transform results. + + Usage: + class GoogleImageGenerationClient(GoogleImagenClient, ImageGenerationClient): + def _parse_content(self, response_data, **parameters): + predictions = super()._parse_content(response_data) + # Extract image from predictions[0]... + """ + + async def _make_request( + self, + request_body: dict[str, Any], + **parameters: Any, + ) -> httpx.Response: + """Make HTTP request to Imagen :predict endpoint.""" + endpoint = config.GoogleImagenEndpoint.CREATE_IMAGE.format( + model_id=self.model.id # type: ignore[attr-defined] + ) + headers = { + **self.auth.get_headers(), # type: ignore[attr-defined] + "Content-Type": ApplicationMimeType.JSON, + } + return await self.http_client.post( # type: ignore[attr-defined,no-any-return] + f"{config.BASE_URL}{endpoint}", + headers=headers, + json_body=request_body, + ) + + def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: + """Imagen API doesn't provide usage metadata.""" + return {} + + def _parse_content(self, response_data: dict[str, Any]) -> Any: + """Parse predictions from response. + + Returns predictions array that capability clients extract from. + """ + predictions = response_data.get("predictions", []) + if not predictions: + msg = "No predictions in response" + raise ValueError(msg) + return predictions + + def _parse_finish_reason(self, response_data: dict[str, Any]) -> FinishReason: + """Imagen API doesn't provide finish reasons.""" + return FinishReason(reason=None) + + def _build_metadata(self, response_data: dict[str, Any]) -> dict[str, Any]: + """Build metadata, filtering out predictions content. + + Keeps metadata like raiFilteredReason, safety attributes. + """ + content_fields = {"predictions"} + filtered_data = { + k: v for k, v in response_data.items() if k not in content_fields + } + return super()._build_metadata(filtered_data) # type: ignore[misc,no-any-return] + + +__all__ = ["GoogleImagenClient"] diff --git a/packages/providers/google/src/celeste_google/imagen/config.py b/packages/providers/google/src/celeste_google/imagen/config.py new file mode 100644 index 00000000..f447051f --- /dev/null +++ b/packages/providers/google/src/celeste_google/imagen/config.py @@ -0,0 +1,12 @@ +"""Configuration for Google Imagen API.""" + +from enum import StrEnum + + +class GoogleImagenEndpoint(StrEnum): + """Endpoints for Imagen API.""" + + CREATE_IMAGE = "/v1beta/models/{model_id}:predict" + + +BASE_URL = "https://generativelanguage.googleapis.com" diff --git a/packages/providers/google/src/celeste_google/imagen/parameters.py b/packages/providers/google/src/celeste_google/imagen/parameters.py new file mode 100644 index 00000000..fd24bbd5 --- /dev/null +++ b/packages/providers/google/src/celeste_google/imagen/parameters.py @@ -0,0 +1,67 @@ +"""Google Imagen API parameter mappers.""" + +from typing import Any + +from celeste.models import Model +from celeste.parameters import ParameterMapper + + +class SampleCountMapper(ParameterMapper): + """Map num_images to Google Imagen parameters.sampleCount field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform num_images into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("parameters", {})["sampleCount"] = validated_value + return request + + +class AspectRatioMapper(ParameterMapper): + """Map aspect_ratio to Google Imagen parameters.aspectRatio field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform aspect_ratio into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("parameters", {})["aspectRatio"] = validated_value + return request + + +class ImageSizeMapper(ParameterMapper): + """Map image_size to Google Imagen parameters.imageSize field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform image_size into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("parameters", {})["imageSize"] = validated_value + return request + + +__all__ = [ + "AspectRatioMapper", + "ImageSizeMapper", + "SampleCountMapper", +] From 0c20b822a527bf25889eadd2a2f9a9db4b274a81 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 15:36:52 +0100 Subject: [PATCH 4/5] feat(image-generation): migrate Google provider to API mixins MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Strategy Pattern - GoogleImageGenerationClient delegates to model-specific strategies - GeminiImageGenerationClient uses GoogleGenerateContentClient mixin - ImagenImageGenerationClient uses GoogleImagenClient mixin - Model ID → Strategy mapping for automatic dispatch ## Gemini Strategy - Extends GoogleGenerateContentClient for multimodal image generation - Uses responseModalities: ['Image'] for image output - Parses inlineData from candidates[].content.parts[] ## Imagen Strategy - Extends GoogleImagenClient for Imagen API - Parses predictions[] array from response - Maps to ImageArtifact with base64 data ## Parameters - Split into gemini_parameters.py and imagen_parameters.py - Extend base API mappers from celeste-google package ## Other Changes - Add celeste-google as workspace dependency - Remove old adapter pattern (gemini_api.py, imagen_api.py, config.py) - Update tests for new strategy pattern --- .../image-generation/pyproject.toml | 1 + .../src/celeste_image_generation/client.py | 10 +- .../src/celeste_image_generation/io.py | 8 +- .../celeste_image_generation/parameters.py | 25 ++-- .../providers/google/__init__.py | 9 +- .../providers/google/client.py | 123 +++++------------- .../providers/google/config.py | 10 -- .../providers/google/gemini.py | 90 +++++++++++++ .../providers/google/gemini_api.py | 83 ------------ .../providers/google/gemini_parameters.py | 35 +++++ .../providers/google/imagen.py | 90 +++++++++++++ .../providers/google/imagen_api.py | 67 ---------- .../providers/google/imagen_parameters.py | 35 +++++ .../providers/google/models.py | 19 +-- .../providers/google/parameters.py | 67 +--------- .../src/celeste_image_generation/streaming.py | 6 +- .../unit_tests/providers/google/__init__.py | 1 + .../providers/google/test_finish_reason.py | 18 +-- 18 files changed, 340 insertions(+), 357 deletions(-) delete mode 100644 packages/capabilities/image-generation/src/celeste_image_generation/providers/google/config.py create mode 100644 packages/capabilities/image-generation/src/celeste_image_generation/providers/google/gemini.py delete mode 100644 packages/capabilities/image-generation/src/celeste_image_generation/providers/google/gemini_api.py create mode 100644 packages/capabilities/image-generation/src/celeste_image_generation/providers/google/gemini_parameters.py create mode 100644 packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen.py delete mode 100644 packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen_api.py create mode 100644 packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen_parameters.py diff --git a/packages/capabilities/image-generation/pyproject.toml b/packages/capabilities/image-generation/pyproject.toml index c50f42dd..d3b01b70 100644 --- a/packages/capabilities/image-generation/pyproject.toml +++ b/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" diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/client.py b/packages/capabilities/image-generation/src/celeste_image_generation/client.py index 301c9cf3..5cf89b19 100644 --- a/packages/capabilities/image-generation/src/celeste_image_generation/client.py +++ b/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)" diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/io.py b/packages/capabilities/image-generation/src/celeste_image_generation/io.py index 90a86ec7..40a3afb3 100644 --- a/packages/capabilities/image-generation/src/celeste_image_generation/io.py +++ b/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 diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/parameters.py b/packages/capabilities/image-generation/src/celeste_image_generation/parameters.py index eebe195a..3d2574d8 100644 --- a/packages/capabilities/image-generation/src/celeste_image_generation/parameters.py +++ b/packages/capabilities/image-generation/src/celeste_image_generation/parameters.py @@ -2,6 +2,7 @@ from enum import StrEnum +from celeste.artifacts import ImageArtifact from celeste.parameters import Parameters @@ -9,27 +10,31 @@ 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 diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/__init__.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/__init__.py index 992b1691..642f7c29 100644 --- a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/__init__.py +++ b/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", +] diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/client.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/client.py index faca14f2..8e1b4978 100644 --- a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/client.py +++ b/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"] diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/config.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/config.py deleted file mode 100644 index c1cb4d28..00000000 --- a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/config.py +++ /dev/null @@ -1,10 +0,0 @@ -"""Google provider configuration for image generation.""" - -# HTTP Configuration -BASE_URL = "https://generativelanguage.googleapis.com/v1beta/models" -IMAGEN_ENDPOINT = "/{model_id}:predict" -GEMINI_ENDPOINT = "/{model_id}:generateContent" - -# Authentication -AUTH_HEADER_NAME = "x-goog-api-key" -AUTH_HEADER_PREFIX = "" # Direct API key, no prefix diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/gemini.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/gemini.py new file mode 100644 index 00000000..42a7ee97 --- /dev/null +++ b/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"] diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/gemini_api.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/gemini_api.py deleted file mode 100644 index 2b371de4..00000000 --- a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/gemini_api.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Gemini API adapter for Google image generation. - -Pure data transformer for Gemini multimodal models (gemini-2.5-flash-image). -Handles request/response structure transformation only. -""" - -from typing import Any - -from celeste_image_generation.io import ImageGenerationUsage - -from . import config - - -class GeminiImageAPIAdapter: - """Adapter for Gemini multimodal API request/response transformation. - - Request format: contents[].parts[] + generationConfig.responseModalities + imageConfig - Response format: candidates[].content.parts[].inlineData (camelCase in REST API) - """ - - def build_request(self, prompt: str, parameters: dict[str, Any]) -> dict[str, Any]: - """Build Gemini API request structure. - - Args: - prompt: Text prompt for image generation. - parameters: Parameter dictionary (aspectRatio, etc.). - - Returns: - Gemini-formatted request with contents[] and generationConfig. - """ - return { - "contents": [{"parts": [{"text": prompt}]}], - "generationConfig": { - "responseModalities": ["Image"], - "imageConfig": parameters, - }, - } - - def parse_response(self, response_data: dict[str, Any]) -> dict[str, Any] | None: - """Parse Gemini API response structure. - - Args: - response_data: Raw API response. - - Returns: - First part containing inlineData with base64 image, or None if blocked. - """ - candidates = response_data.get("candidates", []) - if not candidates: - return None - - candidate = candidates[0] - if candidate.get("finishReason") != "STOP": - return None - return candidate["content"]["parts"][0]["inlineData"] - - def parse_usage(self, response_data: dict[str, Any]) -> ImageGenerationUsage: - """Parse usage from Gemini API response. - - Args: - response_data: Raw API response. - - Returns: - ImageGenerationUsage with token counts and generated_images count. - """ - usage_metadata = response_data.get("usageMetadata", {}) - candidates = response_data.get("candidates", []) - - return ImageGenerationUsage( - input_tokens=usage_metadata.get("promptTokenCount"), - output_tokens=usage_metadata.get("candidatesTokenCount"), - total_tokens=usage_metadata.get("totalTokenCount"), - reasoning_tokens=usage_metadata.get("thoughtsTokenCount"), - generated_images=len(candidates), - ) - - @staticmethod - def endpoint(model_id: str) -> str: - """Get endpoint for model.""" - return config.GEMINI_ENDPOINT.format(model_id=model_id) - - -__all__ = ["GeminiImageAPIAdapter"] diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/gemini_parameters.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/gemini_parameters.py new file mode 100644 index 00000000..9a9acf5c --- /dev/null +++ b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/gemini_parameters.py @@ -0,0 +1,35 @@ +"""Google Gemini parameter mappers for image generation.""" + +from celeste_google.generate_content.parameters import ( + AspectRatioMapper as _AspectRatioMapper, +) +from celeste_google.generate_content.parameters import ( + ImageSizeMapper as _ImageSizeMapper, +) +from celeste_google.generate_content.parameters import ( + MediaContentMapper as _MediaContentMapper, +) + +from celeste.parameters import ParameterMapper +from celeste_image_generation.parameters import ImageGenerationParameter + + +class AspectRatioMapper(_AspectRatioMapper): + name = ImageGenerationParameter.ASPECT_RATIO + + +class QualityMapper(_ImageSizeMapper): + name = ImageGenerationParameter.QUALITY + + +class ReferenceImagesMapper(_MediaContentMapper): + name = ImageGenerationParameter.REFERENCE_IMAGES + + +GEMINI_PARAMETER_MAPPERS: list[ParameterMapper] = [ + AspectRatioMapper(), + QualityMapper(), + ReferenceImagesMapper(), +] + +__all__ = ["GEMINI_PARAMETER_MAPPERS"] diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen.py new file mode 100644 index 00000000..3fd3f74c --- /dev/null +++ b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen.py @@ -0,0 +1,90 @@ +"""Imagen client for Google image generation.""" + +import base64 +from typing import Any, Unpack + +from celeste_google.imagen.client import GoogleImagenClient + +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 + +from .imagen_parameters import IMAGEN_PARAMETER_MAPPERS + + +class ImagenImageGenerationClient(GoogleImagenClient, ImageGenerationClient): + """Google Imagen client for image generation. + + Uses Imagen API format: instances[].prompt → predictions[]. + For Imagen models (imagen-3.x, imagen-4.x). + """ + + @classmethod + def parameter_mappers(cls) -> list[ParameterMapper]: + return IMAGEN_PARAMETER_MAPPERS + + def _init_request(self, inputs: ImageGenerationInput) -> dict[str, Any]: + """Initialize request for Imagen API.""" + return { + "instances": [{"prompt": inputs.prompt}], + "parameters": {}, + } + + def _parse_usage(self, response_data: dict[str, Any]) -> ImageGenerationUsage: + """Parse usage from response.""" + predictions = response_data.get("predictions", []) + return ImageGenerationUsage(num_images=len(predictions)) + + def _parse_content( + self, + response_data: dict[str, Any], + **parameters: Unpack[ImageGenerationParameters], + ) -> ImageArtifact | list[ImageArtifact]: + """Parse content from response. + + Returns ImageArtifact for single image, list[ImageArtifact] for multiple. + """ + predictions = response_data.get("predictions", []) + if not predictions: + return ImageArtifact() + + images: list[ImageArtifact] = [] + for prediction in predictions: + base64_data = prediction.get("bytesBase64Encoded") + if base64_data: + mime_type = ImageMimeType(prediction.get("mimeType", "image/png")) + image_bytes = base64.b64decode(base64_data) + images.append(ImageArtifact(data=image_bytes, mime_type=mime_type)) + + # Return type logic: + # - num_images=1 explicitly → single ImageArtifact + # - num_images>1 explicitly → list (even if fewer returned) + # - num_images=None (not specified) → based on actual count returned + num_images_requested = parameters.get("num_images") + if num_images_requested == 1: + return images[0] if images else ImageArtifact() + if num_images_requested is not None and num_images_requested > 1: + return images if images else [] + # Not specified: return based on what provider actually returned + if len(images) == 1: + return images[0] + return images if images else ImageArtifact() + + def _parse_finish_reason( + self, response_data: dict[str, Any] + ) -> ImageGenerationFinishReason: + """Parse finish reason from response. + + Imagen API doesn't provide finish reasons. + """ + return ImageGenerationFinishReason(reason=None) + + +__all__ = ["ImagenImageGenerationClient"] diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen_api.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen_api.py deleted file mode 100644 index 3db7b85a..00000000 --- a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen_api.py +++ /dev/null @@ -1,67 +0,0 @@ -"""Imagen API adapter for Google image generation. - -Pure data transformer for Imagen models (imagen-3.x, imagen-4.x). -Handles request/response structure transformation only. -""" - -from typing import Any - -from celeste_image_generation.io import ImageGenerationUsage - -from . import config - - -class ImagenAPIAdapter: - """Adapter for Imagen API request/response transformation. - - Request format: instances[].prompt + parameters - Response format: predictions[].bytesBase64Encoded - """ - - def build_request(self, prompt: str, parameters: dict[str, Any]) -> dict[str, Any]: - """Build Imagen API request structure. - - Args: - prompt: Text prompt for image generation. - parameters: Parameter dictionary (aspectRatio, imageSize, etc.). - - Returns: - Imagen-formatted request with instances[] and parameters. - """ - return { - "instances": [{"prompt": prompt}], - "parameters": parameters, - } - - def parse_response(self, response_data: dict[str, Any]) -> dict[str, Any]: - """Parse Imagen API response structure. - - Args: - response_data: Raw API response. - - Returns: - First prediction containing bytesBase64Encoded and mimeType. - """ - return response_data["predictions"][0] - - def parse_usage(self, response_data: dict[str, Any]) -> ImageGenerationUsage: - """Parse usage from Imagen API response. - - Args: - response_data: Raw API response. - - Returns: - ImageGenerationUsage with generated_images count from predictions array. - """ - predictions = response_data.get("predictions", []) - return ImageGenerationUsage( - generated_images=len(predictions), - ) - - @staticmethod - def endpoint(model_id: str) -> str: - """Get endpoint for model.""" - return config.IMAGEN_ENDPOINT.format(model_id=model_id) - - -__all__ = ["ImagenAPIAdapter"] diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen_parameters.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen_parameters.py new file mode 100644 index 00000000..ca06d001 --- /dev/null +++ b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/imagen_parameters.py @@ -0,0 +1,35 @@ +"""Google Imagen parameter mappers for image generation.""" + +from celeste_google.imagen.parameters import ( + AspectRatioMapper as _AspectRatioMapper, +) +from celeste_google.imagen.parameters import ( + ImageSizeMapper as _ImageSizeMapper, +) +from celeste_google.imagen.parameters import ( + SampleCountMapper as _SampleCountMapper, +) + +from celeste.parameters import ParameterMapper +from celeste_image_generation.parameters import ImageGenerationParameter + + +class AspectRatioMapper(_AspectRatioMapper): + name = ImageGenerationParameter.ASPECT_RATIO + + +class QualityMapper(_ImageSizeMapper): + name = ImageGenerationParameter.QUALITY + + +class NumImagesMapper(_SampleCountMapper): + name = ImageGenerationParameter.NUM_IMAGES + + +IMAGEN_PARAMETER_MAPPERS: list[ParameterMapper] = [ + AspectRatioMapper(), + QualityMapper(), + NumImagesMapper(), +] + +__all__ = ["IMAGEN_PARAMETER_MAPPERS"] diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/models.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/models.py index f3e596bd..67aec2b4 100644 --- a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/models.py +++ b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/models.py @@ -1,7 +1,7 @@ """Google models for image generation.""" from celeste import Model, Provider -from celeste.constraints import Choice +from celeste.constraints import Choice, ImagesConstraint, Range from celeste_image_generation.parameters import ImageGenerationParameter # Imagen API models (instances[].prompt → predictions[]) @@ -12,6 +12,7 @@ provider=Provider.GOOGLE, display_name="Imagen 4", parameter_constraints={ + ImageGenerationParameter.NUM_IMAGES: Range(min=1, max=4), ImageGenerationParameter.ASPECT_RATIO: Choice( options=["1:1", "3:4", "4:3", "9:16", "16:9"] ), @@ -23,6 +24,7 @@ provider=Provider.GOOGLE, display_name="Imagen 4 Fast", parameter_constraints={ + ImageGenerationParameter.NUM_IMAGES: Range(min=1, max=4), ImageGenerationParameter.ASPECT_RATIO: Choice( options=["1:1", "3:4", "4:3", "9:16", "16:9"] ), @@ -34,24 +36,13 @@ provider=Provider.GOOGLE, display_name="Imagen 4 Ultra", parameter_constraints={ + ImageGenerationParameter.NUM_IMAGES: Range(min=1, max=4), ImageGenerationParameter.ASPECT_RATIO: Choice( options=["1:1", "3:4", "4:3", "9:16", "16:9"] ), ImageGenerationParameter.QUALITY: Choice(options=["1K", "2K"]), }, ), - # Imagen 3 models (deprecated June 24, 2025) - Support for backwards compatibility - Model( - id="imagen-3.0-generate-002", - provider=Provider.GOOGLE, - display_name="Imagen 3", - parameter_constraints={ - ImageGenerationParameter.ASPECT_RATIO: Choice( - options=["1:1", "3:4", "4:3", "9:16", "16:9"] - ), - ImageGenerationParameter.QUALITY: Choice(options=["1K"]), - }, - ), ] # Gemini API models (contents[].parts[] → candidates[]) @@ -75,6 +66,7 @@ "21:9", ] ), + ImageGenerationParameter.REFERENCE_IMAGES: ImagesConstraint(max_count=3), }, ), Model( @@ -97,6 +89,7 @@ ] ), ImageGenerationParameter.QUALITY: Choice(options=["1K", "2K", "4K"]), + ImageGenerationParameter.REFERENCE_IMAGES: ImagesConstraint(max_count=14), }, ), ] diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/parameters.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/parameters.py index b8201192..c4e37f1f 100644 --- a/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/parameters.py +++ b/packages/capabilities/image-generation/src/celeste_image_generation/providers/google/parameters.py @@ -1,67 +1,12 @@ -"""Google parameter mappers for image generation.""" +"""Google Gemini and Imagen parameter mappers for image generation.""" -from typing import Any - -from celeste import Model from celeste.parameters import ParameterMapper -from celeste_image_generation.parameters import ImageGenerationParameter - - -class AspectRatioMapper(ParameterMapper): - """Map aspect_ratio parameter with validation.""" - - name = ImageGenerationParameter.ASPECT_RATIO - - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform aspect_ratio into provider request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - if "generationConfig" in request: - request.setdefault("generationConfig", {}).setdefault("imageConfig", {})[ - "aspectRatio" - ] = validated_value - else: - request.setdefault("parameters", {})["aspectRatio"] = validated_value - - return request - - -class QualityMapper(ParameterMapper): - """Map quality parameter to imageSize.""" - - name = ImageGenerationParameter.QUALITY - - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform quality into provider imageSize request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - if "generationConfig" in request: - request.setdefault("generationConfig", {}).setdefault("imageConfig", {})[ - "imageSize" - ] = validated_value - else: - request.setdefault("parameters", {})["imageSize"] = validated_value - - return request +from .gemini_parameters import GEMINI_PARAMETER_MAPPERS +from .imagen_parameters import IMAGEN_PARAMETER_MAPPERS -GOOGLE_PARAMETER_MAPPERS: list[ParameterMapper] = [ - AspectRatioMapper(), - QualityMapper(), -] +GOOGLE_PARAMETER_MAPPERS: list[ParameterMapper] = ( + GEMINI_PARAMETER_MAPPERS + IMAGEN_PARAMETER_MAPPERS +) __all__ = ["GOOGLE_PARAMETER_MAPPERS"] diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/streaming.py b/packages/capabilities/image-generation/src/celeste_image_generation/streaming.py index fda1bfe0..48e96103 100644 --- a/packages/capabilities/image-generation/src/celeste_image_generation/streaming.py +++ b/packages/capabilities/image-generation/src/celeste_image_generation/streaming.py @@ -1,7 +1,7 @@ """Streaming for image generation.""" from abc import abstractmethod -from typing import Unpack +from typing import Any, Unpack from celeste.streaming import Stream from celeste_image_generation.io import ( @@ -17,6 +17,10 @@ class ImageGenerationStream( ): """Streaming for image generation.""" + @abstractmethod + def _parse_chunk(self, event: dict[str, Any]) -> ImageGenerationChunk | None: + """Parse SSE event into Chunk (provider-specific).""" + def _parse_output( # type: ignore[override] self, chunks: list[ImageGenerationChunk], diff --git a/packages/capabilities/image-generation/tests/unit_tests/providers/google/__init__.py b/packages/capabilities/image-generation/tests/unit_tests/providers/google/__init__.py index e69de29b..2470e91c 100644 --- a/packages/capabilities/image-generation/tests/unit_tests/providers/google/__init__.py +++ b/packages/capabilities/image-generation/tests/unit_tests/providers/google/__init__.py @@ -0,0 +1 @@ +"""Google provider unit tests for image-generation.""" diff --git a/packages/capabilities/image-generation/tests/unit_tests/providers/google/test_finish_reason.py b/packages/capabilities/image-generation/tests/unit_tests/providers/google/test_finish_reason.py index 7943df14..7ac0ccc7 100644 --- a/packages/capabilities/image-generation/tests/unit_tests/providers/google/test_finish_reason.py +++ b/packages/capabilities/image-generation/tests/unit_tests/providers/google/test_finish_reason.py @@ -3,7 +3,7 @@ from typing import Any import pytest -from celeste_image_generation.providers.google.client import GoogleImageGenerationClient +from celeste_image_generation.providers.google.gemini import GeminiImageGenerationClient from pydantic import SecretStr from celeste.auth import APIKey @@ -12,12 +12,12 @@ class TestParseFinishReason: - """Test _parse_finish_reason method for Google image generation client.""" + """Test _parse_finish_reason method for Gemini image generation client.""" @pytest.fixture - def client(self) -> GoogleImageGenerationClient: - """Create a Google image generation client for testing.""" - return GoogleImageGenerationClient( + def client(self) -> GeminiImageGenerationClient: + """Create a Gemini image generation client for testing.""" + return GeminiImageGenerationClient( model=Model( id="gemini-2.5-flash-image", provider=Provider.GOOGLE, @@ -58,7 +58,7 @@ def client(self) -> GoogleImageGenerationClient: ) def test_parse_finish_reason_with_valid_reason( self, - client: GoogleImageGenerationClient, + client: GeminiImageGenerationClient, finish_reason: str, finish_message: str | None, expected_reason: str, @@ -114,7 +114,7 @@ def test_parse_finish_reason_with_valid_reason( ) def test_parse_finish_reason_returns_none_for_invalid_input( self, - client: GoogleImageGenerationClient, + client: GeminiImageGenerationClient, response_data: dict[str, Any], ) -> None: """Test parsing finish reason returns None for invalid/missing input.""" @@ -125,7 +125,7 @@ def test_parse_finish_reason_returns_none_for_invalid_input( assert result is None def test_parse_finish_reason_empty_string_finish_reason( - self, client: GoogleImageGenerationClient + self, client: GeminiImageGenerationClient ) -> None: """Test parsing finish reason when finishReason is empty string.""" # Arrange @@ -142,7 +142,7 @@ def test_parse_finish_reason_empty_string_finish_reason( assert result is None def test_parse_finish_reason_empty_string_message( - self, client: GoogleImageGenerationClient + self, client: GeminiImageGenerationClient ) -> None: """Test parsing finish reason when finishMessage is empty string.""" # Arrange From ff3765e0006d9f270cca15f09626f77f62ed3c85 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 15:39:48 +0100 Subject: [PATCH 5/5] fix(image-generation): add num_images=1 to Google Imagen test Without num_images parameter, Imagen may return multiple images (list) instead of a single ImageArtifact, causing test assertion to fail. --- .../integration_tests/test_image_generation/test_generate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/capabilities/image-generation/tests/integration_tests/test_image_generation/test_generate.py b/packages/capabilities/image-generation/tests/integration_tests/test_image_generation/test_generate.py index 0f783cf7..78565b65 100644 --- a/packages/capabilities/image-generation/tests/integration_tests/test_image_generation/test_generate.py +++ b/packages/capabilities/image-generation/tests/integration_tests/test_image_generation/test_generate.py @@ -10,7 +10,7 @@ [ (Provider.BFL, "flux-2-pro", {"aspect_ratio": "1024x1024"}), (Provider.OPENAI, "dall-e-2", {}), - (Provider.GOOGLE, "imagen-4.0-fast-generate-001", {}), + (Provider.GOOGLE, "imagen-4.0-fast-generate-001", {"num_images": 1}), (Provider.BYTEDANCE, "seedream-4-0-250828", {}), ], )