diff --git a/.env.example b/.env.example index 87b3885f..0729bccd 100644 --- a/.env.example +++ b/.env.example @@ -39,3 +39,6 @@ TOPAZLABS_API_KEY=your-topazlabs-api-key-here # Perplexity PERPLEXITY_API_KEY=your-perplexity-api-key-here + +# Gradium +GRADIUM_API_KEY=your-gradium-api-key-here diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 79437ce1..54ed205c 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -136,6 +136,7 @@ jobs: BYTEPLUS_API_KEY: ${{ secrets.BYTEPLUS_API_KEY }} XAI_API_KEY: ${{ secrets.XAI_API_KEY }} ELEVENLABS_API_KEY: ${{ secrets.ELEVENLABS_API_KEY }} + GRADIUM_API_KEY: ${{ secrets.GRADIUM_API_KEY }} run: uv run pytest packages/capabilities/${{ matrix.package }}/tests/integration_tests/ -m integration -v build: diff --git a/.gitignore b/.gitignore index 730293e9..4ad08e72 100644 --- a/.gitignore +++ b/.gitignore @@ -157,3 +157,4 @@ uv.lock # Security reports bandit-report.json +mureka.md diff --git a/packages/capabilities/speech-generation/pyproject.toml b/packages/capabilities/speech-generation/pyproject.toml index 921f45fa..a904ff6d 100644 --- a/packages/capabilities/speech-generation/pyproject.toml +++ b/packages/capabilities/speech-generation/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "celeste-speech-generation" -version = "0.3.0" +version = "0.3.1" description = "Speech generation package for Celeste AI. Unified interface for all providers" authors = [{name = "Kamilbenkirane", email = "kamil@withceleste.ai"}] readme = "README.md" @@ -17,7 +17,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", "Typing :: Typed", ] -keywords = ["ai", "speech-generation", "tts", "text-to-speech", "openai", "google", "elevenlabs", "audio-ai"] +keywords = ["ai", "speech-generation", "tts", "text-to-speech", "openai", "google", "elevenlabs", "gradium", "audio-ai"] [project.urls] Homepage = "https://withceleste.ai" @@ -28,6 +28,8 @@ Issues = "https://github.com/withceleste/celeste-python/issues" [tool.uv.sources] celeste-ai = { workspace = true } celeste-elevenlabs = { workspace = true } +celeste-google = { workspace = true } +celeste-gradium = { workspace = true } celeste-openai = { workspace = true } [project.entry-points."celeste.packages"] diff --git a/packages/capabilities/speech-generation/src/celeste_speech_generation/__init__.py b/packages/capabilities/speech-generation/src/celeste_speech_generation/__init__.py index b71207c0..c6dad719 100644 --- a/packages/capabilities/speech-generation/src/celeste_speech_generation/__init__.py +++ b/packages/capabilities/speech-generation/src/celeste_speech_generation/__init__.py @@ -30,6 +30,9 @@ def register_package() -> None: from celeste_speech_generation.providers.google.voices import ( # noqa: E402 GOOGLE_VOICES, ) +from celeste_speech_generation.providers.gradium.voices import ( # noqa: E402 + GRADIUM_VOICES, +) from celeste_speech_generation.providers.openai.voices import ( # noqa: E402 OPENAI_VOICES, ) @@ -37,9 +40,10 @@ def register_package() -> None: from celeste_speech_generation.voices import Voice # noqa: E402 VOICES: list[Voice] = [ + *ELEVENLABS_VOICES, *GOOGLE_VOICES, + *GRADIUM_VOICES, *OPENAI_VOICES, - *ELEVENLABS_VOICES, ] __all__ = [ diff --git a/packages/capabilities/speech-generation/src/celeste_speech_generation/models.py b/packages/capabilities/speech-generation/src/celeste_speech_generation/models.py index c95d7cd4..1cd16031 100644 --- a/packages/capabilities/speech-generation/src/celeste_speech_generation/models.py +++ b/packages/capabilities/speech-generation/src/celeste_speech_generation/models.py @@ -5,10 +5,12 @@ MODELS as ELEVENLABS_MODELS, ) from celeste_speech_generation.providers.google.models import MODELS as GOOGLE_MODELS +from celeste_speech_generation.providers.gradium.models import MODELS as GRADIUM_MODELS from celeste_speech_generation.providers.openai.models import MODELS as OPENAI_MODELS MODELS: list[Model] = [ *GOOGLE_MODELS, *OPENAI_MODELS, *ELEVENLABS_MODELS, + *GRADIUM_MODELS, ] diff --git a/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/__init__.py b/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/__init__.py index aa8f14f6..9fdb918d 100644 --- a/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/__init__.py +++ b/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/__init__.py @@ -14,6 +14,9 @@ def _get_providers() -> list[tuple[Provider, type[Client]]]: from celeste_speech_generation.providers.google.client import ( GoogleSpeechGenerationClient, ) + from celeste_speech_generation.providers.gradium.client import ( + GradiumSpeechGenerationClient, + ) from celeste_speech_generation.providers.openai.client import ( OpenAISpeechGenerationClient, ) @@ -22,6 +25,7 @@ def _get_providers() -> list[tuple[Provider, type[Client]]]: (Provider.GOOGLE, GoogleSpeechGenerationClient), (Provider.OPENAI, OpenAISpeechGenerationClient), (Provider.ELEVENLABS, ElevenLabsSpeechGenerationClient), + (Provider.GRADIUM, GradiumSpeechGenerationClient), ] diff --git a/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/gradium/__init__.py b/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/gradium/__init__.py new file mode 100644 index 00000000..bc01c62b --- /dev/null +++ b/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/gradium/__init__.py @@ -0,0 +1,9 @@ +"""Gradium provider for speech generation.""" + +from .client import GradiumSpeechGenerationClient +from .models import MODELS + +__all__ = [ + "MODELS", + "GradiumSpeechGenerationClient", +] diff --git a/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/gradium/client.py b/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/gradium/client.py new file mode 100644 index 00000000..d2dafb12 --- /dev/null +++ b/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/gradium/client.py @@ -0,0 +1,79 @@ +"""Gradium client implementation for speech generation.""" + +from typing import Any, Unpack + +from celeste_gradium.text_to_speech.client import GradiumTextToSpeechClient + +from celeste.artifacts import AudioArtifact +from celeste.parameters import ParameterMapper +from celeste_speech_generation.client import SpeechGenerationClient +from celeste_speech_generation.io import ( + SpeechGenerationInput, + SpeechGenerationOutput, + SpeechGenerationUsage, +) +from celeste_speech_generation.parameters import SpeechGenerationParameters + +from .parameters import GRADIUM_PARAMETER_MAPPERS + + +class GradiumSpeechGenerationClient(GradiumTextToSpeechClient, SpeechGenerationClient): + """Gradium client for speech generation.""" + + @classmethod + def parameter_mappers(cls) -> list[ParameterMapper]: + return GRADIUM_PARAMETER_MAPPERS + + def _init_request(self, inputs: SpeechGenerationInput) -> dict[str, Any]: + """Initialize request from Gradium API format.""" + return {"text": inputs.text} + + def _parse_usage(self, response_data: dict[str, Any]) -> SpeechGenerationUsage: + """Parse usage from response.""" + usage = super()._parse_usage(response_data) + return SpeechGenerationUsage(**usage) + + def _parse_content( + self, + response_data: dict[str, Any], + **parameters: Unpack[SpeechGenerationParameters], + ) -> AudioArtifact: + """Parse content from response. + + Note: This method is not used for Gradium TTS since we override generate() + to handle WebSocket responses. Kept for interface compliance. + """ + msg = "Gradium TTS uses WebSocket, use generate() override" + raise NotImplementedError(msg) + + async def generate( + self, + *args: str, + **parameters: Unpack[SpeechGenerationParameters], + ) -> SpeechGenerationOutput: + """Generate speech from text. + + Override base generate() to use WebSocket instead of HTTP. + """ + inputs = self._create_inputs(*args, **parameters) + inputs, parameters = self._validate_artifacts(inputs, **parameters) + request_body = self._build_request(inputs, **parameters) + + # Use WebSocket TTS flow + audio_bytes, output_format = await self._websocket_tts(request_body) + + if not audio_bytes: + msg = "No audio data in response" + raise ValueError(msg) + + # Determine MIME type from output_format + mime_type = self._map_output_format_to_mime_type(output_format) + + return self._output_class()( + content=AudioArtifact(data=audio_bytes, mime_type=mime_type), + usage=SpeechGenerationUsage(), + metadata=self._build_metadata({}), + ) + + +__all__ = ["GradiumSpeechGenerationClient"] diff --git a/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/gradium/models.py b/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/gradium/models.py new file mode 100644 index 00000000..8dd86ac1 --- /dev/null +++ b/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/gradium/models.py @@ -0,0 +1,34 @@ +"""Gradium model definitions for speech generation.""" + +from celeste import Model, Provider +from celeste.constraints import Choice, Range +from celeste_speech_generation.constraints import VoiceConstraint +from celeste_speech_generation.parameters import SpeechGenerationParameter + +from .voices import GRADIUM_VOICES + +MODELS: list[Model] = [ + Model( + id="default", + provider=Provider.GRADIUM, + display_name="Gradium Default TTS", + streaming=False, + parameter_constraints={ + SpeechGenerationParameter.VOICE: VoiceConstraint(voices=GRADIUM_VOICES), + SpeechGenerationParameter.OUTPUT_FORMAT: Choice( + options=[ + "wav", + "pcm", + "opus", + "ulaw_8000", + "alaw_8000", + "pcm_16000", + "pcm_24000", + ] + ), + SpeechGenerationParameter.SPEED: Range(min=0.25, max=4.0), + }, + ), +] + +__all__ = ["MODELS"] diff --git a/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/gradium/parameters.py b/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/gradium/parameters.py new file mode 100644 index 00000000..e66e934b --- /dev/null +++ b/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/gradium/parameters.py @@ -0,0 +1,60 @@ +"""Gradium Text-to-Speech parameter mappers for speech generation.""" + +from typing import Any + +from celeste_gradium.text_to_speech.parameters import ( + OutputFormatMapper as _OutputFormatMapper, +) +from celeste_gradium.text_to_speech.parameters import ( + PaddingBonusMapper as _PaddingBonusMapper, +) +from celeste_gradium.text_to_speech.parameters import ( + VoiceMapper as _VoiceMapper, +) + +from celeste.models import Model +from celeste.parameters import ParameterMapper +from celeste_speech_generation.parameters import SpeechGenerationParameter + + +class VoiceMapper(_VoiceMapper): + name = SpeechGenerationParameter.VOICE + + +class OutputFormatMapper(_OutputFormatMapper): + name = SpeechGenerationParameter.OUTPUT_FORMAT + + +class SpeedMapper(_PaddingBonusMapper): + """Translate unified speed to Gradium padding_bonus. + + speed 1.0 → padding_bonus 0 (normal) + speed 0.5 → padding_bonus 2.0 (slower) + speed 2.0 → padding_bonus -4.0 (faster) + """ + + name = SpeechGenerationParameter.SPEED + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform speed into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + # Translate: speed → padding_bonus + padding_bonus = (1.0 - validated_value) * 4.0 + return super().map(request, padding_bonus, model) + + +GRADIUM_PARAMETER_MAPPERS: list[ParameterMapper] = [ + VoiceMapper(), + OutputFormatMapper(), + SpeedMapper(), +] + +__all__ = ["GRADIUM_PARAMETER_MAPPERS"] diff --git a/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/gradium/voices.py b/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/gradium/voices.py new file mode 100644 index 00000000..de4a1d84 --- /dev/null +++ b/packages/capabilities/speech-generation/src/celeste_speech_generation/providers/gradium/voices.py @@ -0,0 +1,101 @@ +"""Gradium voice definitions for speech generation.""" + +from celeste import Provider +from celeste_speech_generation.voices import Voice + +# Gradium flagship voices +# Full list at https://gradium.ai/api_dpocs.html +GRADIUM_VOICES = [ + # English (US) + Voice( + id="YTpq7expH9539ERJ", + provider=Provider.GRADIUM, + name="Emma", + languages={"en"}, + ), + Voice( + id="LFZvm12tW_z0xfGo", + provider=Provider.GRADIUM, + name="Kent", + languages={"en"}, + ), + Voice( + id="jtEKaLYNn6iif5PR", + provider=Provider.GRADIUM, + name="Sydney", + languages={"en"}, + ), + Voice( + id="KWJiFWu2O9nMPYcR", + provider=Provider.GRADIUM, + name="John", + languages={"en"}, + ), + # English (GB) + Voice( + id="ubuXFxVQwVYnZQhy", + provider=Provider.GRADIUM, + name="Eva", + languages={"en"}, + ), + Voice( + id="m86j6D7UZpGzHsNu", + provider=Provider.GRADIUM, + name="Jack", + languages={"en"}, + ), + # French + Voice( + id="b35yykvVppLXyw_l", + provider=Provider.GRADIUM, + name="Elise", + languages={"fr"}, + ), + Voice( + id="axlOaUiFyOZhy4nv", + provider=Provider.GRADIUM, + name="Leo", + languages={"fr"}, + ), + # German + Voice( + id="-uP9MuGtBqAvEyxI", + provider=Provider.GRADIUM, + name="Mia", + languages={"de"}, + ), + Voice( + id="0y1VZjPabOBU3rWy", + provider=Provider.GRADIUM, + name="Maximilian", + languages={"de"}, + ), + # Spanish + Voice( + id="B36pbz5_UoWn4BDl", + provider=Provider.GRADIUM, + name="Valentina", + languages={"es"}, + ), + Voice( + id="xu7iJ_fn2ElcWp2s", + provider=Provider.GRADIUM, + name="Sergio", + languages={"es"}, + ), + # Portuguese (Brazilian) + Voice( + id="pYcGZz9VOo4n2ynh", + provider=Provider.GRADIUM, + name="Alice", + languages={"pt"}, + ), + Voice( + id="M-FvVo9c-jGR4PgP", + provider=Provider.GRADIUM, + name="Davi", + languages={"pt"}, + ), +] + +__all__ = ["GRADIUM_VOICES"] diff --git a/packages/capabilities/speech-generation/tests/integration_tests/test_speech_generation/test_generate.py b/packages/capabilities/speech-generation/tests/integration_tests/test_speech_generation/test_generate.py index 20061909..55252691 100644 --- a/packages/capabilities/speech-generation/tests/integration_tests/test_speech_generation/test_generate.py +++ b/packages/capabilities/speech-generation/tests/integration_tests/test_speech_generation/test_generate.py @@ -19,6 +19,7 @@ "eleven_flash_v2_5", {"voice": "Rachel", "output_format": "mp3_44100_128"}, ), + (Provider.GRADIUM, "default", {}), ], ) @pytest.mark.integration diff --git a/packages/providers/gradium/pyproject.toml b/packages/providers/gradium/pyproject.toml new file mode 100644 index 00000000..91d0f1e7 --- /dev/null +++ b/packages/providers/gradium/pyproject.toml @@ -0,0 +1,18 @@ +[project] +name = "celeste-gradium" +version = "0.3.1" +description = "Gradium provider package for Celeste AI" +authors = [{name = "Kamilbenkirane", email = "kamil@withceleste.ai"}] +license = {text = "Apache-2.0"} +requires-python = ">=3.12" +dependencies = ["celeste-ai", "websockets"] + +[tool.uv.sources] +celeste-ai = { workspace = true } + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/celeste_gradium"] diff --git a/packages/providers/gradium/src/celeste_gradium/__init__.py b/packages/providers/gradium/src/celeste_gradium/__init__.py new file mode 100644 index 00000000..a39dfaa4 --- /dev/null +++ b/packages/providers/gradium/src/celeste_gradium/__init__.py @@ -0,0 +1,3 @@ +"""Gradium provider package for Celeste AI.""" + +__all__: list[str] = [] diff --git a/packages/providers/gradium/src/celeste_gradium/py.typed b/packages/providers/gradium/src/celeste_gradium/py.typed new file mode 100644 index 00000000..e16c76df --- /dev/null +++ b/packages/providers/gradium/src/celeste_gradium/py.typed @@ -0,0 +1 @@ +"" diff --git a/packages/providers/gradium/src/celeste_gradium/text_to_speech/__init__.py b/packages/providers/gradium/src/celeste_gradium/text_to_speech/__init__.py new file mode 100644 index 00000000..701cfd7a --- /dev/null +++ b/packages/providers/gradium/src/celeste_gradium/text_to_speech/__init__.py @@ -0,0 +1 @@ +"""Gradium Text to Speech API provider package.""" diff --git a/packages/providers/gradium/src/celeste_gradium/text_to_speech/client.py b/packages/providers/gradium/src/celeste_gradium/text_to_speech/client.py new file mode 100644 index 00000000..49fe03a4 --- /dev/null +++ b/packages/providers/gradium/src/celeste_gradium/text_to_speech/client.py @@ -0,0 +1,150 @@ +"""Gradium Text-to-Speech API client with shared implementation.""" + +import base64 +import json +from typing import Any + +from websockets.asyncio.client import connect as ws_connect + +from celeste.client import APIMixin +from celeste.mime_types import AudioMimeType + +from . import config + + +class GradiumTextToSpeechClient(APIMixin): + """Mixin for Gradium Text-to-Speech API. + + Provides shared implementation for speech generation via WebSocket: + - _websocket_tts() - Execute WebSocket TTS flow + - _parse_usage() - Returns empty dict (TTS doesn't return usage) + - _map_output_format_to_mime_type() - Map format string to AudioMimeType + + The Gradium Text-to-Speech API uses WebSocket instead of HTTP REST: + 1. Connect to wss://{region}.api.gradium.ai/api/speech/tts + 2. Send setup message with model, voice, format + 3. Receive ready confirmation + 4. Send text to synthesize + 5. Receive audio chunks (base64 encoded) + 6. Receive end message + + Subclasses override generate() to use _websocket_tts() instead of HTTP. + + Usage: + class GradiumSpeechGenerationClient(GradiumTextToSpeechClient, SpeechGenerationClient): + async def generate(self, *args, **parameters): + # Build request... + audio_data, format = await self._websocket_tts(request_body) + # Return SpeechGenerationOutput... + """ + + async def _websocket_tts( + self, + request_body: dict[str, Any], + ) -> tuple[bytes, str]: + """Execute WebSocket TTS flow. + + Args: + request_body: Request with text, voice_id, output_format, json_config. + + Returns: + Tuple of (audio_bytes, output_format). + + Raises: + ValueError: If connection fails or error received. + """ + voice_id = request_body.get("voice_id", config.DEFAULT_VOICE_ID) + output_format = request_body.get("output_format", "wav") + text = request_body.get("text", "") + json_config = request_body.get("json_config") + + url = f"{config.BASE_URL}{config.GradiumTextToSpeechEndpoint.CREATE_SPEECH}" + headers = self.auth.get_headers() + + async with ws_connect(url, additional_headers=headers) as ws: + # 1. Send setup message + setup_msg: dict[str, Any] = { + "type": "setup", + "model_name": self.model.id, + "voice_id": voice_id, + "output_format": output_format, + } + if json_config is not None: + setup_msg["json_config"] = json_config + + await ws.send(json.dumps(setup_msg)) + + # 2. Wait for ready + ready_msg = await ws.recv() + ready = json.loads(ready_msg) + if ready.get("type") != "ready": + msg = f"Expected ready message, got: {ready}" + raise ValueError(msg) + + # 3. Send text + await ws.send(json.dumps({"type": "text", "text": text})) + + # 4. Signal end of input + await ws.send(json.dumps({"type": "end_of_stream"})) + + # 5. Collect audio chunks + audio_chunks: list[bytes] = [] + async for message in ws: + if isinstance(message, bytes): + data = json.loads(message.decode("utf-8")) + else: + data = json.loads(message) + + if data["type"] == "audio": + audio_chunks.append(base64.b64decode(data["audio"])) + elif data["type"] == "end_of_stream": + break + elif data["type"] == "error": + error_msg = data.get("message", "Unknown error") + msg = f"Gradium TTS error: {error_msg}" + raise ValueError(msg) + + return b"".join(audio_chunks), output_format + + def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, Any]: + """Extract usage data from Gradium API response. + + Gradium TTS doesn't provide usage metadata. + Returns empty dict for capability clients to wrap in Usage type. + """ + return {} + + async def _make_request( + self, + request_body: dict[str, Any], + **parameters: Any, + ) -> Any: + """Make HTTP request - not used for Gradium (uses WebSocket). + + Gradium TTS uses WebSocket via _websocket_tts(). + This method satisfies the abstract interface but should not be called. + """ + msg = "Gradium TTS uses WebSocket, use _websocket_tts() instead" + raise NotImplementedError(msg) + + def _map_output_format_to_mime_type( + self, + output_format: str | None, + ) -> AudioMimeType: + """Map Gradium output_format to AudioMimeType. + + Supported formats: wav, pcm, opus, ulaw_8000, alaw_8000, pcm_16000, pcm_24000. + """ + format_map: dict[str, AudioMimeType] = { + "wav": AudioMimeType.WAV, + "pcm": AudioMimeType.WAV, # PCM raw, closest match + "opus": AudioMimeType.OGG, # Opus in OGG container + "ulaw_8000": AudioMimeType.WAV, + "alaw_8000": AudioMimeType.WAV, + "pcm_16000": AudioMimeType.WAV, + "pcm_24000": AudioMimeType.WAV, + } + return format_map.get(output_format or "", AudioMimeType.WAV) + + +__all__ = ["GradiumTextToSpeechClient"] diff --git a/packages/providers/gradium/src/celeste_gradium/text_to_speech/config.py b/packages/providers/gradium/src/celeste_gradium/text_to_speech/config.py new file mode 100644 index 00000000..7db027f6 --- /dev/null +++ b/packages/providers/gradium/src/celeste_gradium/text_to_speech/config.py @@ -0,0 +1,15 @@ +"""Configuration for Gradium Text-to-Speech API.""" + +from enum import StrEnum + + +class GradiumTextToSpeechEndpoint(StrEnum): + """Endpoints for Text-to-Speech API.""" + + CREATE_SPEECH = "/speech/tts" + + +BASE_URL = "wss://eu.api.gradium.ai/api" + +# Default voice ID (Emma - English female) +DEFAULT_VOICE_ID = "YTpq7expH9539ERJ" diff --git a/packages/providers/gradium/src/celeste_gradium/text_to_speech/parameters.py b/packages/providers/gradium/src/celeste_gradium/text_to_speech/parameters.py new file mode 100644 index 00000000..71f79dec --- /dev/null +++ b/packages/providers/gradium/src/celeste_gradium/text_to_speech/parameters.py @@ -0,0 +1,67 @@ +"""Gradium TTS API parameter mappers.""" + +from typing import Any + +from celeste.models import Model +from celeste.parameters import ParameterMapper + + +class VoiceMapper(ParameterMapper): + """Map voice to Gradium voice_id field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform voice into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["voice_id"] = validated_value + return request + + +class PaddingBonusMapper(ParameterMapper): + """Map padding_bonus to Gradium json_config.padding_bonus field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform padding_bonus into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("json_config", {})["padding_bonus"] = validated_value + return request + + +class OutputFormatMapper(ParameterMapper): + """Map output_format to Gradium output_format field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform output_format into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["output_format"] = validated_value + return request + + +__all__ = [ + "OutputFormatMapper", + "PaddingBonusMapper", + "VoiceMapper", +] diff --git a/pyproject.toml b/pyproject.toml index 60bd048f..188f89f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "celeste-ai" -version = "0.3.0" +version = "0.3.1" description = "Open source, type-safe primitives for multi-modal AI. All capabilities, all providers, one interface" authors = [{name = "Kamilbenkirane", email = "kamil@withceleste.ai"}] readme = "README.md" @@ -24,6 +24,7 @@ dependencies = [ "httpx>=0.27.0", "httpx-sse>=0.4.0", "python-dotenv>=1.0.0", + "websockets>=13.0", ] [project.urls] diff --git a/src/celeste/__init__.py b/src/celeste/__init__.py index aa2a5a81..d3dc96c4 100644 --- a/src/celeste/__init__.py +++ b/src/celeste/__init__.py @@ -32,6 +32,7 @@ ) from celeste.types import JsonValue from celeste.utils import image_to_data_uri +from celeste.websocket import WebSocketClient, WebSocketConnection, close_all_ws_clients logger = logging.getLogger(__name__) @@ -143,7 +144,10 @@ def create_client( "Usage", "UsageField", "ValidationError", + "WebSocketClient", + "WebSocketConnection", "close_all_http_clients", + "close_all_ws_clients", "create_client", "get_client_class", "get_model", diff --git a/src/celeste/core.py b/src/celeste/core.py index 92c1d1bf..f663f339 100644 --- a/src/celeste/core.py +++ b/src/celeste/core.py @@ -21,6 +21,7 @@ class Provider(StrEnum): PERPLEXITY = "perplexity" BYTEPLUS = "byteplus" ELEVENLABS = "elevenlabs" + GRADIUM = "gradium" class Capability(StrEnum): diff --git a/src/celeste/credentials.py b/src/celeste/credentials.py index 5a3fe24e..ff7260be 100644 --- a/src/celeste/credentials.py +++ b/src/celeste/credentials.py @@ -26,6 +26,7 @@ Provider.BYTEPLUS: ("celeste_byteplus", "Authorization", "Bearer "), Provider.ELEVENLABS: ("celeste_elevenlabs", "xi-api-key", ""), Provider.BFL: ("celeste_bfl", "x-key", ""), + Provider.GRADIUM: ("celeste_gradium", "x-api-key", ""), } # Provider to credential field mapping @@ -45,6 +46,7 @@ Provider.BYTEPLUS: "byteplus_api_key", Provider.ELEVENLABS: "elevenlabs_api_key", Provider.BFL: "bfl_api_key", + Provider.GRADIUM: "gradium_api_key", } @@ -66,6 +68,7 @@ class Credentials(BaseSettings): byteplus_api_key: SecretStr | None = Field(None, alias="BYTEPLUS_API_KEY") elevenlabs_api_key: SecretStr | None = Field(None, alias="ELEVENLABS_API_KEY") bfl_api_key: SecretStr | None = Field(None, alias="BFL_API_KEY") + gradium_api_key: SecretStr | None = Field(None, alias="GRADIUM_API_KEY") model_config = { "env_file": find_dotenv(), diff --git a/src/celeste/websocket.py b/src/celeste/websocket.py new file mode 100644 index 00000000..c5adabf4 --- /dev/null +++ b/src/celeste/websocket.py @@ -0,0 +1,134 @@ +"""WebSocket client for AI provider APIs using persistent connections.""" + +import json +import logging +from collections.abc import AsyncIterator +from typing import Any + +from websockets.asyncio.client import ClientConnection +from websockets.asyncio.client import connect as ws_connect + +from celeste.core import Capability, Provider + +logger = logging.getLogger(__name__) + +DEFAULT_TIMEOUT = 60.0 + + +class WebSocketClient: + """Async WebSocket client wrapper. + + Provides a simple interface for WebSocket connections with: + - Context manager pattern for automatic cleanup + - JSON message helpers + - Connection state management + """ + + async def connect( + self, + url: str, + headers: dict[str, str] | None = None, + ) -> "WebSocketConnection": + """Open WebSocket connection. + + Args: + url: WebSocket URL (wss:// or ws://). + headers: Optional headers (e.g., authentication). + + Returns: + WebSocketConnection for send/receive operations. + """ + extra_headers = headers or {} + ws = await ws_connect(url, additional_headers=extra_headers) + return WebSocketConnection(ws) + + +class WebSocketConnection: + """Active WebSocket connection wrapper.""" + + def __init__(self, ws: ClientConnection) -> None: + self._ws = ws + + async def send(self, message: str) -> None: + """Send text message.""" + await self._ws.send(message) + + async def send_json(self, data: dict[str, Any]) -> None: + """Send JSON message.""" + await self._ws.send(json.dumps(data)) + + async def recv(self) -> str: + """Receive text message.""" + message = await self._ws.recv() + if isinstance(message, bytes): + return message.decode("utf-8") + return message + + async def recv_json(self) -> dict[str, Any]: + """Receive and parse JSON message.""" + message = await self.recv() + return json.loads(message) # type: ignore[no-any-return] + + def __aiter__(self) -> AsyncIterator[str]: + """Iterate over incoming messages.""" + return self._message_iterator() + + async def _message_iterator(self) -> AsyncIterator[str]: + """Async iterator for incoming messages.""" + async for message in self._ws: + if isinstance(message, bytes): + yield message.decode("utf-8") + else: + yield message + + async def close(self) -> None: + """Close connection.""" + await self._ws.close() + + async def __aenter__(self) -> "WebSocketConnection": + """Enter async context manager.""" + return self + + async def __aexit__(self, *args: object) -> None: + """Exit async context manager and close connection.""" + await self.close() + + +# Module-level registry (mirrors http.py pattern) +_ws_clients: dict[tuple[Provider, Capability], WebSocketClient] = {} + + +def get_ws_client(provider: Provider, capability: Capability) -> WebSocketClient: + """Get or create shared WebSocket client for provider/capability. + + Args: + provider: The AI provider. + capability: The capability being used. + + Returns: + Shared WebSocketClient instance for this provider and capability. + """ + key = (provider, capability) + if key not in _ws_clients: + _ws_clients[key] = WebSocketClient() + return _ws_clients[key] + + +async def close_all_ws_clients() -> None: + """Close all WebSocket clients and clear registry.""" + _ws_clients.clear() + + +def clear_ws_clients() -> None: + """Clear WebSocket client registry.""" + _ws_clients.clear() + + +__all__ = [ + "DEFAULT_TIMEOUT", + "WebSocketClient", + "WebSocketConnection", + "clear_ws_clients", + "close_all_ws_clients", + "get_ws_client", +] diff --git a/tests/unit_tests/test_constraints.py b/tests/unit_tests/test_constraints.py index aeae756a..0748a106 100644 --- a/tests/unit_tests/test_constraints.py +++ b/tests/unit_tests/test_constraints.py @@ -1,9 +1,23 @@ """Tests for constraint validation models.""" import pytest - -from celeste.constraints import Bool, Choice, Float, Int, Pattern, Range, Str +from pydantic import BaseModel + +from celeste.artifacts import ImageArtifact +from celeste.constraints import ( + Bool, + Choice, + Float, + ImageConstraint, + ImagesConstraint, + Int, + Pattern, + Range, + Schema, + Str, +) from celeste.exceptions import ConstraintViolationError +from celeste.mime_types import ImageMimeType class TestChoice: @@ -262,17 +276,37 @@ def test_validates_integer_value(self) -> None: assert result == 42 - def test_accepts_boolean_as_int(self) -> None: - """Test that bool is accepted as int (True=1, False=0).""" + def test_converts_whole_float_to_int(self) -> None: + """Test that whole float is converted to int.""" + constraint = Int() + + result = constraint(42.0) + + assert result == 42 + assert isinstance(result, int) + + def test_rejects_non_whole_float(self) -> None: + """Test that non-whole float raises ConstraintViolationError.""" constraint = Int() + + with pytest.raises(ConstraintViolationError, match=r"Must be int, got 42\.5"): + constraint(42.5) + + def test_accepts_boolean_value(self) -> None: + """Test that bool is accepted (True=1, False=0) since bool is subclass of int.""" + constraint = Int() + assert constraint(True) == 1 assert constraint(False) == 0 - def test_accepts_valid_string(self) -> None: - """Test that valid integer string is converted.""" + def test_converts_valid_string_to_int(self) -> None: + """Test that valid integer string is converted to int.""" constraint = Int() - assert constraint("42") == 42 - assert constraint("-10") == -10 + + result = constraint("42") + + assert result == 42 + assert isinstance(result, int) def test_rejects_invalid_string(self) -> None: """Test that non-integer string raises ConstraintViolationError.""" @@ -281,19 +315,6 @@ def test_rejects_invalid_string(self) -> None: with pytest.raises(ConstraintViolationError, match=r"Must be int, got 'abc'"): constraint("abc") - def test_accepts_whole_float(self) -> None: - """Test that whole number float is converted.""" - constraint = Int() - assert constraint(1.0) == 1 - assert constraint(-5.0) == -5 - - def test_rejects_fractional_float(self) -> None: - """Test that fractional float raises ConstraintViolationError.""" - constraint = Int() - - with pytest.raises(ConstraintViolationError, match=r"Must be int, got 1.1"): - constraint(1.1) - class TestFloat: """Test Float constraint validation.""" @@ -359,3 +380,211 @@ def test_rejects_string_value(self) -> None: with pytest.raises(ConstraintViolationError, match=r"Must be bool, got str"): constraint("true") # type: ignore[arg-type] + + +class TestRangeSpecialValues: + """Test Range constraint special_values bypass behavior.""" + + def test_special_values_bypass_bounds(self) -> None: + """Values in special_values bypass min/max validation.""" + constraint = Range(min=0, max=10, special_values=[-1, 100]) + + # These are outside bounds but allowed via special_values + assert constraint(-1) == -1 + assert constraint(100) == 100 + + def test_non_special_values_still_validated(self) -> None: + """Values not in special_values get normal bounds validation.""" + constraint = Range(min=0, max=10, special_values=[-1]) + + # -2 is not in special_values, so bounds apply + with pytest.raises(ConstraintViolationError, match=r"Must be between 0 and 10"): + constraint(-2) + + def test_normal_values_with_special_values_defined(self) -> None: + """Normal values within bounds still work when special_values defined.""" + constraint = Range(min=0, max=10, special_values=[-1, 100]) + + assert constraint(5) == 5 + assert constraint(0) == 0 + assert constraint(10) == 10 + + +class TestSchema: + """Test Schema constraint validation.""" + + def test_validates_basemodel_subclass(self) -> None: + """Schema accepts valid BaseModel subclass.""" + + class MyModel(BaseModel): + name: str + + constraint = Schema() + result = constraint(MyModel) + + assert result is MyModel + + def test_rejects_non_basemodel_type(self) -> None: + """Schema rejects types that aren't BaseModel subclasses.""" + constraint = Schema() + + with pytest.raises(ConstraintViolationError, match=r"Must be BaseModel"): + constraint(str) # type: ignore[arg-type] + + def test_validates_list_basemodel_type_hint(self) -> None: + """Schema accepts list[BaseModel] type hints.""" + + class MyModel(BaseModel): + value: int + + constraint = Schema() + result = constraint(list[MyModel]) # type: ignore[arg-type] + + assert result == list[MyModel] # type: ignore[comparison-overlap] + + def test_rejects_list_of_non_models(self) -> None: + """Schema rejects list[T] where T is not a BaseModel subclass.""" + constraint = Schema() + + with pytest.raises( + ConstraintViolationError, match=r"List type must be BaseModel" + ): + constraint(list[str]) # type: ignore[arg-type] + + +class TestImageConstraint: + """Test ImageConstraint validation for single image artifacts.""" + + def test_rejects_list_input(self) -> None: + """ImageConstraint requires single artifact, not a list.""" + constraint = ImageConstraint() + artifact = ImageArtifact(data=b"test") + + with pytest.raises( + ConstraintViolationError, + match=r"requires a single ImageArtifact, not a list", + ): + constraint([artifact]) # type: ignore[arg-type] + + def test_validates_image_artifact_type(self) -> None: + """ImageConstraint rejects non-ImageArtifact types.""" + constraint = ImageConstraint() + + with pytest.raises(ConstraintViolationError, match=r"Must be ImageArtifact"): + constraint("not an artifact") # type: ignore[arg-type] + + def test_accepts_valid_artifact(self) -> None: + """Valid ImageArtifact passes through unchanged.""" + constraint = ImageConstraint() + artifact = ImageArtifact(data=b"test image data") + + result = constraint(artifact) + + assert result is artifact + + def test_filters_mime_types_when_specified(self) -> None: + """ImageConstraint rejects unsupported MIME types.""" + constraint = ImageConstraint(supported_mime_types=[ImageMimeType.PNG]) + jpeg_artifact = ImageArtifact(data=b"test", mime_type=ImageMimeType.JPEG) + + with pytest.raises(ConstraintViolationError, match=r"mime_type must be one of"): + constraint(jpeg_artifact) + + def test_accepts_supported_mime_type(self) -> None: + """ImageConstraint accepts artifact with supported MIME type.""" + constraint = ImageConstraint(supported_mime_types=[ImageMimeType.PNG]) + png_artifact = ImageArtifact(data=b"test", mime_type=ImageMimeType.PNG) + + result = constraint(png_artifact) + + assert result is png_artifact + + def test_accepts_any_mime_when_none_specified(self) -> None: + """No MIME filtering when supported_mime_types is None.""" + constraint = ImageConstraint(supported_mime_types=None) + artifact = ImageArtifact(data=b"test", mime_type=ImageMimeType.JPEG) + + result = constraint(artifact) + + assert result is artifact + + +class TestImagesConstraint: + """Test ImagesConstraint validation for image artifact lists.""" + + def test_normalizes_single_artifact_to_list(self) -> None: + """Single ImageArtifact is wrapped in a list.""" + constraint = ImagesConstraint() + artifact = ImageArtifact(data=b"test") + + result = constraint(artifact) + + assert isinstance(result, list) + assert len(result) == 1 + assert result[0] is artifact + + def test_accepts_list_of_artifacts(self) -> None: + """List of ImageArtifacts passes through.""" + constraint = ImagesConstraint() + artifacts = [ImageArtifact(data=b"1"), ImageArtifact(data=b"2")] + + result = constraint(artifacts) + + assert result == artifacts + + def test_enforces_max_count(self) -> None: + """ImagesConstraint rejects when count exceeds max_count.""" + constraint = ImagesConstraint(max_count=2) + artifacts = [ImageArtifact(data=b"x") for _ in range(3)] + + with pytest.raises(ConstraintViolationError, match=r"at most 2 image"): + constraint(artifacts) + + def test_validates_each_artifact_type(self) -> None: + """ImagesConstraint reports index (1-indexed) of invalid artifact.""" + constraint = ImagesConstraint(supported_mime_types=[ImageMimeType.PNG]) + artifacts = [ + ImageArtifact(data=b"1", mime_type=ImageMimeType.PNG), + "not an artifact", # Invalid at index 2 + ImageArtifact(data=b"3", mime_type=ImageMimeType.PNG), + ] + + with pytest.raises( + ConstraintViolationError, match=r"Image 2.*Must be ImageArtifact" + ): + constraint(artifacts) # type: ignore[arg-type] + + def test_filters_mime_types_per_image(self) -> None: + """Each image's MIME type is validated against supported types.""" + constraint = ImagesConstraint(supported_mime_types=[ImageMimeType.PNG]) + artifacts = [ + ImageArtifact(data=b"1", mime_type=ImageMimeType.PNG), + ImageArtifact(data=b"2", mime_type=ImageMimeType.JPEG), # Invalid + ] + + with pytest.raises( + ConstraintViolationError, match=r"Image 2.*mime_type must be one of" + ): + constraint(artifacts) + + def test_handles_empty_list(self) -> None: + """Empty list is valid (no images to validate).""" + constraint = ImagesConstraint() + + result = constraint([]) + + assert result == [] + + def test_accepts_all_valid_images(self) -> None: + """All images with valid MIME types pass validation.""" + constraint = ImagesConstraint( + supported_mime_types=[ImageMimeType.PNG, ImageMimeType.JPEG] + ) + artifacts = [ + ImageArtifact(data=b"1", mime_type=ImageMimeType.PNG), + ImageArtifact(data=b"2", mime_type=ImageMimeType.JPEG), + ] + + result = constraint(artifacts) + + assert result == artifacts diff --git a/tests/unit_tests/test_parameters.py b/tests/unit_tests/test_parameters.py index de41061d..13c7cd03 100644 --- a/tests/unit_tests/test_parameters.py +++ b/tests/unit_tests/test_parameters.py @@ -7,6 +7,7 @@ from celeste.core import Parameter from celeste.models import Model +from celeste.types import StructuredOutput class DefaultParseOutputMapper: @@ -20,7 +21,9 @@ def map(self, request: dict[str, Any], value: Any, model: Model) -> dict[str, An request["temperature"] = value return request - def parse_output(self, content: object, value: object | None) -> object: + def parse_output( + self, content: StructuredOutput, value: object | None + ) -> StructuredOutput: """Default implementation: return content unchanged.""" return content @@ -41,7 +44,7 @@ class TestParameterMapperProtocol: ], ) def test_parse_output_returns_content_unchanged( - self, content: object, value: object | None + self, content: StructuredOutput, value: object | None ) -> None: """Default parse_output implementation returns content unchanged.