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

autobyteus/memory/models/memory_types.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class MemoryType(str, Enum):
+    RAW_TRACE = "raw_trace"
+    EPISODIC = "episodic"
+    SEMANTIC = "semantic"

autobyteus/memory/models/raw_trace_item.py

@@ -0,0 +1,79 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+from .memory_types import MemoryType
+
+
+@dataclass
+class RawTraceItem:
+    id: str
+    ts: float
+    turn_id: str
+    seq: int
+    trace_type: str
+    content: str
+    source_event: str
+    media: Optional[Dict[str, List[str]]] = None
+    tool_name: Optional[str] = None
+    tool_call_id: Optional[str] = None
+    tool_args: Optional[Dict[str, Any]] = None
+    tool_result: Optional[Any] = None
+    tool_error: Optional[str] = None
+    correlation_id: Optional[str] = None
+    tags: List[str] = field(default_factory=list)
+    tool_result_ref: Optional[str] = None
+
+    @property
+    def memory_type(self) -> MemoryType:
+        return MemoryType.RAW_TRACE
+
+    def to_dict(self) -> Dict[str, Any]:
+        data: Dict[str, Any] = {
+            "id": self.id,
+            "ts": self.ts,
+            "turn_id": self.turn_id,
+            "seq": self.seq,
+            "trace_type": self.trace_type,
+            "content": self.content,
+            "source_event": self.source_event,
+        }
+        if self.media is not None:
+            data["media"] = self.media
+        if self.tool_name is not None:
+            data["tool_name"] = self.tool_name
+        if self.tool_call_id is not None:
+            data["tool_call_id"] = self.tool_call_id
+        if self.tool_args is not None:
+            data["tool_args"] = self.tool_args
+        if self.tool_result is not None:
+            data["tool_result"] = self.tool_result
+        if self.tool_error is not None:
+            data["tool_error"] = self.tool_error
+        if self.correlation_id is not None:
+            data["correlation_id"] = self.correlation_id
+        if self.tags:
+            data["tags"] = self.tags
+        if self.tool_result_ref is not None:
+            data["tool_result_ref"] = self.tool_result_ref
+        return data
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "RawTraceItem":
+        return cls(
+            id=data["id"],
+            ts=data["ts"],
+            turn_id=data["turn_id"],
+            seq=data["seq"],
+            trace_type=data["trace_type"],
+            content=data.get("content", ""),
+            source_event=data["source_event"],
+            media=data.get("media"),
+            tool_name=data.get("tool_name"),
+            tool_call_id=data.get("tool_call_id"),
+            tool_args=data.get("tool_args"),
+            tool_result=data.get("tool_result"),
+            tool_error=data.get("tool_error"),
+            correlation_id=data.get("correlation_id"),
+            tags=data.get("tags", []),
+            tool_result_ref=data.get("tool_result_ref"),
+        )

autobyteus/memory/models/semantic_item.py

@@ -0,0 +1,41 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, List
+
+from .memory_types import MemoryType
+
+
+@dataclass
+class SemanticItem:
+    id: str
+    ts: float
+    fact: str
+    tags: List[str] = field(default_factory=list)
+    confidence: float = 0.0
+    salience: float = 0.0
+
+    @property
+    def memory_type(self) -> MemoryType:
+        return MemoryType.SEMANTIC
+
+    def to_dict(self) -> Dict[str, Any]:
+        data: Dict[str, Any] = {
+            "id": self.id,
+            "ts": self.ts,
+            "fact": self.fact,
+            "confidence": self.confidence,
+            "salience": self.salience,
+        }
+        if self.tags:
+            data["tags"] = self.tags
+        return data
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "SemanticItem":
+        return cls(
+            id=data["id"],
+            ts=data["ts"],
+            fact=data.get("fact", ""),
+            tags=data.get("tags", []),
+            confidence=data.get("confidence", 0.0),
+            salience=data.get("salience", 0.0),
+        )

autobyteus/memory/models/tool_interaction.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Optional, Dict
+
+
+class ToolInteractionStatus(Enum):
+    PENDING = "pending"
+    SUCCESS = "success"
+    ERROR = "error"
+
+
+@dataclass
+class ToolInteraction:
+    tool_call_id: str
+    turn_id: Optional[str]
+    tool_name: Optional[str]
+    arguments: Optional[Dict[str, Any]]
+    result: Optional[Any]
+    error: Optional[str]
+    status: ToolInteractionStatus

autobyteus/memory/policies/__init__.py

@@ -0,0 +1,5 @@
+from autobyteus.memory.policies.compaction_policy import CompactionPolicy
+
+__all__ = [
+    "CompactionPolicy",
+]

autobyteus/memory/policies/compaction_policy.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class CompactionPolicy:
+    trigger_ratio: float = 0.8
+    raw_tail_turns: int = 4
+    max_item_chars: int = 2000
+    safety_margin_tokens: int = 256
+
+    def should_compact(self, prompt_tokens: int, input_budget: int) -> bool:
+        if input_budget <= 0:
+            return True
+        if prompt_tokens >= input_budget:
+            return True
+        return prompt_tokens >= int(self.trigger_ratio * input_budget)

autobyteus/memory/retrieval/__init__.py

@@ -0,0 +1,7 @@
+from autobyteus.memory.retrieval.memory_bundle import MemoryBundle
+from autobyteus.memory.retrieval.retriever import Retriever
+
+__all__ = [
+    "MemoryBundle",
+    "Retriever",
+]

autobyteus/memory/retrieval/memory_bundle.py

@@ -0,0 +1,11 @@
+from dataclasses import dataclass, field
+from typing import List
+
+from autobyteus.memory.models.episodic_item import EpisodicItem
+from autobyteus.memory.models.semantic_item import SemanticItem
+
+
+@dataclass
+class MemoryBundle:
+    episodic: List[EpisodicItem] = field(default_factory=list)
+    semantic: List[SemanticItem] = field(default_factory=list)

autobyteus/memory/retrieval/retriever.py

@@ -0,0 +1,13 @@
+from autobyteus.memory.models.memory_types import MemoryType
+from autobyteus.memory.retrieval.memory_bundle import MemoryBundle
+from autobyteus.memory.store.base_store import MemoryStore
+
+
+class Retriever:
+    def __init__(self, store: MemoryStore):
+        self.store = store
+
+    def retrieve(self, max_episodic: int, max_semantic: int) -> MemoryBundle:
+        episodic = self.store.list(MemoryType.EPISODIC, limit=max_episodic)
+        semantic = self.store.list(MemoryType.SEMANTIC, limit=max_semantic)
+        return MemoryBundle(episodic=episodic, semantic=semantic)

autobyteus/memory/store/__init__.py

@@ -0,0 +1,7 @@
+from autobyteus.memory.store.base_store import MemoryStore
+from autobyteus.memory.store.file_store import FileMemoryStore
+
+__all__ = [
+    "MemoryStore",
+    "FileMemoryStore",
+]

autobyteus/memory/store/base_store.py

@@ -0,0 +1,14 @@
+from abc import ABC, abstractmethod
+from typing import Iterable, List, Optional
+
+from autobyteus.memory.models.memory_types import MemoryType
+
+
+class MemoryStore(ABC):
+    @abstractmethod
+    def add(self, items: Iterable[object]) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def list(self, memory_type: MemoryType, limit: Optional[int] = None) -> List[object]:
+        raise NotImplementedError

autobyteus/memory/store/file_store.py

@@ -0,0 +1,98 @@
+import json
+from pathlib import Path
+from typing import Iterable, List, Optional, Union
+
+from autobyteus.memory.models.memory_types import MemoryType
+from autobyteus.memory.models.raw_trace_item import RawTraceItem
+from autobyteus.memory.models.episodic_item import EpisodicItem
+from autobyteus.memory.models.semantic_item import SemanticItem
+from autobyteus.memory.store.base_store import MemoryStore
+
+
+class FileMemoryStore(MemoryStore):
+    def __init__(self, base_dir: Union[str, Path], agent_id: str):
+        self.base_dir = Path(base_dir)
+        self.agent_id = agent_id
+        self.agent_dir = self.base_dir / "agents" / agent_id
+        self.agent_dir.mkdir(parents=True, exist_ok=True)
+
+    def add(self, items: Iterable[object]) -> None:
+        for item in items:
+            memory_type = getattr(item, "memory_type", None)
+            if memory_type is None:
+                raise ValueError("Memory item missing memory_type")
+            path = self._get_file_path(memory_type)
+            record = item.to_dict() if hasattr(item, "to_dict") else item
+            with path.open("a", encoding="utf-8") as handle:
+                handle.write(json.dumps(record) + "\n")
+
+    def list(self, memory_type: MemoryType, limit: Optional[int] = None) -> List[object]:
+        path = self._get_file_path(memory_type)
+        if not path.exists():
+            return []
+        with path.open("r", encoding="utf-8") as handle:
+            lines = [line.strip() for line in handle.readlines() if line.strip()]
+        if limit is not None:
+            lines = lines[-limit:]
+        return [self._deserialize(memory_type, json.loads(line)) for line in lines]
+
+    def list_raw_trace_dicts(self) -> List[dict]:
+        path = self._get_file_path(MemoryType.RAW_TRACE)
+        if not path.exists():
+            return []
+        with path.open("r", encoding="utf-8") as handle:
+            return [json.loads(line) for line in handle if line.strip()]
+
+    def read_archive_raw_traces(self) -> List[dict]:
+        path = self._get_archive_path()
+        if not path.exists():
+            return []
+        with path.open("r", encoding="utf-8") as handle:
+            return [json.loads(line) for line in handle if line.strip()]
+
+    def prune_raw_traces(self, keep_turn_ids: set[str], archive: bool = True) -> None:
+        raw_items = self.list_raw_trace_dicts()
+        if not raw_items:
+            return
+
+        keep = []
+        removed = []
+        for item in raw_items:
+            if item.get("turn_id") in keep_turn_ids:
+                keep.append(item)
+            else:
+                removed.append(item)
+
+        raw_path = self._get_file_path(MemoryType.RAW_TRACE)
+        tmp_path = raw_path.with_suffix(".jsonl.tmp")
+        with tmp_path.open("w", encoding="utf-8") as handle:
+            for item in keep:
+                handle.write(json.dumps(item) + "\n")
+        tmp_path.replace(raw_path)
+
+        if archive and removed:
+            archive_path = self._get_archive_path()
+            with archive_path.open("a", encoding="utf-8") as handle:
+                for item in removed:
+                    handle.write(json.dumps(item) + "\n")
+
+    def _get_file_path(self, memory_type: MemoryType) -> Path:
+        if memory_type == MemoryType.RAW_TRACE:
+            return self.agent_dir / "raw_traces.jsonl"
+        if memory_type == MemoryType.EPISODIC:
+            return self.agent_dir / "episodic.jsonl"
+        if memory_type == MemoryType.SEMANTIC:
+            return self.agent_dir / "semantic.jsonl"
+        raise ValueError(f"Unknown memory type: {memory_type}")
+
+    def _deserialize(self, memory_type: MemoryType, data: dict) -> object:
+        if memory_type == MemoryType.RAW_TRACE:
+            return RawTraceItem.from_dict(data)
+        if memory_type == MemoryType.EPISODIC:
+            return EpisodicItem.from_dict(data)
+        if memory_type == MemoryType.SEMANTIC:
+            return SemanticItem.from_dict(data)
+        raise ValueError(f"Unknown memory type: {memory_type}")
+
+    def _get_archive_path(self) -> Path:
+        return self.agent_dir / "raw_traces_archive.jsonl"

autobyteus/memory/tool_interaction_builder.py

@@ -0,0 +1,46 @@
+from typing import Dict, List
+
+from autobyteus.memory.models.raw_trace_item import RawTraceItem
+from autobyteus.memory.models.tool_interaction import ToolInteraction, ToolInteractionStatus
+
+
+def build_tool_interactions(raw_traces: List[RawTraceItem]) -> List[ToolInteraction]:
+    interactions: Dict[str, ToolInteraction] = {}
+
+    for trace in raw_traces:
+        if trace.trace_type not in {"tool_call", "tool_result"}:
+            continue
+
+        tool_call_id = trace.tool_call_id
+        if not tool_call_id:
+            continue
+
+        interaction = interactions.get(tool_call_id)
+        if interaction is None:
+            interaction = ToolInteraction(
+                tool_call_id=tool_call_id,
+                turn_id=trace.turn_id,
+                tool_name=trace.tool_name,
+                arguments=None,
+                result=None,
+                error=None,
+                status=ToolInteractionStatus.PENDING,
+            )
+            interactions[tool_call_id] = interaction
+
+        if trace.trace_type == "tool_call":
+            interaction.tool_name = trace.tool_name
+            interaction.arguments = trace.tool_args
+            if interaction.status == ToolInteractionStatus.PENDING and interaction.error:
+                interaction.status = ToolInteractionStatus.ERROR
+            continue
+
+        if trace.trace_type == "tool_result":
+            interaction.tool_name = trace.tool_name or interaction.tool_name
+            interaction.result = trace.tool_result
+            interaction.error = trace.tool_error
+            interaction.status = (
+                ToolInteractionStatus.ERROR if trace.tool_error else ToolInteractionStatus.SUCCESS
+            )
+
+    return list(interactions.values())

autobyteus/memory/turn_tracker.py

@@ -0,0 +1,9 @@
+class TurnTracker:
+    def __init__(self, start: int = 1):
+        if start < 1:
+            raise ValueError("start must be >= 1")
+        self._counter = start - 1
+
+    def next_turn_id(self) -> str:
+        self._counter += 1
+        return f"turn_{self._counter:04d}"

autobyteus/multimedia/audio/api/__init__.py

@@ -1,4 +1,5 @@
-from .gemini_audio_client import GeminiAudioClient
 from .autobyteus_audio_client import AutobyteusAudioClient
+from .gemini_audio_client import GeminiAudioClient
+from .openai_audio_client import OpenAIAudioClient
 
-__all__ = ["GeminiAudioClient", "AutobyteusAudioClient"]
+__all__ = ["AutobyteusAudioClient", "GeminiAudioClient", "OpenAIAudioClient"]

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,62 @@ 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.  **AI Studio Mode:** Set `GEMINI_API_KEY`.
+    2.  **Vertex AI Mode:** Set `VERTEX_AI_PROJECT` and `VERTEX_AI_LOCATION`.
     """
 
     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 +123,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 +177,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 +197,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])