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/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:*)' - 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 }} 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..098fe7ca 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": {}, @@ -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() 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: diff --git a/src/celeste/modalities/videos/providers/google/client.py b/src/celeste/modalities/videos/providers/google/client.py index 744bef7b..c8a24e75 100644 --- a/src/celeste/modalities/videos/providers/google/client.py +++ b/src/celeste/modalities/videos/providers/google/client.py @@ -3,6 +3,7 @@ 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,6 +52,12 @@ def _parse_content( ) -> VideoArtifact: """Parse content from response.""" video_data = super()._parse_content(response_data) + # Handle inline base64 response (Vertex can return bytesBase64Encoded) + if "bytesBase64Encoded" in video_data: + mime_type = video_data.get("mimeType", VideoMimeType.MP4) + 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: @@ -62,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/anthropic/messages/client.py b/src/celeste/providers/anthropic/messages/client.py index a80c377e..4253710e 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,23 @@ 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.""" + if isinstance(self.auth, GoogleADC): + 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]: """Build headers with beta features extracted from request.""" beta_features: list[str] = request_body.pop("_beta_features", []) @@ -85,7 +107,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 +133,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..6d86adeb 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,22 @@ 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.""" + if isinstance(self.auth, GoogleADC): + return self.auth.build_url(self._get_vertex_endpoint(endpoint)) + return f"{config.BASE_URL}{endpoint}" + def _build_request( self, inputs: Any, @@ -64,7 +81,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 +106,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 +117,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/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/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..acc5a619 100644 --- a/src/celeste/providers/google/embeddings/client.py +++ b/src/celeste/providers/google/embeddings/client.py @@ -8,10 +8,27 @@ from celeste.io import FinishReason from celeste.mime_types import ApplicationMimeType +from ..auth import GoogleADC from . import config 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], @@ -30,17 +47,24 @@ 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) - - 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): + 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( self, @@ -50,6 +74,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 @@ -58,7 +93,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 +100,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, ) @@ -78,16 +112,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/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..61d6e58d 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,26 @@ 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.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: + 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): + 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( self, request_body: dict[str, Any], @@ -42,14 +67,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 +91,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..2ff96198 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 VertexGenerateContentEndpoint(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..778b285f 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,24 @@ 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): + 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( self, request_body: dict[str, Any], @@ -41,14 +64,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..9348c8f0 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,65 @@ 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): + 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.""" + if isinstance(self.auth, GoogleADC): + return self.auth.build_url( + config.VertexVeoEndpoint.FETCH_OPERATION, model_id=self.model.id + ) + 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 +114,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 +123,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 +135,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 +159,22 @@ 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) + 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", [] ) - 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..d1617e9e 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,23 @@ 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.""" + if isinstance(self.auth, GoogleADC): + 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( self, inputs: Any, @@ -64,7 +82,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 +107,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..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 @@ -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. @@ -20,15 +23,42 @@ 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) - _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): + # 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( self, inputs: Any, @@ -101,6 +131,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/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" 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) 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