path: root/__init__.py

"""MiMo V2.5 TTS Provider Plugin for Hermes Agent.

Integrates Xiaomi's MiMo V2.5 Text-to-Speech API as a Hermes TTS provider.
Supports three models:
  - mimo-v2.5-tts          — built-in high-quality voices
  - mimo-v2.5-tts-voicedesign — custom voice from text description
  - mimo-v2.5-tts-voiceclone  — replicate voice from audio samples

Requires: XIAOMI_API_KEY environment variable.
"""

from __future__ import annotations

import base64
import json
import logging
import os
from typing import Any, Dict, List, Optional

from agent.tts_provider import TTSProvider, DEFAULT_OUTPUT_FORMAT, VALID_OUTPUT_FORMATS

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_API_BASE_DEFAULT = "https://api.xiaomimimo.com/v1"


def _get_api_base() -> str:
    """Return the MiMo API base URL.

    Checks XIAOMI_BASE_URL env var first (matches the LLM provider config),
    then falls back to the public endpoint.
    """
    return os.environ.get("XIAOMI_BASE_URL", _API_BASE_DEFAULT).rstrip("/")


_MODEL_DEFAULT = "mimo-v2.5-tts"
_VOICE_DEFAULT = "mimo_default"

# Format mapping: Hermes format → MiMo API format
_FORMAT_MAP = {
    "mp3": "wav",   # MiMo doesn't output mp3; fall back to wav
    "wav": "wav",
    "ogg": "wav",
    "opus": "wav",
    "flac": "wav",
}

_BUILTIN_VOICES: List[Dict[str, Any]] = [
    {"id": "mimo_default", "display": "MiMo Default (auto-select cluster default)", "language": "auto", "gender": ""},
    {"id": "冰糖", "display": "Bingtang (冰糖)", "language": "zh", "gender": "female"},
    {"id": "茉莉", "display": "Moli (茉莉)", "language": "zh", "gender": "female"},
    {"id": "苏打", "display": "Suda (苏打)", "language": "zh", "gender": "male"},
    {"id": "白桦", "display": "Baihua (白桦)", "language": "zh", "gender": "male"},
    {"id": "Mia", "display": "Mia", "language": "en", "gender": "female"},
    {"id": "Chloe", "display": "Chloe", "language": "en", "gender": "female"},
    {"id": "Milo", "display": "Milo", "language": "en", "gender": "male"},
    {"id": "Dean", "display": "Dean", "language": "en", "gender": "male"},
]

_MODELS: List[Dict[str, Any]] = [
    {"id": "mimo-v2.5-tts", "display": "MiMo V2.5 TTS (built-in voices)", "max_text_length": 10000},
    {"id": "mimo-v2.5-tts-voicedesign", "display": "MiMo V2.5 TTS Voice Design", "max_text_length": 10000},
    {"id": "mimo-v2.5-tts-voiceclone", "display": "MiMo V2.5 TTS Voice Clone", "max_text_length": 10000},
]


# ---------------------------------------------------------------------------
# Provider
# ---------------------------------------------------------------------------

class MiMoTTSProvider(TTSProvider):
    """Xiaomi MiMo V2.5 Text-to-Speech provider."""

    @property
    def name(self) -> str:
        return "mimo"

    @property
    def display_name(self) -> str:
        return "MiMo TTS (Xiaomi)"

    def is_available(self) -> bool:
        return bool(os.environ.get("XIAOMI_API_KEY"))

    def list_voices(self) -> List[Dict[str, Any]]:
        return _BUILTIN_VOICES

    def list_models(self) -> List[Dict[str, Any]]:
        return _MODELS

    def default_model(self) -> Optional[str]:
        return _MODEL_DEFAULT

    def default_voice(self) -> Optional[str]:
        return _VOICE_DEFAULT

    def get_setup_schema(self) -> Dict[str, Any]:
        return {
            "name": "MiMo TTS (Xiaomi)",
            "badge": "paid",
            "tag": "V2.5 — built-in voices, voice design, voice cloning",
            "env_vars": [
                {
                    "key": "XIAOMI_API_KEY",
                    "prompt": "Xiaomi MiMo API key",
                    "url": "https://mimo.mi.com",
                },
            ],
        }

    @property
    def voice_compatible(self) -> bool:
        return True

    def synthesize(
        self,
        text: str,
        output_path: str,
        *,
        voice: Optional[str] = None,
        model: Optional[str] = None,
        speed: Optional[float] = None,
        format: str = DEFAULT_OUTPUT_FORMAT,
        **extra: Any,
    ) -> str:
        """Synthesize text via MiMo V2.5 TTS API.

        The text is placed in the ``assistant`` message (per MiMo API spec).
        Style instructions (from ``extra["style"]``) go in the ``user`` message.
        """
        import urllib.request

        api_key = os.environ.get("XIAOMI_API_KEY")
        if not api_key:
            raise RuntimeError(
                "XIAOMI_API_KEY environment variable is not set. "
                "Get your key at https://mimo.mi.com"
            )

        model_id = model or _MODEL_DEFAULT
        voice_id = voice or _VOICE_DEFAULT
        mimo_format = _FORMAT_MAP.get(format, "wav")

        # Build messages
        messages: list = []

        # Style instruction in user message (optional)
        style = extra.get("style")
        if style:
            messages.append({"role": "user", "content": str(style)})

        # Text to synthesize in assistant message
        messages.append({"role": "assistant", "content": text})

        # Build request body
        # voice-design and voice-clone models don't accept the voice field
        audio_config: Dict[str, Any] = {"format": mimo_format}
        if "voicedesign" not in model_id and "voiceclone" not in model_id:
            audio_config["voice"] = voice_id
        body: Dict[str, Any] = {
            "model": model_id,
            "messages": messages,
            "audio": audio_config,
        }

        # Speed control via style instructions if speed is set
        # (MiMo doesn't have a direct speed param; we prepend to style)
        if speed is not None and speed != 1.0 and not style:
            speed_desc = "faster" if speed > 1.0 else "slower"
            body["messages"].insert(0, {
                "role": "user",
                "content": f"Speak at a slightly {speed_desc} pace than normal.",
            })

        data = json.dumps(body).encode("utf-8")

        req = urllib.request.Request(
            f"{_get_api_base()}/chat/completions",
            data=data,
            headers={
                "Content-Type": "application/json",
                "api-key": api_key,
            },
            method="POST",
        )

        logger.info("MiMo TTS: synthesizing %d chars with model=%s voice=%s", len(text), model_id, voice_id)

        try:
            with urllib.request.urlopen(req, timeout=60) as resp:
                resp_data = json.loads(resp.read().decode("utf-8"))
        except Exception as exc:
            raise RuntimeError(f"MiMo TTS API request failed: {exc}") from exc

        # Extract audio from response
        try:
            audio_b64 = resp_data["choices"][0]["message"]["audio"]["data"]
        except (KeyError, IndexError, TypeError) as exc:
            raise RuntimeError(
                f"MiMo TTS API returned unexpected response structure: {exc}\n"
                f"Response: {json.dumps(resp_data, ensure_ascii=False)[:500]}"
            ) from exc

        audio_bytes = base64.b64decode(audio_b64)

        # Ensure output has correct extension for the actual format
        # MiMo always returns wav regardless of what we asked
        if not output_path.endswith(".wav"):
            output_path = os.path.splitext(output_path)[0] + ".wav"

        os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
        with open(output_path, "wb") as f:
            f.write(audio_bytes)

        logger.info("MiMo TTS: wrote %d bytes to %s", len(audio_bytes), output_path)
        return output_path

    def stream(
        self,
        text: str,
        *,
        voice: Optional[str] = None,
        model: Optional[str] = None,
        format: str = "opus",
        **extra: Any,
    ):
        """Stream synthesized audio chunks via MiMo V2.5 TTS streaming API."""
        import urllib.request

        api_key = os.environ.get("XIAOMI_API_KEY")
        if not api_key:
            raise RuntimeError("XIAOMI_API_KEY environment variable is not set.")

        model_id = model or _MODEL_DEFAULT
        voice_id = voice or _VOICE_DEFAULT

        messages: list = []
        style = extra.get("style")
        if style:
            messages.append({"role": "user", "content": str(style)})
        messages.append({"role": "assistant", "content": text})

        audio_config: Dict[str, Any] = {"format": "pcm16"}
        if "voicedesign" not in model_id and "voiceclone" not in model_id:
            audio_config["voice"] = voice_id
        body = {
            "model": model_id,
            "messages": messages,
            "audio": audio_config,
            "stream": True,
        }

        data = json.dumps(body).encode("utf-8")
        req = urllib.request.Request(
            f"{_get_api_base()}/chat/completions",
            data=data,
            headers={
                "Content-Type": "application/json",
                "api-key": api_key,
            },
            method="POST",
        )

        logger.info("MiMo TTS streaming: synthesizing %d chars", len(text))

        with urllib.request.urlopen(req, timeout=120) as resp:
            buffer = b""
            for line in resp:
                line = line.strip()
                if not line:
                    continue
                # SSE format: "data: {...}"
                if line.startswith(b"data: "):
                    payload = line[6:]
                    if payload == b"[DONE]":
                        break
                    try:
                        chunk = json.loads(payload)
                        delta = chunk.get("choices", [{}])[0].get("delta", {})
                        audio_data = delta.get("audio", {})
                        if isinstance(audio_data, dict) and "data" in audio_data:
                            pcm_bytes = base64.b64decode(audio_data["data"])
                            yield pcm_bytes
                    except (json.JSONDecodeError, KeyError, IndexError):
                        continue


# ---------------------------------------------------------------------------
# Plugin entry point
# ---------------------------------------------------------------------------

def register(ctx):
    """Hermes plugin entry point — register the MiMo TTS provider."""
    ctx.register_tts_provider(MiMoTTSProvider())