"""P2: TTS 流式合成服务 — WebSocket 实时音频推送。 通过 WebSocket 将合成音频以二进制帧实时推送给客户端。 - 短文本(≤500 字):合成完整音频后分块推送 - 长文本(>500 字):分段并发合成,逐段推送音频 """ from __future__ import annotations import asyncio import logging from typing import Any, Optional import httpx 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)