mytoolkit/bin/commands/voice.py

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

import asyncio
import json
import os
import uuid
from pathlib import Path

import click

from bin.config import config
from bin.utils import handle_errors

_RESOURCE_ID = "volc.service_type.10029"
_ENDPOINT = "wss://openspeech.bytedance.com/api/v3/tts/unidirectional/stream"

_VOICE_CHOICES = [
    "zh_female_cancan_mars_bigtts",
    "zh_female_shuangkuaisisi_moon_bigtts",
    "zh_male_wennuanahu_moon_bigtts",
    "zh_male_sunwukong_moon_bigtts",
]

_FORMAT_CHOICES = ["mp3", "wav", "pcm"]


def _get_credentials():
    """Get Volcengine credentials from config."""
    appid = config.get("volc_appid")
    access_token = config.get("volc_access_token")
    if not appid or not access_token:
        return None, None
    return appid, access_token


def _synthesize_sync(
    text: str,
    speaker: str,
    fmt: str,
    speech_rate: int,
    output_path: str,
) -> str:
    """Synchronous wrapper around async WebSocket TTS."""
    try:
        from bin.commands.protocols import (
            EventType,
            MsgType,
            full_client_request,
            receive_message,
        )
        import websockets
    except ImportError as e:
        raise RuntimeError(f"Missing dependency: {e}. Run: pip install websockets")

    appid, access_token = _get_credentials()
    if not appid or not access_token:
        raise click.UsageError(
            "Volcengine credentials not set. Run:\n"
            "  mytoolkit env set volc_appid <appid>\n"
            "  mytoolkit env set volc_access_token <token>"
        )

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

    async def _run() -> bytearray:
        try:
            ws = await websockets.connect(
                _ENDPOINT,
                additional_headers=headers,
                max_size=10 * 1024 * 1024,
                open_timeout=30,
            )
        except Exception as e:
            raise RuntimeError(f"WebSocket connection failed: {e}") from e

        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")

    Path(output_path).write_bytes(audio_data)
    return output_path


@click.group()
def voice():
    """Voice synthesis (TTS) commands."""
    pass


@voice.command("tts")
@click.argument("text", required=False)
@click.option("-f", "--file", type=click.Path(exists=True), help="从文件读取文本")
@click.option(
    "-v",
    "--voice",
    "speaker",
    type=click.Choice(_VOICE_CHOICES),
    default="zh_male_wennuanahu_moon_bigtts",
    help="音色",
)
@click.option("--format", "fmt", type=click.Choice(_FORMAT_CHOICES), default="mp3", help="输出格式")
@click.option(
    "--speed",
    type=click.IntRange(-50, 100),
    default=0,
    help="语速 (-50=0.5倍, 0=正常, 100=2倍)",
)
@click.option("-o", "--output", help="输出文件路径")
@handle_errors
def tts_cmd(text, file, speaker, fmt, speed, output):
    """使用豆包/火山引擎语音 API 将文本转为语音。"""
    if file:
        text = Path(file).read_text(encoding="utf-8").strip()
    elif not text:
        raise click.UsageError("必须提供 text 或使用 -f/--file 从文件读取")

    click.echo(f"Synthesizing: {text[:40]}{'...' if len(text) > 40 else ''}")

    save_path = output or f"output.{fmt}"
    result = _synthesize_sync(text, speaker, fmt, speed, save_path)
    click.echo(os.path.abspath(result))