mytoolkit/mytoolkit/api/voice.py

"""Voice / TTS API client via Volcengine (ByteDance) WebSocket API."""

import asyncio
import json
import uuid
from pathlib import Path

# TTS 1.0 (moon/mars) vs 2.0 (uranus / ICL_uranus)
_RESOURCE_TTS_1 = "volc.service_type.10029"
_RESOURCE_TTS_2 = "seed-tts-2.0"
_ENDPOINT = "wss://openspeech.bytedance.com/api/v3/tts/unidirectional/stream"
# Newer HTTP endpoint: X-Api-Key auth (no AppID/Token), streamed base64.
# Note: the legacy appid/token path is WebSocket (_ENDPOINT); the api-key path
# is this HTTP endpoint (the /plan/ variant is Agent-Plan-only and rejects it).
_ENDPOINT_HTTP = "https://openspeech.bytedance.com/api/v3/tts/unidirectional"

# Product labels: 经典男声 / 经典女声 / 阳光甜妹
_DEFAULT_VOICE = "zh_female_xiaohe_uranus_bigtts"


def resource_id_for_speaker(speaker: str) -> str:
    """Pick Volcengine resource id for a speaker."""
    s = speaker or ""
    if "uranus" in s or s.startswith("ICL_uranus_"):
        return _RESOURCE_TTS_2
    return _RESOURCE_TTS_1


def synthesize_tts(
    appid: str,
    token: str,
    text: str,
    speaker: str = _DEFAULT_VOICE,
    fmt: str = "mp3",
    speech_rate: int = 0,
    output_path: str = "output.mp3",
    api_key: str | None = None,
) -> str:
    """Synchronous TTS synthesis via Volcengine API.

    Prefers the newer plan HTTP endpoint (``api_key``, X-Api-Key auth); falls
    back to the legacy WebSocket endpoint (appid + access_token). Returns the
    absolute path to the output audio file.
    """
    if api_key:
        return synthesize_tts_http(
            api_key,
            text=text,
            speaker=speaker,
            fmt=fmt,
            speech_rate=speech_rate,
            output_path=output_path,
        )

    # Delayed import to avoid circular dependency when bin.commands is not yet loaded
    from mytoolkit.protocols import (
        EventType,
        MsgType,
        full_client_request,
        receive_message,
    )

    try:
        import websockets
    except ImportError as e:
        raise RuntimeError(
            f"Missing dependency: {e}. Run: pip install websockets"
        ) from e

    headers = {
        "X-Api-App-Id": appid,
        "X-Api-Access-Key": token,
        "X-Api-Resource-Id": resource_id_for_speaker(speaker),
        "X-Api-Connect-Id": str(uuid.uuid4()),
    }

    async def _run() -> bytearray:
        ws = await websockets.connect(
            _ENDPOINT,
            additional_headers=headers,
            max_size=10 * 1024 * 1024,
            open_timeout=30,
        )

        request = {
            "user": {"uid": str(uuid.uuid4())},
            "req_params": {
                "speaker": speaker,
                "audio_params": {
                    "format": fmt,
                    "sample_rate": 24000,
                    "speech_rate": speech_rate,
                },
                "text": text,
            },
        }

        await full_client_request(ws, json.dumps(request).encode())

        audio_data = bytearray()
        while True:
            msg = await receive_message(ws)
            if (
                msg.type == MsgType.FullServerResponse
                and msg.event == EventType.SessionFinished
            ):
                break
            elif msg.type == MsgType.AudioOnlyServer:
                audio_data.extend(msg.payload)
            elif msg.type == MsgType.Error:
                await ws.close()
                raise RuntimeError(
                    f"TTS service error: {msg.payload.decode('utf-8', errors='replace')}"
                )

        await ws.close()
        return audio_data

    audio_data = asyncio.run(_run())
    if not audio_data:
        raise RuntimeError("No audio data received from TTS service")

    out = Path(output_path)
    out.write_bytes(audio_data)
    return str(out.resolve())


def synthesize_tts_http(
    api_key: str,
    text: str,
    *,
    speaker: str = _DEFAULT_VOICE,
    fmt: str = "mp3",
    speech_rate: int = 0,
    output_path: str = "output.mp3",
    timeout_sec: float = 60.0,
) -> str:
    """TTS via Volcengine v3 plan HTTP endpoint (X-Api-Key auth).

    One POST, streamed newline-delimited JSON; each line carries a base64
    audio chunk plus a terminal ``code == 20000000`` line. No AppID/Token.
    """
    import base64
    import json
    import urllib.error
    import urllib.request

    payload = {
        "req_params": {
            "text": text,
            "speaker": speaker,
            "audio_params": {
                "format": fmt,
                "sample_rate": 24000,
                "speech_rate": speech_rate,
            },
        }
    }
    req = urllib.request.Request(
        _ENDPOINT_HTTP,
        data=json.dumps(payload).encode("utf-8"),
        method="POST",
    )
    req.add_header("X-Api-Key", api_key)
    req.add_header("X-Api-Resource-Id", _RESOURCE_TTS_2)
    req.add_header("Content-Type", "application/json")
    req.add_header("Connection", "keep-alive")
    req.add_header("X-Control-Require-Usage-Tokens-Return", "*")

    try:
        resp = urllib.request.urlopen(req, timeout=timeout_sec)
    except urllib.error.HTTPError as exc:
        detail = exc.read()[:300].decode("utf-8", errors="replace")
        raise RuntimeError(
            f"TTS service error: HTTP {exc.code} {detail}"
        ) from exc
    except OSError as exc:
        raise RuntimeError(f"TTS service error: {exc}") from exc

    audio_data = bytearray()
    for raw in resp:
        line = raw.decode("utf-8", errors="replace").strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
        except json.JSONDecodeError:
            continue
        code = obj.get("code", 0)
        if obj.get("data"):
            audio_data.extend(base64.b64decode(obj["data"]))
        if code == 20000000:  # terminal success
            break
        if code > 0:
            raise RuntimeError(
                f"TTS service error: {json.dumps(obj, ensure_ascii=False)[:300]}"
            )
    if not audio_data:
        raise RuntimeError("No audio data received from TTS service")

    out = Path(output_path)
    out.write_bytes(audio_data)
    return str(out.resolve())