497 lines
18 KiB
Python
497 lines
18 KiB
Python
import uuid
|
||
import websockets
|
||
import time
|
||
import fastrand
|
||
import json
|
||
import asyncio
|
||
import os
|
||
import aiofiles
|
||
from datetime import datetime
|
||
from typing import Dict, Any, Optional as OptionalType
|
||
|
||
from app.constants.tts import APP_ID, TOKEN, SPEAKER
|
||
|
||
# 协议常量
|
||
PROTOCOL_VERSION = 0b0001
|
||
DEFAULT_HEADER_SIZE = 0b0001
|
||
FULL_CLIENT_REQUEST = 0b0001
|
||
AUDIO_ONLY_RESPONSE = 0b1011
|
||
FULL_SERVER_RESPONSE = 0b1001
|
||
ERROR_INFORMATION = 0b1111
|
||
MsgTypeFlagWithEvent = 0b100
|
||
NO_SERIALIZATION = 0b0000
|
||
JSON = 0b0001
|
||
COMPRESSION_NO = 0b0000
|
||
|
||
# 事件类型
|
||
EVENT_NONE = 0
|
||
EVENT_Start_Connection = 1
|
||
EVENT_FinishConnection = 2
|
||
EVENT_ConnectionStarted = 50
|
||
EVENT_StartSession = 100
|
||
EVENT_FinishSession = 102
|
||
EVENT_SessionStarted = 150
|
||
EVENT_SessionFinished = 152
|
||
EVENT_TaskRequest = 200
|
||
EVENT_TTSSentenceEnd = 351
|
||
EVENT_TTSResponse = 352
|
||
|
||
# 音频文件保存目录
|
||
TEMP_AUDIO_DIR = "./temp_audio"
|
||
|
||
|
||
# 确保音频目录存在
|
||
async def ensure_audio_dir():
|
||
"""异步创建音频目录"""
|
||
if not os.path.exists(TEMP_AUDIO_DIR):
|
||
os.makedirs(TEMP_AUDIO_DIR, exist_ok=True)
|
||
|
||
|
||
# 生成时间戳文件名
|
||
def generate_audio_filename() -> str:
|
||
"""生成基于时间戳的音频文件名"""
|
||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3] # 精确到毫秒
|
||
return f"{timestamp}.mp3"
|
||
|
||
|
||
# ... 保留所有原有的类定义和工具函数 ...
|
||
|
||
class Header:
|
||
def __init__(self,
|
||
protocol_version=PROTOCOL_VERSION,
|
||
header_size=DEFAULT_HEADER_SIZE,
|
||
message_type=0,
|
||
message_type_specific_flags=0,
|
||
serial_method=NO_SERIALIZATION,
|
||
compression_type=COMPRESSION_NO,
|
||
reserved_data=0):
|
||
self.header_size = header_size
|
||
self.protocol_version = protocol_version
|
||
self.message_type = message_type
|
||
self.message_type_specific_flags = message_type_specific_flags
|
||
self.serial_method = serial_method
|
||
self.compression_type = compression_type
|
||
self.reserved_data = reserved_data
|
||
|
||
def as_bytes(self) -> bytes:
|
||
return bytes([
|
||
(self.protocol_version << 4) | self.header_size,
|
||
(self.message_type << 4) | self.message_type_specific_flags,
|
||
(self.serial_method << 4) | self.compression_type,
|
||
self.reserved_data
|
||
])
|
||
|
||
|
||
class Optional:
|
||
def __init__(self, event=EVENT_NONE, sessionId=None, sequence=None):
|
||
self.event = event
|
||
self.sessionId = sessionId
|
||
self.errorCode = 0
|
||
self.connectionId = None
|
||
self.response_meta_json = None
|
||
self.sequence = sequence
|
||
|
||
def as_bytes(self) -> bytes:
|
||
option_bytes = bytearray()
|
||
if self.event != EVENT_NONE:
|
||
option_bytes.extend(self.event.to_bytes(4, "big", signed=True))
|
||
if self.sessionId is not None:
|
||
session_id_bytes = str.encode(self.sessionId)
|
||
size = len(session_id_bytes).to_bytes(4, "big", signed=True)
|
||
option_bytes.extend(size)
|
||
option_bytes.extend(session_id_bytes)
|
||
if self.sequence is not None:
|
||
option_bytes.extend(self.sequence.to_bytes(4, "big", signed=True))
|
||
return option_bytes
|
||
|
||
|
||
class Response:
|
||
def __init__(self, header, optional):
|
||
self.optional = optional
|
||
self.header = header
|
||
self.payload = None
|
||
self.payload_json = None
|
||
|
||
|
||
# 工具函数
|
||
def gen_log_id():
|
||
"""生成logID"""
|
||
ts = int(time.time() * 1000)
|
||
r = fastrand.pcg32bounded(1 << 24) + (1 << 20)
|
||
local_ip = "00000000000000000000000000000000"
|
||
return f"02{ts}{local_ip}{r:08x}"
|
||
|
||
|
||
def get_payload_bytes(uid='1234', event=EVENT_NONE, text='', speaker='', audio_format='mp3',
|
||
audio_sample_rate=24000):
|
||
return str.encode(json.dumps({
|
||
"user": {"uid": uid},
|
||
"event": event,
|
||
"namespace": "BidirectionalTTS",
|
||
"req_params": {
|
||
"text": text,
|
||
"speaker": speaker,
|
||
"audio_params": {
|
||
"format": audio_format,
|
||
"sample_rate": audio_sample_rate,
|
||
"enable_timestamp": True,
|
||
}
|
||
}
|
||
}))
|
||
|
||
|
||
def read_res_content(res, offset):
|
||
content_size = int.from_bytes(res[offset: offset + 4], "big", signed=True)
|
||
offset += 4
|
||
content = str(res[offset: offset + content_size], encoding='utf8')
|
||
offset += content_size
|
||
return content, offset
|
||
|
||
|
||
def read_res_payload(res, offset):
|
||
payload_size = int.from_bytes(res[offset: offset + 4], "big", signed=True)
|
||
offset += 4
|
||
payload = res[offset: offset + payload_size]
|
||
offset += payload_size
|
||
return payload, offset
|
||
|
||
|
||
def parser_response(res) -> Response:
|
||
if isinstance(res, str):
|
||
raise RuntimeError(res)
|
||
response = Response(Header(), Optional())
|
||
# 解析结果
|
||
header = response.header
|
||
num = 0b00001111
|
||
header.protocol_version = res[0] >> 4 & num
|
||
header.header_size = res[0] & 0x0f
|
||
header.message_type = (res[1] >> 4) & num
|
||
header.message_type_specific_flags = res[1] & 0x0f
|
||
header.serial_method = res[2] >> num
|
||
header.message_compression = res[2] & 0x0f
|
||
header.reserved_data = res[3]
|
||
|
||
offset = 4
|
||
optional = response.optional
|
||
if header.message_type == FULL_SERVER_RESPONSE or AUDIO_ONLY_RESPONSE:
|
||
# read event
|
||
if header.message_type_specific_flags == MsgTypeFlagWithEvent:
|
||
optional.event = int.from_bytes(res[offset:offset + 4], "big", signed=True)
|
||
offset += 4
|
||
if optional.event == EVENT_NONE:
|
||
return response
|
||
# read connectionId
|
||
elif optional.event == EVENT_ConnectionStarted:
|
||
optional.connectionId, offset = read_res_content(res, offset)
|
||
elif optional.event == EVENT_SessionStarted or optional.event == EVENT_SessionFinished:
|
||
optional.sessionId, offset = read_res_content(res, offset)
|
||
optional.response_meta_json, offset = read_res_content(res, offset)
|
||
elif optional.event == EVENT_TTSResponse:
|
||
optional.sessionId, offset = read_res_content(res, offset)
|
||
response.payload, offset = read_res_payload(res, offset)
|
||
elif optional.event == EVENT_TTSSentenceEnd:
|
||
optional.sessionId, offset = read_res_content(res, offset)
|
||
response.payload_json, offset = read_res_content(res, offset)
|
||
|
||
elif header.message_type == ERROR_INFORMATION:
|
||
optional.errorCode = int.from_bytes(res[offset:offset + 4], "big", signed=True)
|
||
offset += 4
|
||
response.payload, offset = read_res_payload(res, offset)
|
||
return response
|
||
|
||
|
||
async def send_event(ws, header, optional=None, payload=None):
|
||
full_client_request = bytearray(header)
|
||
if optional is not None:
|
||
full_client_request.extend(optional)
|
||
if payload is not None:
|
||
payload_size = len(payload).to_bytes(4, 'big', signed=True)
|
||
full_client_request.extend(payload_size)
|
||
full_client_request.extend(payload)
|
||
await ws.send(full_client_request)
|
||
|
||
|
||
# TTS状态管理类，添加消息ID和任务追踪
|
||
class TTSState:
|
||
def __init__(self, message_id: str):
|
||
self.message_id = message_id
|
||
self.volc_ws: OptionalType[websockets.WebSocketServerProtocol] = None
|
||
self.session_id: OptionalType[str] = None
|
||
self.task: OptionalType[asyncio.Task] = None # 用于追踪异步任务
|
||
self.is_processing = False
|
||
self.audio_data = bytearray() # 用于收集音频数据
|
||
self.audio_filename = None # 保存的文件名
|
||
|
||
|
||
# 全局状态管理
|
||
class TTSManager:
|
||
def __init__(self):
|
||
# WebSocket -> 消息ID -> TTS状态
|
||
self.connections: Dict[any, Dict[str, TTSState]] = {}
|
||
# 会话ID -> 消息ID 的映射，用于路由响应
|
||
self.session_to_message: Dict[str, str] = {}
|
||
# 消息ID -> WebSocket 的映射
|
||
self.message_to_websocket: Dict[str, any] = {}
|
||
|
||
def get_connection_states(self, websocket) -> Dict[str, TTSState]:
|
||
"""获取WebSocket连接的所有TTS状态"""
|
||
if websocket not in self.connections:
|
||
self.connections[websocket] = {}
|
||
return self.connections[websocket]
|
||
|
||
def add_tts_state(self, websocket, message_id: str) -> TTSState:
|
||
"""添加新的TTS状态"""
|
||
states = self.get_connection_states(websocket)
|
||
if message_id in states:
|
||
# 如果已存在，先清理旧的
|
||
self.cleanup_message_state(websocket, message_id)
|
||
|
||
tts_state = TTSState(message_id)
|
||
states[message_id] = tts_state
|
||
self.message_to_websocket[message_id] = websocket
|
||
return tts_state
|
||
|
||
def get_tts_state(self, websocket, message_id: str) -> OptionalType[TTSState]:
|
||
"""获取指定的TTS状态"""
|
||
states = self.get_connection_states(websocket)
|
||
return states.get(message_id)
|
||
|
||
def register_session(self, session_id: str, message_id: str):
|
||
"""注册会话ID和消息ID的映射"""
|
||
self.session_to_message[session_id] = message_id
|
||
|
||
def get_message_by_session(self, session_id: str) -> OptionalType[str]:
|
||
"""根据会话ID获取消息ID"""
|
||
return self.session_to_message.get(session_id)
|
||
|
||
def get_websocket_by_message(self, message_id: str):
|
||
"""根据消息ID获取WebSocket"""
|
||
return self.message_to_websocket.get(message_id)
|
||
|
||
def cleanup_message_state(self, websocket, message_id: str):
|
||
"""清理指定消息的状态"""
|
||
states = self.get_connection_states(websocket)
|
||
if message_id in states:
|
||
tts_state = states[message_id]
|
||
# 取消任务
|
||
if tts_state.task and not tts_state.task.done():
|
||
tts_state.task.cancel()
|
||
# 清理映射
|
||
if tts_state.session_id and tts_state.session_id in self.session_to_message:
|
||
del self.session_to_message[tts_state.session_id]
|
||
if message_id in self.message_to_websocket:
|
||
del self.message_to_websocket[message_id]
|
||
# 删除状态
|
||
del states[message_id]
|
||
|
||
def cleanup_connection(self, websocket):
|
||
"""清理整个连接的状态"""
|
||
if websocket in self.connections:
|
||
states = self.connections[websocket]
|
||
for message_id in list(states.keys()):
|
||
self.cleanup_message_state(websocket, message_id)
|
||
del self.connections[websocket]
|
||
|
||
|
||
# 全局TTS管理器实例
|
||
tts_manager = TTSManager()
|
||
|
||
|
||
# 初始化独立的TTS连接
|
||
async def create_tts_connection() -> websockets.WebSocketServerProtocol:
|
||
"""创建独立的TTS连接"""
|
||
log_id = gen_log_id()
|
||
ws_header = {
|
||
"X-Api-App-Key": APP_ID,
|
||
"X-Api-Access-Key": TOKEN,
|
||
"X-Api-Resource-Id": 'volc.service_type.10029',
|
||
"X-Api-Connect-Id": str(uuid.uuid4()),
|
||
"X-Tt-Logid": log_id,
|
||
}
|
||
url = 'wss://openspeech.bytedance.com/api/v3/tts/bidirection'
|
||
volc_ws = await websockets.connect(url, additional_headers=ws_header, max_size=1000000000)
|
||
|
||
# 启动连接
|
||
header = Header(message_type=FULL_CLIENT_REQUEST,
|
||
message_type_specific_flags=MsgTypeFlagWithEvent).as_bytes()
|
||
optional = Optional(event=EVENT_Start_Connection).as_bytes()
|
||
payload = str.encode("{}")
|
||
await send_event(volc_ws, header, optional, payload)
|
||
|
||
# 等待连接确认
|
||
raw_data = await volc_ws.recv()
|
||
res = parser_response(raw_data)
|
||
if res.optional.event != EVENT_ConnectionStarted:
|
||
raise Exception("TTS连接失败")
|
||
|
||
return volc_ws
|
||
|
||
|
||
# 异步保存音频文件
|
||
async def save_audio_file(audio_data: bytes, filename: str) -> str:
|
||
"""异步保存音频文件"""
|
||
await ensure_audio_dir()
|
||
file_path = os.path.join(TEMP_AUDIO_DIR, filename)
|
||
|
||
async with aiofiles.open(file_path, 'wb') as f:
|
||
await f.write(audio_data)
|
||
|
||
return file_path
|
||
|
||
|
||
# 处理单个TTS任务
|
||
async def process_tts_task(websocket, message_id: str, text: str, speaker: str = None):
|
||
"""处理单个TTS任务（独立协程）"""
|
||
tts_state = None
|
||
# 使用传入的speaker，如果没有则使用默认的
|
||
selected_speaker = speaker if speaker else SPEAKER
|
||
|
||
try:
|
||
print(f"开始处理TTS任务 [{message_id}]: {text}, 使用说话人: {selected_speaker}")
|
||
# 获取TTS状态
|
||
tts_state = tts_manager.get_tts_state(websocket, message_id)
|
||
if not tts_state:
|
||
raise Exception(f"找不到TTS状态: {message_id}")
|
||
tts_state.is_processing = True
|
||
# 生成音频文件名
|
||
tts_state.audio_filename = generate_audio_filename()
|
||
# 创建独立的TTS连接
|
||
tts_state.volc_ws = await create_tts_connection()
|
||
# 创建会话
|
||
tts_state.session_id = uuid.uuid4().__str__().replace('-', '')
|
||
tts_manager.register_session(tts_state.session_id, message_id)
|
||
print(f"创建TTS会话 [{message_id}]: {tts_state.session_id}")
|
||
header = Header(message_type=FULL_CLIENT_REQUEST,
|
||
message_type_specific_flags=MsgTypeFlagWithEvent,
|
||
serial_method=JSON).as_bytes()
|
||
optional = Optional(event=EVENT_StartSession, sessionId=tts_state.session_id).as_bytes()
|
||
# 使用选择的speaker
|
||
payload = get_payload_bytes(event=EVENT_StartSession, speaker=selected_speaker)
|
||
await send_event(tts_state.volc_ws, header, optional, payload)
|
||
raw_data = await tts_state.volc_ws.recv()
|
||
res = parser_response(raw_data)
|
||
if res.optional.event != EVENT_SessionStarted:
|
||
raise Exception("TTS会话启动失败")
|
||
print(f"TTS会话创建成功 [{message_id}]: {tts_state.session_id}")
|
||
# 发送文本到TTS服务
|
||
print(f"发送文本到TTS服务 [{message_id}]...")
|
||
header = Header(message_type=FULL_CLIENT_REQUEST,
|
||
message_type_specific_flags=MsgTypeFlagWithEvent,
|
||
serial_method=JSON).as_bytes()
|
||
optional = Optional(event=EVENT_TaskRequest, sessionId=tts_state.session_id).as_bytes()
|
||
# 使用选择的speaker
|
||
payload = get_payload_bytes(event=EVENT_TaskRequest, text=text, speaker=selected_speaker)
|
||
await send_event(tts_state.volc_ws, header, optional, payload)
|
||
# 接收TTS响应并发送到前端
|
||
print(f"开始接收TTS响应 [{message_id}]...")
|
||
audio_count = 0
|
||
try:
|
||
while True:
|
||
raw_data = await asyncio.wait_for(
|
||
tts_state.volc_ws.recv(),
|
||
timeout=30
|
||
)
|
||
res = parser_response(raw_data)
|
||
print(f"收到TTS事件 [{message_id}]: {res.optional.event}")
|
||
if res.optional.event == EVENT_TTSSentenceEnd:
|
||
print(f"句子结束事件 [{message_id}] - 直接完成")
|
||
break
|
||
elif res.optional.event == EVENT_SessionFinished:
|
||
print(f"收到会话结束事件 [{message_id}]")
|
||
break
|
||
elif res.optional.event == EVENT_TTSResponse:
|
||
audio_count += 1
|
||
print(f"收到音频数据 [{message_id}] #{audio_count}，大小: {len(res.payload)}")
|
||
|
||
# 收集音频数据
|
||
tts_state.audio_data.extend(res.payload)
|
||
|
||
# 发送音频数据到前端
|
||
await websocket.send_json({
|
||
"id": audio_count,
|
||
"type": "tts_audio_data",
|
||
"messageId": message_id,
|
||
"audioData": res.payload.hex() # 转为hex字符串
|
||
})
|
||
else:
|
||
print(f"未知TTS事件 [{message_id}]: {res.optional.event}")
|
||
except asyncio.TimeoutError:
|
||
print(f"TTS响应超时 [{message_id}]，强制结束")
|
||
# 异步保存音频文件
|
||
if tts_state.audio_data:
|
||
file_path = await save_audio_file(
|
||
bytes(tts_state.audio_data),
|
||
tts_state.audio_filename
|
||
)
|
||
print(f"音频文件已保存 [{message_id}]: {file_path}")
|
||
|
||
# 发送完成消息，包含文件路径和使用的speaker
|
||
await websocket.send_json({
|
||
"type": "tts_audio_complete",
|
||
"messageId": message_id,
|
||
"audioFile": tts_state.audio_filename,
|
||
"audioPath": os.path.join(TEMP_AUDIO_DIR, tts_state.audio_filename) if tts_state.audio_data else None,
|
||
"speaker": selected_speaker
|
||
})
|
||
print(f"TTS处理完成 [{message_id}]，共发送 {audio_count} 个音频包，使用说话人: {selected_speaker}")
|
||
except asyncio.CancelledError:
|
||
print(f"TTS任务被取消 [{message_id}]")
|
||
await websocket.send_json({
|
||
"type": "tts_error",
|
||
"messageId": message_id,
|
||
"message": "TTS任务被取消"
|
||
})
|
||
except Exception as e:
|
||
print(f"TTS处理异常 [{message_id}]: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
await websocket.send_json({
|
||
"type": "tts_error",
|
||
"messageId": message_id,
|
||
"message": f"TTS处理失败: {str(e)}"
|
||
})
|
||
finally:
|
||
# 清理资源
|
||
if tts_state:
|
||
tts_state.is_processing = False
|
||
if tts_state.volc_ws:
|
||
try:
|
||
await tts_state.volc_ws.close()
|
||
except:
|
||
pass
|
||
# 清理状态
|
||
tts_manager.cleanup_message_state(websocket, message_id)
|
||
|
||
|
||
# 启动TTS文本转换
|
||
async def handle_tts_text(websocket, message_id: str, text: str, speaker: str = None):
|
||
"""启动TTS文本转换"""
|
||
# 创建新的TTS状态
|
||
print(speaker)
|
||
tts_state = tts_manager.add_tts_state(websocket, message_id)
|
||
# 启动异步任务，传入speaker参数
|
||
tts_state.task = asyncio.create_task(
|
||
process_tts_task(websocket, message_id, text, speaker)
|
||
)
|
||
|
||
|
||
# 取消TTS任务
|
||
async def handle_tts_cancel(websocket, message_id: str):
|
||
"""取消TTS任务"""
|
||
tts_state = tts_manager.get_tts_state(websocket, message_id)
|
||
if tts_state and tts_state.task and not tts_state.task.done():
|
||
tts_state.task.cancel()
|
||
await websocket.send_json({
|
||
"type": "tts_complete",
|
||
"messageId": message_id
|
||
})
|
||
tts_manager.cleanup_message_state(websocket, message_id)
|
||
|
||
|
||
# 清理连接的所有TTS资源
|
||
async def cleanup_connection_tts(websocket):
|
||
"""清理连接的所有TTS资源"""
|
||
print(f"清理连接的TTS资源...")
|
||
tts_manager.cleanup_connection(websocket)
|
||
print("TTS资源清理完成")
|