xiaoxia-saas/packages/application/tts_job/streaming_service.py

"""P2: TTS 流式合成服务 — WebSocket 实时音频推送。

通过 WebSocket 将合成音频以二进制帧实时推送给客户端。
- 短文本（≤500 字）：合成完整音频后分块推送
- 长文本（>500 字）：分段并发合成，逐段推送音频
"""

from __future__ import annotations

import asyncio
import logging
from typing import Any, Optional

from packages.application.cosyvoice_service import CosyVoiceError, CosyVoiceService
from packages.application.tts_job.text_splitter import split_text
from packages.shared.url_security import ALLOWED_AUDIO_MIME_TYPES, safe_download_bytes

logger = logging.getLogger(__name__)

# WebSocket 二进制帧块大小（4KB）
_AUDIO_CHUNK_SIZE = 4096
# 分段并发上限
_MAX_STREAMING_SEGMENT_WORKERS = 5
# 长文本分段阈值
_SEGMENT_THRESHOLD = 500
# WebSocket 最大文本长度
_MAX_TEXT_LENGTH = 10000


class TTSStreamingError(Exception):
    """TTS 流式合成异常。"""

    pass


class TTSStreamingService:
    """TTS 流式合成服务。

    通过 WebSocket 实时推送合成音频。
    使用 CosyVoiceService（同步 REST API）合成，
    通过 asyncio.to_thread 桥接到异步 WebSocket。
    """

    def __init__(self, cosyvoice_service: CosyVoiceService) -> None:
        self._cosyvoice = cosyvoice_service

    async def synthesize_and_stream(self, websocket: Any, params: dict) -> None:
        """根据文本长度选择流式合成策略。

        Args:
            websocket: FastAPI WebSocket 连接
            params: 合成参数（text, voice_id, sample_rate, format, speed）
        """
        text = params.get("text", "")
        if not text:
            await self._send_json(websocket, {"type": "error", "message": "文本不能为空"})
            return

        if len(text) > _MAX_TEXT_LENGTH:
            await self._send_json(
                websocket,
                {"type": "error", "message": f"文本过长，最大 {_MAX_TEXT_LENGTH} 字"},
            )
            return

        if len(text) <= _SEGMENT_THRESHOLD:
            await self._stream_short_text(websocket, params)
        else:
            await self._stream_long_text(websocket, params)

    # ── 短文本流式合成 ────────────────────────────────────────

    async def _stream_short_text(self, websocket: Any, params: dict) -> None:
        """短文本：合成完整音频后分块推送。"""
        text = params["text"]
        voice_id = params.get("voice_id", "")
        sample_rate = params.get("sample_rate", 0)
        audio_format = params.get("format", "mp3")
        speed = params.get("speed", 1.0)

        await self._send_json(
            websocket,
            {"type": "started", "segment_count": 1, "total_segments": 1},
        )

        # 在线程池中执行同步合成
        try:
            result = await asyncio.to_thread(
                self._cosyvoice.submit_synthesize_task,
                text=text,
                voice_id=voice_id,
                sample_rate=sample_rate,
                format=audio_format,
                speed=speed,
            )
        except CosyVoiceError as e:
            logger.error(f"流式合成失败: {e}")
            await self._send_json(websocket, {"type": "error", "message": str(e)})
            return
        except Exception as e:
            logger.error(f"流式合成意外错误: {e}")
            await self._send_json(websocket, {"type": "error", "message": f"合成失败: {e}"})
            return

        audio_url = result.get("audio_url", "")
        if not audio_url:
            await self._send_json(websocket, {"type": "error", "message": "合成未返回音频 URL"})
            return

        # 下载并流式推送音频
        try:
            audio_data = await asyncio.to_thread(self._download_audio, audio_url)
            total_bytes = await self._stream_audio_chunks(websocket, audio_data)

            await self._send_json(
                websocket,
                {
                    "type": "done",
                    "duration": result.get("duration", 0.0),
                    "file_size": total_bytes,
                    "format": audio_format,
                },
            )
        except Exception as e:
            logger.error(f"音频流式推送失败: {e}")
            await self._send_json(websocket, {"type": "error", "message": f"音频推送失败: {e}"})

    # ── 长文本分段流式合成 ────────────────────────────────────

    async def _stream_long_text(self, websocket: Any, params: dict) -> None:
        """长文本：分段并发合成，逐段推送音频。"""
        text = params["text"]
        voice_id = params.get("voice_id", "")
        sample_rate = params.get("sample_rate", 0)
        audio_format = params.get("format", "mp3")
        speed = params.get("speed", 1.0)

        segments = split_text(text, max_chars=_SEGMENT_THRESHOLD)
        segment_count = len(segments)

        logger.info(f"流式分段合成: 原文={len(text)}字, 段数={segment_count}")

        await self._send_json(
            websocket,
            {"type": "started", "segment_count": segment_count, "total_segments": segment_count},
        )

        # 并发合成所有分段，按顺序流式推送
        queue: asyncio.Queue[tuple[int, Optional[bytes], Optional[str]]] = asyncio.Queue()

        async def _synthesize_one(idx: int, seg_text: str) -> None:
            """合成单个分段并放入队列。"""
            try:
                result = await asyncio.to_thread(
                    self._cosyvoice.submit_synthesize_task,
                    text=seg_text,
                    voice_id=voice_id,
                    sample_rate=sample_rate,
                    format=audio_format,
                    speed=speed,
                )
                audio_url = result.get("audio_url", "")
                if audio_url:
                    audio_data = await asyncio.to_thread(self._download_audio, audio_url)
                    await queue.put((idx, audio_data, None))
                else:
                    await queue.put((idx, None, "合成未返回音频 URL"))
            except Exception as e:
                await queue.put((idx, None, str(e)))

        # 启动并发合成任务
        workers = [asyncio.create_task(_synthesize_one(idx, seg)) for idx, seg in enumerate(segments)]

        # 按顺序消费队列，流式推送
        total_bytes = 0
        total_duration = 0.0
        consumed = 0

        try:
            while consumed < segment_count:
                idx, audio_data, error = await queue.get()
                consumed += 1

                if error:
                    logger.error(f"分段 {idx + 1} 合成失败: {error}")
                    await self._send_json(
                        websocket,
                        {"type": "error", "message": f"分段 {idx + 1} 合成失败: {error}"},
                    )
                    # 取消剩余 worker
                    for w in workers:
                        w.cancel()
                    return

                if audio_data:
                    seg_bytes = await self._stream_audio_chunks(websocket, audio_data)
                    total_bytes += seg_bytes

                await self._send_json(
                    websocket,
                    {"type": "segment_done", "segment": idx + 1, "total": segment_count},
                )

            # 等待所有 worker 完成
            await asyncio.gather(*workers, return_exceptions=True)

            await self._send_json(
                websocket,
                {
                    "type": "done",
                    "duration": total_duration,
                    "file_size": total_bytes,
                    "format": audio_format,
                },
            )

        except Exception as e:
            logger.error(f"流式分段推送失败: {e}")
            await self._send_json(websocket, {"type": "error", "message": f"推送失败: {e}"})
            for w in workers:
                w.cancel()

    # ── 工具方法 ────────────────────────────────────────────

    def _download_audio(self, url: str) -> bytes:
        """下载音频数据（含 SSRF 防护 + 大小限制 + 重定向校验）。"""
        return safe_download_bytes(
            url,
            purpose="tts_streaming_download",
            allowed_mime_types=ALLOWED_AUDIO_MIME_TYPES,
            timeout=60.0,
        )

    async def _stream_audio_chunks(self, websocket: Any, audio_data: bytes) -> int:
        """将音频数据分块通过 WebSocket 推送。

        Returns:
            推送的总字节数
        """
        total = 0
        for offset in range(0, len(audio_data), _AUDIO_CHUNK_SIZE):
            chunk = audio_data[offset : offset + _AUDIO_CHUNK_SIZE]
            await websocket.send_bytes(chunk)
            total += len(chunk)
        return total

    async def _send_json(self, websocket: Any, data: dict) -> None:
        """发送 JSON 帧，失败时记录日志。"""
        try:
            await websocket.send_json(data)
        except Exception as e:
            logger.warning("WebSocket JSON 发送失败: type=%s error=%s", data.get("type", "?"), e)