autobyteus 1.1.8__py3-none-any.whl → 1.2.0__py3-none-any.whl

autobyteus/multimedia/audio/api/gemini_audio_client.py

@@ -1,13 +1,11 @@
-import asyncio
-import base64
 import logging
 import os
+import base64
 import uuid
 import wave
-from typing import Optional, Dict, Any, TYPE_CHECKING, List
-
-# Old/legacy Gemini SDK (as requested)
-import google.generativeai as genai
+from typing import Optional, List, Dict, Any, TYPE_CHECKING
+from google import genai
+from google.genai import types as genai_types
 
 from autobyteus.multimedia.audio.base_audio_client import BaseAudioClient
 from autobyteus.multimedia.utils.response_types import SpeechGenerationResponse
@@ -19,77 +17,31 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-def _save_audio_bytes_to_wav(
-    pcm_bytes: bytes,
-    channels: int = 1,
-    rate: int = 24000,
-    sample_width: int = 2,
-) -> str:
-    """
-    Save raw PCM (s16le) audio bytes to a temporary WAV file and return the file path.
-
-    Gemini TTS models output mono, 24 kHz, 16-bit PCM by default.
-    """
+def _save_audio_bytes_to_wav(pcm_bytes: bytes, channels=1, rate=24000, sample_width=2) -> str:
+    """Saves PCM audio bytes to a temporary WAV file and returns the path."""
     temp_dir = "/tmp/autobyteus_audio"
     os.makedirs(temp_dir, exist_ok=True)
     file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.wav")
-
+    
     try:
         with wave.open(file_path, "wb") as wf:
             wf.setnchannels(channels)
-            wf.setsampwidth(sample_width)  # 2 bytes => 16-bit
+            wf.setsampwidth(sample_width)
             wf.setframerate(rate)
             wf.writeframes(pcm_bytes)
-        logger.info("Successfully saved generated audio to %s", file_path)
+        logger.info(f"Successfully saved generated audio to {file_path}")
         return file_path
     except Exception as e:
-        logger.error("Failed to save audio to WAV file at %s: %s", file_path, e)
-        raise
-
-
-def _extract_inline_audio_bytes(response) -> bytes:
-    """
-    Extract inline audio bytes from a google.generativeai response.
-
-    The legacy SDK returns a Response object with candidates -> content -> parts[0].inline_data.data.
-    Depending on version, `.data` can be bytes or base64-encoded str.
-    """
-    try:
-        # Access the first candidate's first part's inline_data
-        part = response.candidates[0].content.parts[0]
-        inline = getattr(part, "inline_data", None)
-        if not inline or not hasattr(inline, "data"):
-            raise ValueError("No inline audio data found in response.")
-        data = inline.data
-        if isinstance(data, bytes):
-            return data
-        if isinstance(data, str):
-            return base64.b64decode(data)
-        raise TypeError(f"Unexpected inline_data.data type: {type(data)}")
-    except Exception as e:
-        logger.error("Failed to extract audio from response: %s", e)
+        logger.error(f"Failed to save audio to WAV file at {file_path}: {e}")
         raise
 
 
 class GeminiAudioClient(BaseAudioClient):
     """
-    An audio client that uses Google's Gemini models for TTS via the *legacy* SDK
-    (`google.generativeai`).
-
-    Usage notes:
-      - Ensure your model value is a TTS-capable model (e.g. "gemini-2.5-flash-preview-tts"
-        or "gemini-2.5-pro-preview-tts").
-      - Single-speaker is default. For simple usage, provide `voice_name` (e.g. "Kore", "Puck")
-        in MultimediaConfig or generation_config.
-      - Multi-speaker preview exists in the API; if you want it, pass:
-          generation_config = {
-              "mode": "multi-speaker",
-              "speakers": [
-                  {"speaker": "Alice", "voice_name": "Kore"},
-                  {"speaker": "Bob",   "voice_name": "Puck"},
-              ]
-          }
-        and make sure your prompt contains lines for each named speaker.
+    An audio client that uses Google's Gemini models for audio tasks.
+
+    **Setup Requirements:**
+    1.  **Authentication:** Set the `GEMINI_API_KEY` environment variable with your API key.
     """
 
     def __init__(self, model: "AudioModel", config: "MultimediaConfig"):
@@ -99,121 +51,100 @@ class GeminiAudioClient(BaseAudioClient):
             raise ValueError("Please set the GEMINI_API_KEY environment variable.")
 
         try:
-            # Legacy library uses a global configure call
-            genai.configure(api_key=api_key)
-            # Create a GenerativeModel handle
-            self._model = genai.GenerativeModel(self.model.value or "gemini-2.5-flash-preview-tts")
-            logger.info("GeminiAudioClient (legacy SDK) configured for model '%s'.", self.model.value)
+            self.client = genai.Client()
+            self.async_client = self.client.aio
+            logger.info(f"GeminiAudioClient initialized for model '{self.model.name}'.")
         except Exception as e:
-            logger.error("Failed to configure Gemini client: %s", e)
+            logger.error(f"Failed to configure Gemini client: {e}")
             raise RuntimeError(f"Failed to configure Gemini client: {e}")
 
-    @staticmethod
-    def _build_single_speaker_generation_config(voice_name: str) -> Dict[str, Any]:
-        """
-        Build generation_config for single-speaker TTS in the legacy SDK.
-        Key bits:
-          - response_mime_type => request audio
-          - speech_config.voice_config.prebuilt_voice_config.voice_name => set the voice
-        """
-        return {
-            "response_mime_type": "audio/pcm",
-            "speech_config": {
-                "voice_config": {
-                    "prebuilt_voice_config": {
-                        "voice_name": voice_name,
-                    }
-                }
-            },
-        }
-
-    @staticmethod
-    def _build_multi_speaker_generation_config(speakers: List[Dict[str, str]]) -> Dict[str, Any]:
-        """
-        Build generation_config for multi-speaker TTS (preview).
-        `speakers` = [{"speaker": "...", "voice_name": "..."}, ...]
-        """
-        speaker_voice_configs = []
-        for s in speakers:
-            spk = s.get("speaker")
-            vname = s.get("voice_name")
-            if not spk or not vname:
-                raise ValueError("Each speaker must include 'speaker' and 'voice_name'.")
-            speaker_voice_configs.append(
-                {
-                    "speaker": spk,
-                    "voice_config": {
-                        "prebuilt_voice_config": {
-                            "voice_name": vname,
-                        }
-                    },
-                }
-            )
-        return {
-            "response_mime_type": "audio/pcm",
-            "speech_config": {
-                "multi_speaker_voice_config": {
-                    "speaker_voice_configs": speaker_voice_configs
-                }
-            },
-        }
 
     async def generate_speech(
         self,
         prompt: str,
-        generation_config: Optional[Dict[str, Any]] = None
+        generation_config: Optional[Dict[str, Any]] = None,
+        **kwargs
     ) -> SpeechGenerationResponse:
         """
-        Generates spoken audio from text using a Gemini TTS model through the legacy SDK.
-
-        Implementation details:
-          - We call `GenerativeModel.generate_content(...)` with a `generation_config`
-            that asks for AUDIO and sets the voice settings.
-          - The legacy SDK call is synchronous; we offload to a worker thread.
+        Generates spoken audio from text using a Gemini TTS model, supporting single-speaker,
+        multi-speaker, and style-controlled generation.
         """
         try:
-            logger.info("Generating speech with Gemini TTS (legacy SDK) model '%s'...", self.model.value)
-
-            # Merge base config with per-call overrides
-            final_cfg = self.config.to_dict().copy()
+            logger.info(f"Generating speech with Gemini TTS model '{self.model.value}'...")
+            
+            final_config = self.config.to_dict().copy()
             if generation_config:
-                final_cfg.update(generation_config or {})
-
-            # Style instructions: prepend if provided
-            style_instructions = final_cfg.get("style_instructions")
+                final_config.update(generation_config)
+            
+            # Handle style instructions by prepending them to the prompt
+            style_instructions = final_config.get("style_instructions")
             final_prompt = f"{style_instructions}: {prompt}" if style_instructions else prompt
-            logger.debug("Final prompt for TTS (truncated): '%s...'", final_prompt[:160])
+            logger.debug(f"Final prompt for TTS: '{final_prompt[:100]}...'")
 
-            # Mode & voice
-            mode = final_cfg.get("mode", "single-speaker")
-            default_voice = final_cfg.get("voice_name", "Kore")
+            speech_config = None
+            mode = final_config.get("mode", "single-speaker")
 
+            # Handle multi-speaker generation
             if mode == "multi-speaker":
-                speakers = final_cfg.get("speakers")
-                if not speakers or not isinstance(speakers, list):
-                    raise ValueError(
-                        "For multi-speaker mode, provide generation_config['speakers'] "
-                        "as a list of {'speaker': <name>, 'voice_name': <prebuilt voice>}."
+                speaker_mapping_list = final_config.get("speaker_mapping")
+                if not speaker_mapping_list or not isinstance(speaker_mapping_list, list):
+                    raise ValueError("Multi-speaker mode requires a 'speaker_mapping' list in generation_config.")
+                
+                logger.info(f"Configuring multi-speaker TTS with mapping: {speaker_mapping_list}")
+                speaker_voice_configs = []
+                for mapping_item in speaker_mapping_list:
+                    speaker = mapping_item.get("speaker")
+                    voice_name = mapping_item.get("voice")
+                    if not speaker or not voice_name:
+                        logger.warning(f"Skipping invalid item in speaker_mapping list: {mapping_item}")
+                        continue
+                    
+                    speaker_voice_configs.append(
+                        genai_types.SpeakerVoiceConfig(
+                            speaker=speaker,
+                            voice_config=genai_types.VoiceConfig(
+                                prebuilt_voice_config=genai_types.PrebuiltVoiceConfig(voice_name=voice_name)
+                            )
+                        )
                     )
-                gen_config = self._build_multi_speaker_generation_config(speakers)
-            else:
-                gen_config = self._build_single_speaker_generation_config(default_voice)
+                
+                if not speaker_voice_configs:
+                    raise ValueError("The 'speaker_mapping' list was empty or contained no valid mappings.")
 
-            # Run the blocking gen call in a thread so this coroutine stays non-blocking
-            response = await asyncio.to_thread(
-                self._model.generate_content,
-                final_prompt,
-                generation_config=gen_config,
-            )
+                speech_config = genai_types.SpeechConfig(
+                    multi_speaker_voice_config=genai_types.MultiSpeakerVoiceConfig(speaker_voice_configs=speaker_voice_configs)
+                )
 
-            audio_pcm = _extract_inline_audio_bytes(response)
+            # Handle single-speaker generation (default)
+            else:
+                voice_name = final_config.get("voice_name", "Kore") # A default voice
+                logger.info(f"Configuring single-speaker TTS with voice: '{voice_name}'")
+                speech_config = genai_types.SpeechConfig(
+                    voice_config=genai_types.VoiceConfig(
+                        prebuilt_voice_config=genai_types.PrebuiltVoiceConfig(voice_name=voice_name)
+                    )
+                )
+
+            # The google-genai library's TTS endpoint uses a synchronous call.
+            resp = self.client.models.generate_content(
+                model=self.model.value,
+                contents=final_prompt,
+                config=genai_types.GenerateContentConfig(
+                    response_modalities=["AUDIO"],
+                    speech_config=speech_config
+                ),
+            )
+            
+            audio_b64 = resp.candidates[0].content.parts[0].inline_data.data
+            audio_pcm = base64.b64decode(audio_b64)
+            
             audio_path = _save_audio_bytes_to_wav(audio_pcm)
 
             return SpeechGenerationResponse(audio_urls=[audio_path])
 
         except Exception as e:
-            logger.error("Error during Google Gemini speech generation (legacy SDK): %s", str(e))
+            logger.error(f"Error during Google Gemini speech generation: {str(e)}")
             raise ValueError(f"Google Gemini speech generation failed: {str(e)}")
 
     async def cleanup(self):
-        logger.debug("GeminiAudioClient cleanup called (legacy SDK; nothing to release).")
+        logger.debug("GeminiAudioClient cleanup called.")

autobyteus/multimedia/audio/audio_client_factory.py

@@ -8,6 +8,7 @@ from autobyteus.multimedia.audio.api.gemini_audio_client import GeminiAudioClien
 from autobyteus.multimedia.audio.autobyteus_audio_provider import AutobyteusAudioModelProvider
 from autobyteus.multimedia.utils.multimedia_config import MultimediaConfig
 from autobyteus.utils.singleton import SingletonMeta
+from autobyteus.utils.parameter_schema import ParameterSchema, ParameterDefinition, ParameterType
 
 logger = logging.getLogger(__name__)
 
@@ -46,34 +47,58 @@ class AudioClientFactory(metaclass=SingletonMeta):
     def _initialize_registry():
         """Initializes the registry with built-in audio models."""
         
+        # --- Define a clear schema for speaker mapping items using ParameterSchema ---
+        speaker_mapping_item_schema = ParameterSchema(parameters=[
+            ParameterDefinition(
+                name="speaker",
+                param_type=ParameterType.STRING,
+                description="The speaker's name as it appears in the prompt (e.g., 'Joe').",
+                required=True
+            ),
+            ParameterDefinition(
+                name="voice",
+                param_type=ParameterType.ENUM,
+                description="The voice to assign to this speaker.",
+                enum_values=GEMINI_TTS_VOICES,
+                required=True
+            )
+        ])
+
         # Google Gemini Audio Models
+        gemini_tts_schema = ParameterSchema(parameters=[
+            ParameterDefinition(
+                name="mode",
+                param_type=ParameterType.ENUM,
+                default_value="single-speaker",
+                enum_values=["single-speaker", "multi-speaker"],
+                description="The speech generation mode. 'single-speaker' for a consistent voice, or 'multi-speaker' to assign different voices to speakers identified in the prompt."
+            ),
+            ParameterDefinition(
+                name="voice_name",
+                param_type=ParameterType.ENUM,
+                default_value="Kore",
+                enum_values=GEMINI_TTS_VOICES,
+                description="The voice to use for single-speaker generation."
+            ),
+            ParameterDefinition(
+                name="style_instructions",
+                param_type=ParameterType.STRING,
+                description="Optional instructions on the style of speech, e.g., 'Say this in a dramatic whisper'."
+            ),
+            ParameterDefinition(
+                name="speaker_mapping",
+                param_type=ParameterType.ARRAY,
+                description="Required for multi-speaker mode. A list of objects, each mapping a speaker name from the prompt to a voice name.",
+                array_item_schema=speaker_mapping_item_schema
+            )
+        ])
+        
         gemini_tts_model = AudioModel(
             name="gemini-2.5-flash-tts",
             value="gemini-2.5-flash-preview-tts",
             provider=MultimediaProvider.GOOGLE,
             client_class=GeminiAudioClient,
-            parameter_schema={
-                "mode": { 
-                    "type": "string", 
-                    "default": "single-speaker", 
-                    "allowed_values": ["single-speaker", "multi-speaker"],
-                    "description": "The speech generation mode. 'single-speaker' for a consistent voice, or 'multi-speaker' to assign different voices to speakers identified in the prompt."
-                },
-                "voice_name": { 
-                    "type": "string", 
-                    "default": "Kore",
-                    "allowed_values": GEMINI_TTS_VOICES,
-                    "description": "The voice to use for single-speaker generation."
-                },
-                "style_instructions": {
-                    "type": "string",
-                    "description": "Optional instructions on the style of speech, e.g., 'Say this in a dramatic whisper'."
-                },
-                "speaker_mapping": {
-                    "type": "object",
-                    "description": "Required for multi-speaker mode. An object mapping speaker names from the prompt (e.g., 'Joe') to a voice name (e.g., 'Puck')."
-                }
-            }
+            parameter_schema=gemini_tts_schema
         )
 
         models_to_register = [

autobyteus/multimedia/audio/audio_model.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 import logging
-from typing import TYPE_CHECKING, Type, Optional, Iterator, Dict, Any
+from typing import TYPE_CHECKING, Type, Optional, Iterator, Dict, Any, Union
 from urllib.parse import urlparse
 
 from autobyteus.multimedia.providers import MultimediaProvider
 from autobyteus.multimedia.runtimes import MultimediaRuntime
 from autobyteus.multimedia.utils.multimedia_config import MultimediaConfig
+from autobyteus.utils.parameter_schema import ParameterSchema
 
 if TYPE_CHECKING:
     from autobyteus.multimedia.audio.base_audio_client import BaseAudioClient
@@ -47,7 +48,7 @@ class AudioModel(metaclass=AudioModelMeta):
         value: str,
         provider: MultimediaProvider,
         client_class: Type["BaseAudioClient"],
-        parameter_schema: Optional[Dict[str, Any]] = None,
+        parameter_schema: Optional[Union[Dict[str, Any], ParameterSchema]] = None,
         runtime: MultimediaRuntime = MultimediaRuntime.API,
         host_url: Optional[str] = None
     ):
@@ -57,13 +58,19 @@ class AudioModel(metaclass=AudioModelMeta):
         self.client_class = client_class
         self.runtime = runtime
         self.host_url = host_url
-        self.parameter_schema = parameter_schema if parameter_schema else {}
+
+        if isinstance(parameter_schema, dict):
+            self.parameter_schema = ParameterSchema.from_dict(parameter_schema)
+        elif parameter_schema is None:
+            self.parameter_schema = ParameterSchema()
+        else:
+            self.parameter_schema = parameter_schema
 
         # Automatically build default_config from the schema's default values
         default_params = {
-            key: meta.get("default")
-            for key, meta in self.parameter_schema.items()
-            if "default" in meta
+            param.name: param.default_value
+            for param in self.parameter_schema.parameters
+            if param.default_value is not None
         }
         self.default_config = MultimediaConfig(params=default_params)
 

autobyteus/multimedia/audio/autobyteus_audio_provider.py

@@ -3,7 +3,7 @@ from typing import List
 import os
 from urllib.parse import urlparse
 
-from autobyteus_llm_client import AutobyteusClient
+from autobyteus.clients import AutobyteusClient
 from autobyteus.multimedia.audio.api.autobyteus_audio_client import AutobyteusAudioClient
 from autobyteus.multimedia.audio.audio_model import AudioModel
 from autobyteus.multimedia.providers import MultimediaProvider
@@ -19,7 +19,9 @@ class AutobyteusAudioModelProvider:
 
     @staticmethod
     def _get_hosts() -> List[str]:
-        """Gets Autobyteus server hosts from env vars."""
+        """
+        Gets Autobyteus server hosts from env vars. Skips discovery if no host is configured.
+        """
         hosts_str = os.getenv('AUTOBYTEUS_LLM_SERVER_HOSTS')
         if hosts_str:
             return [host.strip() for host in hosts_str.split(',')]
@@ -28,7 +30,7 @@ class AutobyteusAudioModelProvider:
         if legacy_host:
             return [legacy_host]
             
-        return [AutobyteusAudioModelProvider.DEFAULT_SERVER_URL]
+        return []
 
     @staticmethod
     def discover_and_register():
@@ -37,6 +39,10 @@ class AutobyteusAudioModelProvider:
             from autobyteus.multimedia.audio.audio_client_factory import AudioClientFactory
 
             hosts = AutobyteusAudioModelProvider._get_hosts()
+            if not hosts:
+                logger.info("No Autobyteus server hosts configured. Skipping Autobyteus audio model discovery.")
+                return
+
             total_registered_count = 0
 
             for host_url in hosts:

autobyteus/multimedia/audio/base_audio_client.py

@@ -20,7 +20,8 @@ class BaseAudioClient(ABC):
     async def generate_speech(
         self,
         prompt: str,
-        generation_config: Optional[Dict[str, Any]] = None
+        generation_config: Optional[Dict[str, Any]] = None,
+        **kwargs
     ) -> SpeechGenerationResponse:
         """
         Generates spoken audio from text (Text-to-Speech).
@@ -29,6 +30,7 @@ class BaseAudioClient(ABC):
             prompt (str): The text to be converted to speech.
             generation_config (Optional[Dict[str, Any]]): Provider-specific parameters
                                                         (e.g., voice_name, speaker_mapping).
+            **kwargs: Additional keyword arguments for extensibility.
 
         Returns:
             SpeechGenerationResponse: An object containing URLs or paths to the generated audio files.

autobyteus/multimedia/image/api/autobyteus_image_client.py

@@ -1,6 +1,6 @@
 import logging
 from typing import Optional, List, Dict, Any, TYPE_CHECKING
-from autobyteus_llm_client import AutobyteusClient
+from autobyteus.clients import AutobyteusClient
 from autobyteus.multimedia.image.base_image_client import BaseImageClient
 from autobyteus.multimedia.utils.response_types import ImageGenerationResponse
 
@@ -27,7 +27,8 @@ class AutobyteusImageClient(BaseImageClient):
         self,
         prompt: str,
         input_image_urls: Optional[List[str]] = None,
-        generation_config: Optional[Dict[str, Any]] = None
+        generation_config: Optional[Dict[str, Any]] = None,
+        **kwargs
     ) -> ImageGenerationResponse:
         """
         Generates an image by calling the generate_image endpoint on the remote Autobyteus server.
@@ -38,7 +39,8 @@ class AutobyteusImageClient(BaseImageClient):
             prompt=prompt,
             input_image_urls=input_image_urls,
             mask_url=None, # Not used in pure generation
-            generation_config=generation_config
+            generation_config=generation_config,
+            **kwargs
         )
 
     async def edit_image(
@@ -46,7 +48,8 @@ class AutobyteusImageClient(BaseImageClient):
         prompt: str,
         input_image_urls: List[str],
         mask_url: Optional[str] = None,
-        generation_config: Optional[Dict[str, Any]] = None
+        generation_config: Optional[Dict[str, Any]] = None,
+        **kwargs
     ) -> ImageGenerationResponse:
         """
         Edits an image by calling the generate_image endpoint on the remote Autobyteus server.
@@ -55,7 +58,8 @@ class AutobyteusImageClient(BaseImageClient):
             prompt=prompt,
             input_image_urls=input_image_urls,
             mask_url=mask_url,
-            generation_config=generation_config
+            generation_config=generation_config,
+            **kwargs
         )
     
     async def _call_remote_generate(
@@ -63,7 +67,8 @@ class AutobyteusImageClient(BaseImageClient):
         prompt: str,
         input_image_urls: Optional[List[str]],
         mask_url: Optional[str],
-        generation_config: Optional[Dict[str, Any]]
+        generation_config: Optional[Dict[str, Any]],
+        **kwargs
     ) -> ImageGenerationResponse:
         """Internal helper to call the remote server."""
         try:
@@ -72,6 +77,8 @@ class AutobyteusImageClient(BaseImageClient):
             # The model name for the remote server is the `value`, not the unique `model_identifier`
             model_name_for_server = self.model.name
 
+            # Note: The underlying autobyteus_client.generate_image does not currently accept **kwargs.
+            # They are accepted here for interface consistency and future-proofing.
             response_data = await self.autobyteus_client.generate_image(
                 model_name=model_name_for_server,
                 prompt=prompt,