# -*- coding: utf-8 -*- import asyncio import io import json import os import struct import uuid from dataclasses import dataclass from enum import IntEnum from typing import Callable, List, Optional import websockets from loguru import logger from pydub import AudioSegment # Protocol definitions (from podcasts_protocols) class MsgType(IntEnum): """Message type enumeration""" Invalid = 0 FullClientRequest = 0b1 AudioOnlyClient = 0b10 FullServerResponse = 0b1001 AudioOnlyServer = 0b1011 FrontEndResultServer = 0b1100 Error = 0b1111 ServerACK = AudioOnlyServer class MsgTypeFlagBits(IntEnum): """Message type flag bits""" NoSeq = 0 PositiveSeq = 0b1 LastNoSeq = 0b10 NegativeSeq = 0b11 WithEvent = 0b100 class VersionBits(IntEnum): """Version bits""" Version1 = 1 class HeaderSizeBits(IntEnum): """Header size bits""" HeaderSize4 = 1 HeaderSize8 = 2 HeaderSize12 = 3 HeaderSize16 = 4 class SerializationBits(IntEnum): """Serialization method bits""" Raw = 0 JSON = 0b1 Thrift = 0b11 Custom = 0b1111 class CompressionBits(IntEnum): """Compression method bits""" None_ = 0 Gzip = 0b1 Custom = 0b1111 class EventType(IntEnum): """Event type enumeration""" None_ = 0 StartConnection = 1 StartTask = 1 FinishConnection = 2 FinishTask = 2 ConnectionStarted = 50 TaskStarted = 50 ConnectionFailed = 51 TaskFailed = 51 ConnectionFinished = 52 TaskFinished = 52 StartSession = 100 CancelSession = 101 FinishSession = 102 SessionStarted = 150 SessionCanceled = 151 SessionFinished = 152 SessionFailed = 153 UsageResponse = 154 ChargeData = 154 TaskRequest = 200 UpdateConfig = 201 AudioMuted = 250 SayHello = 300 TTSSentenceStart = 350 TTSSentenceEnd = 351 TTSResponse = 352 TTSEnded = 359 PodcastRoundStart = 360 PodcastRoundResponse = 361 PodcastRoundEnd = 362 PodcastEnd = 363 @dataclass class Message: """Message object""" version: VersionBits = VersionBits.Version1 header_size: HeaderSizeBits = HeaderSizeBits.HeaderSize4 type: MsgType = MsgType.Invalid flag: MsgTypeFlagBits = MsgTypeFlagBits.NoSeq serialization: SerializationBits = SerializationBits.JSON compression: CompressionBits = CompressionBits.None_ event: EventType = EventType.None_ session_id: str = "" connect_id: str = "" sequence: int = 0 error_code: int = 0 payload: bytes = b"" @classmethod def from_bytes(cls, data: bytes) -> "Message": """Create message object from bytes""" if len(data) < 3: raise ValueError(f"Data too short: expected at least 3 bytes, got {len(data)}") type_and_flag = data[1] msg_type = MsgType(type_and_flag >> 4) flag = MsgTypeFlagBits(type_and_flag & 0b00001111) msg = cls(type=msg_type, flag=flag) msg.unmarshal(data) return msg def marshal(self) -> bytes: """Serialize message to bytes""" buffer = io.BytesIO() header = [ (self.version << 4) | self.header_size, (self.type << 4) | self.flag, (self.serialization << 4) | self.compression, ] header_size = 4 * self.header_size if padding := header_size - len(header): header.extend([0] * padding) buffer.write(bytes(header)) writers = self._get_writers() for writer in writers: writer(buffer) return buffer.getvalue() def unmarshal(self, data: bytes) -> None: """Deserialize message from bytes""" buffer = io.BytesIO(data) version_and_header_size = buffer.read(1)[0] self.version = VersionBits(version_and_header_size >> 4) self.header_size = HeaderSizeBits(version_and_header_size & 0b00001111) buffer.read(1) serialization_compression = buffer.read(1)[0] self.serialization = SerializationBits(serialization_compression >> 4) self.compression = CompressionBits(serialization_compression & 0b00001111) header_size = 4 * self.header_size read_size = 3 if padding_size := header_size - read_size: buffer.read(padding_size) readers = self._get_readers() for reader in readers: reader(buffer) remaining = buffer.read() if remaining: raise ValueError(f"Unexpected data after message: {remaining}") def _get_writers(self) -> List[Callable[[io.BytesIO], None]]: """Get list of writer functions""" writers = [] if self.flag == MsgTypeFlagBits.WithEvent: writers.extend([self._write_event, self._write_session_id]) if self.type in [MsgType.FullClientRequest, MsgType.FullServerResponse, MsgType.FrontEndResultServer, MsgType.AudioOnlyClient, MsgType.AudioOnlyServer]: if self.flag in [MsgTypeFlagBits.PositiveSeq, MsgTypeFlagBits.NegativeSeq]: writers.append(self._write_sequence) elif self.type == MsgType.Error: writers.append(self._write_error_code) else: raise ValueError(f"Unsupported message type: {self.type}") writers.append(self._write_payload) return writers def _get_readers(self) -> List[Callable[[io.BytesIO], None]]: """Get list of reader functions""" readers = [] if self.type in [MsgType.FullClientRequest, MsgType.FullServerResponse, MsgType.FrontEndResultServer, MsgType.AudioOnlyClient, MsgType.AudioOnlyServer]: if self.flag in [MsgTypeFlagBits.PositiveSeq, MsgTypeFlagBits.NegativeSeq]: readers.append(self._read_sequence) elif self.type == MsgType.Error: readers.append(self._read_error_code) if self.flag == MsgTypeFlagBits.WithEvent: readers.extend([self._read_event, self._read_session_id, self._read_connect_id]) readers.append(self._read_payload) return readers def _write_event(self, buffer: io.BytesIO) -> None: buffer.write(struct.pack(">i", self.event)) def _write_session_id(self, buffer: io.BytesIO) -> None: if self.event in [EventType.StartConnection, EventType.FinishConnection, EventType.ConnectionStarted, EventType.ConnectionFailed]: return session_id_bytes = self.session_id.encode("utf-8") size = len(session_id_bytes) if size > 0xFFFFFFFF: raise ValueError(f"Session ID size ({size}) exceeds max(uint32)") buffer.write(struct.pack(">I", size)) if size > 0: buffer.write(session_id_bytes) def _write_sequence(self, buffer: io.BytesIO) -> None: buffer.write(struct.pack(">i", self.sequence)) def _write_error_code(self, buffer: io.BytesIO) -> None: buffer.write(struct.pack(">I", self.error_code)) def _write_payload(self, buffer: io.BytesIO) -> None: size = len(self.payload) if size > 0xFFFFFFFF: raise ValueError(f"Payload size ({size}) exceeds max(uint32)") buffer.write(struct.pack(">I", size)) buffer.write(self.payload) def _read_event(self, buffer: io.BytesIO) -> None: event_bytes = buffer.read(4) if event_bytes: self.event = EventType(struct.unpack(">i", event_bytes)[0]) def _read_session_id(self, buffer: io.BytesIO) -> None: if self.event in [EventType.StartConnection, EventType.FinishConnection, EventType.ConnectionStarted, EventType.ConnectionFailed, EventType.ConnectionFinished]: return size_bytes = buffer.read(4) if size_bytes: size = struct.unpack(">I", size_bytes)[0] if size > 0: session_id_bytes = buffer.read(size) if len(session_id_bytes) == size: self.session_id = session_id_bytes.decode("utf-8") def _read_connect_id(self, buffer: io.BytesIO) -> None: if self.event in [EventType.ConnectionStarted, EventType.ConnectionFailed, EventType.ConnectionFinished]: size_bytes = buffer.read(4) if size_bytes: size = struct.unpack(">I", size_bytes)[0] if size > 0: self.connect_id = buffer.read(size).decode("utf-8") def _read_sequence(self, buffer: io.BytesIO) -> None: sequence_bytes = buffer.read(4) if sequence_bytes: self.sequence = struct.unpack(">i", sequence_bytes)[0] def _read_error_code(self, buffer: io.BytesIO) -> None: error_code_bytes = buffer.read(4) if error_code_bytes: self.error_code = struct.unpack(">I", error_code_bytes)[0] def _read_payload(self, buffer: io.BytesIO) -> None: size_bytes = buffer.read(4) if size_bytes: size = struct.unpack(">I", size_bytes)[0] if size > 0: self.payload = buffer.read(size) async def receive_message(websocket: websockets.WebSocketClientProtocol) -> Message: """Receive message from websocket""" try: data = await websocket.recv() if isinstance(data, str): raise ValueError(f"Unexpected text message: {data}") elif isinstance(data, bytes): msg = Message.from_bytes(data) # logger.debug(f"Received: {msg}") return msg else: raise ValueError(f"Unexpected message type: {type(data)}") except Exception as e: logger.error(f"Failed to receive message: {e}") raise async def wait_for_event(websocket: websockets.WebSocketClientProtocol, msg_type: MsgType, event_type: EventType) -> Message: """Wait for specific event""" while True: msg = await receive_message(websocket) if msg.type != msg_type or msg.event != event_type: raise ValueError(f"Unexpected message: {msg}") if msg.type == msg_type and msg.event == event_type: return msg async def start_connection(websocket: websockets.WebSocketClientProtocol) -> None: """Start connection""" msg = Message(type=MsgType.FullClientRequest, flag=MsgTypeFlagBits.WithEvent) msg.event = EventType.StartConnection msg.payload = b"{}" logger.debug(f"Sending: {msg}") await websocket.send(msg.marshal()) async def finish_connection(websocket: websockets.WebSocketClientProtocol) -> None: """Finish connection""" msg = Message(type=MsgType.FullClientRequest, flag=MsgTypeFlagBits.WithEvent) msg.event = EventType.FinishConnection msg.payload = b"{}" logger.debug(f"Sending: {msg}") await websocket.send(msg.marshal()) async def start_session(websocket: websockets.WebSocketClientProtocol, payload: bytes, session_id: str) -> None: """Start session""" msg = Message(type=MsgType.FullClientRequest, flag=MsgTypeFlagBits.WithEvent) msg.event = EventType.StartSession msg.session_id = session_id msg.payload = payload logger.debug(f"Sending: {msg}") await websocket.send(msg.marshal()) async def finish_session(websocket: websockets.WebSocketClientProtocol, session_id: str) -> None: """Finish session""" msg = Message(type=MsgType.FullClientRequest, flag=MsgTypeFlagBits.WithEvent) msg.event = EventType.FinishSession msg.session_id = session_id msg.payload = b"{}" logger.debug(f"Sending: {msg}") await websocket.send(msg.marshal()) class PodcastRoundPostProcessor: def __init__(self, session_id, data_manager): self.session_id = session_id self.data_manager = data_manager self.temp_merged_audio_name = "merged_audio.mp3" self.output_merged_audio_name = f"{session_id}-merged_audio.mp3" self.subtitle_timestamps = [] # 记录字幕时间戳 self.current_audio_duration = 0.0 # 当前音频时长 self.merged_audio = None # 用于存储合并的音频对象 self.merged_audio_bytes = None async def init(self): if self.data_manager: await self.data_manager.create_podcast_temp_session_dir(self.session_id) async def postprocess_round(self, current_round, voice, audio, podcast_texts): text = "" if podcast_texts: text = podcast_texts[-1].get("text", "") logger.debug(f"Processing round: {current_round}, voice: {voice}, text: {text}, audio: {len(audio)} bytes") new_segment = AudioSegment.from_mp3(io.BytesIO(bytes(audio))) round_duration = len(new_segment) / 1000.0 if self.merged_audio is None: self.merged_audio = new_segment else: self.merged_audio = self.merged_audio + new_segment # 保存合并后的音频到临时文件（用于前端实时访问） merged_io = io.BytesIO() self.merged_audio.export(merged_io, format="mp3") self.merged_audio_bytes = merged_io.getvalue() if self.data_manager: await self.data_manager.save_podcast_temp_session_file(self.session_id, self.temp_merged_audio_name, self.merged_audio_bytes) merged_file_size = len(self.merged_audio_bytes) # 记录字幕时间戳 self.subtitle_timestamps.append( { "start": self.current_audio_duration, "end": self.current_audio_duration + round_duration, "text": text, "speaker": voice, } ) self.current_audio_duration += round_duration logger.debug(f"Merged audio updated: {merged_file_size} bytes, duration: {self.current_audio_duration:.2f}s") return { "url": f"/api/v1/podcast/audio?session_id={self.session_id}&filename={self.temp_merged_audio_name}", "size": merged_file_size, "duration": self.current_audio_duration, "round": current_round, "text": text, "speaker": voice, } async def postprocess_final(self): if self.data_manager: await self.data_manager.save_podcast_output_file(self.output_merged_audio_name, self.merged_audio_bytes) return { "subtitles": self.subtitle_timestamps, "audio_name": self.output_merged_audio_name, } async def cleanup(self): if self.data_manager: await self.data_manager.clear_podcast_temp_session_dir(self.session_id) self.data_manager = None class VolcEnginePodcastClient: """ VolcEngine Podcast客户端 支持多种播客类型: - action=0: 文本转播客 - action=3: NLP文本转播客 - action=4: 提示词生成播客 """ def __init__(self): self.endpoint = "wss://openspeech.bytedance.com/api/v3/sami/podcasttts" self.appid = os.getenv("VOLCENGINE_PODCAST_APPID") self.access_token = os.getenv("VOLCENGINE_PODCAST_ACCESS_TOKEN") self.app_key = "aGjiRDfUWi" self.proxy = os.getenv("HTTPS_PROXY", None) if self.proxy: logger.info(f"volcengine podcast use proxy: {self.proxy}") async def podcast_request( self, session_id: str, data_manager=None, text: str = "", input_url: str = "", prompt_text: str = "", nlp_texts: str = "", action: int = 0, resource_id: str = "volc.service_type.10050", encoding: str = "mp3", input_id: str = "test_podcast", speaker_info: str = '{"random_order":false}', use_head_music: bool = False, use_tail_music: bool = False, only_nlp_text: bool = False, return_audio_url: bool = False, skip_round_audio_save: bool = False, on_round_complete: Optional[Callable] = None, ): """ 执行播客请求 Args: text: 输入文本 (action=0时使用) input_url: Web URL或文件URL (action=0时使用) prompt_text: 提示词文本 (action=4时必须) nlp_texts: NLP文本 (action=3时必须) action: 播客类型 (0/3/4) resource_id: 音频资源ID encoding: 音频格式 (mp3/wav) input_id: 唯一输入标识 speaker_info: 播客说话人信息 use_head_music: 是否使用开头音乐 use_tail_music: 是否使用结尾音乐 only_nlp_text: 是否只返回播客文本 return_audio_url: 是否返回音频URL skip_round_audio_save: 是否跳过单轮音频保存 output_dir: 输出目录 on_round_complete: 轮次完成回调函数 """ if not self.appid or not self.access_token: logger.error("APP ID or Access Key is required") return None, None headers = { "X-Api-App-Id": self.appid, "X-Api-App-Key": self.app_key, "X-Api-Access-Key": self.access_token, "X-Api-Resource-Id": resource_id, "X-Api-Connect-Id": str(uuid.uuid4()), } is_podcast_round_end = True audio_received = False last_round_id = -1 task_id = "" websocket = None retry_num = 5 audio = bytearray() voice = "" current_round = 0 podcast_texts = [] post_processor = PodcastRoundPostProcessor(session_id, data_manager) await post_processor.init() try: while retry_num > 0: # 建立WebSocket连接 websocket = await websockets.connect(self.endpoint, additional_headers=headers) logger.debug(f"WebSocket connected: {websocket.response.headers}") # 构建请求参数 if input_url: req_params = { "input_id": input_id, "nlp_texts": json.loads(nlp_texts) if nlp_texts else None, "prompt_text": prompt_text, "action": action, "use_head_music": use_head_music, "use_tail_music": use_tail_music, "input_info": { "input_url": input_url, "return_audio_url": return_audio_url, "only_nlp_text": only_nlp_text, }, "speaker_info": json.loads(speaker_info) if speaker_info else None, "audio_config": {"format": encoding, "sample_rate": 24000, "speech_rate": 0}, } else: req_params = { "input_id": input_id, "input_text": text, "nlp_texts": json.loads(nlp_texts) if nlp_texts else None, "prompt_text": prompt_text, "action": action, "use_head_music": use_head_music, "use_tail_music": use_tail_music, "input_info": { "input_url": input_url, "return_audio_url": return_audio_url, "only_nlp_text": only_nlp_text, }, "speaker_info": json.loads(speaker_info) if speaker_info else None, "audio_config": {"format": encoding, "sample_rate": 24000, "speech_rate": 0}, } logger.debug(f"Request params: {json.dumps(req_params, indent=2, ensure_ascii=False)}") if not is_podcast_round_end: req_params["retry_info"] = {"retry_task_id": task_id, "last_finished_round_id": last_round_id} # Start connection await start_connection(websocket) await wait_for_event(websocket, MsgType.FullServerResponse, EventType.ConnectionStarted) session_id = str(uuid.uuid4()) if not task_id: task_id = session_id # Start session await start_session(websocket, json.dumps(req_params).encode(), session_id) await wait_for_event(websocket, MsgType.FullServerResponse, EventType.SessionStarted) # Finish session await finish_session(websocket, session_id) while True: msg = await receive_message(websocket) # 音频数据块 if msg.type == MsgType.AudioOnlyServer and msg.event == EventType.PodcastRoundResponse: if not audio_received and audio: audio_received = True audio.extend(msg.payload) # 错误信息 elif msg.type == MsgType.Error: raise RuntimeError(f"Server error: {msg.payload.decode()}") elif msg.type == MsgType.FullServerResponse: # 播客 round 开始 if msg.event == EventType.PodcastRoundStart: data = json.loads(msg.payload.decode()) if data.get("text"): filtered_payload = {"text": data.get("text"), "speaker": data.get("speaker")} podcast_texts.append(filtered_payload) voice = data.get("speaker") current_round = data.get("round_id") if current_round == -1: voice = "head_music" if current_round == 9999: voice = "tail_music" is_podcast_round_end = False logger.debug(f"New round started: {data}") # 播客 round 结束 if msg.event == EventType.PodcastRoundEnd: data = json.loads(msg.payload.decode()) logger.debug(f"Podcast round end: {data}") if data.get("is_error"): break is_podcast_round_end = True last_round_id = current_round if audio: round_info = await post_processor.postprocess_round(current_round, voice, audio, podcast_texts) if on_round_complete: await on_round_complete(round_info) audio.clear() # 播客结束 if msg.event == EventType.PodcastEnd: data = json.loads(msg.payload.decode()) logger.info(f"Podcast end: {data}") # 会话结束 if msg.event == EventType.SessionFinished: break if not audio_received and not only_nlp_text: raise RuntimeError("No audio data received") # 保持连接 await finish_connection(websocket) await wait_for_event(websocket, MsgType.FullServerResponse, EventType.ConnectionFinished) # 播客结束, 保存最终音频文件 if is_podcast_round_end: podcast_info = await post_processor.postprocess_final() return podcast_info else: logger.error(f"Current podcast not finished, resuming from round {last_round_id}") retry_num -= 1 await asyncio.sleep(1) if websocket: await websocket.close() finally: await post_processor.cleanup() if websocket: await websocket.close() return None async def test(args): """ Podcast测试函数 Args: args: dict, 包含所有podcast参数 """ client = VolcEnginePodcastClient() # 设置默认参数 params = { "text": "", "input_url": "https://zhuanlan.zhihu.com/p/607822576", "prompt_text": "", "nlp_texts": "", "action": 0, "resource_id": "volc.service_type.10050", "encoding": "mp3", "input_id": "test_podcast", "speaker_info": '{"random_order":false}', "use_head_music": False, "use_tail_music": False, "only_nlp_text": False, "return_audio_url": True, "skip_round_audio_save": False, "output_dir": "output", } # 覆盖默认参数 if args: params.update(args) await client.podcast_request(**params) if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() parser.add_argument("--text", default="", help="Input text Use when action in [0]") parser.add_argument("--input_url", default="", help="Web url or file url Use when action in [0]") parser.add_argument("--prompt_text", default="", help="Input Prompt Text must not empty when action in [4]") parser.add_argument("--nlp_texts", default="", help="Input NLP Texts must not empty when action in [3]") parser.add_argument("--resource_id", default="volc.service_type.10050", help="Audio Resource ID") parser.add_argument("--encoding", default="mp3", choices=["mp3", "wav"], help="Audio format") parser.add_argument("--input_id", default="test_podcast", help="Unique input identifier") parser.add_argument("--speaker_info", default='{"random_order":false}', help="Podcast Speaker Info") parser.add_argument("--use_head_music", default=False, action="store_true", help="Enable head music") parser.add_argument("--use_tail_music", default=False, action="store_true", help="Enable tail music") parser.add_argument("--only_nlp_text", default=False, action="store_true", help="Enable only podcast text when action in [0, 4]") parser.add_argument("--return_audio_url", default=False, action="store_true", help="Enable return audio url that can download") parser.add_argument("--action", default=0, type=int, choices=[0, 3, 4], help="different podcast type") parser.add_argument("--skip_round_audio_save", default=False, action="store_true", help="skip round audio save") parser.add_argument("--output_dir", default="output", help="Output directory") args = parser.parse_args() kwargs = {k: v for k, v in vars(args).items() if v is not None and not (isinstance(v, bool) and not v)} asyncio.run(test(kwargs))