diff --git a/packages/providers/google/pyproject.toml b/packages/providers/google/pyproject.toml new file mode 100644 index 00000000..daafb511 --- /dev/null +++ b/packages/providers/google/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "celeste-google" +version = "0.1.0" +description = "Google (Gemini API) provider package for Celeste AI" +authors = [{name = "Kamilbenkirane", email = "kamil@withceleste.ai"}] +license = {text = "Apache-2.0"} +requires-python = ">=3.12" +dependencies = ["celeste-ai", "httpx", "google-auth"] + +[tool.uv.sources] +celeste-ai = { workspace = true } + +[project.entry-points."celeste.providers"] +google = "celeste_google:register_provider" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/celeste_google"] diff --git a/packages/providers/google/src/celeste_google/__init__.py b/packages/providers/google/src/celeste_google/__init__.py new file mode 100644 index 00000000..ec93ff5d --- /dev/null +++ b/packages/providers/google/src/celeste_google/__init__.py @@ -0,0 +1 @@ +"""Google provider package for Celeste AI.""" diff --git a/packages/providers/google/src/celeste_google/generate_content/__init__.py b/packages/providers/google/src/celeste_google/generate_content/__init__.py new file mode 100644 index 00000000..a18780ee --- /dev/null +++ b/packages/providers/google/src/celeste_google/generate_content/__init__.py @@ -0,0 +1 @@ +"""Google GenerateContent API provider package.""" diff --git a/packages/providers/google/src/celeste_google/generate_content/client.py b/packages/providers/google/src/celeste_google/generate_content/client.py new file mode 100644 index 00000000..eaa078b3 --- /dev/null +++ b/packages/providers/google/src/celeste_google/generate_content/client.py @@ -0,0 +1,124 @@ +"""Google GenerateContent API client with shared implementation.""" + +from collections.abc import AsyncIterator +from typing import Any + +import httpx + +from celeste.core import UsageField +from celeste.io import FinishReason +from celeste.mime_types import ApplicationMimeType + +from . import config + + +class GoogleGenerateContentClient: + """Mixin for GenerateContent API capabilities. + + Provides shared implementation for all capabilities using the GenerateContent API: + - _make_request() - HTTP POST to generateContent + - _make_stream_request() - HTTP streaming to streamGenerateContent + - _parse_usage() - Extract usage dict from usageMetadata + - _parse_content() - Extract parts array from first candidate + - _parse_finish_reason() - Extract finish reason string from candidates + - _build_metadata() - Filter content fields + + Capability clients extend parsing methods via super() to wrap/transform results. + + Usage: + class GoogleTextGenerationClient(GoogleGenerateContentClient, TextGenerationClient): + def _parse_content(self, response_data, **parameters): + parts = super()._parse_content(response_data) + text = parts[0].get("text") or "" + return self._transform_output(text, **parameters) + """ + + async def _make_request( + self, + request_body: dict[str, Any], + **parameters: Any, + ) -> httpx.Response: + """Make HTTP request to generateContent endpoint.""" + endpoint = config.GoogleGenerateContentEndpoint.GENERATE_CONTENT.format( + model_id=self.model.id # type: ignore[attr-defined] + ) + headers = { + **self.auth.get_headers(), # type: ignore[attr-defined] + "Content-Type": ApplicationMimeType.JSON, + } + return await self.http_client.post( # type: ignore[attr-defined,no-any-return] + f"{config.BASE_URL}{endpoint}", + headers=headers, + json_body=request_body, + ) + + def _make_stream_request( + self, + request_body: dict[str, Any], + **parameters: Any, + ) -> AsyncIterator[dict[str, Any]]: + """Make streaming request to streamGenerateContent endpoint.""" + endpoint = config.GoogleGenerateContentEndpoint.STREAM_GENERATE_CONTENT.format( + model_id=self.model.id # type: ignore[attr-defined] + ) + headers = { + **self.auth.get_headers(), # type: ignore[attr-defined] + "Content-Type": ApplicationMimeType.JSON, + } + return self.http_client.stream_post( # type: ignore[attr-defined,no-any-return] + f"{config.BASE_URL}{endpoint}", + headers=headers, + json_body=request_body, + ) + + @staticmethod + def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | None]: + """Map Google Gemini usage fields to unified names. + + Shared by client and streaming across all capabilities. + """ + return { + UsageField.INPUT_TOKENS: usage_data.get("promptTokenCount"), + UsageField.OUTPUT_TOKENS: usage_data.get("candidatesTokenCount"), + UsageField.TOTAL_TOKENS: usage_data.get("totalTokenCount"), + UsageField.REASONING_TOKENS: usage_data.get("thoughtsTokenCount"), + } + + def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: + """Extract usage data from Gemini usageMetadata.""" + usage_metadata = response_data.get("usageMetadata", {}) + return self.map_usage_fields(usage_metadata) + + def _parse_content(self, response_data: dict[str, Any]) -> Any: + """Return all candidates from response. + + Returns list of candidate objects that capability clients extract content from. + """ + candidates = response_data.get("candidates", []) + if not candidates: + msg = "No candidates in response" + raise ValueError(msg) + return candidates + + def _parse_finish_reason(self, response_data: dict[str, Any]) -> FinishReason: + """Extract finish reason from Gemini candidates. + + Returns FinishReason that capability clients wrap in their specific type. + """ + candidates = response_data.get("candidates", []) + if not candidates: + reason = None + else: + reason = candidates[0].get("finishReason") + return FinishReason(reason=reason) + + def _build_metadata(self, response_data: dict[str, Any]) -> dict[str, Any]: + """Build metadata dictionary, filtering out content fields.""" + content_fields = {"candidates"} + filtered_data = { + k: v for k, v in response_data.items() if k not in content_fields + } + return super()._build_metadata(filtered_data) # type: ignore[misc,no-any-return] + + +__all__ = ["GoogleGenerateContentClient"] diff --git a/packages/providers/google/src/celeste_google/generate_content/config.py b/packages/providers/google/src/celeste_google/generate_content/config.py new file mode 100644 index 00000000..1075ea7a --- /dev/null +++ b/packages/providers/google/src/celeste_google/generate_content/config.py @@ -0,0 +1,23 @@ +"""Configuration for Google GenerateContent API.""" + +from enum import StrEnum + + +class GoogleGenerateContentEndpoint(StrEnum): + """Endpoints for GenerateContent API.""" + + GENERATE_CONTENT = "/v1beta/models/{model_id}:generateContent" + STREAM_GENERATE_CONTENT = "/v1beta/models/{model_id}:streamGenerateContent?alt=sse" + COUNT_TOKENS = "/v1beta/models/{model_id}:countTokens" + EMBED_CONTENT = "/v1beta/models/{model_id}:embedContent" + BATCH_EMBED_CONTENTS = "/v1beta/models/{model_id}:batchEmbedContents" + LIST_MODELS = "/v1beta/models" + GET_MODEL = "/v1beta/models/{model_id}" + UPLOAD_FILE = "/upload/v1beta/files" + LIST_FILES = "/v1beta/files" + GET_FILE = "/v1beta/files/{file_id}" + DELETE_FILE = "/v1beta/files/{file_id}" + BATCH_GENERATE_CONTENT = "/v1beta/models/{model_id}:batchGenerateContent" + + +BASE_URL = "https://generativelanguage.googleapis.com" diff --git a/packages/providers/google/src/celeste_google/generate_content/parameters.py b/packages/providers/google/src/celeste_google/generate_content/parameters.py new file mode 100644 index 00000000..831cf883 --- /dev/null +++ b/packages/providers/google/src/celeste_google/generate_content/parameters.py @@ -0,0 +1,310 @@ +"""Google Gemini API parameter mappers.""" + +import base64 +import json +from typing import Any, get_args, get_origin + +from pydantic import BaseModel, TypeAdapter + +from celeste.artifacts import ImageArtifact +from celeste.mime_types import ApplicationMimeType, ImageMimeType +from celeste.models import Model +from celeste.parameters import ParameterMapper +from celeste.types import StructuredOutput + + +class TemperatureMapper(ParameterMapper): + """Map temperature to Google generationConfig.temperature field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform temperature into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("generationConfig", {})["temperature"] = validated_value + return request + + +class MaxOutputTokensMapper(ParameterMapper): + """Map max_tokens to Google generationConfig.maxOutputTokens field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform max_tokens into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("generationConfig", {})["maxOutputTokens"] = validated_value + return request + + +class ThinkingBudgetMapper(ParameterMapper): + """Map thinkingBudget to Google generationConfig.thinkingConfig.thinkingBudget field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform thinkingBudget into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("generationConfig", {}).setdefault("thinkingConfig", {})[ + "thinkingBudget" + ] = validated_value + return request + + +class ThinkingLevelMapper(ParameterMapper): + """Map thinkingLevel to Google generationConfig.thinkingConfig.thinkingLevel field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform thinkingLevel into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("generationConfig", {}).setdefault("thinkingConfig", {})[ + "thinkingLevel" + ] = validated_value + return request + + +class AspectRatioMapper(ParameterMapper): + """Map aspect_ratio to Google generationConfig.imageConfig.aspectRatio field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform aspect_ratio into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("generationConfig", {}).setdefault("imageConfig", {})[ + "aspectRatio" + ] = validated_value + return request + + +class ImageSizeMapper(ParameterMapper): + """Map image_size to Google generationConfig.imageConfig.imageSize field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform image_size into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("generationConfig", {}).setdefault("imageConfig", {})[ + "imageSize" + ] = validated_value + return request + + +class MediaContentMapper(ParameterMapper): + """Map reference_images to Google contents.parts field.""" + + def _build_image_part(self, image: ImageArtifact) -> dict[str, Any]: + """Build inline_data part from image artifact.""" + if image.url: + return {"file_data": {"file_uri": image.url}} + elif image.data: + base64_data = ( + base64.b64encode(image.data).decode("utf-8") + if isinstance(image.data, bytes) + else image.data + ) + return { + "inline_data": { + "mime_type": image.mime_type or ImageMimeType.JPEG, + "data": base64_data, + } + } + elif image.path: + with open(image.path, "rb") as f: + image_bytes = f.read() + base64_data = base64.b64encode(image_bytes).decode("utf-8") + return { + "inline_data": { + "mime_type": image.mime_type or ImageMimeType.JPEG, + "data": base64_data, + } + } + else: + msg = "ImageArtifact must have url, data, or path" + raise ValueError(msg) + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform reference_images into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + # Convert list of ImageArtifact to list of image parts + image_parts = [self._build_image_part(img) for img in validated_value] + + # Add image parts before text in contents[0].parts + if "contents" in request and len(request["contents"]) > 0: + parts = request["contents"][0].get("parts", []) + # Find text part and insert images before it + text_index = next( + (i for i, part in enumerate(parts) if "text" in part), len(parts) + ) + # Insert image parts before text + parts[text_index:text_index] = image_parts + request["contents"][0]["parts"] = parts + + return request + + +class WebSearchMapper(ParameterMapper): + """Map web_search to Google tools field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform web_search into provider request.""" + validated_value = self._validate_value(value, model) + if not validated_value: + return request + + request["tools"] = [{"google_search": {}}] + return request + + +class OutputSchemaMapper(ParameterMapper): + """Map output_schema to Google generationConfig.responseJsonSchema field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform output_schema into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + schema = self._convert_to_google_schema(validated_value) + + config = request.setdefault("generationConfig", {}) + config["responseJsonSchema"] = schema + config["responseMimeType"] = ApplicationMimeType.JSON + + return request + + def parse_output( + self, content: StructuredOutput, value: object | None + ) -> StructuredOutput: + """Parse JSON to BaseModel using Pydantic's TypeAdapter.""" + if value is None: + return content + + # If content is already a BaseModel, return it unchanged + if isinstance(content, BaseModel): + return content + if isinstance(content, list) and content and isinstance(content[0], BaseModel): + return content + + parsed = json.loads(content) if isinstance(content, str) else content + + # For list[T], handle various formats Google might return + origin = get_origin(value) + if origin is list and isinstance(parsed, dict): + # Google returns arrays directly in JSON schema, but might wrap in object + if "items" in parsed: + parsed = parsed["items"] + else: + # If it's a dict but not wrapped, try to extract array values + parsed = list(parsed.values()) if parsed else [] + + return TypeAdapter(value).validate_python(parsed) + + def _convert_to_google_schema(self, output_schema: Any) -> dict[str, Any]: # noqa: ANN401 + """Convert Pydantic BaseModel or list[BaseModel] to Google JSON Schema format.""" + origin = get_origin(output_schema) + if origin is list: + inner_type = get_args(output_schema)[0] + items_schema = inner_type.model_json_schema() + defs = items_schema.get("$defs", {}) + items_schema_clean = {k: v for k, v in items_schema.items() if k != "$defs"} + json_schema = {"type": "array", "items": items_schema_clean} + if defs: + json_schema["$defs"] = defs + else: + json_schema = output_schema.model_json_schema() + + json_schema = self._remove_unsupported_fields(json_schema) + return json_schema + + def _remove_unsupported_fields(self, schema: dict[str, Any]) -> dict[str, Any]: + """Remove unsupported metadata fields from schema.""" + result: dict[str, Any] = {} + + for key, value in schema.items(): + if key == "title": + continue + + if isinstance(value, dict): + result[key] = self._remove_unsupported_fields(value) + elif isinstance(value, list): + result[key] = [ + self._remove_unsupported_fields(item) + if isinstance(item, dict) + else item + for item in value + ] + else: + result[key] = value + + return result + + +__all__ = [ + "AspectRatioMapper", + "ImageSizeMapper", + "MaxOutputTokensMapper", + "MediaContentMapper", + "OutputSchemaMapper", + "TemperatureMapper", + "ThinkingBudgetMapper", + "ThinkingLevelMapper", + "WebSearchMapper", +] diff --git a/packages/providers/google/src/celeste_google/generate_content/streaming.py b/packages/providers/google/src/celeste_google/generate_content/streaming.py new file mode 100644 index 00000000..4d71d666 --- /dev/null +++ b/packages/providers/google/src/celeste_google/generate_content/streaming.py @@ -0,0 +1,54 @@ +"""Google GenerateContent SSE parsing for streaming.""" + +from typing import Any + +from .client import GoogleGenerateContentClient + + +class GoogleGenerateContentStream: + """Mixin for GenerateContent SSE parsing. + + Provides shared implementation for all capabilities using GenerateContent streaming: + - _parse_chunk() - Parse SSE event into raw chunk dict + + Capability streams extend via super() to wrap results in typed Chunks. + + Usage: + class GoogleTextGenerationStream(GoogleGenerateContentStream, TextGenerationStream): + def _parse_chunk(self, event): + raw = super()._parse_chunk(event) + if not raw: + return None + return TextGenerationChunk(...) + """ + + def _parse_chunk(self, event: dict[str, Any]) -> dict[str, Any] | None: + """Parse SSE event into raw chunk data.""" + candidates = event.get("candidates", []) + if not candidates: + return None + + candidate = candidates[0] + content = candidate.get("content", {}) + parts = content.get("parts", []) + + text_delta = parts[0].get("text") if parts else None + finish_reason = candidate.get("finishReason") + + usage = None + usage_data = event.get("usageMetadata") + if usage_data: + usage = GoogleGenerateContentClient.map_usage_fields(usage_data) + + if not text_delta and not finish_reason: + return None + + return { + "content": text_delta or "", + "finish_reason": finish_reason, + "usage": usage, + "raw_event": event, + } + + +__all__ = ["GoogleGenerateContentStream"] diff --git a/packages/providers/google/src/celeste_google/py.typed b/packages/providers/google/src/celeste_google/py.typed new file mode 100644 index 00000000..e69de29b diff --git a/src/celeste/auth.py b/src/celeste/auth.py index 3affe19f..14acadc1 100644 --- a/src/celeste/auth.py +++ b/src/celeste/auth.py @@ -66,9 +66,9 @@ def get_auth_class(auth_type: str) -> type[Authentication]: Raises: ValueError: If auth type is not registered. """ - from celeste.registry import _load_from_entry_points + from celeste.registry import _load_providers_from_entry_points - _load_from_entry_points() + _load_providers_from_entry_points() if auth_type not in _auth_classes: msg = f"Unknown auth type: {auth_type}. Available: {list(_auth_classes.keys())}" diff --git a/src/celeste/client.py b/src/celeste/client.py index 92ee479d..c9d5eb98 100644 --- a/src/celeste/client.py +++ b/src/celeste/client.py @@ -20,6 +20,7 @@ from celeste.models import Model from celeste.parameters import ParameterMapper, Parameters from celeste.streaming import Stream +from celeste.types import StructuredOutput class Client[In: Input, Out: Output, Params: Parameters](ABC, BaseModel): @@ -123,7 +124,7 @@ def _parse_content( self, response_data: dict[str, Any], **parameters: Unpack[Params], # type: ignore[misc] - ) -> object: + ) -> StructuredOutput: """Parse content from provider response.""" ... @@ -205,9 +206,9 @@ def _handle_error_response(self, response: httpx.Response) -> None: def _transform_output( self, - content: object, + content: StructuredOutput, **parameters: Unpack[Params], # type: ignore[misc] - ) -> object: + ) -> StructuredOutput: """Transform content using parameter mapper output transformations.""" for mapper in self.parameter_mappers(): value = parameters.get(mapper.name) diff --git a/src/celeste/http.py b/src/celeste/http.py index 2b0361b7..c267d87e 100644 --- a/src/celeste/http.py +++ b/src/celeste/http.py @@ -14,7 +14,7 @@ MAX_CONNECTIONS = 20 MAX_KEEPALIVE_CONNECTIONS = 10 -DEFAULT_TIMEOUT = 60.0 +DEFAULT_TIMEOUT = 180.0 class HTTPClient: diff --git a/src/celeste/parameters.py b/src/celeste/parameters.py index 2c2b7915..d904bdc8 100644 --- a/src/celeste/parameters.py +++ b/src/celeste/parameters.py @@ -6,6 +6,7 @@ from celeste.exceptions import UnsupportedParameterError from celeste.models import Model +from celeste.types import StructuredOutput class Parameters(TypedDict, total=False): @@ -32,7 +33,9 @@ def map(self, request: dict[str, Any], value: Any, model: Model) -> dict[str, An """ ... - def parse_output(self, content: Any, value: object | None) -> object: # noqa: ANN401 + def parse_output( + self, content: StructuredOutput, value: object | None + ) -> StructuredOutput: """Optionally transform parsed content based on parameter value (default: return unchanged).""" return content diff --git a/src/celeste/registry.py b/src/celeste/registry.py index 48690d9f..514b51b9 100644 --- a/src/celeste/registry.py +++ b/src/celeste/registry.py @@ -3,6 +3,7 @@ import importlib.metadata _loaded_packages: set[str] = set() +_loaded_providers: set[str] = set() def _load_from_entry_points() -> None: @@ -22,3 +23,22 @@ def _load_from_entry_points() -> None: # The function should register models and clients when called register_func() _loaded_packages.add(ep.name) + + +def _load_providers_from_entry_points() -> None: + """Load auth from installed provider packages via entry points.""" + + entry_points = importlib.metadata.entry_points(group="celeste.providers") + + # Early return if all providers are already loaded + entry_point_names = {ep.name for ep in entry_points} + if entry_point_names.issubset(_loaded_providers): + return + + for ep in entry_points: + if ep.name in _loaded_providers: + continue + register_func = ep.load() + # The function should register auth types when called + register_func() + _loaded_providers.add(ep.name) diff --git a/tests/unit_tests/test_client.py b/tests/unit_tests/test_client.py index 48f76060..7bc24905 100644 --- a/tests/unit_tests/test_client.py +++ b/tests/unit_tests/test_client.py @@ -20,6 +20,7 @@ from celeste.models import Model from celeste.parameters import ParameterMapper, Parameters from celeste.streaming import Stream +from celeste.types import StructuredOutput class ParamEnum(StrEnum): @@ -92,7 +93,9 @@ def map( request[actual_map_key] = value return request - def parse_output(self, content: object, value: object | None) -> object: + def parse_output( + self, content: StructuredOutput, value: object | None + ) -> StructuredOutput: return content return TestMapperClass() @@ -120,7 +123,9 @@ def map( request[actual_map_key] = value return request - def parse_output(self, content: object, value: object | None) -> object: + def parse_output( + self, content: StructuredOutput, value: object | None + ) -> StructuredOutput: if value is not None: return f"{content}_transformed_with_{value}" return content diff --git a/tests/unit_tests/test_http.py b/tests/unit_tests/test_http.py index 6fbd219e..415d0a4e 100644 --- a/tests/unit_tests/test_http.py +++ b/tests/unit_tests/test_http.py @@ -296,7 +296,7 @@ async def test_post_uses_default_timeout_when_not_specified( # Assert - Verify default timeout was used mock_httpx_client.post.assert_called_once() call_kwargs = mock_httpx_client.post.call_args[1] - assert call_kwargs["timeout"] == 60.0 + assert call_kwargs["timeout"] == 180.0 async def test_get_uses_default_timeout_when_not_specified( self, mock_httpx_client: AsyncMock @@ -316,7 +316,7 @@ async def test_get_uses_default_timeout_when_not_specified( # Assert - Verify default timeout was used mock_httpx_client.get.assert_called_once() call_kwargs = mock_httpx_client.get.call_args[1] - assert call_kwargs["timeout"] == 60.0 + assert call_kwargs["timeout"] == 180.0 async def test_custom_timeout_passed_to_httpx( self, mock_httpx_client: AsyncMock