autobyteus 1.2.1__py3-none-any.whl → 1.3.0__py3-none-any.whl

autobyteus/memory/working_context_snapshot_serializer.py

@@ -0,0 +1,135 @@
+import json
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+from autobyteus.llm.utils.messages import (
+    Message,
+    MessageRole,
+    ToolCallPayload,
+    ToolCallSpec,
+    ToolResultPayload,
+)
+from autobyteus.memory.working_context_snapshot import WorkingContextSnapshot
+
+
+class WorkingContextSnapshotSerializer:
+    @staticmethod
+    def serialize(working_context_snapshot: WorkingContextSnapshot, metadata: Dict[str, Any]) -> Dict[str, Any]:
+        payload = {
+            "schema_version": metadata.get("schema_version", 1),
+            "agent_id": metadata.get("agent_id"),
+            "epoch_id": metadata.get("epoch_id", working_context_snapshot.epoch_id),
+            "last_compaction_ts": metadata.get("last_compaction_ts", working_context_snapshot.last_compaction_ts),
+            "messages": [WorkingContextSnapshotSerializer._serialize_message(msg) for msg in working_context_snapshot.build_messages()],
+        }
+        return payload
+
+    @staticmethod
+    def deserialize(payload: Dict[str, Any]) -> Tuple[WorkingContextSnapshot, Dict[str, Any]]:
+        messages = [
+            WorkingContextSnapshotSerializer._deserialize_message(msg)
+            for msg in payload.get("messages", [])
+            if isinstance(msg, dict)
+        ]
+        snapshot = WorkingContextSnapshot(initial_messages=messages)
+        metadata = {
+            "schema_version": payload.get("schema_version"),
+            "agent_id": payload.get("agent_id"),
+            "epoch_id": payload.get("epoch_id"),
+            "last_compaction_ts": payload.get("last_compaction_ts"),
+        }
+        if isinstance(metadata["epoch_id"], int):
+            snapshot.epoch_id = metadata["epoch_id"]
+        if metadata["last_compaction_ts"] is not None:
+            snapshot.last_compaction_ts = metadata["last_compaction_ts"]
+        return snapshot, metadata
+
+    @staticmethod
+    def validate(payload: Dict[str, Any]) -> bool:
+        if not isinstance(payload, dict):
+            return False
+        if not isinstance(payload.get("schema_version"), int):
+            return False
+        if not isinstance(payload.get("agent_id"), str):
+            return False
+        messages = payload.get("messages")
+        if not isinstance(messages, list):
+            return False
+        for msg in messages:
+            if not isinstance(msg, dict):
+                return False
+            if not isinstance(msg.get("role"), str):
+                return False
+        return True
+
+    @staticmethod
+    def _serialize_message(message: Message) -> Dict[str, Any]:
+        base = message.to_dict()
+        if base.get("tool_payload"):
+            base["tool_payload"] = WorkingContextSnapshotSerializer._normalize_tool_payload(base["tool_payload"])
+        return base
+
+    @staticmethod
+    def _deserialize_message(data: Dict[str, Any]) -> Message:
+        role = MessageRole(data.get("role"))
+        tool_payload = WorkingContextSnapshotSerializer._deserialize_tool_payload(data.get("tool_payload"))
+        return Message(
+            role=role,
+            content=data.get("content"),
+            reasoning_content=data.get("reasoning_content"),
+            image_urls=data.get("image_urls") or [],
+            audio_urls=data.get("audio_urls") or [],
+            video_urls=data.get("video_urls") or [],
+            tool_payload=tool_payload,
+        )
+
+    @staticmethod
+    def _normalize_tool_payload(payload: Dict[str, Any]) -> Dict[str, Any]:
+        if "tool_calls" in payload:
+            return {
+                "tool_calls": [
+                    {
+                        "id": call.get("id"),
+                        "name": call.get("name"),
+                        "arguments": WorkingContextSnapshotSerializer._safe_json_value(call.get("arguments")),
+                    }
+                    for call in payload.get("tool_calls", [])
+                ]
+            }
+        return {
+            "tool_call_id": payload.get("tool_call_id"),
+            "tool_name": payload.get("tool_name"),
+            "tool_result": WorkingContextSnapshotSerializer._safe_json_value(payload.get("tool_result")),
+            "tool_error": payload.get("tool_error"),
+        }
+
+    @staticmethod
+    def _deserialize_tool_payload(payload: Optional[Dict[str, Any]]) -> Optional[Any]:
+        if not payload:
+            return None
+        if "tool_calls" in payload:
+            calls = []
+            for call in payload.get("tool_calls", []) or []:
+                calls.append(
+                    ToolCallSpec(
+                        id=str(call.get("id")),
+                        name=str(call.get("name")),
+                        arguments=call.get("arguments") or {},
+                    )
+                )
+            return ToolCallPayload(tool_calls=calls)
+        if "tool_call_id" in payload:
+            return ToolResultPayload(
+                tool_call_id=str(payload.get("tool_call_id")),
+                tool_name=str(payload.get("tool_name")),
+                tool_result=payload.get("tool_result"),
+                tool_error=payload.get("tool_error"),
+            )
+        return None
+
+    @staticmethod
+    def _safe_json_value(value: Any) -> Any:
+        try:
+            json.dumps(value)
+            return value
+        except TypeError:
+            return str(value)

autobyteus/multimedia/audio/api/autobyteus_audio_client.py

@@ -1,4 +1,5 @@
 import logging
+import uuid
 from typing import Optional, List, Dict, Any, TYPE_CHECKING
 from autobyteus.clients import AutobyteusClient
 from autobyteus.multimedia.audio.base_audio_client import BaseAudioClient
@@ -13,6 +14,7 @@ logger = logging.getLogger(__name__)
 class AutobyteusAudioClient(BaseAudioClient):
     """
     An audio client that connects to an Autobyteus server instance for audio tasks.
+    Maintains a persistent session ID for stateful interactions.
     """
 
     def __init__(self, model: "AudioModel", config: "MultimediaConfig"):
@@ -21,7 +23,9 @@ class AutobyteusAudioClient(BaseAudioClient):
             raise ValueError("AutobyteusAudioClient requires a host_url in its AudioModel.")
         
         self.autobyteus_client = AutobyteusClient(server_url=model.host_url)
-        logger.info(f"AutobyteusAudioClient initialized for model '{model.name}' on host '{model.host_url}'.")
+        self.session_id = str(uuid.uuid4())
+        logger.info(f"AutobyteusAudioClient initialized for model '{model.name}' "
+                    f"on host '{model.host_url}' with session_id '{self.session_id}'.")
 
     async def generate_speech(
         self,
@@ -33,7 +37,7 @@ class AutobyteusAudioClient(BaseAudioClient):
         Generates speech by calling the generate_speech endpoint on the remote Autobyteus server.
         """
         try:
-            logger.info(f"Sending speech generation request for model '{self.model.name}' to {self.model.host_url}")
+            logger.info(f"Sending speech generation request for model '{self.model.name}' to {self.model.host_url} (Session: {self.session_id})")
             
             model_name_for_server = self.model.name
 
@@ -42,7 +46,8 @@ class AutobyteusAudioClient(BaseAudioClient):
             response_data = await self.autobyteus_client.generate_speech(
                 model_name=model_name_for_server,
                 prompt=prompt,
-                generation_config=generation_config
+                generation_config=generation_config,
+                session_id=self.session_id
             )
             
             audio_urls = response_data.get("audio_urls", [])
@@ -56,7 +61,16 @@ class AutobyteusAudioClient(BaseAudioClient):
             raise
 
     async def cleanup(self):
-        """Closes the underlying AutobyteusClient."""
+        """
+        Notifies the server to cleanup the session, then closes the underlying HTTP client.
+        """
         if self.autobyteus_client:
-            await self.autobyteus_client.close()
+            try:
+                logger.info(f"Notifying server to cleanup audio session '{self.session_id}'...")
+                await self.autobyteus_client.cleanup_audio_session(self.session_id)
+            except Exception as e:
+                logger.error(f"Failed to cleanup remote audio session '{self.session_id}': {e}")
+            finally:
+                await self.autobyteus_client.close()
+
         logger.debug("AutobyteusAudioClient cleaned up.")

autobyteus/multimedia/audio/api/gemini_audio_client.py

@@ -9,6 +9,8 @@ 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
+from autobyteus.utils.gemini_helper import initialize_gemini_client_with_runtime
+from autobyteus.utils.gemini_model_mapping import resolve_model_for_runtime
 
 if TYPE_CHECKING:
     from autobyteus.multimedia.audio.audio_model import AudioModel
@@ -17,11 +19,22 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+_AUDIO_TEMP_DIR = "/tmp/autobyteus_audio"
+
+_AUDIO_MIME_EXTENSION_MAP = {
+    "audio/wav": "wav",
+    "audio/x-wav": "wav",
+    "audio/mpeg": "mp3",
+    "audio/mp3": "mp3",
+    "audio/ogg": "ogg",
+    "audio/webm": "webm",
+}
+
+
 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")
+    os.makedirs(_AUDIO_TEMP_DIR, exist_ok=True)
+    file_path = os.path.join(_AUDIO_TEMP_DIR, f"{uuid.uuid4()}.wav")
     
     try:
         with wave.open(file_path, "wb") as wf:
@@ -36,22 +49,63 @@ def _save_audio_bytes_to_wav(pcm_bytes: bytes, channels=1, rate=24000, sample_wi
         raise
 
 
+def _save_audio_bytes(audio_bytes: bytes, extension: Optional[str]) -> str:
+    """Saves audio bytes to a temporary file and returns the path."""
+    os.makedirs(_AUDIO_TEMP_DIR, exist_ok=True)
+    suffix = (extension or "bin").lstrip(".")
+    file_path = os.path.join(_AUDIO_TEMP_DIR, f"{uuid.uuid4()}.{suffix}")
+    try:
+        with open(file_path, "wb") as audio_file:
+            audio_file.write(audio_bytes)
+        logger.info(f"Successfully saved generated audio to {file_path}")
+        return file_path
+    except Exception as e:
+        logger.error(f"Failed to save audio to file at {file_path}: {e}")
+        raise
+
+
+def _parse_mime_type(mime_type: Optional[str]) -> tuple[str, Dict[str, str]]:
+    if not mime_type:
+        return "", {}
+    parts = [part.strip() for part in mime_type.split(";") if part.strip()]
+    base = parts[0].lower() if parts else ""
+    params: Dict[str, str] = {}
+    for part in parts[1:]:
+        if "=" in part:
+            key, value = part.split("=", 1)
+            params[key.strip().lower()] = value.strip()
+    return base, params
+
+
+def _coerce_audio_bytes(audio_data: Any) -> bytes:
+    if audio_data is None:
+        return b""
+    if isinstance(audio_data, bytes):
+        return audio_data
+    if isinstance(audio_data, bytearray):
+        return bytes(audio_data)
+    if isinstance(audio_data, memoryview):
+        return audio_data.tobytes()
+    if isinstance(audio_data, str):
+        return base64.b64decode(audio_data)
+    return bytes(audio_data)
+
+
 class GeminiAudioClient(BaseAudioClient):
     """
     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.
+    1.  **Vertex AI Express Mode:** Set `VERTEX_AI_API_KEY`.
+    2.  **Vertex AI Mode:** Set `VERTEX_AI_PROJECT` and `VERTEX_AI_LOCATION`.
+    3.  **AI Studio Mode:** Set `GEMINI_API_KEY`.
     """
 
     def __init__(self, model: "AudioModel", config: "MultimediaConfig"):
         super().__init__(model, config)
-        api_key = os.getenv("GEMINI_API_KEY")
-        if not api_key:
-            raise ValueError("Please set the GEMINI_API_KEY environment variable.")
-
+        
         try:
-            self.client = genai.Client()
+            self.client, self.runtime_info = initialize_gemini_client_with_runtime()
             self.async_client = self.client.aio
             logger.info(f"GeminiAudioClient initialized for model '{self.model.name}'.")
         except Exception as e:
@@ -70,8 +124,6 @@ class GeminiAudioClient(BaseAudioClient):
         multi-speaker, and style-controlled generation.
         """
         try:
-            logger.info(f"Generating speech with Gemini TTS model '{self.model.value}'...")
-            
             final_config = self.config.to_dict().copy()
             if generation_config:
                 final_config.update(generation_config)
@@ -126,8 +178,19 @@ class GeminiAudioClient(BaseAudioClient):
                 )
 
             # The google-genai library's TTS endpoint uses a synchronous call.
+            # FIX: Ensure no 'models/' prefix is used here.
+            runtime_adjusted_model = resolve_model_for_runtime(
+                self.model.value,
+                modality="tts",
+                runtime=getattr(self, "runtime_info", None) and self.runtime_info.runtime,
+            )
+            logger.info(
+                "Generating speech with Gemini TTS model '%s' (requested '%s').",
+                runtime_adjusted_model,
+                self.model.value,
+                )
             resp = self.client.models.generate_content(
-                model=self.model.value,
+                model=runtime_adjusted_model,
                 contents=final_prompt,
                 config=genai_types.GenerateContentConfig(
                     response_modalities=["AUDIO"],
@@ -135,10 +198,40 @@ class GeminiAudioClient(BaseAudioClient):
                 ),
             )
             
-            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)
+            part = resp.candidates[0].content.parts[0]
+            inline_data = part.inline_data
+            if not inline_data or not inline_data.data:
+                raise ValueError("Gemini TTS response did not include audio data.")
+
+            mime_type, mime_params = _parse_mime_type(inline_data.mime_type)
+            audio_bytes = _coerce_audio_bytes(inline_data.data)
+            if not audio_bytes:
+                raise ValueError("Gemini TTS returned empty audio data.")
+
+            logger.info(
+                "Received Gemini TTS audio payload (mime_type='%s', bytes=%d).",
+                mime_type or "unknown",
+                len(audio_bytes),
+            )
+
+            if not mime_type or mime_type.startswith("audio/pcm") or mime_type == "audio/l16":
+                rate = 24000
+                channels = 1
+                if "rate" in mime_params:
+                    try:
+                        rate = int(mime_params["rate"])
+                    except ValueError:
+                        logger.warning("Invalid sample rate in mime_type '%s'; using default 24000.", inline_data.mime_type)
+                if "channels" in mime_params:
+                    try:
+                        channels = int(mime_params["channels"])
+                    except ValueError:
+                        logger.warning("Invalid channel count in mime_type '%s'; using default 1.", inline_data.mime_type)
+
+                audio_path = _save_audio_bytes_to_wav(audio_bytes, channels=channels, rate=rate, sample_width=2)
+            else:
+                extension = _AUDIO_MIME_EXTENSION_MAP.get(mime_type, "bin")
+                audio_path = _save_audio_bytes(audio_bytes, extension)
 
             return SpeechGenerationResponse(audio_urls=[audio_path])
 

autobyteus/multimedia/audio/audio_client_factory.py

@@ -13,13 +13,51 @@ from autobyteus.utils.parameter_schema import ParameterSchema, ParameterDefiniti
 
 logger = logging.getLogger(__name__)
 
-GEMINI_TTS_VOICES = [
-    "Zephyr", "Puck", "Charon", "Kore", "Fenrir", "Leda", "Orus", "Aoede", 
-    "Callirrhoe", "Autonoe", "Enceladus", "Iapetus", "Umbriel", "Algieba", 
-    "Despina", "Erinome", "Algenib", "Rasalgethi", "Laomedeia", "Achernar", 
-    "Alnilam", "Schedar", "Gacrux", "Pulcherrima", "Achird", "Zubenelgenubi", 
-    "Vindemiatrix", "Sadachbia", "Sadaltager", "Sulafat"
+# Enhanced metadata for Google Gemini TTS voices, including gender and description.
+GEMINI_VOICE_DETAILS = {
+    "Zephyr": {"gender": "female", "description": "Bright, Higher pitch"},
+    "Puck": {"gender": "male", "description": "Upbeat, Middle pitch"},
+    "Charon": {"gender": "male", "description": "Informative, Lower pitch"},
+    "Kore": {"gender": "female", "description": "Firm, Middle pitch"},
+    "Fenrir": {"gender": "male", "description": "Excitable, Lower middle pitch"},
+    "Leda": {"gender": "female", "description": "Youthful, Higher pitch"},
+    "Orus": {"gender": "male", "description": "Firm, Lower middle pitch"},
+    "Aoede": {"gender": "female", "description": "Breezy, Middle pitch"},
+    "Callirrhoe": {"gender": "female", "description": "Easy-going, Middle pitch"},
+    "Autonoe": {"gender": "female", "description": "Bright, Middle pitch"},
+    "Enceladus": {"gender": "male", "description": "Breathy, Lower pitch"},
+    "Iapetus": {"gender": "male", "description": "Clear, Lower middle pitch"},
+    "Umbriel": {"gender": "male", "description": "Easy-going, Lower middle pitch"},
+    "Algieba": {"gender": "male", "description": "Smooth, Lower pitch"},
+    "Despina": {"gender": "female", "description": "Smooth, Middle pitch"},
+    "Erinome": {"gender": "female", "description": "Clear, Middle pitch"},
+    "Algenib": {"gender": "male", "description": "Gravelly, Lower pitch"},
+    "Rasalgethi": {"gender": "male", "description": "Informative, Middle pitch"},
+    "Laomedeia": {"gender": "female", "description": "Upbeat, Higher pitch"},
+    "Achernar": {"gender": "female", "description": "Soft, Higher pitch"},
+    "Alnilam": {"gender": "male", "description": "Firm, Lower middle pitch"},
+    "Schedar": {"gender": "male", "description": "Even, Lower middle pitch"},
+    "Gacrux": {"gender": "female", "description": "Mature, Middle pitch"},
+    "Pulcherrima": {"gender": "female", "description": "Forward, Middle pitch"},
+    "Achird": {"gender": "male", "description": "Friendly, Lower middle pitch"},
+    "Zubenelgenubi": {"gender": "male", "description": "Casual, Lower middle pitch"},
+    "Vindemiatrix": {"gender": "female", "description": "Gentle, Middle pitch"},
+    "Sadachbia": {"gender": "male", "description": "Lively, Lower pitch"},
+    "Sadaltager": {"gender": "male", "description": "Knowledgeable, Middle pitch"},
+    "Sulafat": {"gender": "female", "description": "Warm, Middle pitch"},
+}
+
+# The list of voice names, derived from the keys of the details dictionary.
+# This is used for the `enum_values` to maintain compatibility.
+GEMINI_TTS_VOICES = list(GEMINI_VOICE_DETAILS.keys())
+
+# Generate a formatted string of voice metadata to be appended to parameter descriptions.
+_voice_descriptions_list = [
+    f"- {name} ({details['gender']}): {details['description']}"
+    for name, details in GEMINI_VOICE_DETAILS.items()
 ]
+GEMINI_VOICE_METADATA_DESC = "\n\nDetailed Voice Options:\n" + "\n".join(_voice_descriptions_list)
+
 
 OPENAI_TTS_VOICES = [
     "alloy", "ash", "ballad", "coral", "echo", "fable", "onyx",
@@ -64,7 +102,7 @@ class AudioClientFactory(metaclass=SingletonMeta):
             ParameterDefinition(
                 name="voice",
                 param_type=ParameterType.ENUM,
-                description="The voice to assign to this speaker.",
+                description="The voice to assign to this speaker." + GEMINI_VOICE_METADATA_DESC,
                 enum_values=GEMINI_TTS_VOICES,
                 required=True
             )
@@ -84,7 +122,7 @@ class AudioClientFactory(metaclass=SingletonMeta):
                 param_type=ParameterType.ENUM,
                 default_value="Kore",
                 enum_values=GEMINI_TTS_VOICES,
-                description="The voice to use for single-speaker generation."
+                description="The voice to use for single-speaker generation." + GEMINI_VOICE_METADATA_DESC
             ),
             ParameterDefinition(
                 name="style_instructions",
@@ -102,7 +140,7 @@ class AudioClientFactory(metaclass=SingletonMeta):
         gemini_tts_model = AudioModel(
             name="gemini-2.5-flash-tts",
             value="gemini-2.5-flash-preview-tts",
-            provider=MultimediaProvider.GOOGLE,
+            provider=MultimediaProvider.GEMINI,
             client_class=GeminiAudioClient,
             parameter_schema=gemini_tts_schema
         )

autobyteus/multimedia/audio/audio_model.py

@@ -79,7 +79,8 @@ class AudioModel(metaclass=AudioModelMeta):
         """Returns the unique identifier for the model."""
         if self.runtime == MultimediaRuntime.AUTOBYTEUS and self.host_url:
             try:
-                host = urlparse(self.host_url).hostname
+                parsed = urlparse(self.host_url)
+                host = parsed.netloc or parsed.hostname or self.host_url
                 return f"{self.name}@{host}"
             except Exception:
                 return f"{self.name}@{self.host_url}" # Fallback

autobyteus/multimedia/image/api/autobyteus_image_client.py

@@ -1,4 +1,5 @@
 import logging
+import uuid
 from typing import Optional, List, Dict, Any, TYPE_CHECKING
 from autobyteus.clients import AutobyteusClient
 from autobyteus.multimedia.image.base_image_client import BaseImageClient
@@ -13,6 +14,7 @@ logger = logging.getLogger(__name__)
 class AutobyteusImageClient(BaseImageClient):
     """
     An image client that connects to an Autobyteus LLM server instance for image tasks.
+    Maintains a persistent session ID for stateful interactions (e.g. conversational editing).
     """
 
     def __init__(self, model: "ImageModel", config: "MultimediaConfig"):
@@ -21,7 +23,9 @@ class AutobyteusImageClient(BaseImageClient):
             raise ValueError("AutobyteusImageClient requires a host_url in its ImageModel.")
         
         self.autobyteus_client = AutobyteusClient(server_url=model.host_url)
-        logger.info(f"AutobyteusImageClient initialized for model '{self.model.name}' on host '{model.host_url}'.")
+        self.session_id = str(uuid.uuid4())
+        logger.info(f"AutobyteusImageClient initialized for model '{self.model.name}' "
+                    f"on host '{model.host_url}' with session_id '{self.session_id}'.")
 
     async def generate_image(
         self,
@@ -72,7 +76,7 @@ class AutobyteusImageClient(BaseImageClient):
     ) -> ImageGenerationResponse:
         """Internal helper to call the remote server."""
         try:
-            logger.info(f"Sending image generation request for model '{self.model.name}' to {self.model.host_url}")
+            logger.info(f"Sending image generation request for model '{self.model.name}' to {self.model.host_url} (Session: {self.session_id})")
             
             # The model name for the remote server is the `value`, not the unique `model_identifier`
             model_name_for_server = self.model.name
@@ -84,7 +88,8 @@ class AutobyteusImageClient(BaseImageClient):
                 prompt=prompt,
                 input_image_urls=input_image_urls,
                 mask_url=mask_url,
-                generation_config=generation_config
+                generation_config=generation_config,
+                session_id=self.session_id
             )
             
             image_urls = response_data.get("image_urls", [])
@@ -98,7 +103,16 @@ class AutobyteusImageClient(BaseImageClient):
             raise
 
     async def cleanup(self):
-        """Closes the underlying AutobyteusClient."""
+        """
+        Notifies the server to cleanup the session, then closes the underlying HTTP client.
+        """
         if self.autobyteus_client:
-            await self.autobyteus_client.close()
+            try:
+                logger.info(f"Notifying server to cleanup image session '{self.session_id}'...")
+                await self.autobyteus_client.cleanup_image_session(self.session_id)
+            except Exception as e:
+                logger.error(f"Failed to cleanup remote image session '{self.session_id}': {e}")
+            finally:
+                await self.autobyteus_client.close()
+        
         logger.debug("AutobyteusImageClient cleaned up.")

autobyteus/multimedia/image/api/gemini_image_client.py

@@ -1,14 +1,13 @@
 import logging
 import base64
-import os
 from typing import Optional, List, Dict, Any, TYPE_CHECKING
-from google import genai
-from PIL import Image
-import requests
+from google.genai import types as genai_types
 
 from autobyteus.multimedia.image.base_image_client import BaseImageClient
 from autobyteus.multimedia.utils.response_types import ImageGenerationResponse
 from autobyteus.multimedia.utils.api_utils import load_image_from_url
+from autobyteus.utils.gemini_helper import initialize_gemini_client_with_runtime
+from autobyteus.utils.gemini_model_mapping import resolve_model_for_runtime
 
 if TYPE_CHECKING:
     from autobyteus.multimedia.image.image_model import ImageModel
@@ -21,17 +20,16 @@ class GeminiImageClient(BaseImageClient):
     An image client that uses Google's Gemini models for image generation tasks.
 
     **Setup Requirements:**
-    1.  **Authentication:** Set the `GEMINI_API_KEY` environment variable with your API key.
+    1.  **Vertex AI Express Mode:** Set `VERTEX_AI_API_KEY`.
+    2.  **Vertex AI Mode:** Set `VERTEX_AI_PROJECT` and `VERTEX_AI_LOCATION`.
+    3.  **AI Studio Mode:** Set `GEMINI_API_KEY`.
     """
 
     def __init__(self, model: "ImageModel", config: "MultimediaConfig"):
         super().__init__(model, config)
-        api_key = os.getenv("GEMINI_API_KEY")
-        if not api_key:
-            raise ValueError("Please set the GEMINI_API_KEY environment variable.")
         
         try:
-            self.client = genai.Client()
+            self.client, self.runtime_info = initialize_gemini_client_with_runtime()
             self.async_client = self.client.aio
             logger.info(f"GeminiImageClient initialized for model '{self.model.name}'.")
         except Exception as e:
@@ -60,16 +58,40 @@ class GeminiImageClient(BaseImageClient):
                     except Exception as e:
                         logger.error(f"Skipping image at '{url}' due to loading error: {e}")
 
-            # Note: The google-genai library uses the synchronous client for the `.generate_content` method on a model
-            # even in an async context, as there isn't a direct async equivalent exposed for this specific call on the model object.
-            # We use the top-level async client for other potential future calls if the library API changes.
-            model_instance = self.client.get_generative_model(model_name=f"models/{self.model.value}")
-            response = await model_instance.generate_content_async(contents=content)
+            config_dict: Dict[str, Any] = {}
+            if self.config and self.config.params:
+                config_dict.update(self.config.params)
+            if generation_config:
+                config_dict.update(generation_config)
+            if "response_modalities" not in config_dict:
+                if getattr(self, "runtime_info", None) and self.runtime_info.runtime == "vertex":
+                    config_dict["response_modalities"] = ["TEXT", "IMAGE"]
+                else:
+                    config_dict["response_modalities"] = ["IMAGE"]
+            config = genai_types.GenerateContentConfig(**config_dict)
+            
+            # FIX: Removed 'models/' prefix from model_name to support Vertex AI
+            runtime_adjusted_model = resolve_model_for_runtime(
+                self.model.value,
+                modality="image",
+                runtime=getattr(self, "runtime_info", None) and self.runtime_info.runtime,
+            )
+            if runtime_adjusted_model != self.model.value:
+                logger.info(
+                    "Using runtime-adjusted Gemini image model '%s' (requested '%s').",
+                    runtime_adjusted_model,
+                    self.model.value,
+                )
+            response = await self.async_client.models.generate_content(
+                model=runtime_adjusted_model,
+                contents=content,
+                config=config,
+            )
 
 
             image_urls = []
-            for part in response.parts:
-                if part.inline_data and "image" in part.inline_data.mime_type:
+            for part in response.parts or []:
+                if part.inline_data and part.inline_data.mime_type and "image" in part.inline_data.mime_type:
                     image_bytes = part.inline_data.data
                     base64_image = base64.b64encode(image_bytes).decode("utf-8")
                     data_uri = f"data:{part.inline_data.mime_type};base64,{base64_image}"
@@ -77,7 +99,7 @@ class GeminiImageClient(BaseImageClient):
             
             if not image_urls:
                 # Check for a safety-related refusal to generate content
-                if response.prompt_feedback.block_reason:
+                if response.prompt_feedback and response.prompt_feedback.block_reason:
                     reason = response.prompt_feedback.block_reason.name
                     logger.error(f"Image generation blocked due to safety settings. Reason: {reason}")
                     raise ValueError(f"Image generation failed due to safety settings: {reason}")