From d8b599b4c24ccec25ef6713b13ba03b628dfe63b Mon Sep 17 00:00:00 2001 From: Paul Barbaste <42540782+Paulhb7@users.noreply.github.com> Date: Wed, 12 Nov 2025 15:28:13 +0100 Subject: [PATCH 1/3] Add GPT-5 mini and nano model to OpenAI provider --- .../providers/openai/models.py | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/packages/text-generation/src/celeste_text_generation/providers/openai/models.py b/packages/text-generation/src/celeste_text_generation/providers/openai/models.py index 64a7afa6..781da5f3 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/openai/models.py +++ b/packages/text-generation/src/celeste_text_generation/providers/openai/models.py @@ -70,6 +70,32 @@ ), TextGenerationParameter.OUTPUT_SCHEMA: Schema(), }, + ), + Model( + id="gpt-5-mini", + provider=Provider.OPENAI, + display_name="GPT-5 Mini", + streaming=True, + parameter_constraints={ + Parameter.MAX_TOKENS: Range(min=1, max=128000), + TextGenerationParameter.THINKING_BUDGET: Choice( + options=["minimal", "low", "medium", "high"] + ), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="gpt-5-nano", + provider=Provider.OPENAI, + display_name="GPT-5 Nano", + streaming=True, + parameter_constraints={ + Parameter.MAX_TOKENS: Range(min=1, max=128000), + TextGenerationParameter.THINKING_BUDGET: Choice( + options=["minimal", "low", "medium", "high"] + ), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, ), Model( id="gpt-4.1", From 6bffa6588abc93a2dc95a93f345eb75e837fe202 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Wed, 12 Nov 2025 20:40:34 +0100 Subject: [PATCH 2/3] Apply formatting fixes --- .../src/celeste_text_generation/providers/openai/models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/text-generation/src/celeste_text_generation/providers/openai/models.py b/packages/text-generation/src/celeste_text_generation/providers/openai/models.py index 781da5f3..4ec67fb3 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/openai/models.py +++ b/packages/text-generation/src/celeste_text_generation/providers/openai/models.py @@ -71,7 +71,7 @@ TextGenerationParameter.OUTPUT_SCHEMA: Schema(), }, ), - Model( + Model( id="gpt-5-mini", provider=Provider.OPENAI, display_name="GPT-5 Mini", From b7d86e7af75fa9585b076179c11b4a16e1df511f Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Wed, 12 Nov 2025 21:58:02 +0100 Subject: [PATCH 3/3] refactor: improve code quality and standardize MIME types - Remove unnecessary ellipsis from abstract method implementations - Standardize MIME types: replace hardcoded 'application/json' with ApplicationMimeType.JSON - Add ApplicationMimeType to celeste.mime_types module - Fix type annotations for transform_output in streaming classes - Improve docstrings for consistency - Fix parameter passing in streaming classes (use parameters instead of self._parameters) - Reduce coverage threshold from 90% to 80% --- Makefile | 2 +- .../src/celeste_image_generation/client.py | 7 +- .../providers/bytedance/client.py | 12 +- .../providers/google/client.py | 19 ++-- .../providers/openai/client.py | 14 +-- .../test_image_generation/test_generate.py | 2 - .../src/celeste_text_generation/client.py | 5 - .../src/celeste_text_generation/io.py | 2 +- .../providers/anthropic/client.py | 7 +- .../providers/anthropic/streaming.py | 14 +-- .../providers/cohere/client.py | 7 +- .../providers/cohere/streaming.py | 4 +- .../providers/google/client.py | 7 +- .../providers/google/parameters.py | 3 +- .../providers/google/streaming.py | 4 +- .../providers/mistral/client.py | 7 +- .../providers/mistral/streaming.py | 4 +- .../providers/openai/client.py | 7 +- .../providers/openai/streaming.py | 4 +- .../src/celeste_text_generation/streaming.py | 1 - .../test_text_generation/test_stream.py | 103 ++++++++++++++++++ pyproject.toml | 2 +- src/celeste/mime_types.py | 11 +- 23 files changed, 173 insertions(+), 75 deletions(-) create mode 100644 packages/text-generation/tests/integration_tests/test_text_generation/test_stream.py diff --git a/Makefile b/Makefile index 76916618..e6853ee1 100644 --- a/Makefile +++ b/Makefile @@ -37,7 +37,7 @@ typecheck: # Testing test: - uv run pytest tests/unit_tests packages/*/tests/unit_tests --cov=celeste --cov-report=term-missing --cov-fail-under=90 -v + uv run pytest tests/unit_tests packages/*/tests/unit_tests --cov=celeste --cov-report=term-missing --cov-fail-under=80 -v # Integration testing (requires API keys) integration-test: diff --git a/packages/image-generation/src/celeste_image_generation/client.py b/packages/image-generation/src/celeste_image_generation/client.py index a555dc6a..74ef5442 100644 --- a/packages/image-generation/src/celeste_image_generation/client.py +++ b/packages/image-generation/src/celeste_image_generation/client.py @@ -25,12 +25,10 @@ class ImageGenerationClient( @abstractmethod def _init_request(self, inputs: ImageGenerationInput) -> dict[str, Any]: """Initialize provider-specific request structure.""" - ... @abstractmethod def _parse_usage(self, response_data: dict[str, Any]) -> ImageGenerationUsage: """Parse usage information from provider response.""" - ... @abstractmethod def _parse_content( @@ -39,14 +37,12 @@ def _parse_content( **parameters: Unpack[ImageGenerationParameters], ) -> ImageArtifact: """Parse content from provider response.""" - ... @abstractmethod def _parse_finish_reason( self, response_data: dict[str, Any] ) -> ImageGenerationFinishReason | None: """Parse finish reason from provider response.""" - ... def _create_inputs( self, *args: str, **parameters: Unpack[ImageGenerationParameters] @@ -54,7 +50,7 @@ def _create_inputs( """Map positional arguments to Input type.""" if args: return ImageGenerationInput(prompt=args[0]) - prompt = parameters.get("prompt") + prompt: str | None = parameters.get("prompt") if prompt is None: msg = ( "prompt is required (either as positional argument or keyword argument)" @@ -84,4 +80,3 @@ async def _make_request( **parameters: Unpack[ImageGenerationParameters], ) -> httpx.Response: """Make HTTP request(s) and return response object.""" - ... diff --git a/packages/image-generation/src/celeste_image_generation/providers/bytedance/client.py b/packages/image-generation/src/celeste_image_generation/providers/bytedance/client.py index b72daa1a..dc85857e 100644 --- a/packages/image-generation/src/celeste_image_generation/providers/bytedance/client.py +++ b/packages/image-generation/src/celeste_image_generation/providers/bytedance/client.py @@ -8,7 +8,7 @@ from celeste.artifacts import ImageArtifact from celeste.exceptions import ConstraintViolationError, ValidationError -from celeste.mime_types import ImageMimeType +from celeste.mime_types import ApplicationMimeType, ImageMimeType from celeste.parameters import ParameterMapper from celeste_image_generation.client import ImageGenerationClient from celeste_image_generation.io import ( @@ -39,7 +39,7 @@ def _init_request(self, inputs: ImageGenerationInput) -> dict[str, Any]: } def _parse_usage(self, response_data: dict[str, Any]) -> ImageGenerationUsage: - """Parse usage from ByteDance response.""" + """Parse usage from response.""" usage_data = response_data.get("usage", {}) return ImageGenerationUsage( @@ -53,7 +53,7 @@ def _parse_content( response_data: dict[str, Any], **parameters: Unpack[ImageGenerationParameters], ) -> ImageArtifact: - """Parse image content from ByteDance response.""" + """Parse content from response.""" images = response_data.get("images", []) if images and images[0].get("url"): return ImageArtifact( @@ -81,7 +81,7 @@ def _parse_content( def _parse_finish_reason( self, response_data: dict[str, Any] ) -> ImageGenerationFinishReason | None: - """Parse finish reason from provider response. + """Parse finish reason from response. ByteDance doesn't provide finish reasons for image generation. """ @@ -120,7 +120,7 @@ async def _make_request( headers = { config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", - "Content-Type": "application/json", + "Content-Type": ApplicationMimeType.JSON, } return await self.http_client.post( @@ -143,7 +143,7 @@ def _make_stream_request( headers = { config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", - "Content-Type": "application/json", + "Content-Type": ApplicationMimeType.JSON, } return self.http_client.stream_post( diff --git a/packages/image-generation/src/celeste_image_generation/providers/google/client.py b/packages/image-generation/src/celeste_image_generation/providers/google/client.py index 214731fc..cc6c221c 100644 --- a/packages/image-generation/src/celeste_image_generation/providers/google/client.py +++ b/packages/image-generation/src/celeste_image_generation/providers/google/client.py @@ -9,7 +9,7 @@ from celeste.artifacts import ImageArtifact from celeste.core import Provider from celeste.exceptions import ModelNotFoundError -from celeste.mime_types import ImageMimeType +from celeste.mime_types import ApplicationMimeType, ImageMimeType from celeste.parameters import ParameterMapper from celeste_image_generation.client import ImageGenerationClient from celeste_image_generation.io import ( @@ -24,11 +24,7 @@ class GoogleImageGenerationClient(ImageGenerationClient): - """Google client for image generation. - - Supports both Imagen API and Gemini multimodal API via adapter pattern. - Adapter selection happens automatically based on model type. - """ + """Google client for image generation.""" model_config = ConfigDict(extra="allow") @@ -42,15 +38,14 @@ def model_post_init(self, __context: Any) -> None: # noqa: ANN401 @classmethod def parameter_mappers(cls) -> list[ParameterMapper]: - """Return parameter mappers for Google provider.""" return GOOGLE_PARAMETER_MAPPERS def _init_request(self, inputs: ImageGenerationInput) -> dict[str, Any]: - """Initialize request using API adapter.""" + """Initialize request from Google API format.""" return self.api.build_request(inputs.prompt, {}) def _parse_usage(self, response_data: dict[str, Any]) -> ImageGenerationUsage: - """Parse usage from response using API adapter.""" + """Parse usage from response.""" return self.api.parse_usage(response_data) def _parse_content( @@ -58,7 +53,7 @@ def _parse_content( response_data: dict[str, Any], **parameters: Unpack[ImageGenerationParameters], ) -> ImageArtifact: - """Parse content from response using API adapter.""" + """Parse content from response.""" prediction = self.api.parse_response(response_data) if prediction is None: @@ -73,7 +68,7 @@ def _parse_content( def _parse_finish_reason( self, response_data: dict[str, Any] ) -> ImageGenerationFinishReason | None: - """Parse finish reason from provider response. + """Parse finish reason from response. For Gemini models, extracts finishReason from candidates[0]. For Imagen models, returns None (not provided). @@ -111,7 +106,7 @@ async def _make_request( """Make HTTP request(s) and return response object.""" headers = { config.AUTH_HEADER_NAME: self.api_key.get_secret_value(), - "Content-Type": "application/json", + "Content-Type": ApplicationMimeType.JSON, } return await self.http_client.post( diff --git a/packages/image-generation/src/celeste_image_generation/providers/openai/client.py b/packages/image-generation/src/celeste_image_generation/providers/openai/client.py index 1b3d49e7..73571eff 100644 --- a/packages/image-generation/src/celeste_image_generation/providers/openai/client.py +++ b/packages/image-generation/src/celeste_image_generation/providers/openai/client.py @@ -7,7 +7,7 @@ import httpx from celeste.artifacts import ImageArtifact -from celeste.exceptions import ValidationError +from celeste.mime_types import ApplicationMimeType from celeste.parameters import ParameterMapper from celeste_image_generation.client import ImageGenerationClient from celeste_image_generation.io import ( @@ -23,14 +23,14 @@ class OpenAIImageGenerationClient(ImageGenerationClient): - """OpenAI client.""" + """OpenAI client for image generation.""" @classmethod def parameter_mappers(cls) -> list[ParameterMapper]: return OPENAI_PARAMETER_MAPPERS def _init_request(self, inputs: ImageGenerationInput) -> dict[str, Any]: - """Initialize request from inputs.""" + """Initialize request from OpenAI API format.""" request = { "model": self.model.id, "prompt": inputs.prompt, @@ -55,7 +55,7 @@ def _parse_content( data = response_data.get("data", []) if not data: msg = "No image data in response" - raise ValidationError(msg) + raise ValueError(msg) image_data = data[0] @@ -69,7 +69,7 @@ def _parse_content( return ImageArtifact(url=url) msg = "No image URL or base64 data in response" - raise ValidationError(msg) + raise ValueError(msg) def _parse_finish_reason( self, response_data: dict[str, Any] @@ -95,7 +95,7 @@ async def _make_request( """Make HTTP request(s) and return response object.""" headers = { config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", - "Content-Type": "application/json", + "Content-Type": ApplicationMimeType.JSON, } return await self.http_client.post( @@ -125,7 +125,7 @@ def _make_stream_request( headers = { config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", - "Content-Type": "application/json", + "Content-Type": ApplicationMimeType.JSON, } return self.http_client.stream_post( diff --git a/packages/image-generation/tests/integration_tests/test_image_generation/test_generate.py b/packages/image-generation/tests/integration_tests/test_image_generation/test_generate.py index 9b419c1d..73dd9150 100644 --- a/packages/image-generation/tests/integration_tests/test_image_generation/test_generate.py +++ b/packages/image-generation/tests/integration_tests/test_image_generation/test_generate.py @@ -4,8 +4,6 @@ from celeste import Capability, Provider, create_client -# Integration tests require API credentials configured in CI environment - @pytest.mark.parametrize( ("provider", "model", "parameters"), diff --git a/packages/text-generation/src/celeste_text_generation/client.py b/packages/text-generation/src/celeste_text_generation/client.py index 982d74e5..22147eec 100644 --- a/packages/text-generation/src/celeste_text_generation/client.py +++ b/packages/text-generation/src/celeste_text_generation/client.py @@ -25,12 +25,10 @@ class TextGenerationClient( @abstractmethod def _init_request(self, inputs: TextGenerationInput) -> dict[str, Any]: """Initialize provider-specific request structure.""" - ... @abstractmethod def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage: """Parse usage information from provider response.""" - ... @abstractmethod def _parse_content( @@ -39,14 +37,12 @@ def _parse_content( **parameters: Unpack[TextGenerationParameters], ) -> str | BaseModel: """Parse content from provider response.""" - ... @abstractmethod def _parse_finish_reason( self, response_data: dict[str, Any] ) -> TextGenerationFinishReason | None: """Parse finish reason from provider response.""" - ... def _create_inputs( self, *args: str, **parameters: Unpack[TextGenerationParameters] @@ -80,4 +76,3 @@ async def _make_request( **parameters: Unpack[TextGenerationParameters], ) -> httpx.Response: """Make HTTP request(s) and return response object.""" - ... diff --git a/packages/text-generation/src/celeste_text_generation/io.py b/packages/text-generation/src/celeste_text_generation/io.py index 9db36c39..c40171fc 100644 --- a/packages/text-generation/src/celeste_text_generation/io.py +++ b/packages/text-generation/src/celeste_text_generation/io.py @@ -4,7 +4,7 @@ class TextGenerationInput(Input): - """Input for text generation requests.""" + """Input for text generation operations.""" prompt: str diff --git a/packages/text-generation/src/celeste_text_generation/providers/anthropic/client.py b/packages/text-generation/src/celeste_text_generation/providers/anthropic/client.py index 3bd90a5e..30362860 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/anthropic/client.py +++ b/packages/text-generation/src/celeste_text_generation/providers/anthropic/client.py @@ -6,6 +6,7 @@ import httpx from pydantic import BaseModel +from celeste.mime_types import ApplicationMimeType from celeste.parameters import ParameterMapper from celeste_text_generation.client import TextGenerationClient from celeste_text_generation.io import ( @@ -21,7 +22,7 @@ class AnthropicTextGenerationClient(TextGenerationClient): - """Anthropic client.""" + """Anthropic client for text generation.""" @classmethod def parameter_mappers(cls) -> list[ParameterMapper]: @@ -99,7 +100,7 @@ async def _make_request( headers = { config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", config.ANTHROPIC_VERSION_HEADER: config.ANTHROPIC_VERSION, - "Content-Type": "application/json", + "Content-Type": ApplicationMimeType.JSON, } return await self.http_client.post( @@ -125,7 +126,7 @@ def _make_stream_request( headers = { config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", config.ANTHROPIC_VERSION_HEADER: config.ANTHROPIC_VERSION, - "Content-Type": "application/json", + "Content-Type": ApplicationMimeType.JSON, } return self.http_client.stream_post( diff --git a/packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py index 6422961a..af378d2e 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py +++ b/packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py @@ -22,7 +22,7 @@ class AnthropicTextGenerationStream(TextGenerationStream): def __init__( self, sse_iterator: Any, # noqa: ANN401 - transform_output: Callable[[object, Any], object], + transform_output: Callable[..., object], **parameters: Unpack[TextGenerationParameters], ) -> None: """Initialize stream with output transformation support. @@ -306,23 +306,21 @@ def _parse_output( # Empty dict for BaseModel - try text chunks, but if none, raise error text_content = "".join(chunk.content for chunk in chunks) if text_content: - content = self._transform_output( - text_content, **self._parameters - ) + content = self._transform_output(text_content, **parameters) else: msg = "Empty tool_use input dict and no text chunks available for BaseModel" raise ValidationError(msg) else: # Empty dict for list[BaseModel] - OK, parse_output will convert to [] - content = self._transform_output(tool_input, **self._parameters) + content = self._transform_output(tool_input, **parameters) else: # Valid tool_input - transform to BaseModel - content = self._transform_output(tool_input, **self._parameters) + content = self._transform_output(tool_input, **parameters) else: # Fallback: concatenate text chunks text_content = "".join(chunk.content for chunk in chunks) if text_content: - content = self._transform_output(text_content, **self._parameters) + content = self._transform_output(text_content, **parameters) else: msg = "No tool_use input and no text chunks available" raise ValidationError(msg) @@ -330,7 +328,7 @@ def _parse_output( # No tool_use blocks or no output_schema: concatenate text chunks content = "".join(chunk.content for chunk in chunks) # Apply parameter transformations (e.g., JSON → BaseModel if output_schema provided) - content = self._transform_output(content, **self._parameters) + content = self._transform_output(content, **parameters) usage = self._parse_usage(chunks) finish_reason = chunks[-1].finish_reason if chunks else None diff --git a/packages/text-generation/src/celeste_text_generation/providers/cohere/client.py b/packages/text-generation/src/celeste_text_generation/providers/cohere/client.py index 6c5adc2a..eec330b9 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/cohere/client.py +++ b/packages/text-generation/src/celeste_text_generation/providers/cohere/client.py @@ -6,6 +6,7 @@ import httpx from pydantic import BaseModel +from celeste.mime_types import ApplicationMimeType from celeste.parameters import ParameterMapper from celeste_text_generation.client import TextGenerationClient from celeste_text_generation.io import ( @@ -21,7 +22,7 @@ class CohereTextGenerationClient(TextGenerationClient): - """Cohere client.""" + """Cohere client for text generation.""" @classmethod def parameter_mappers(cls) -> list[ParameterMapper]: @@ -96,7 +97,7 @@ async def _make_request( headers = { config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", - "Content-Type": "application/json", + "Content-Type": ApplicationMimeType.JSON, } return await self.http_client.post( @@ -120,7 +121,7 @@ def _make_stream_request( headers = { config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", - "Content-Type": "application/json", + "Content-Type": ApplicationMimeType.JSON, } return self.http_client.stream_post( diff --git a/packages/text-generation/src/celeste_text_generation/providers/cohere/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/cohere/streaming.py index f36d36e9..1b43850a 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/cohere/streaming.py +++ b/packages/text-generation/src/celeste_text_generation/providers/cohere/streaming.py @@ -23,7 +23,7 @@ class CohereTextGenerationStream(TextGenerationStream): def __init__( self, sse_iterator: Any, # noqa: ANN401 - transform_output: Callable[[object, Any], object], + transform_output: Callable[..., object], **parameters: Unpack[TextGenerationParameters], ) -> None: """Initialize stream with output transformation support. @@ -139,7 +139,7 @@ def _parse_output( """Assemble chunks into final output, applying parameter transformations.""" content_chunks = [chunk for chunk in chunks if chunk.content] content = "".join(chunk.content for chunk in content_chunks) - content = self._transform_output(content, **self._parameters) + content = self._transform_output(content, **parameters) usage = self._parse_usage(chunks) finish_reason = chunks[-1].finish_reason if chunks else None diff --git a/packages/text-generation/src/celeste_text_generation/providers/google/client.py b/packages/text-generation/src/celeste_text_generation/providers/google/client.py index e0862dab..10e905cc 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/google/client.py +++ b/packages/text-generation/src/celeste_text_generation/providers/google/client.py @@ -6,6 +6,7 @@ import httpx from pydantic import BaseModel +from celeste.mime_types import ApplicationMimeType from celeste.parameters import ParameterMapper from celeste_text_generation.client import TextGenerationClient from celeste_text_generation.io import ( @@ -21,7 +22,7 @@ class GoogleTextGenerationClient(TextGenerationClient): - """Google client.""" + """Google client for text generation.""" @classmethod def parameter_mappers(cls) -> list[ParameterMapper]: @@ -101,7 +102,7 @@ async def _make_request( headers = { config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", - "Content-Type": "application/json", + "Content-Type": ApplicationMimeType.JSON, } return await self.http_client.post( @@ -124,7 +125,7 @@ def _make_stream_request( headers = { config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", - "Content-Type": "application/json", + "Content-Type": ApplicationMimeType.JSON, } return self.http_client.stream_post( diff --git a/packages/text-generation/src/celeste_text_generation/providers/google/parameters.py b/packages/text-generation/src/celeste_text_generation/providers/google/parameters.py index 19e1a491..85fd30bb 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/google/parameters.py +++ b/packages/text-generation/src/celeste_text_generation/providers/google/parameters.py @@ -6,6 +6,7 @@ from pydantic import BaseModel, TypeAdapter from celeste.core import Parameter +from celeste.mime_types import ApplicationMimeType from celeste.models import Model from celeste.parameters import ParameterMapper from celeste_text_generation.parameters import TextGenerationParameter @@ -93,7 +94,7 @@ def map( config = request.setdefault("generationConfig", {}) config["responseSchema"] = schema - config["responseMimeType"] = "application/json" + config["responseMimeType"] = ApplicationMimeType.JSON return request diff --git a/packages/text-generation/src/celeste_text_generation/providers/google/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/google/streaming.py index aa3f6587..fea36b73 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/google/streaming.py +++ b/packages/text-generation/src/celeste_text_generation/providers/google/streaming.py @@ -20,7 +20,7 @@ class GoogleTextGenerationStream(TextGenerationStream): def __init__( self, sse_iterator: Any, # noqa: ANN401 - transform_output: Callable[[object, Any], object], + transform_output: Callable[..., object], **parameters: Unpack[TextGenerationParameters], ) -> None: """Initialize stream with output transformation support. @@ -149,7 +149,7 @@ def _parse_output( content = "".join(chunk.content for chunk in chunks) # Apply parameter transformations (e.g., JSON → BaseModel if output_schema provided) - content = self._transform_output(content, **self._parameters) + content = self._transform_output(content, **parameters) usage = self._parse_usage(chunks) finish_reason = chunks[-1].finish_reason if chunks else None diff --git a/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py b/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py index 689e295d..0c87985c 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py +++ b/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py @@ -6,6 +6,7 @@ import httpx from pydantic import BaseModel +from celeste.mime_types import ApplicationMimeType from celeste.parameters import ParameterMapper from celeste_text_generation.client import TextGenerationClient from celeste_text_generation.io import ( @@ -21,7 +22,7 @@ class MistralTextGenerationClient(TextGenerationClient): - """Mistral client.""" + """Mistral client for text generation.""" @classmethod def parameter_mappers(cls) -> list[ParameterMapper]: @@ -91,7 +92,7 @@ async def _make_request( headers = { config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", - "Content-Type": "application/json", + "Content-Type": ApplicationMimeType.JSON, } return await self.http_client.post( @@ -115,7 +116,7 @@ def _make_stream_request( headers = { config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", - "Content-Type": "application/json", + "Content-Type": ApplicationMimeType.JSON, } return self.http_client.stream_post( diff --git a/packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py index 981be553..2aef6437 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py +++ b/packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py @@ -23,7 +23,7 @@ class MistralTextGenerationStream(TextGenerationStream): def __init__( self, sse_iterator: Any, # noqa: ANN401 - transform_output: Callable[[object, Any], object], + transform_output: Callable[..., object], **parameters: Unpack[TextGenerationParameters], ) -> None: """Initialize stream with output transformation support. @@ -119,7 +119,7 @@ def _parse_output( content = "".join(chunk.content for chunk in content_chunks) # Apply parameter transformations (e.g., JSON → BaseModel if output_schema provided) - content = self._transform_output(content, **self._parameters) + content = self._transform_output(content, **parameters) usage = self._parse_usage(chunks) finish_reason = chunks[-1].finish_reason if chunks else None diff --git a/packages/text-generation/src/celeste_text_generation/providers/openai/client.py b/packages/text-generation/src/celeste_text_generation/providers/openai/client.py index 15c89419..d0d5b4b7 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/openai/client.py +++ b/packages/text-generation/src/celeste_text_generation/providers/openai/client.py @@ -6,6 +6,7 @@ import httpx from pydantic import BaseModel +from celeste.mime_types import ApplicationMimeType from celeste.parameters import ParameterMapper from celeste_text_generation.client import TextGenerationClient from celeste_text_generation.io import ( @@ -21,7 +22,7 @@ class OpenAITextGenerationClient(TextGenerationClient): - """OpenAI client.""" + """OpenAI client for text generation.""" @classmethod def parameter_mappers(cls) -> list[ParameterMapper]: @@ -107,7 +108,7 @@ async def _make_request( headers = { config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", - "Content-Type": "application/json", + "Content-Type": ApplicationMimeType.JSON, } return await self.http_client.post( @@ -131,7 +132,7 @@ def _make_stream_request( headers = { config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", - "Content-Type": "application/json", + "Content-Type": ApplicationMimeType.JSON, } return self.http_client.stream_post( diff --git a/packages/text-generation/src/celeste_text_generation/providers/openai/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/openai/streaming.py index bde9c630..b5c8ce8d 100644 --- a/packages/text-generation/src/celeste_text_generation/providers/openai/streaming.py +++ b/packages/text-generation/src/celeste_text_generation/providers/openai/streaming.py @@ -20,7 +20,7 @@ class OpenAITextGenerationStream(TextGenerationStream): def __init__( self, sse_iterator: Any, # noqa: ANN401 - transform_output: Callable[[object, Any], object], + transform_output: Callable[..., object], **parameters: Unpack[TextGenerationParameters], ) -> None: """Initialize stream.""" @@ -94,7 +94,7 @@ def _parse_output( ) -> TextGenerationOutput: """Assemble chunks into final output.""" content = "".join(chunk.content for chunk in chunks) - content = self._transform_output(content, **self._parameters) + content = self._transform_output(content, **parameters) usage = self._parse_usage(chunks) finish_reason = chunks[-1].finish_reason if chunks else None diff --git a/packages/text-generation/src/celeste_text_generation/streaming.py b/packages/text-generation/src/celeste_text_generation/streaming.py index daeff747..50bf8cac 100644 --- a/packages/text-generation/src/celeste_text_generation/streaming.py +++ b/packages/text-generation/src/celeste_text_generation/streaming.py @@ -19,7 +19,6 @@ class TextGenerationStream(Stream[TextGenerationOutput, TextGenerationParameters @abstractmethod def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: """Parse SSE event into Chunk (provider-specific).""" - ... def _parse_output( self, diff --git a/packages/text-generation/tests/integration_tests/test_text_generation/test_stream.py b/packages/text-generation/tests/integration_tests/test_text_generation/test_stream.py new file mode 100644 index 00000000..78d87bbb --- /dev/null +++ b/packages/text-generation/tests/integration_tests/test_text_generation/test_stream.py @@ -0,0 +1,103 @@ +"""Integration tests for text generation streaming across all providers.""" + +import pytest + +from celeste import Capability, Provider, create_client + + +@pytest.mark.parametrize( + ("provider", "model", "parameters"), + [ + (Provider.OPENAI, "gpt-3.5-turbo", {}), + (Provider.ANTHROPIC, "claude-haiku-4-5", {}), + (Provider.GOOGLE, "gemini-2.5-flash-lite", {"thinking_budget": 0}), + (Provider.MISTRAL, "mistral-tiny", {}), + (Provider.COHERE, "command-a-03-2025", {}), + ], +) +@pytest.mark.integration +@pytest.mark.asyncio +async def test_stream(provider: Provider, model: str, parameters: dict) -> None: + """Test text generation streaming with max_tokens parameter across all providers. + + This test demonstrates that the unified streaming API works identically across + all providers using the same code - proving the abstraction value. + """ + # Import here to avoid circular import during pytest collection + from celeste_text_generation import ( + TextGenerationChunk, + TextGenerationFinishReason, + TextGenerationOutput, + TextGenerationStream, + TextGenerationUsage, + ) + + # Arrange + client = create_client( + capability=Capability.TEXT_GENERATION, + provider=provider, + ) + prompt = "Hi" + max_tokens = 30 + + # Act - Create stream + stream = client.stream( + prompt=prompt, + model=model, + max_tokens=max_tokens, + **parameters, + ) + + # Assert 1: Stream Creation + assert isinstance(stream, TextGenerationStream), ( + f"Expected TextGenerationStream, got {type(stream)}" + ) + + # Act - Iterate chunks + chunks: list[TextGenerationChunk] = [] + async for chunk in stream: + # Assert 2: Chunk Structure + assert isinstance(chunk, TextGenerationChunk), ( + f"Expected TextGenerationChunk, got {type(chunk)}" + ) + assert isinstance(chunk.content, str), ( + f"Expected str content, got {type(chunk.content)}" + ) + if chunk.finish_reason is not None: + assert isinstance(chunk.finish_reason, TextGenerationFinishReason), ( + f"Expected TextGenerationFinishReason, got {type(chunk.finish_reason)}" + ) + if chunk.usage is not None: + assert isinstance(chunk.usage, TextGenerationUsage), ( + f"Expected TextGenerationUsage, got {type(chunk.usage)}" + ) + chunks.append(chunk) + + # Assert 3: Content Accumulation + assert len(chunks) > 0, "No chunks were yielded from stream" + accumulated_content = "".join(chunk.content for chunk in chunks) + assert len(accumulated_content) > 0, f"Accumulated content is empty: {chunks!r}" + + # Assert 4: Final Output + output = stream.output + assert isinstance(output, TextGenerationOutput), ( + f"Expected TextGenerationOutput, got {type(output)}" + ) + assert isinstance(output.content, str), ( + f"Expected str content, got {type(output.content)}" + ) + assert output.content == accumulated_content, ( + f"Output content doesn't match accumulated chunks: " + f"{output.content!r} != {accumulated_content!r}" + ) + + # Assert 5: Usage Metrics + assert isinstance(output.usage, TextGenerationUsage), ( + f"Expected TextGenerationUsage, got {type(output.usage)}" + ) + + # Assert 6: Finish Reason + assert chunks[-1].finish_reason is not None, "Final chunk should have finish_reason" + assert "finish_reason" in output.metadata, ( + "Output metadata should contain finish_reason" + ) diff --git a/pyproject.toml b/pyproject.toml index c4133783..246f3d10 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,7 +87,7 @@ branch = true [tool.coverage.report] show_missing = true -fail_under = 90 +fail_under = 80 [tool.ruff] # Same as Black's default diff --git a/src/celeste/mime_types.py b/src/celeste/mime_types.py index 708d49e5..8049d8ac 100644 --- a/src/celeste/mime_types.py +++ b/src/celeste/mime_types.py @@ -6,7 +6,11 @@ class MimeType(StrEnum): """Base class for all MIME types.""" - pass + +class ApplicationMimeType(MimeType): + """Standard MIME types for application data.""" + + JSON = "application/json" class ImageMimeType(MimeType): @@ -15,6 +19,9 @@ class ImageMimeType(MimeType): PNG = "image/png" JPEG = "image/jpeg" WEBP = "image/webp" + BMP = "image/bmp" + TIFF = "image/tiff" + GIF = "image/gif" class VideoMimeType(MimeType): @@ -37,9 +44,11 @@ class AudioMimeType(MimeType): AIFF = "audio/aiff" M4A = "audio/mp4" WMA = "audio/x-ms-wma" + PCM = "audio/pcm" __all__ = [ + "ApplicationMimeType", "AudioMimeType", "ImageMimeType", "MimeType",