From 232e18fdb265296797c4f99382b85c4354a66915 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Fri, 6 Feb 2026 17:49:46 +0100 Subject: [PATCH 1/9] feat: add Vertex AI support for all providers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Route requests through Vertex AI when GoogleADC auth is provided. Supports Google, Anthropic, Mistral, and DeepSeek providers across text, images, and videos modalities. Includes Veo polling fix (fetchPredictOperation), error handler hardening, Gemini image role fix, and DeepSeek usage parser fix. WIP: Veo Vertex inline video (bytesBase64Encoded) parsing not yet handled — needs base64 decoding or storageUri in request. Co-Authored-By: Claude Opus 4.6 --- src/celeste/client.py | 5 +- .../images/providers/google/gemini.py | 2 +- .../videos/providers/google/client.py | 4 +- .../providers/anthropic/messages/client.py | 38 +- .../providers/anthropic/messages/config.py | 7 + src/celeste/providers/deepseek/chat/client.py | 39 +- src/celeste/providers/deepseek/chat/config.py | 6 + .../providers/google/cloud_tts/client.py | 3 +- .../providers/google/embeddings/client.py | 35 +- .../providers/google/embeddings/config.py | 7 + .../google/generate_content/client.py | 44 +- .../google/generate_content/config.py | 8 + src/celeste/providers/google/imagen/client.py | 34 +- src/celeste/providers/google/imagen/config.py | 6 + src/celeste/providers/google/veo/client.py | 104 +++- src/celeste/providers/google/veo/config.py | 7 + src/celeste/providers/mistral/chat/client.py | 33 +- src/celeste/providers/mistral/chat/config.py | 7 + .../{api_slug}/client.py.template | 21 + .../integration_tests/images/test_generate.py | 29 + tests/integration_tests/text/test_generate.py | 29 + .../text/test_stream_generate.py | 31 + .../integration_tests/videos/test_generate.py | 23 + tests/unit_tests/test_vertex_routing.py | 557 ++++++++++++++++++ 24 files changed, 1040 insertions(+), 39 deletions(-) create mode 100644 tests/unit_tests/test_vertex_routing.py diff --git a/src/celeste/client.py b/src/celeste/client.py index 481fd6b8..6df75c6a 100644 --- a/src/celeste/client.py +++ b/src/celeste/client.py @@ -296,9 +296,8 @@ def _handle_error_response(self, response: httpx.Response) -> None: """Handle error responses from provider APIs.""" if not response.is_success: try: - error_data = response.json() - error_msg = error_data.get("error", {}).get("message", response.text) - except JSONDecodeError: + error_msg = response.json()["error"]["message"] + except (JSONDecodeError, KeyError, TypeError, IndexError): error_msg = response.text or f"HTTP {response.status_code}" raise httpx.HTTPStatusError( diff --git a/src/celeste/modalities/images/providers/google/gemini.py b/src/celeste/modalities/images/providers/google/gemini.py index e7143226..276d87f7 100644 --- a/src/celeste/modalities/images/providers/google/gemini.py +++ b/src/celeste/modalities/images/providers/google/gemini.py @@ -82,7 +82,7 @@ def _init_request(self, inputs: ImageInput) -> dict[str, Any]: parts.append({"text": inputs.prompt}) return { - "contents": [{"parts": parts}], + "contents": [{"role": "user", "parts": parts}], "generationConfig": { "responseModalities": ["TEXT", "IMAGE"], "imageConfig": {}, diff --git a/src/celeste/modalities/videos/providers/google/client.py b/src/celeste/modalities/videos/providers/google/client.py index 744bef7b..fda19da7 100644 --- a/src/celeste/modalities/videos/providers/google/client.py +++ b/src/celeste/modalities/videos/providers/google/client.py @@ -51,7 +51,9 @@ def _parse_content( ) -> VideoArtifact: """Parse content from response.""" video_data = super()._parse_content(response_data) - return VideoArtifact(url=video_data.get("uri")) + return VideoArtifact( + url=video_data.get("uri") or video_data.get("gcsUri"), + ) def _parse_finish_reason(self, response_data: dict[str, Any]) -> VideoFinishReason: """Parse finish reason from response.""" diff --git a/src/celeste/providers/anthropic/messages/client.py b/src/celeste/providers/anthropic/messages/client.py index a80c377e..1919cfd1 100644 --- a/src/celeste/providers/anthropic/messages/client.py +++ b/src/celeste/providers/anthropic/messages/client.py @@ -7,6 +7,7 @@ from celeste.core import UsageField from celeste.io import FinishReason from celeste.mime_types import ApplicationMimeType +from celeste.providers.google.auth import GoogleADC from . import config @@ -22,6 +23,10 @@ class AnthropicMessagesClient(APIMixin): - _parse_finish_reason() - Extract finish reason from response - _build_metadata() - Filter content fields + Auth-based endpoint selection: + - GoogleADC auth -> Vertex AI endpoints (Claude on Google Cloud) + - API key auth -> Anthropic API endpoints + Usage: class AnthropicTextGenerationClient(AnthropicMessagesClient, TextGenerationClient): def _parse_content(self, response_data, **parameters): @@ -32,6 +37,35 @@ def _parse_content(self, response_data, **parameters): return "" """ + def _get_vertex_endpoint( + self, anthropic_endpoint: str, streaming: bool = False + ) -> str: + """Map Anthropic endpoint to Vertex AI endpoint.""" + if streaming: + return config.VertexAnthropicEndpoint.STREAM_MESSAGE + return config.VertexAnthropicEndpoint.CREATE_MESSAGE + + def _build_url(self, endpoint: str, streaming: bool = False) -> str: + """Build full URL based on auth type. + + - GoogleADC auth -> Vertex AI endpoints + - API key auth -> Anthropic API endpoints + """ + if isinstance(self.auth, GoogleADC): + project_id = self.auth.resolved_project_id + if project_id is None: + raise ValueError( + "Vertex AI requires a project_id. " + "Pass project_id to GoogleADC() or ensure credentials have a project." + ) + + vertex_endpoint = self._get_vertex_endpoint(endpoint, streaming=streaming) + base_url = self.auth.get_vertex_base_url() + return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location, model_id=self.model.id)}" + + # Default: Anthropic API + return f"{config.BASE_URL}{endpoint}" + def _build_headers(self, request_body: dict[str, Any]) -> dict[str, str]: """Build headers with beta features extracted from request.""" beta_features: list[str] = request_body.pop("_beta_features", []) @@ -85,7 +119,7 @@ async def _make_request( endpoint = config.AnthropicMessagesEndpoint.CREATE_MESSAGE response = await self.http_client.post( - f"{config.BASE_URL}{endpoint}", + url=self._build_url(endpoint, streaming=False), headers=headers, json_body=request_body, ) @@ -111,7 +145,7 @@ def _make_stream_request( endpoint = config.AnthropicMessagesEndpoint.CREATE_MESSAGE return self.http_client.stream_post( - f"{config.BASE_URL}{endpoint}", + url=self._build_url(endpoint, streaming=True), headers=headers, json_body=request_body, ) diff --git a/src/celeste/providers/anthropic/messages/config.py b/src/celeste/providers/anthropic/messages/config.py index 3199bed8..7c3e6f40 100644 --- a/src/celeste/providers/anthropic/messages/config.py +++ b/src/celeste/providers/anthropic/messages/config.py @@ -12,6 +12,13 @@ class AnthropicMessagesEndpoint(StrEnum): GET_MODEL = "/v1/models/{model_id}" +class VertexAnthropicEndpoint(StrEnum): + """Endpoints for Anthropic on Vertex AI.""" + + CREATE_MESSAGE = "/v1/projects/{project_id}/locations/{location}/publishers/anthropic/models/{model_id}:rawPredict" + STREAM_MESSAGE = "/v1/projects/{project_id}/locations/{location}/publishers/anthropic/models/{model_id}:streamRawPredict" + + BASE_URL = "https://api.anthropic.com" # Required diff --git a/src/celeste/providers/deepseek/chat/client.py b/src/celeste/providers/deepseek/chat/client.py index f8547520..19819290 100644 --- a/src/celeste/providers/deepseek/chat/client.py +++ b/src/celeste/providers/deepseek/chat/client.py @@ -7,6 +7,7 @@ from celeste.core import UsageField from celeste.io import FinishReason from celeste.mime_types import ApplicationMimeType +from celeste.providers.google.auth import GoogleADC from . import config @@ -31,6 +32,36 @@ def _parse_content(self, response_data, **parameters): return self._transform_output(content, **parameters) """ + def _get_vertex_endpoint(self, deepseek_endpoint: str) -> str: + """Map DeepSeek endpoint to Vertex AI endpoint.""" + mapping: dict[str, str] = { + config.DeepSeekChatEndpoint.CREATE_CHAT: config.VertexDeepSeekEndpoint.CREATE_CHAT, + } + vertex_endpoint = mapping.get(deepseek_endpoint) + if vertex_endpoint is None: + raise ValueError(f"No Vertex AI endpoint mapping for: {deepseek_endpoint}") + return vertex_endpoint + + def _build_url(self, endpoint: str) -> str: + """Build full URL based on auth type. + + - GoogleADC auth -> Vertex AI endpoints (OpenAI-compatible) + - API key auth -> DeepSeek API endpoints + """ + if isinstance(self.auth, GoogleADC): + project_id = self.auth.resolved_project_id + if project_id is None: + raise ValueError( + "Vertex AI requires a project_id. " + "Pass project_id to GoogleADC() or ensure credentials have a project." + ) + + vertex_endpoint = self._get_vertex_endpoint(endpoint) + base_url = self.auth.get_vertex_base_url() + return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location)}" + + return f"{config.BASE_URL}{endpoint}" + def _build_request( self, inputs: Any, @@ -64,7 +95,7 @@ async def _make_request( } response = await self.http_client.post( - f"{config.BASE_URL}{endpoint}", + self._build_url(endpoint), headers=headers, json_body=request_body, ) @@ -89,7 +120,7 @@ def _make_stream_request( } return self.http_client.stream_post( - f"{config.BASE_URL}{endpoint}", + self._build_url(endpoint), headers=headers, json_body=request_body, ) @@ -100,8 +131,8 @@ def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | float | None Shared by client and streaming across all capabilities. """ - prompt_tokens_details = usage_data.get("prompt_tokens_details", {}) - completion_tokens_details = usage_data.get("completion_tokens_details", {}) + prompt_tokens_details = usage_data.get("prompt_tokens_details") or {} + completion_tokens_details = usage_data.get("completion_tokens_details") or {} return { UsageField.INPUT_TOKENS: usage_data.get("prompt_tokens"), UsageField.OUTPUT_TOKENS: usage_data.get("completion_tokens"), diff --git a/src/celeste/providers/deepseek/chat/config.py b/src/celeste/providers/deepseek/chat/config.py index 53c4a127..90ce02e5 100644 --- a/src/celeste/providers/deepseek/chat/config.py +++ b/src/celeste/providers/deepseek/chat/config.py @@ -10,4 +10,10 @@ class DeepSeekChatEndpoint(StrEnum): LIST_MODELS = "/models" +class VertexDeepSeekEndpoint(StrEnum): + """Endpoints for DeepSeek on Vertex AI (OpenAI-compatible).""" + + CREATE_CHAT = "/v1/projects/{project_id}/locations/{location}/endpoints/openapi/chat/completions" + + BASE_URL = "https://api.deepseek.com" diff --git a/src/celeste/providers/google/cloud_tts/client.py b/src/celeste/providers/google/cloud_tts/client.py index 86854b50..8d1aaa58 100644 --- a/src/celeste/providers/google/cloud_tts/client.py +++ b/src/celeste/providers/google/cloud_tts/client.py @@ -51,7 +51,8 @@ def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | float | None def model_post_init(self, _context: Any) -> None: """Override auth to use ADC for Cloud TTS (not API key like Gemini).""" super().model_post_init(_context) # type: ignore[misc] - object.__setattr__(self, "auth", GoogleADC()) + if not isinstance(self.auth, GoogleADC): + object.__setattr__(self, "auth", GoogleADC()) async def _make_request( self, diff --git a/src/celeste/providers/google/embeddings/client.py b/src/celeste/providers/google/embeddings/client.py index db267f4b..76270e17 100644 --- a/src/celeste/providers/google/embeddings/client.py +++ b/src/celeste/providers/google/embeddings/client.py @@ -8,6 +8,7 @@ from celeste.io import FinishReason from celeste.mime_types import ApplicationMimeType +from ..auth import GoogleADC from . import config @@ -36,12 +37,43 @@ def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | float | None - _make_request() - HTTP POST to embedContent or batchEmbedContents endpoint - _parse_content() - Extract embedding vectors (generic) + Auth-based endpoint selection: + - GoogleADC auth -> Vertex AI endpoints + - API key auth -> Gemini API endpoints + Capability clients extend via super(): class GoogleEmbeddingsClient(GoogleEmbeddingsClient, EmbeddingsClient): def _parse_content(self, response_data, **params): return super()._parse_content(response_data) # No transformation needed """ + def _get_vertex_endpoint(self, gemini_endpoint: str) -> str: + """Map Gemini Embeddings endpoint to Vertex AI endpoint.""" + mapping: dict[str, str] = { + config.GoogleEmbeddingsEndpoint.EMBED_CONTENT: config.VertexEmbeddingsEndpoint.EMBED_CONTENT, + config.GoogleEmbeddingsEndpoint.BATCH_EMBED_CONTENTS: config.VertexEmbeddingsEndpoint.BATCH_EMBED_CONTENTS, + } + vertex_endpoint = mapping.get(gemini_endpoint) + if vertex_endpoint is None: + raise ValueError(f"No Vertex AI endpoint mapping for: {gemini_endpoint}") + return vertex_endpoint + + def _build_url(self, endpoint: str) -> str: + """Build full URL based on auth type.""" + if isinstance(self.auth, GoogleADC): + project_id = self.auth.resolved_project_id + if project_id is None: + raise ValueError( + "Vertex AI requires a project_id. " + "Pass project_id to GoogleADC() or ensure credentials have a project." + ) + + vertex_endpoint = self._get_vertex_endpoint(endpoint) + base_url = self.auth.get_vertex_base_url() + return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location, model_id=self.model.id)}" + + return f"{config.BASE_URL}{endpoint.format(model_id=self.model.id)}" + async def _make_request( self, request_body: dict[str, Any], @@ -58,7 +90,6 @@ async def _make_request( ) if endpoint is None: endpoint = endpoint_template - endpoint = endpoint.format(model_id=self.model.id) headers = { **self.auth.get_headers(), @@ -66,7 +97,7 @@ async def _make_request( } response = await self.http_client.post( - f"{config.BASE_URL}{endpoint}", + self._build_url(endpoint), headers=headers, json_body=request_body, ) diff --git a/src/celeste/providers/google/embeddings/config.py b/src/celeste/providers/google/embeddings/config.py index 33822fa0..188c884b 100644 --- a/src/celeste/providers/google/embeddings/config.py +++ b/src/celeste/providers/google/embeddings/config.py @@ -10,4 +10,11 @@ class GoogleEmbeddingsEndpoint(StrEnum): BATCH_EMBED_CONTENTS = "/v1beta/models/{model_id}:batchEmbedContents" +class VertexEmbeddingsEndpoint(StrEnum): + """Endpoints for Embeddings on Vertex AI.""" + + EMBED_CONTENT = "/v1/projects/{project_id}/locations/{location}/publishers/google/models/{model_id}:predict" + BATCH_EMBED_CONTENTS = "/v1/projects/{project_id}/locations/{location}/publishers/google/models/{model_id}:predict" + + BASE_URL = "https://generativelanguage.googleapis.com" diff --git a/src/celeste/providers/google/generate_content/client.py b/src/celeste/providers/google/generate_content/client.py index ef270342..d04d70a4 100644 --- a/src/celeste/providers/google/generate_content/client.py +++ b/src/celeste/providers/google/generate_content/client.py @@ -8,6 +8,7 @@ from celeste.io import FinishReason from celeste.mime_types import ApplicationMimeType +from ..auth import GoogleADC from . import config @@ -22,6 +23,10 @@ class GoogleGenerateContentClient(APIMixin): - _parse_finish_reason() - Extract finish reason string from candidates - _build_metadata() - Filter content fields + Auth-based endpoint selection: + - GoogleADC auth -> Vertex AI endpoints + - API key auth -> Gemini API endpoints + Capability clients extend parsing methods via super() to wrap/transform results. Usage: @@ -32,6 +37,39 @@ def _parse_content(self, response_data, **parameters): return self._transform_output(text, **parameters) """ + def _get_vertex_endpoint(self, gemini_endpoint: str) -> str: + """Map Gemini endpoint to Vertex AI endpoint.""" + mapping: dict[str, str] = { + config.GoogleGenerateContentEndpoint.GENERATE_CONTENT: config.VertexEndpoint.GENERATE_CONTENT, + config.GoogleGenerateContentEndpoint.STREAM_GENERATE_CONTENT: config.VertexEndpoint.STREAM_GENERATE_CONTENT, + config.GoogleGenerateContentEndpoint.COUNT_TOKENS: config.VertexEndpoint.COUNT_TOKENS, + } + vertex_endpoint = mapping.get(gemini_endpoint) + if vertex_endpoint is None: + raise ValueError(f"No Vertex AI endpoint mapping for: {gemini_endpoint}") + return vertex_endpoint + + def _build_url(self, endpoint: str) -> str: + """Build full URL based on auth type. + + - GoogleADC auth -> Vertex AI endpoints + - API key auth -> Gemini API endpoints + """ + if isinstance(self.auth, GoogleADC): + project_id = self.auth.resolved_project_id + if project_id is None: + raise ValueError( + "Vertex AI requires a project_id. " + "Pass project_id to GoogleADC() or ensure credentials have a project." + ) + + vertex_endpoint = self._get_vertex_endpoint(endpoint) + base_url = self.auth.get_vertex_base_url() + return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location, model_id=self.model.id)}" + + # Default: Gemini API + return f"{config.BASE_URL}{endpoint.format(model_id=self.model.id)}" + async def _make_request( self, request_body: dict[str, Any], @@ -42,14 +80,13 @@ async def _make_request( """Make HTTP request to generateContent endpoint.""" if endpoint is None: endpoint = config.GoogleGenerateContentEndpoint.GENERATE_CONTENT - endpoint = endpoint.format(model_id=self.model.id) headers = { **self.auth.get_headers(), "Content-Type": ApplicationMimeType.JSON, } response = await self.http_client.post( - f"{config.BASE_URL}{endpoint}", + url=self._build_url(endpoint), headers=headers, json_body=request_body, ) @@ -67,14 +104,13 @@ def _make_stream_request( """Make streaming request to streamGenerateContent endpoint.""" if endpoint is None: endpoint = config.GoogleGenerateContentEndpoint.STREAM_GENERATE_CONTENT - endpoint = 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}{endpoint}", + url=self._build_url(endpoint), headers=headers, json_body=request_body, ) diff --git a/src/celeste/providers/google/generate_content/config.py b/src/celeste/providers/google/generate_content/config.py index 9787afc2..a5fe0181 100644 --- a/src/celeste/providers/google/generate_content/config.py +++ b/src/celeste/providers/google/generate_content/config.py @@ -20,4 +20,12 @@ class GoogleGenerateContentEndpoint(StrEnum): BATCH_GENERATE_CONTENT = "/v1beta/models/{model_id}:batchGenerateContent" +class VertexEndpoint(StrEnum): + """Endpoints for Vertex AI (when using GoogleADC auth).""" + + GENERATE_CONTENT = "/v1/projects/{project_id}/locations/{location}/publishers/google/models/{model_id}:generateContent" + STREAM_GENERATE_CONTENT = "/v1/projects/{project_id}/locations/{location}/publishers/google/models/{model_id}:streamGenerateContent?alt=sse" + COUNT_TOKENS = "/v1/projects/{project_id}/locations/{location}/publishers/google/models/{model_id}:countTokens" + + BASE_URL = "https://generativelanguage.googleapis.com" diff --git a/src/celeste/providers/google/imagen/client.py b/src/celeste/providers/google/imagen/client.py index ba7fb5f8..5dd94039 100644 --- a/src/celeste/providers/google/imagen/client.py +++ b/src/celeste/providers/google/imagen/client.py @@ -9,6 +9,7 @@ from celeste.io import FinishReason from celeste.mime_types import ApplicationMimeType +from ..auth import GoogleADC from . import config @@ -22,6 +23,10 @@ class GoogleImagenClient(APIMixin): - _parse_finish_reason() - Returns None (Imagen doesn't provide finish reasons) - _build_metadata() - Filter out predictions content + Auth-based endpoint selection: + - GoogleADC auth -> Vertex AI endpoints + - API key auth -> Gemini API endpoints + Capability clients extend parsing methods via super() to wrap/transform results. Usage: @@ -31,6 +36,32 @@ def _parse_content(self, response_data, **parameters): # Extract image from predictions[0]... """ + def _get_vertex_endpoint(self, gemini_endpoint: str) -> str: + """Map Gemini Imagen endpoint to Vertex AI endpoint.""" + mapping: dict[str, str] = { + config.GoogleImagenEndpoint.CREATE_IMAGE: config.VertexImagenEndpoint.CREATE_IMAGE, + } + vertex_endpoint = mapping.get(gemini_endpoint) + if vertex_endpoint is None: + raise ValueError(f"No Vertex AI endpoint mapping for: {gemini_endpoint}") + return vertex_endpoint + + def _build_url(self, endpoint: str) -> str: + """Build full URL based on auth type.""" + if isinstance(self.auth, GoogleADC): + project_id = self.auth.resolved_project_id + if project_id is None: + raise ValueError( + "Vertex AI requires a project_id. " + "Pass project_id to GoogleADC() or ensure credentials have a project." + ) + + vertex_endpoint = self._get_vertex_endpoint(endpoint) + base_url = self.auth.get_vertex_base_url() + return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location, model_id=self.model.id)}" + + return f"{config.BASE_URL}{endpoint.format(model_id=self.model.id)}" + async def _make_request( self, request_body: dict[str, Any], @@ -41,14 +72,13 @@ async def _make_request( """Make HTTP request to Imagen :predict endpoint.""" if endpoint is None: endpoint = config.GoogleImagenEndpoint.CREATE_IMAGE - endpoint = endpoint.format(model_id=self.model.id) headers = { **self.auth.get_headers(), "Content-Type": ApplicationMimeType.JSON, } response = await self.http_client.post( - f"{config.BASE_URL}{endpoint}", + self._build_url(endpoint), headers=headers, json_body=request_body, ) diff --git a/src/celeste/providers/google/imagen/config.py b/src/celeste/providers/google/imagen/config.py index f2d16c48..b8981096 100644 --- a/src/celeste/providers/google/imagen/config.py +++ b/src/celeste/providers/google/imagen/config.py @@ -9,4 +9,10 @@ class GoogleImagenEndpoint(StrEnum): CREATE_IMAGE = "/v1beta/models/{model_id}:predict" +class VertexImagenEndpoint(StrEnum): + """Endpoints for Imagen on Vertex AI.""" + + CREATE_IMAGE = "/v1/projects/{project_id}/locations/{location}/publishers/google/models/{model_id}:predict" + + BASE_URL = "https://generativelanguage.googleapis.com" diff --git a/src/celeste/providers/google/veo/client.py b/src/celeste/providers/google/veo/client.py index 7c29643c..f3441fd2 100644 --- a/src/celeste/providers/google/veo/client.py +++ b/src/celeste/providers/google/veo/client.py @@ -10,6 +10,7 @@ from celeste.io import FinishReason from celeste.mime_types import ApplicationMimeType +from ..auth import GoogleADC from . import config logger = logging.getLogger(__name__) @@ -34,6 +35,79 @@ async def download_content(self, artifact: VideoArtifact) -> VideoArtifact: return VideoArtifact(data=video_bytes, mime_type=VideoMimeType.MP4, ...) """ + def _get_vertex_endpoint(self, gemini_endpoint: str) -> str: + """Map Gemini Veo endpoint to Vertex AI endpoint.""" + mapping: dict[str, str] = { + config.GoogleVeoEndpoint.CREATE_VIDEO: config.VertexVeoEndpoint.CREATE_VIDEO, + } + vertex_endpoint = mapping.get(gemini_endpoint) + if vertex_endpoint is None: + raise ValueError(f"No Vertex AI endpoint mapping for: {gemini_endpoint}") + return vertex_endpoint + + def _build_url(self, endpoint: str) -> str: + """Build full URL based on auth type.""" + if isinstance(self.auth, GoogleADC): + project_id = self.auth.resolved_project_id + if project_id is None: + raise ValueError( + "Vertex AI requires a project_id. " + "Pass project_id to GoogleADC() or ensure credentials have a project." + ) + + vertex_endpoint = self._get_vertex_endpoint(endpoint) + base_url = self.auth.get_vertex_base_url() + return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location, model_id=self.model.id)}" + + return f"{config.BASE_URL}{endpoint.format(model_id=self.model.id)}" + + def _build_poll_url(self, operation_name: str) -> str: + """Build polling URL for long-running operations based on auth type.""" + if isinstance(self.auth, GoogleADC): + project_id = self.auth.resolved_project_id + base_url = self.auth.get_vertex_base_url() + poll_path = config.VertexVeoEndpoint.FETCH_OPERATION.format( + project_id=project_id, + location=self.auth.location, + model_id=self.model.id, + ) + return f"{base_url}{poll_path}" + + poll_path = config.GoogleVeoEndpoint.GET_OPERATION.format( + operation_name=operation_name + ) + return f"{config.BASE_URL}{poll_path}" + + async def _make_poll_request(self, operation_name: str) -> dict[str, Any]: + """Poll a long-running operation. + + Vertex AI uses POST to fetchPredictOperation with operationName in body. + AI Studio uses GET to /v1beta/{operation_name}. + """ + headers = { + **self.auth.get_headers(), + "Content-Type": ApplicationMimeType.JSON, + } + poll_url = self._build_poll_url(operation_name) + + if isinstance(self.auth, GoogleADC): + response = await self.http_client.post( + poll_url, + headers=headers, + json_body={"operationName": operation_name}, + timeout=config.DEFAULT_TIMEOUT, + ) + else: + response = await self.http_client.get( + poll_url, + headers=headers, + timeout=config.DEFAULT_TIMEOUT, + ) + + self._handle_error_response(response) + data: dict[str, Any] = response.json() + return data + def _make_stream_request( self, request_body: dict[str, Any], @@ -54,8 +128,6 @@ async def _make_request( """Make HTTP request with async polling for Veo video generation.""" if endpoint is None: endpoint = config.GoogleVeoEndpoint.CREATE_VIDEO - endpoint = endpoint.format(model_id=self.model.id) - url = f"{config.BASE_URL}{endpoint}" auth_headers = self.auth.get_headers() headers = { @@ -65,7 +137,7 @@ async def _make_request( logger.info(f"Initiating video generation with model {self.model.id}") response = await self.http_client.post( - url, + self._build_url(endpoint), headers=headers, json_body=request_body, timeout=config.DEFAULT_TIMEOUT, @@ -77,21 +149,11 @@ async def _make_request( operation_name = operation_data["name"] logger.info(f"Video generation started: {operation_name}") - poll_url = f"{config.BASE_URL}{config.GoogleVeoEndpoint.GET_OPERATION.format(operation_name=operation_name)}" - poll_headers = auth_headers - while True: await asyncio.sleep(config.POLL_INTERVAL) logger.debug(f"Polling operation status: {operation_name}") - poll_response = await self.http_client.get( - poll_url, - headers=poll_headers, - timeout=config.DEFAULT_TIMEOUT, - ) - - self._handle_error_response(poll_response) - operation_data = poll_response.json() + operation_data = await self._make_poll_request(operation_name) if operation_data.get("done"): if "error" in operation_data: @@ -111,10 +173,18 @@ def _parse_content(self, response_data: dict[str, Any]) -> Any: Returns generic dict with video data that capability clients wrap in artifacts. """ - generate_response = response_data.get("response", {}).get( - "generateVideoResponse", {} + response = response_data.get("response", {}) + + if isinstance(self.auth, GoogleADC): + videos = response.get("videos", []) + if not videos: + msg = "No videos in response" + raise ValueError(msg) + return videos[0] + + generated_samples = response.get("generateVideoResponse", {}).get( + "generatedSamples", [] ) - generated_samples = generate_response.get("generatedSamples", []) if not generated_samples: msg = "No generated samples in response" raise ValueError(msg) diff --git a/src/celeste/providers/google/veo/config.py b/src/celeste/providers/google/veo/config.py index 4c773fea..ceb8d3f9 100644 --- a/src/celeste/providers/google/veo/config.py +++ b/src/celeste/providers/google/veo/config.py @@ -10,6 +10,13 @@ class GoogleVeoEndpoint(StrEnum): GET_OPERATION = "/v1beta/{operation_name}" +class VertexVeoEndpoint(StrEnum): + """Endpoints for Veo on Vertex AI.""" + + CREATE_VIDEO = "/v1/projects/{project_id}/locations/{location}/publishers/google/models/{model_id}:predictLongRunning" + FETCH_OPERATION = "/v1/projects/{project_id}/locations/{location}/publishers/google/models/{model_id}:fetchPredictOperation" + + BASE_URL = "https://generativelanguage.googleapis.com" # Polling Configuration diff --git a/src/celeste/providers/mistral/chat/client.py b/src/celeste/providers/mistral/chat/client.py index db23247b..f9e6e85a 100644 --- a/src/celeste/providers/mistral/chat/client.py +++ b/src/celeste/providers/mistral/chat/client.py @@ -7,6 +7,7 @@ from celeste.core import UsageField from celeste.io import FinishReason from celeste.mime_types import ApplicationMimeType +from celeste.providers.google.auth import GoogleADC from . import config @@ -31,6 +32,34 @@ def _parse_content(self, response_data, **parameters): return self._transform_output(content, **parameters) """ + def _get_vertex_endpoint( + self, mistral_endpoint: str, streaming: bool = False + ) -> str: + """Map Mistral endpoint to Vertex AI endpoint.""" + if streaming: + return config.VertexMistralEndpoint.STREAM_CHAT + return config.VertexMistralEndpoint.CREATE_CHAT + + def _build_url(self, endpoint: str, streaming: bool = False) -> str: + """Build full URL based on auth type. + + - GoogleADC auth -> Vertex AI endpoints + - API key auth -> Mistral API endpoints + """ + if isinstance(self.auth, GoogleADC): + project_id = self.auth.resolved_project_id + if project_id is None: + raise ValueError( + "Vertex AI requires a project_id. " + "Pass project_id to GoogleADC() or ensure credentials have a project." + ) + + vertex_endpoint = self._get_vertex_endpoint(endpoint, streaming=streaming) + base_url = self.auth.get_vertex_base_url() + return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location, model_id=self.model.id)}" + + return f"{config.BASE_URL}{endpoint}" + def _build_request( self, inputs: Any, @@ -64,7 +93,7 @@ async def _make_request( } response = await self.http_client.post( - f"{config.BASE_URL}{endpoint}", + self._build_url(endpoint), headers=headers, json_body=request_body, ) @@ -89,7 +118,7 @@ def _make_stream_request( } return self.http_client.stream_post( - f"{config.BASE_URL}{endpoint}", + self._build_url(endpoint, streaming=True), headers=headers, json_body=request_body, ) diff --git a/src/celeste/providers/mistral/chat/config.py b/src/celeste/providers/mistral/chat/config.py index 4b32ff4a..0c323cc7 100644 --- a/src/celeste/providers/mistral/chat/config.py +++ b/src/celeste/providers/mistral/chat/config.py @@ -36,6 +36,13 @@ class MistralChatEndpoint(StrEnum): CREATE_CHAT_MODERATION = "/v1/chat/moderations" +class VertexMistralEndpoint(StrEnum): + """Endpoints for Mistral on Vertex AI.""" + + CREATE_CHAT = "/v1/projects/{project_id}/locations/{location}/publishers/mistralai/models/{model_id}:rawPredict" + STREAM_CHAT = "/v1/projects/{project_id}/locations/{location}/publishers/mistralai/models/{model_id}:streamRawPredict" + + BASE_URL = "https://api.mistral.ai" # Alternative diff --git a/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/client.py.template b/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/client.py.template index aaa75fe1..8885d1f9 100644 --- a/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/client.py.template +++ b/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/client.py.template @@ -20,6 +20,7 @@ class {Provider}{Api}Client(APIMixin): Provides shared HTTP implementation: - _make_request(endpoint=...) - HTTP POST to specified endpoint - _make_stream_request() - HTTP streaming (if supported, otherwise raises StreamingNotSupportedError) + - _make_poll_request() - Poll long-running operations (if supported, otherwise remove) - _parse_usage() - Extract usage dict from response - _parse_content() - Extract content from response - _parse_finish_reason() - Extract finish reason (if provided) @@ -101,6 +102,26 @@ class {Provider}{Api}Client(APIMixin): json_body=request_body, ) + async def _make_poll_request( + self, + operation_name: str, + ) -> dict[str, Any]: + """Poll a long-running operation. + + If this API does not use long-running operations, remove this method. + Override for Vertex AI support (POST to fetchPredictOperation). + """ + headers = self.auth.get_headers() + poll_url = f"{config.BASE_URL}{config.{Provider}{Api}Endpoint.GET_OPERATION.format(operation_name=operation_name)}" + + response = await self.http_client.get( + poll_url, + headers=headers, + ) + self._handle_error_response(response) + data: dict[str, Any] = response.json() + return data + @staticmethod def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | float | None]: """Map {Provider} {Api} usage fields to unified names.""" diff --git a/tests/integration_tests/images/test_generate.py b/tests/integration_tests/images/test_generate.py index 705b2985..9c46d72b 100644 --- a/tests/integration_tests/images/test_generate.py +++ b/tests/integration_tests/images/test_generate.py @@ -14,6 +14,7 @@ from celeste import Modality, Provider, create_client # noqa: E402 from celeste.artifacts import ImageArtifact # noqa: E402 from celeste.modalities.images import ImageOutput, ImageUsage # noqa: E402 +from celeste.providers.google.auth import GoogleADC # noqa: E402 @pytest.mark.parametrize( @@ -79,3 +80,31 @@ def test_sync_generate() -> None: ) assert isinstance(content, ImageArtifact) assert content.has_content + + +@pytest.mark.parametrize( + ("model", "parameters"), + [ + ("imagen-4.0-fast-generate-001", {"num_images": 1}), + ("gemini-2.5-flash-image", {}), + ], +) +@pytest.mark.integration +@pytest.mark.asyncio +async def test_vertex_generate(model: str, parameters: dict) -> None: + """Test image generation via Vertex AI.""" + client = create_client( + modality=Modality.IMAGES, + provider=Provider.GOOGLE, + model=model, + auth=GoogleADC(), + ) + + response = await client.generate( + prompt="A red apple on a white background", + **parameters, + ) + + assert isinstance(response, ImageOutput) + assert isinstance(response.content, ImageArtifact) + assert response.content.has_content diff --git a/tests/integration_tests/text/test_generate.py b/tests/integration_tests/text/test_generate.py index 26eaafd5..8be78590 100644 --- a/tests/integration_tests/text/test_generate.py +++ b/tests/integration_tests/text/test_generate.py @@ -19,6 +19,7 @@ list_models, ) from celeste.modalities.text import TextOutput, TextUsage # noqa: E402 +from celeste.providers.google.auth import GoogleADC # noqa: E402 TEST_MAX_TOKENS = 200 @@ -87,3 +88,31 @@ def test_sync_generate() -> None: assert isinstance(response, TextOutput) assert response.content or response.finish_reason is not None + + +@pytest.mark.parametrize( + ("provider", "model", "location"), + [ + ("google", "gemini-2.5-flash", "global"), + ("anthropic", "claude-haiku-4-5", "us-east5"), + ("mistral", "mistral-small-2503", "us-central1"), + ("deepseek", "deepseek-ai/deepseek-v3.2-maas", "global"), + ], +) +@pytest.mark.integration +@pytest.mark.asyncio +async def test_vertex_generate(provider: str, model: str, location: str) -> None: + """Test text generation via Vertex AI - one model per provider.""" + client = create_client( + modality=Modality.TEXT, + provider=provider, + model=model, + auth=GoogleADC(location=location), + ) + + response = await client.generate(prompt="Hi", max_tokens=TEST_MAX_TOKENS) + + assert isinstance(response, TextOutput) + if not response.content: + assert response.finish_reason is not None + assert isinstance(response.usage, TextUsage) diff --git a/tests/integration_tests/text/test_stream_generate.py b/tests/integration_tests/text/test_stream_generate.py index c08c6b6a..e3e1e25d 100644 --- a/tests/integration_tests/text/test_stream_generate.py +++ b/tests/integration_tests/text/test_stream_generate.py @@ -19,6 +19,7 @@ list_models, ) from celeste.modalities.text import TextChunk, TextUsage # noqa: E402 +from celeste.providers.google.auth import GoogleADC # noqa: E402 TEST_MAX_TOKENS = 200 @@ -91,3 +92,33 @@ def test_sync_stream_generate() -> None: for _chunk in client.sync.stream.generate(prompt="Hi", max_tokens=TEST_MAX_TOKENS): pass # Just exhaust the stream + + +@pytest.mark.parametrize( + ("provider", "model", "location"), + [ + ("google", "gemini-2.5-flash", "global"), + ("anthropic", "claude-haiku-4-5", "us-east5"), + ("mistral", "mistral-small-2503", "us-central1"), + ("deepseek", "deepseek-ai/deepseek-v3.2-maas", "global"), + ], +) +@pytest.mark.integration +@pytest.mark.asyncio +async def test_vertex_stream_generate(provider: str, model: str, location: str) -> None: + """Test streaming text generation via Vertex AI - one model per provider.""" + client = create_client( + modality=Modality.TEXT, + provider=provider, + model=model, + auth=GoogleADC(location=location), + ) + + chunks: list[TextChunk] = [] + async for chunk in client.stream.generate(prompt="Hi", max_tokens=TEST_MAX_TOKENS): + chunks.append(chunk) + + if not chunks: + return + + assert all(isinstance(c, TextChunk) for c in chunks) diff --git a/tests/integration_tests/videos/test_generate.py b/tests/integration_tests/videos/test_generate.py index 9e9ef834..e7dc9c03 100644 --- a/tests/integration_tests/videos/test_generate.py +++ b/tests/integration_tests/videos/test_generate.py @@ -14,6 +14,7 @@ from celeste import Modality, Provider, create_client # noqa: E402 from celeste.artifacts import VideoArtifact # noqa: E402 from celeste.modalities.videos import VideoOutput, VideoUsage # noqa: E402 +from celeste.providers.google.auth import GoogleADC # noqa: E402 @pytest.mark.parametrize( @@ -90,3 +91,25 @@ def test_sync_generate() -> None: assert isinstance(response, VideoOutput) assert isinstance(response.content, VideoArtifact) assert response.content.has_content + + +@pytest.mark.integration +@pytest.mark.slow +@pytest.mark.asyncio +async def test_vertex_generate() -> None: + """Test video generation via Vertex AI.""" + client = create_client( + modality=Modality.VIDEOS, + provider=Provider.GOOGLE, + model="veo-3.0-fast-generate-001", + auth=GoogleADC(location="us-central1"), + ) + + response = await client.generate( + prompt="A cat walking on the beach", + duration=4, + ) + + assert isinstance(response, VideoOutput) + assert isinstance(response.content, VideoArtifact) + assert response.content.has_content diff --git a/tests/unit_tests/test_vertex_routing.py b/tests/unit_tests/test_vertex_routing.py new file mode 100644 index 00000000..31f51f78 --- /dev/null +++ b/tests/unit_tests/test_vertex_routing.py @@ -0,0 +1,557 @@ +"""Tests for Vertex AI URL routing across all Vertex-enabled providers.""" + +from unittest.mock import MagicMock, patch + +import pytest +from pydantic import SecretStr + +from celeste.auth import AuthHeader +from celeste.models import Model +from celeste.providers.google.auth import ( + VERTEX_GLOBAL_BASE_URL, + GoogleADC, +) + +# --- GoogleADC helpers --- + + +class TestGoogleADCResolvedProjectId: + """Test resolved_project_id property.""" + + @patch("google.auth.default") + @patch("google.auth.transport.requests.Request") + def test_returns_explicit_project_id( + self, mock_request: MagicMock, mock_default: MagicMock + ) -> None: + """Explicit project_id takes priority over ADC-inferred project.""" + mock_creds = MagicMock() + mock_creds.valid = True + mock_creds.token = "fake-token" # nosec B105 + mock_default.return_value = (mock_creds, "adc-project") + + auth = GoogleADC(project_id="explicit-project") + assert auth.resolved_project_id == "explicit-project" + + @patch("google.auth.default") + @patch("google.auth.transport.requests.Request") + def test_falls_back_to_adc_project( + self, mock_request: MagicMock, mock_default: MagicMock + ) -> None: + """Falls back to project from ADC credentials when no explicit project_id.""" + mock_creds = MagicMock() + mock_creds.valid = True + mock_creds.token = "fake-token" # nosec B105 + mock_default.return_value = (mock_creds, "adc-project") + + auth = GoogleADC() + assert auth.resolved_project_id == "adc-project" + + @patch("google.auth.default") + @patch("google.auth.transport.requests.Request") + def test_returns_none_when_no_project( + self, mock_request: MagicMock, mock_default: MagicMock + ) -> None: + """Returns None when neither explicit nor ADC project is available.""" + mock_creds = MagicMock() + mock_creds.valid = True + mock_creds.token = "fake-token" # nosec B105 + mock_default.return_value = (mock_creds, None) + + auth = GoogleADC() + assert auth.resolved_project_id is None + + +class TestGoogleADCGetVertexBaseUrl: + """Test get_vertex_base_url method.""" + + def test_global_location(self) -> None: + """Global location returns the global Vertex AI endpoint.""" + auth = GoogleADC(location="global") + assert auth.get_vertex_base_url() == VERTEX_GLOBAL_BASE_URL + + def test_regional_location(self) -> None: + """Regional location returns location-prefixed endpoint.""" + auth = GoogleADC(location="us-central1") + assert ( + auth.get_vertex_base_url() + == "https://us-central1-aiplatform.googleapis.com" + ) + + def test_europe_location(self) -> None: + """European location returns correct regional endpoint.""" + auth = GoogleADC(location="europe-west4") + assert ( + auth.get_vertex_base_url() + == "https://europe-west4-aiplatform.googleapis.com" + ) + + +# --- Helpers for client URL testing --- + + +def _make_mock_model(model_id: str = "gemini-2.0-flash") -> Model: + """Create a minimal Model for testing.""" + return Model(id=model_id, provider="google", display_name="Test Model") + + +def _make_adc( + project_id: str | None = "test-project", + location: str = "us-central1", +) -> GoogleADC: + """Create a GoogleADC with pre-loaded credentials to avoid real auth calls.""" + auth = GoogleADC(project_id=project_id, location=location) + # Pre-set private attrs so resolved_project_id doesn't trigger real google.auth + object.__setattr__(auth, "_credentials", MagicMock(valid=True, token="fake")) # nosec B106 + object.__setattr__(auth, "_project", "adc-fallback-project") + return auth + + +def _make_api_key() -> AuthHeader: + """Create a simple API key auth.""" + return AuthHeader( + secret=SecretStr("test-api-key"), + header="x-goog-api-key", + prefix="", + ) + + +# --- GenerateContent routing --- + + +class TestGenerateContentRouting: + """Test _build_url in GoogleGenerateContentClient.""" + + def _make_client(self, auth: AuthHeader | GoogleADC) -> MagicMock: + from celeste.providers.google.generate_content.client import ( + GoogleGenerateContentClient, + ) + + client = MagicMock(spec=GoogleGenerateContentClient) + client.auth = auth + client.model = _make_mock_model() + # Bind the real methods + client._build_url = GoogleGenerateContentClient._build_url.__get__(client) + client._get_vertex_endpoint = ( + GoogleGenerateContentClient._get_vertex_endpoint.__get__(client) + ) + return client + + def test_api_key_uses_gemini_endpoint(self) -> None: + from celeste.providers.google.generate_content import config + + client = self._make_client(_make_api_key()) + url = client._build_url(config.GoogleGenerateContentEndpoint.GENERATE_CONTENT) + assert ( + url + == "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent" + ) + + def test_adc_uses_vertex_endpoint(self) -> None: + client = self._make_client(_make_adc()) + from celeste.providers.google.generate_content import config + + url = client._build_url(config.GoogleGenerateContentEndpoint.GENERATE_CONTENT) + assert "us-central1-aiplatform.googleapis.com" in url + assert "projects/test-project" in url + assert "locations/us-central1" in url + assert "gemini-2.0-flash:generateContent" in url + + def test_adc_global_location(self) -> None: + from celeste.providers.google.generate_content import config + + client = self._make_client(_make_adc(location="global")) + url = client._build_url(config.GoogleGenerateContentEndpoint.GENERATE_CONTENT) + assert url.startswith("https://aiplatform.googleapis.com") + + def test_adc_stream_endpoint(self) -> None: + from celeste.providers.google.generate_content import config + + client = self._make_client(_make_adc()) + url = client._build_url( + config.GoogleGenerateContentEndpoint.STREAM_GENERATE_CONTENT + ) + assert "streamGenerateContent" in url + assert "us-central1-aiplatform.googleapis.com" in url + + def test_adc_no_project_raises(self) -> None: + client = self._make_client(_make_adc(project_id=None)) + # Override the fallback too + object.__setattr__(client.auth, "_project", None) + from celeste.providers.google.generate_content import config + + with pytest.raises(ValueError, match="Vertex AI requires a project_id"): + client._build_url(config.GoogleGenerateContentEndpoint.GENERATE_CONTENT) + + +# --- Anthropic Messages routing --- + + +class TestAnthropicMessagesRouting: + """Test _build_url in AnthropicMessagesClient.""" + + def _make_client(self, auth: AuthHeader | GoogleADC) -> MagicMock: + from celeste.providers.anthropic.messages.client import AnthropicMessagesClient + + client = MagicMock(spec=AnthropicMessagesClient) + client.auth = auth + client.model = _make_mock_model("claude-sonnet-4-5-20250929") + client._build_url = AnthropicMessagesClient._build_url.__get__(client) + client._get_vertex_endpoint = ( + AnthropicMessagesClient._get_vertex_endpoint.__get__(client) + ) + return client + + def test_api_key_uses_anthropic_endpoint(self) -> None: + from celeste.providers.anthropic.messages import config + + client = self._make_client(_make_api_key()) + url = client._build_url(config.AnthropicMessagesEndpoint.CREATE_MESSAGE) + assert url == "https://api.anthropic.com/v1/messages" + + def test_adc_uses_vertex_rawpredict(self) -> None: + from celeste.providers.anthropic.messages import config + + client = self._make_client(_make_adc()) + url = client._build_url( + config.AnthropicMessagesEndpoint.CREATE_MESSAGE, streaming=False + ) + assert "us-central1-aiplatform.googleapis.com" in url + assert "projects/test-project" in url + assert "publishers/anthropic" in url + assert "rawPredict" in url + + def test_adc_streaming_uses_stream_rawpredict(self) -> None: + from celeste.providers.anthropic.messages import config + + client = self._make_client(_make_adc()) + url = client._build_url( + config.AnthropicMessagesEndpoint.CREATE_MESSAGE, streaming=True + ) + assert "streamRawPredict" in url + + def test_adc_no_project_raises(self) -> None: + from celeste.providers.anthropic.messages import config + + client = self._make_client(_make_adc(project_id=None)) + object.__setattr__(client.auth, "_project", None) + with pytest.raises(ValueError, match="Vertex AI requires a project_id"): + client._build_url(config.AnthropicMessagesEndpoint.CREATE_MESSAGE) + + +# --- Imagen routing --- + + +class TestImagenRouting: + """Test _build_url in GoogleImagenClient.""" + + def _make_client(self, auth: AuthHeader | GoogleADC) -> MagicMock: + from celeste.providers.google.imagen.client import GoogleImagenClient + + client = MagicMock(spec=GoogleImagenClient) + client.auth = auth + client.model = _make_mock_model("imagen-3.0-generate-002") + client._build_url = GoogleImagenClient._build_url.__get__(client) + client._get_vertex_endpoint = GoogleImagenClient._get_vertex_endpoint.__get__( + client + ) + return client + + def test_api_key_uses_gemini_endpoint(self) -> None: + from celeste.providers.google.imagen import config + + client = self._make_client(_make_api_key()) + url = client._build_url(config.GoogleImagenEndpoint.CREATE_IMAGE) + assert ( + url + == "https://generativelanguage.googleapis.com/v1beta/models/imagen-3.0-generate-002:predict" + ) + + def test_adc_uses_vertex_endpoint(self) -> None: + from celeste.providers.google.imagen import config + + client = self._make_client(_make_adc()) + url = client._build_url(config.GoogleImagenEndpoint.CREATE_IMAGE) + assert "us-central1-aiplatform.googleapis.com" in url + assert "projects/test-project" in url + assert "publishers/google" in url + assert "imagen-3.0-generate-002:predict" in url + + def test_adc_no_project_raises(self) -> None: + from celeste.providers.google.imagen import config + + client = self._make_client(_make_adc(project_id=None)) + object.__setattr__(client.auth, "_project", None) + with pytest.raises(ValueError, match="Vertex AI requires a project_id"): + client._build_url(config.GoogleImagenEndpoint.CREATE_IMAGE) + + +# --- Embeddings routing --- + + +class TestEmbeddingsRouting: + """Test _build_url in GoogleEmbeddingsClient.""" + + def _make_client(self, auth: AuthHeader | GoogleADC) -> MagicMock: + from celeste.providers.google.embeddings.client import GoogleEmbeddingsClient + + client = MagicMock(spec=GoogleEmbeddingsClient) + client.auth = auth + client.model = _make_mock_model("text-embedding-004") + client._build_url = GoogleEmbeddingsClient._build_url.__get__(client) + client._get_vertex_endpoint = ( + GoogleEmbeddingsClient._get_vertex_endpoint.__get__(client) + ) + return client + + def test_api_key_uses_gemini_endpoint(self) -> None: + from celeste.providers.google.embeddings import config + + client = self._make_client(_make_api_key()) + url = client._build_url(config.GoogleEmbeddingsEndpoint.EMBED_CONTENT) + assert ( + url + == "https://generativelanguage.googleapis.com/v1beta/models/text-embedding-004:embedContent" + ) + + def test_adc_uses_vertex_endpoint(self) -> None: + from celeste.providers.google.embeddings import config + + client = self._make_client(_make_adc()) + url = client._build_url(config.GoogleEmbeddingsEndpoint.EMBED_CONTENT) + assert "us-central1-aiplatform.googleapis.com" in url + assert "projects/test-project" in url + assert "publishers/google" in url + assert "text-embedding-004:predict" in url + + def test_adc_batch_uses_vertex_endpoint(self) -> None: + from celeste.providers.google.embeddings import config + + client = self._make_client(_make_adc()) + url = client._build_url(config.GoogleEmbeddingsEndpoint.BATCH_EMBED_CONTENTS) + assert "us-central1-aiplatform.googleapis.com" in url + assert "text-embedding-004:predict" in url + + def test_adc_no_project_raises(self) -> None: + from celeste.providers.google.embeddings import config + + client = self._make_client(_make_adc(project_id=None)) + object.__setattr__(client.auth, "_project", None) + with pytest.raises(ValueError, match="Vertex AI requires a project_id"): + client._build_url(config.GoogleEmbeddingsEndpoint.EMBED_CONTENT) + + +# --- Veo routing --- + + +class TestVeoRouting: + """Test _build_url and _build_poll_url in GoogleVeoClient.""" + + def _make_client(self, auth: AuthHeader | GoogleADC) -> MagicMock: + from celeste.providers.google.veo.client import GoogleVeoClient + + client = MagicMock(spec=GoogleVeoClient) + client.auth = auth + client.model = _make_mock_model("veo-2.0-generate-001") + client._build_url = GoogleVeoClient._build_url.__get__(client) + client._get_vertex_endpoint = GoogleVeoClient._get_vertex_endpoint.__get__( + client + ) + client._build_poll_url = GoogleVeoClient._build_poll_url.__get__(client) + return client + + def test_api_key_uses_gemini_endpoint(self) -> None: + from celeste.providers.google.veo import config + + client = self._make_client(_make_api_key()) + url = client._build_url(config.GoogleVeoEndpoint.CREATE_VIDEO) + assert ( + url + == "https://generativelanguage.googleapis.com/v1beta/models/veo-2.0-generate-001:predictLongRunning" + ) + + def test_adc_uses_vertex_endpoint(self) -> None: + from celeste.providers.google.veo import config + + client = self._make_client(_make_adc()) + url = client._build_url(config.GoogleVeoEndpoint.CREATE_VIDEO) + assert "us-central1-aiplatform.googleapis.com" in url + assert "projects/test-project" in url + assert "publishers/google" in url + assert "veo-2.0-generate-001:predictLongRunning" in url + + def test_api_key_poll_url(self) -> None: + client = self._make_client(_make_api_key()) + url = client._build_poll_url("operations/abc123") + assert ( + url == "https://generativelanguage.googleapis.com/v1beta/operations/abc123" + ) + + def test_adc_poll_url(self) -> None: + client = self._make_client(_make_adc()) + url = client._build_poll_url( + "projects/test-project/locations/us-central1/operations/abc123" + ) + assert "us-central1-aiplatform.googleapis.com" in url + assert "fetchPredictOperation" in url + + def test_adc_no_project_raises(self) -> None: + from celeste.providers.google.veo import config + + client = self._make_client(_make_adc(project_id=None)) + object.__setattr__(client.auth, "_project", None) + with pytest.raises(ValueError, match="Vertex AI requires a project_id"): + client._build_url(config.GoogleVeoEndpoint.CREATE_VIDEO) + + +# --- Mistral routing --- + + +class TestMistralRouting: + """Test _build_url in MistralChatClient.""" + + def _make_client(self, auth: AuthHeader | GoogleADC) -> MagicMock: + from celeste.providers.mistral.chat.client import MistralChatClient + + client = MagicMock(spec=MistralChatClient) + client.auth = auth + client.model = _make_mock_model("mistral-large-2411") + client._build_url = MistralChatClient._build_url.__get__(client) + client._get_vertex_endpoint = MistralChatClient._get_vertex_endpoint.__get__( + client + ) + return client + + def test_api_key_uses_mistral_endpoint(self) -> None: + from celeste.providers.mistral.chat import config + + client = self._make_client(_make_api_key()) + url = client._build_url(config.MistralChatEndpoint.CREATE_CHAT_COMPLETION) + assert url == "https://api.mistral.ai/v1/chat/completions" + + def test_adc_uses_vertex_rawpredict(self) -> None: + from celeste.providers.mistral.chat import config + + client = self._make_client(_make_adc()) + url = client._build_url(config.MistralChatEndpoint.CREATE_CHAT_COMPLETION) + assert "us-central1-aiplatform.googleapis.com" in url + assert "projects/test-project" in url + assert "publishers/mistralai" in url + assert "mistral-large-2411:rawPredict" in url + + def test_adc_streaming_uses_stream_rawpredict(self) -> None: + from celeste.providers.mistral.chat import config + + client = self._make_client(_make_adc()) + url = client._build_url( + config.MistralChatEndpoint.CREATE_CHAT_COMPLETION, streaming=True + ) + assert "streamRawPredict" in url + assert "publishers/mistralai" in url + + def test_adc_global_location(self) -> None: + from celeste.providers.mistral.chat import config + + client = self._make_client(_make_adc(location="global")) + url = client._build_url(config.MistralChatEndpoint.CREATE_CHAT_COMPLETION) + assert url.startswith("https://aiplatform.googleapis.com") + + def test_adc_no_project_raises(self) -> None: + from celeste.providers.mistral.chat import config + + client = self._make_client(_make_adc(project_id=None)) + object.__setattr__(client.auth, "_project", None) + with pytest.raises(ValueError, match="Vertex AI requires a project_id"): + client._build_url(config.MistralChatEndpoint.CREATE_CHAT_COMPLETION) + + +# --- DeepSeek routing --- + + +class TestDeepSeekRouting: + """Test _build_url in DeepSeekChatClient.""" + + def _make_client(self, auth: AuthHeader | GoogleADC) -> MagicMock: + from celeste.providers.deepseek.chat.client import DeepSeekChatClient + + client = MagicMock(spec=DeepSeekChatClient) + client.auth = auth + client.model = _make_mock_model("deepseek-ai/deepseek-v3-0324") + client._build_url = DeepSeekChatClient._build_url.__get__(client) + client._get_vertex_endpoint = DeepSeekChatClient._get_vertex_endpoint.__get__( + client + ) + return client + + def test_api_key_uses_deepseek_endpoint(self) -> None: + from celeste.providers.deepseek.chat import config + + client = self._make_client(_make_api_key()) + url = client._build_url(config.DeepSeekChatEndpoint.CREATE_CHAT) + assert url == "https://api.deepseek.com/v1/chat/completions" + + def test_adc_uses_vertex_openai_compatible(self) -> None: + from celeste.providers.deepseek.chat import config + + client = self._make_client(_make_adc()) + url = client._build_url(config.DeepSeekChatEndpoint.CREATE_CHAT) + assert "us-central1-aiplatform.googleapis.com" in url + assert "projects/test-project" in url + assert "endpoints/openapi/chat/completions" in url + + def test_adc_url_has_no_model_id(self) -> None: + """Model ID should be in request body only, not in the URL.""" + from celeste.providers.deepseek.chat import config + + client = self._make_client(_make_adc()) + url = client._build_url(config.DeepSeekChatEndpoint.CREATE_CHAT) + assert "deepseek" not in url.lower().split("endpoints")[1] + + def test_adc_global_location(self) -> None: + from celeste.providers.deepseek.chat import config + + client = self._make_client(_make_adc(location="global")) + url = client._build_url(config.DeepSeekChatEndpoint.CREATE_CHAT) + assert url.startswith("https://aiplatform.googleapis.com") + + def test_adc_no_project_raises(self) -> None: + from celeste.providers.deepseek.chat import config + + client = self._make_client(_make_adc(project_id=None)) + object.__setattr__(client.auth, "_project", None) + with pytest.raises(ValueError, match="Vertex AI requires a project_id"): + client._build_url(config.DeepSeekChatEndpoint.CREATE_CHAT) + + +# --- Cloud TTS auth preservation --- + + +class TestCloudTTSAuthPreservation: + """Test that Cloud TTS preserves user-provided GoogleADC.""" + + def test_preserves_user_provided_adc(self) -> None: + """User-provided GoogleADC should not be overwritten by model_post_init logic.""" + # Test the conditional logic: isinstance check should preserve existing GoogleADC + user_adc = GoogleADC(project_id="my-project", location="europe-west4") + assert isinstance(user_adc, GoogleADC) + + # Verify the source code has the isinstance guard + import inspect + + from celeste.providers.google.cloud_tts.client import GoogleCloudTTSClient + + source = inspect.getsource(GoogleCloudTTSClient.model_post_init) + assert "isinstance(self.auth, GoogleADC)" in source + + def test_non_adc_auth_gets_replaced(self) -> None: + """Non-GoogleADC auth should trigger replacement (verified via source inspection).""" + api_key = _make_api_key() + assert not isinstance(api_key, GoogleADC) + + import inspect + + from celeste.providers.google.cloud_tts.client import GoogleCloudTTSClient + + source = inspect.getsource(GoogleCloudTTSClient.model_post_init) + # Ensure the conditional creates new GoogleADC when auth is not GoogleADC + assert "if not isinstance(self.auth, GoogleADC)" in source + assert 'object.__setattr__(self, "auth", GoogleADC())' in source From d1e42c532fe844ae4d27b45d6e1f1ac83c8b8113 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 9 Feb 2026 11:54:37 +0100 Subject: [PATCH 2/9] fix(ci): install gcp extra so Vertex routing tests can import google-auth The google-auth package is optional under [gcp], but unit tests in test_vertex_routing.py and Vertex integration tests need it importable. Co-Authored-By: Claude Opus 4.6 --- .github/workflows/ci.yml | 1 + .github/workflows/publish.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ad0caba2..36b7abae 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -85,6 +85,7 @@ jobs: - uses: ./.github/actions/setup-python-uv with: python-version: ${{ matrix.python-version }} + - run: uv sync --extra gcp - run: uv run pytest tests/unit_tests -v --cov=celeste --cov-report=term-missing --cov-report=xml --cov-report=html --cov-fail-under=80 - uses: codecov/codecov-action@v4 if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.12' diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 8cf7b8f3..ab94de10 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -58,6 +58,7 @@ jobs: workload_identity_provider: ${{ secrets.GCP_WORKLOAD_IDENTITY_PROVIDER }} service_account: ${{ secrets.GCP_SERVICE_ACCOUNT }} - uses: ./.github/actions/setup-python-uv + - run: uv sync --extra gcp - name: Run ${{ matrix.package }} integration tests env: OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} From ff8b03c7b04b24a09134491d56308fab78b0583f Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 9 Feb 2026 11:54:50 +0100 Subject: [PATCH 3/9] fix(veo): handle Vertex inline bytesBase64Encoded and videoGcsUri key mismatch Vertex Veo responses use videoGcsUri (not uri/gcsUri) and can return inline base64 instead of a GCS URL. Normalize the key and decode inline responses directly into VideoArtifact. Co-Authored-By: Claude Opus 4.6 --- .../videos/providers/google/client.py | 18 +++++++++++++----- src/celeste/providers/google/veo/client.py | 6 +++++- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/src/celeste/modalities/videos/providers/google/client.py b/src/celeste/modalities/videos/providers/google/client.py index fda19da7..194c7609 100644 --- a/src/celeste/modalities/videos/providers/google/client.py +++ b/src/celeste/modalities/videos/providers/google/client.py @@ -1,8 +1,10 @@ """Google videos client.""" +import base64 from typing import Any, Unpack from celeste.artifacts import VideoArtifact +from celeste.mime_types import VideoMimeType from celeste.parameters import ParameterMapper from celeste.providers.google.veo import config from celeste.providers.google.veo.client import GoogleVeoClient as GoogleVeoMixin @@ -51,9 +53,12 @@ def _parse_content( ) -> VideoArtifact: """Parse content from response.""" video_data = super()._parse_content(response_data) - return VideoArtifact( - url=video_data.get("uri") or video_data.get("gcsUri"), - ) + # Handle inline base64 response (Vertex can return bytesBase64Encoded) + if "bytesBase64Encoded" in video_data: + video_bytes = base64.b64decode(video_data["bytesBase64Encoded"]) + mime_type = video_data.get("mimeType", VideoMimeType.MP4) + return VideoArtifact(data=video_bytes, mime_type=mime_type) + return VideoArtifact(url=video_data.get("uri")) def _parse_finish_reason(self, response_data: dict[str, Any]) -> VideoFinishReason: """Parse finish reason from response.""" @@ -64,13 +69,16 @@ async def download_content(self, artifact: VideoArtifact) -> VideoArtifact: """Download video content from GCS URL. Args: - artifact: VideoArtifact with URL to download. + artifact: VideoArtifact with URL or inline data to download. Returns: VideoArtifact with downloaded bytes data. """ + if artifact.data is not None: + return artifact + if artifact.url is None: - msg = "Artifact has no URL to download" + msg = "Artifact has no URL or data to download" raise ValueError(msg) video_bytes = await super().download_content(artifact.url) diff --git a/src/celeste/providers/google/veo/client.py b/src/celeste/providers/google/veo/client.py index f3441fd2..20dfb66b 100644 --- a/src/celeste/providers/google/veo/client.py +++ b/src/celeste/providers/google/veo/client.py @@ -180,7 +180,11 @@ def _parse_content(self, response_data: dict[str, Any]) -> Any: if not videos: msg = "No videos in response" raise ValueError(msg) - return videos[0] + video = videos[0] + # Normalize Vertex key "videoGcsUri" to "uri" for consistency + if "videoGcsUri" in video: + video["uri"] = video.pop("videoGcsUri") + return video generated_samples = response.get("generateVideoResponse", {}).get( "generatedSamples", [] From 24acae0de1b60c446c56a24b0737f0fdf4b8f9a5 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 9 Feb 2026 11:55:52 +0100 Subject: [PATCH 4/9] fix(embeddings): adapt request/response format for Vertex :predict endpoint Vertex embeddings uses :predict with instances format, not :embedContent. Build correct request body in _init_request when auth is GoogleADC, and parse predictions response format in _parse_content. Add integration test. Co-Authored-By: Claude Opus 4.6 --- .../embeddings/providers/google/client.py | 5 +++++ .../providers/google/embeddings/client.py | 13 +++++++++--- .../embeddings/test_embed.py | 21 +++++++++++++++++++ 3 files changed, 36 insertions(+), 3 deletions(-) diff --git a/src/celeste/modalities/embeddings/providers/google/client.py b/src/celeste/modalities/embeddings/providers/google/client.py index d09cc43f..a24f21e9 100644 --- a/src/celeste/modalities/embeddings/providers/google/client.py +++ b/src/celeste/modalities/embeddings/providers/google/client.py @@ -3,6 +3,7 @@ from typing import Any, Unpack from celeste.parameters import ParameterMapper +from celeste.providers.google.auth import GoogleADC from celeste.providers.google.embeddings.client import ( GoogleEmbeddingsClient as GoogleEmbeddingsMixin, ) @@ -30,6 +31,10 @@ def _init_request(self, inputs: EmbeddingsInput) -> dict[str, Any]: """Build Google embeddings request from inputs.""" texts = inputs.text if isinstance(inputs.text, list) else [inputs.text] + # Vertex :predict endpoint uses instances format + if isinstance(self.auth, GoogleADC): + return {"instances": [{"content": text} for text in texts]} + if len(texts) == 1: return {"content": {"parts": [{"text": texts[0]}]}} else: diff --git a/src/celeste/providers/google/embeddings/client.py b/src/celeste/providers/google/embeddings/client.py index 76270e17..e22010dd 100644 --- a/src/celeste/providers/google/embeddings/client.py +++ b/src/celeste/providers/google/embeddings/client.py @@ -109,16 +109,23 @@ def _parse_content(self, response_data: dict[str, Any]) -> list[list[float]]: """Extract embedding vectors from response. Returns list of embedding vectors (already generic - no artifacts needed). + Handles both Gemini API and Vertex AI :predict response formats. """ - # Single embedding response + # Vertex :predict response + if "predictions" in response_data: + return [ + pred["embeddings"]["values"] for pred in response_data["predictions"] + ] + + # Gemini single embedding response if "embedding" in response_data: return [response_data["embedding"]["values"]] - # Batch embedding response + # Gemini batch embedding response if "embeddings" in response_data: return [emb["values"] for emb in response_data["embeddings"]] - msg = "Unexpected response format: missing 'embedding' or 'embeddings' field" + msg = "Unexpected response format: missing 'embedding', 'embeddings', or 'predictions' field" raise ValueError(msg) def _parse_usage( diff --git a/tests/integration_tests/embeddings/test_embed.py b/tests/integration_tests/embeddings/test_embed.py index c558c04f..1d79b0c7 100644 --- a/tests/integration_tests/embeddings/test_embed.py +++ b/tests/integration_tests/embeddings/test_embed.py @@ -22,6 +22,7 @@ EmbeddingsOutput, EmbeddingsUsage, ) +from celeste.providers.google.auth import GoogleADC # noqa: E402 @pytest.mark.parametrize( @@ -105,3 +106,23 @@ def test_sync_embed() -> None: assert isinstance(response, EmbeddingsOutput) assert response.content is not None + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_vertex_embed() -> None: + """Test embeddings via Vertex AI.""" + client = create_client( + modality=Modality.EMBEDDINGS, + provider="google", + model="gemini-embedding-001", + auth=GoogleADC(), + ) + + response = await client.embed("Hello world") + + assert isinstance(response, EmbeddingsOutput) + assert response.content is not None + assert isinstance(response.content, list) + assert len(response.content) > 0 + assert isinstance(response.content[0], float) From a43d0273547240b31433bc8c3eac6884775b4ba2 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 9 Feb 2026 11:57:22 +0100 Subject: [PATCH 5/9] chore: fix trailing newlines in workflow files Co-Authored-By: Claude Opus 4.6 --- .github/workflows/claude-code-review.yml | 1 - .github/workflows/claude.yml | 1 - 2 files changed, 2 deletions(-) diff --git a/.github/workflows/claude-code-review.yml b/.github/workflows/claude-code-review.yml index b5e8cfd4..25f4ad18 100644 --- a/.github/workflows/claude-code-review.yml +++ b/.github/workflows/claude-code-review.yml @@ -41,4 +41,3 @@ jobs: prompt: '/code-review:code-review ${{ github.repository }}/pull/${{ github.event.pull_request.number }}' # See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md # or https://code.claude.com/docs/en/cli-reference for available options - diff --git a/.github/workflows/claude.yml b/.github/workflows/claude.yml index d300267f..9471a059 100644 --- a/.github/workflows/claude.yml +++ b/.github/workflows/claude.yml @@ -47,4 +47,3 @@ jobs: # See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md # or https://code.claude.com/docs/en/cli-reference for available options # claude_args: '--allowed-tools Bash(gh pr:*)' - From a9b51ad42508a90c5a37da033e0ada62a82625cc Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 9 Feb 2026 12:15:22 +0100 Subject: [PATCH 6/9] refactor(vertex): move embeddings auth to provider mixin, rename VertexEndpoint, update templates - Move isinstance(self.auth, GoogleADC) check from modality _init_request() to provider mixin _make_request() for embeddings, keeping auth logic in provider layer - Fix misplaced class docstring in GoogleEmbeddingsClient mixin - Rename VertexEndpoint to VertexGenerateContentEndpoint for consistency with VertexImagenEndpoint, VertexEmbeddingsEndpoint, etc. - Add Vertex AI routing patterns (commented) to provider templates Co-Authored-By: Claude Opus 4.6 --- .../embeddings/providers/google/client.py | 5 --- .../providers/google/embeddings/client.py | 43 ++++++++++++------- .../google/generate_content/client.py | 6 +-- .../google/generate_content/config.py | 2 +- .../{api_slug}/client.py.template | 35 +++++++++++++++ .../{api_slug}/config.py.template | 9 ++++ 6 files changed, 75 insertions(+), 25 deletions(-) diff --git a/src/celeste/modalities/embeddings/providers/google/client.py b/src/celeste/modalities/embeddings/providers/google/client.py index a24f21e9..d09cc43f 100644 --- a/src/celeste/modalities/embeddings/providers/google/client.py +++ b/src/celeste/modalities/embeddings/providers/google/client.py @@ -3,7 +3,6 @@ from typing import Any, Unpack from celeste.parameters import ParameterMapper -from celeste.providers.google.auth import GoogleADC from celeste.providers.google.embeddings.client import ( GoogleEmbeddingsClient as GoogleEmbeddingsMixin, ) @@ -31,10 +30,6 @@ def _init_request(self, inputs: EmbeddingsInput) -> dict[str, Any]: """Build Google embeddings request from inputs.""" texts = inputs.text if isinstance(inputs.text, list) else [inputs.text] - # Vertex :predict endpoint uses instances format - if isinstance(self.auth, GoogleADC): - return {"instances": [{"content": text} for text in texts]} - if len(texts) == 1: return {"content": {"parts": [{"text": texts[0]}]}} else: diff --git a/src/celeste/providers/google/embeddings/client.py b/src/celeste/providers/google/embeddings/client.py index e22010dd..ce5d343c 100644 --- a/src/celeste/providers/google/embeddings/client.py +++ b/src/celeste/providers/google/embeddings/client.py @@ -13,6 +13,22 @@ class GoogleEmbeddingsClient(APIMixin): + """Mixin for Embeddings API capabilities. + + Provides shared implementation for embeddings using the Embeddings API: + - _make_request() - HTTP POST to embedContent or batchEmbedContents endpoint + - _parse_content() - Extract embedding vectors (generic) + + Auth-based endpoint selection: + - GoogleADC auth -> Vertex AI endpoints + - API key auth -> Gemini API endpoints + + Capability clients extend via super(): + class GoogleEmbeddingsClient(GoogleEmbeddingsClient, EmbeddingsClient): + def _parse_content(self, response_data, **params): + return super()._parse_content(response_data) # No transformation needed + """ + def _make_stream_request( self, request_body: dict[str, Any], @@ -31,22 +47,6 @@ def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | float | None """ return {} - """Mixin for Embeddings API capabilities. - - Provides shared implementation for embeddings using the Embeddings API: - - _make_request() - HTTP POST to embedContent or batchEmbedContents endpoint - - _parse_content() - Extract embedding vectors (generic) - - Auth-based endpoint selection: - - GoogleADC auth -> Vertex AI endpoints - - API key auth -> Gemini API endpoints - - Capability clients extend via super(): - class GoogleEmbeddingsClient(GoogleEmbeddingsClient, EmbeddingsClient): - def _parse_content(self, response_data, **params): - return super()._parse_content(response_data) # No transformation needed - """ - def _get_vertex_endpoint(self, gemini_endpoint: str) -> str: """Map Gemini Embeddings endpoint to Vertex AI endpoint.""" mapping: dict[str, str] = { @@ -82,6 +82,17 @@ async def _make_request( **parameters: Any, ) -> dict[str, Any]: """Make HTTP request to embeddings endpoint.""" + # Vertex :predict expects {"instances": [{"content": "..."}]} format + if isinstance(self.auth, GoogleADC): + if "requests" in request_body: + texts = [ + req["content"]["parts"][0]["text"] + for req in request_body["requests"] + ] + else: + texts = [request_body["content"]["parts"][0]["text"]] + request_body = {"instances": [{"content": text} for text in texts]} + is_batch = "requests" in request_body endpoint_template = ( config.GoogleEmbeddingsEndpoint.BATCH_EMBED_CONTENTS diff --git a/src/celeste/providers/google/generate_content/client.py b/src/celeste/providers/google/generate_content/client.py index d04d70a4..bcdc3f53 100644 --- a/src/celeste/providers/google/generate_content/client.py +++ b/src/celeste/providers/google/generate_content/client.py @@ -40,9 +40,9 @@ def _parse_content(self, response_data, **parameters): def _get_vertex_endpoint(self, gemini_endpoint: str) -> str: """Map Gemini endpoint to Vertex AI endpoint.""" mapping: dict[str, str] = { - config.GoogleGenerateContentEndpoint.GENERATE_CONTENT: config.VertexEndpoint.GENERATE_CONTENT, - config.GoogleGenerateContentEndpoint.STREAM_GENERATE_CONTENT: config.VertexEndpoint.STREAM_GENERATE_CONTENT, - config.GoogleGenerateContentEndpoint.COUNT_TOKENS: config.VertexEndpoint.COUNT_TOKENS, + config.GoogleGenerateContentEndpoint.GENERATE_CONTENT: config.VertexGenerateContentEndpoint.GENERATE_CONTENT, + config.GoogleGenerateContentEndpoint.STREAM_GENERATE_CONTENT: config.VertexGenerateContentEndpoint.STREAM_GENERATE_CONTENT, + config.GoogleGenerateContentEndpoint.COUNT_TOKENS: config.VertexGenerateContentEndpoint.COUNT_TOKENS, } vertex_endpoint = mapping.get(gemini_endpoint) if vertex_endpoint is None: diff --git a/src/celeste/providers/google/generate_content/config.py b/src/celeste/providers/google/generate_content/config.py index a5fe0181..2ff96198 100644 --- a/src/celeste/providers/google/generate_content/config.py +++ b/src/celeste/providers/google/generate_content/config.py @@ -20,7 +20,7 @@ class GoogleGenerateContentEndpoint(StrEnum): BATCH_GENERATE_CONTENT = "/v1beta/models/{model_id}:batchGenerateContent" -class VertexEndpoint(StrEnum): +class VertexGenerateContentEndpoint(StrEnum): """Endpoints for Vertex AI (when using GoogleADC auth).""" GENERATE_CONTENT = "/v1/projects/{project_id}/locations/{location}/publishers/google/models/{model_id}:generateContent" diff --git a/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/client.py.template b/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/client.py.template index 8885d1f9..d900a819 100644 --- a/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/client.py.template +++ b/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/client.py.template @@ -13,6 +13,9 @@ from celeste.mime_types import ApplicationMimeType from . import config +# Optional: Import for Vertex AI support (when using GoogleADC auth). +# from ..auth import GoogleADC + class {Provider}{Api}Client(APIMixin): """Mixin for {Provider} {Api} API. @@ -26,10 +29,42 @@ class {Provider}{Api}Client(APIMixin): - _parse_finish_reason() - Extract finish reason (if provided) - _build_metadata() - Filter content fields + Auth-based endpoint selection (optional, for Google providers): + - GoogleADC auth -> Vertex AI endpoints (via _build_url + _get_vertex_endpoint) + - API key auth -> native API endpoints + Modality clients pass endpoint parameter to route operations: await self._predict(inputs, endpoint=config.{Provider}{Api}Endpoint.CREATE_..., **parameters) """ + # Optional: Vertex AI endpoint routing. Add these two methods when the provider + # supports Vertex AI. They replace the inline f"{config.BASE_URL}{endpoint}" URL + # construction in _make_request() / _make_stream_request() with self._build_url(endpoint). + # + # def _get_vertex_endpoint(self, native_endpoint: str) -> str: + # """Map native endpoint to Vertex AI endpoint.""" + # mapping: dict[str, str] = { + # config.{Provider}{Api}Endpoint.CREATE_...: config.Vertex{Api}Endpoint.CREATE_..., + # } + # vertex_endpoint = mapping.get(native_endpoint) + # if vertex_endpoint is None: + # raise ValueError(f"No Vertex AI endpoint mapping for: {native_endpoint}") + # return vertex_endpoint + # + # def _build_url(self, endpoint: str) -> str: + # """Build full URL based on auth type.""" + # if isinstance(self.auth, GoogleADC): + # project_id = self.auth.resolved_project_id + # if project_id is None: + # raise ValueError( + # "Vertex AI requires a project_id. " + # "Pass project_id to GoogleADC() or ensure credentials have a project." + # ) + # vertex_endpoint = self._get_vertex_endpoint(endpoint) + # base_url = self.auth.get_vertex_base_url() + # return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location, model_id=self.model.id)}" + # return f"{config.BASE_URL}{endpoint.format(model_id=self.model.id)}" + def _build_request( self, inputs: Any, diff --git a/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/config.py.template b/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/config.py.template index 4560a1d1..d7b09b7f 100644 --- a/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/config.py.template +++ b/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/config.py.template @@ -17,4 +17,13 @@ class {Provider}{Api}Endpoint(StrEnum): pass +# Optional: Vertex AI endpoints (when using GoogleADC auth). +# Mirror the same endpoint names as the native class above. +# +# class Vertex{Api}Endpoint(StrEnum): +# """Endpoints for {Api} on Vertex AI.""" +# +# CREATE_... = "/v1/projects/{project_id}/locations/{location}/publishers/google/models/{model_id}:predict" + + BASE_URL = "https://api.{provider}.com" From ecdb5f55943814acf5f53790e126bdb587515ea0 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 9 Feb 2026 12:53:04 +0100 Subject: [PATCH 7/9] refactor(vertex): centralize URL building in GoogleADC.build_url() Move duplicated project_id validation, base URL resolution, and endpoint formatting from 7 provider _build_url() methods into GoogleADC.build_url(). Also remove manual base64.b64decode from video client (Artifact validator handles it). Co-Authored-By: Claude Opus 4.6 --- .../videos/providers/google/client.py | 6 ++--- .../providers/anthropic/messages/client.py | 22 ++++------------ src/celeste/providers/deepseek/chat/client.py | 18 ++----------- src/celeste/providers/google/auth.py | 25 ++++++++++++++++++ .../providers/google/embeddings/client.py | 14 +++------- .../google/generate_content/client.py | 21 +++------------ src/celeste/providers/google/imagen/client.py | 14 +++------- src/celeste/providers/google/veo/client.py | 26 +++++-------------- src/celeste/providers/mistral/chat/client.py | 21 ++++----------- .../{api_slug}/client.py.template | 12 +++------ 10 files changed, 59 insertions(+), 120 deletions(-) diff --git a/src/celeste/modalities/videos/providers/google/client.py b/src/celeste/modalities/videos/providers/google/client.py index 194c7609..c8a24e75 100644 --- a/src/celeste/modalities/videos/providers/google/client.py +++ b/src/celeste/modalities/videos/providers/google/client.py @@ -1,6 +1,5 @@ """Google videos client.""" -import base64 from typing import Any, Unpack from celeste.artifacts import VideoArtifact @@ -55,9 +54,10 @@ def _parse_content( video_data = super()._parse_content(response_data) # Handle inline base64 response (Vertex can return bytesBase64Encoded) if "bytesBase64Encoded" in video_data: - video_bytes = base64.b64decode(video_data["bytesBase64Encoded"]) mime_type = video_data.get("mimeType", VideoMimeType.MP4) - return VideoArtifact(data=video_bytes, mime_type=mime_type) + return VideoArtifact( + data=video_data["bytesBase64Encoded"], mime_type=mime_type + ) return VideoArtifact(url=video_data.get("uri")) def _parse_finish_reason(self, response_data: dict[str, Any]) -> VideoFinishReason: diff --git a/src/celeste/providers/anthropic/messages/client.py b/src/celeste/providers/anthropic/messages/client.py index 1919cfd1..4253710e 100644 --- a/src/celeste/providers/anthropic/messages/client.py +++ b/src/celeste/providers/anthropic/messages/client.py @@ -46,24 +46,12 @@ def _get_vertex_endpoint( return config.VertexAnthropicEndpoint.CREATE_MESSAGE def _build_url(self, endpoint: str, streaming: bool = False) -> str: - """Build full URL based on auth type. - - - GoogleADC auth -> Vertex AI endpoints - - API key auth -> Anthropic API endpoints - """ + """Build full URL based on auth type.""" if isinstance(self.auth, GoogleADC): - project_id = self.auth.resolved_project_id - if project_id is None: - raise ValueError( - "Vertex AI requires a project_id. " - "Pass project_id to GoogleADC() or ensure credentials have a project." - ) - - vertex_endpoint = self._get_vertex_endpoint(endpoint, streaming=streaming) - base_url = self.auth.get_vertex_base_url() - return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location, model_id=self.model.id)}" - - # Default: Anthropic API + return self.auth.build_url( + self._get_vertex_endpoint(endpoint, streaming=streaming), + model_id=self.model.id, + ) return f"{config.BASE_URL}{endpoint}" def _build_headers(self, request_body: dict[str, Any]) -> dict[str, str]: diff --git a/src/celeste/providers/deepseek/chat/client.py b/src/celeste/providers/deepseek/chat/client.py index 19819290..6d86adeb 100644 --- a/src/celeste/providers/deepseek/chat/client.py +++ b/src/celeste/providers/deepseek/chat/client.py @@ -43,23 +43,9 @@ def _get_vertex_endpoint(self, deepseek_endpoint: str) -> str: return vertex_endpoint def _build_url(self, endpoint: str) -> str: - """Build full URL based on auth type. - - - GoogleADC auth -> Vertex AI endpoints (OpenAI-compatible) - - API key auth -> DeepSeek API endpoints - """ + """Build full URL based on auth type.""" if isinstance(self.auth, GoogleADC): - project_id = self.auth.resolved_project_id - if project_id is None: - raise ValueError( - "Vertex AI requires a project_id. " - "Pass project_id to GoogleADC() or ensure credentials have a project." - ) - - vertex_endpoint = self._get_vertex_endpoint(endpoint) - base_url = self.auth.get_vertex_base_url() - return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location)}" - + return self.auth.build_url(self._get_vertex_endpoint(endpoint)) return f"{config.BASE_URL}{endpoint}" def _build_request( diff --git a/src/celeste/providers/google/auth.py b/src/celeste/providers/google/auth.py index 24eba8d8..df4e4cfb 100644 --- a/src/celeste/providers/google/auth.py +++ b/src/celeste/providers/google/auth.py @@ -68,5 +68,30 @@ def get_vertex_base_url(self) -> str: return VERTEX_GLOBAL_BASE_URL return VERTEX_BASE_URL.format(location=self.location) + def build_url(self, vertex_endpoint: str, model_id: str | None = None) -> str: + """Build a complete Vertex AI URL from an endpoint template. + + Args: + vertex_endpoint: Endpoint template with {project_id}, {location}, and + optionally {model_id} placeholders. + model_id: Model identifier for endpoints that include it in the path. + """ + project_id = self.resolved_project_id + if project_id is None: + raise ValueError( + "Vertex AI requires a project_id. " + "Pass project_id to GoogleADC() or ensure credentials have a project." + ) + base_url = self.get_vertex_base_url() + if model_id is not None: + endpoint = vertex_endpoint.format( + project_id=project_id, location=self.location, model_id=model_id + ) + else: + endpoint = vertex_endpoint.format( + project_id=project_id, location=self.location + ) + return f"{base_url}{endpoint}" + __all__ = ["GoogleADC"] diff --git a/src/celeste/providers/google/embeddings/client.py b/src/celeste/providers/google/embeddings/client.py index ce5d343c..acc5a619 100644 --- a/src/celeste/providers/google/embeddings/client.py +++ b/src/celeste/providers/google/embeddings/client.py @@ -61,17 +61,9 @@ def _get_vertex_endpoint(self, gemini_endpoint: str) -> str: def _build_url(self, endpoint: str) -> str: """Build full URL based on auth type.""" if isinstance(self.auth, GoogleADC): - project_id = self.auth.resolved_project_id - if project_id is None: - raise ValueError( - "Vertex AI requires a project_id. " - "Pass project_id to GoogleADC() or ensure credentials have a project." - ) - - vertex_endpoint = self._get_vertex_endpoint(endpoint) - base_url = self.auth.get_vertex_base_url() - return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location, model_id=self.model.id)}" - + return self.auth.build_url( + self._get_vertex_endpoint(endpoint), model_id=self.model.id + ) return f"{config.BASE_URL}{endpoint.format(model_id=self.model.id)}" async def _make_request( diff --git a/src/celeste/providers/google/generate_content/client.py b/src/celeste/providers/google/generate_content/client.py index bcdc3f53..61d6e58d 100644 --- a/src/celeste/providers/google/generate_content/client.py +++ b/src/celeste/providers/google/generate_content/client.py @@ -50,24 +50,11 @@ def _get_vertex_endpoint(self, gemini_endpoint: str) -> str: return vertex_endpoint def _build_url(self, endpoint: str) -> str: - """Build full URL based on auth type. - - - GoogleADC auth -> Vertex AI endpoints - - API key auth -> Gemini API endpoints - """ + """Build full URL based on auth type.""" if isinstance(self.auth, GoogleADC): - project_id = self.auth.resolved_project_id - if project_id is None: - raise ValueError( - "Vertex AI requires a project_id. " - "Pass project_id to GoogleADC() or ensure credentials have a project." - ) - - vertex_endpoint = self._get_vertex_endpoint(endpoint) - base_url = self.auth.get_vertex_base_url() - return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location, model_id=self.model.id)}" - - # Default: Gemini API + return self.auth.build_url( + self._get_vertex_endpoint(endpoint), model_id=self.model.id + ) return f"{config.BASE_URL}{endpoint.format(model_id=self.model.id)}" async def _make_request( diff --git a/src/celeste/providers/google/imagen/client.py b/src/celeste/providers/google/imagen/client.py index 5dd94039..778b285f 100644 --- a/src/celeste/providers/google/imagen/client.py +++ b/src/celeste/providers/google/imagen/client.py @@ -49,17 +49,9 @@ def _get_vertex_endpoint(self, gemini_endpoint: str) -> str: def _build_url(self, endpoint: str) -> str: """Build full URL based on auth type.""" if isinstance(self.auth, GoogleADC): - project_id = self.auth.resolved_project_id - if project_id is None: - raise ValueError( - "Vertex AI requires a project_id. " - "Pass project_id to GoogleADC() or ensure credentials have a project." - ) - - vertex_endpoint = self._get_vertex_endpoint(endpoint) - base_url = self.auth.get_vertex_base_url() - return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location, model_id=self.model.id)}" - + return self.auth.build_url( + self._get_vertex_endpoint(endpoint), model_id=self.model.id + ) return f"{config.BASE_URL}{endpoint.format(model_id=self.model.id)}" async def _make_request( diff --git a/src/celeste/providers/google/veo/client.py b/src/celeste/providers/google/veo/client.py index 20dfb66b..9348c8f0 100644 --- a/src/celeste/providers/google/veo/client.py +++ b/src/celeste/providers/google/veo/client.py @@ -48,31 +48,17 @@ def _get_vertex_endpoint(self, gemini_endpoint: str) -> str: def _build_url(self, endpoint: str) -> str: """Build full URL based on auth type.""" if isinstance(self.auth, GoogleADC): - project_id = self.auth.resolved_project_id - if project_id is None: - raise ValueError( - "Vertex AI requires a project_id. " - "Pass project_id to GoogleADC() or ensure credentials have a project." - ) - - vertex_endpoint = self._get_vertex_endpoint(endpoint) - base_url = self.auth.get_vertex_base_url() - return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location, model_id=self.model.id)}" - + return self.auth.build_url( + self._get_vertex_endpoint(endpoint), model_id=self.model.id + ) return f"{config.BASE_URL}{endpoint.format(model_id=self.model.id)}" def _build_poll_url(self, operation_name: str) -> str: - """Build polling URL for long-running operations based on auth type.""" + """Build polling URL for long-running operations.""" if isinstance(self.auth, GoogleADC): - project_id = self.auth.resolved_project_id - base_url = self.auth.get_vertex_base_url() - poll_path = config.VertexVeoEndpoint.FETCH_OPERATION.format( - project_id=project_id, - location=self.auth.location, - model_id=self.model.id, + return self.auth.build_url( + config.VertexVeoEndpoint.FETCH_OPERATION, model_id=self.model.id ) - return f"{base_url}{poll_path}" - poll_path = config.GoogleVeoEndpoint.GET_OPERATION.format( operation_name=operation_name ) diff --git a/src/celeste/providers/mistral/chat/client.py b/src/celeste/providers/mistral/chat/client.py index f9e6e85a..d1617e9e 100644 --- a/src/celeste/providers/mistral/chat/client.py +++ b/src/celeste/providers/mistral/chat/client.py @@ -41,23 +41,12 @@ def _get_vertex_endpoint( return config.VertexMistralEndpoint.CREATE_CHAT def _build_url(self, endpoint: str, streaming: bool = False) -> str: - """Build full URL based on auth type. - - - GoogleADC auth -> Vertex AI endpoints - - API key auth -> Mistral API endpoints - """ + """Build full URL based on auth type.""" if isinstance(self.auth, GoogleADC): - project_id = self.auth.resolved_project_id - if project_id is None: - raise ValueError( - "Vertex AI requires a project_id. " - "Pass project_id to GoogleADC() or ensure credentials have a project." - ) - - vertex_endpoint = self._get_vertex_endpoint(endpoint, streaming=streaming) - base_url = self.auth.get_vertex_base_url() - return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location, model_id=self.model.id)}" - + return self.auth.build_url( + self._get_vertex_endpoint(endpoint, streaming=streaming), + model_id=self.model.id, + ) return f"{config.BASE_URL}{endpoint}" def _build_request( diff --git a/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/client.py.template b/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/client.py.template index d900a819..b828b76b 100644 --- a/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/client.py.template +++ b/templates/providers/{provider_slug}/src/celeste_{provider_slug}/{api_slug}/client.py.template @@ -54,15 +54,9 @@ class {Provider}{Api}Client(APIMixin): # def _build_url(self, endpoint: str) -> str: # """Build full URL based on auth type.""" # if isinstance(self.auth, GoogleADC): - # project_id = self.auth.resolved_project_id - # if project_id is None: - # raise ValueError( - # "Vertex AI requires a project_id. " - # "Pass project_id to GoogleADC() or ensure credentials have a project." - # ) - # vertex_endpoint = self._get_vertex_endpoint(endpoint) - # base_url = self.auth.get_vertex_base_url() - # return f"{base_url}{vertex_endpoint.format(project_id=project_id, location=self.auth.location, model_id=self.model.id)}" + # return self.auth.build_url( + # self._get_vertex_endpoint(endpoint), model_id=self.model.id + # ) # return f"{config.BASE_URL}{endpoint.format(model_id=self.model.id)}" def _build_request( From 81f9352c57de2e5052d2875d2954efe1ac158d64 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 9 Feb 2026 12:58:34 +0100 Subject: [PATCH 8/9] fix(images): let Artifact validator handle base64 decoding in Gemini images Same pattern as the video client fix - pass base64 string directly to ImageArtifact(data=...) instead of manual base64.b64decode(). Co-Authored-By: Claude Opus 4.6 --- src/celeste/modalities/images/providers/google/gemini.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/celeste/modalities/images/providers/google/gemini.py b/src/celeste/modalities/images/providers/google/gemini.py index 276d87f7..098fe7ca 100644 --- a/src/celeste/modalities/images/providers/google/gemini.py +++ b/src/celeste/modalities/images/providers/google/gemini.py @@ -113,8 +113,7 @@ def _parse_content( if not base64_data: continue 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)) + artifacts.append(ImageArtifact(data=base64_data, mime_type=mime_type)) if not artifacts: return ImageArtifact() From 3b4d1c1b2462adebcb72f4333c554c99c7fc6d96 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 9 Feb 2026 13:02:38 +0100 Subject: [PATCH 9/9] fix(images): let Artifact validator handle base64 decoding in Imagen Same pattern as Gemini images and Veo video fixes - pass base64 string directly to ImageArtifact(data=...) instead of manual base64.b64decode(). Co-Authored-By: Claude Opus 4.6 --- src/celeste/modalities/images/providers/google/imagen.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/celeste/modalities/images/providers/google/imagen.py b/src/celeste/modalities/images/providers/google/imagen.py index bd7a1b11..fe95a311 100644 --- a/src/celeste/modalities/images/providers/google/imagen.py +++ b/src/celeste/modalities/images/providers/google/imagen.py @@ -1,6 +1,5 @@ """Imagen client for Google images modality.""" -import base64 from typing import Any, Unpack from celeste.artifacts import ImageArtifact @@ -61,8 +60,7 @@ def _parse_content( if not base64_data: continue mime_type = ImageMimeType(prediction.get("mimeType", "image/png")) - image_bytes = base64.b64decode(base64_data) - images.append(ImageArtifact(data=image_bytes, mime_type=mime_type)) + images.append(ImageArtifact(data=base64_data, mime_type=mime_type)) num_images_requested = parameters.get("num_images") if num_images_requested == 1: