470 lines
17 KiB
Python
470 lines
17 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
豆包音频处理模块
|
||
简化版WebSocket API，支持音频文件上传和返回音频处理
|
||
"""
|
||
|
||
import asyncio
|
||
import gzip
|
||
import json
|
||
import uuid
|
||
import wave
|
||
import struct
|
||
import time
|
||
import os
|
||
from typing import Dict, Any, Optional
|
||
import websockets
|
||
|
||
|
||
class DoubaoConfig:
|
||
"""豆包配置"""
|
||
|
||
def __init__(self):
|
||
self.base_url = "wss://openspeech.bytedance.com/api/v3/realtime/dialogue"
|
||
self.app_id = "8718217928"
|
||
self.access_key = "ynJMX-5ix1FsJvswC9KTNlGUdubcchqc"
|
||
self.app_key = "PlgvMymc7f3tQnJ6"
|
||
self.resource_id = "volc.speech.dialog"
|
||
|
||
def get_headers(self) -> Dict[str, str]:
|
||
"""获取请求头"""
|
||
return {
|
||
"X-Api-App-ID": self.app_id,
|
||
"X-Api-Access-Key": self.access_key,
|
||
"X-Api-Resource-Id": self.resource_id,
|
||
"X-Api-App-Key": self.app_key,
|
||
"X-Api-Connect-Id": str(uuid.uuid4()),
|
||
}
|
||
|
||
|
||
class DoubaoProtocol:
|
||
"""豆包协议处理"""
|
||
|
||
# 协议常量
|
||
PROTOCOL_VERSION = 0b0001
|
||
CLIENT_FULL_REQUEST = 0b0001
|
||
CLIENT_AUDIO_ONLY_REQUEST = 0b0010
|
||
SERVER_FULL_RESPONSE = 0b1001
|
||
SERVER_ACK = 0b1011
|
||
SERVER_ERROR_RESPONSE = 0b1111
|
||
|
||
NO_SEQUENCE = 0b0000
|
||
POS_SEQUENCE = 0b0001
|
||
MSG_WITH_EVENT = 0b0100
|
||
JSON = 0b0001
|
||
NO_SERIALIZATION = 0b0000
|
||
GZIP = 0b0001
|
||
|
||
@classmethod
|
||
def generate_header(cls, message_type=CLIENT_FULL_REQUEST,
|
||
message_type_specific_flags=MSG_WITH_EVENT,
|
||
serial_method=JSON, compression_type=GZIP) -> bytes:
|
||
"""生成协议头"""
|
||
header = bytearray()
|
||
header.append((cls.PROTOCOL_VERSION << 4) | 1) # version + header_size
|
||
header.append((message_type << 4) | message_type_specific_flags)
|
||
header.append((serial_method << 4) | compression_type)
|
||
header.append(0x00) # reserved
|
||
return bytes(header)
|
||
|
||
@classmethod
|
||
def parse_response(cls, res: bytes) -> Dict[str, Any]:
|
||
"""解析响应"""
|
||
if isinstance(res, str):
|
||
return {}
|
||
|
||
protocol_version = res[0] >> 4
|
||
header_size = res[0] & 0x0f
|
||
message_type = res[1] >> 4
|
||
message_type_specific_flags = res[1] & 0x0f
|
||
serialization_method = res[2] >> 4
|
||
message_compression = res[2] & 0x0f
|
||
|
||
payload = res[header_size * 4:]
|
||
result = {}
|
||
|
||
if message_type == cls.SERVER_FULL_RESPONSE or message_type == cls.SERVER_ACK:
|
||
result['message_type'] = 'SERVER_FULL_RESPONSE'
|
||
if message_type == cls.SERVER_ACK:
|
||
result['message_type'] = 'SERVER_ACK'
|
||
|
||
start = 0
|
||
if message_type_specific_flags & cls.MSG_WITH_EVENT:
|
||
result['event'] = int.from_bytes(payload[:4], "big", signed=False)
|
||
start += 4
|
||
|
||
payload = payload[start:]
|
||
session_id_size = int.from_bytes(payload[:4], "big", signed=True)
|
||
session_id = payload[4:session_id_size+4]
|
||
result['session_id'] = str(session_id)
|
||
payload = payload[4 + session_id_size:]
|
||
|
||
payload_size = int.from_bytes(payload[:4], "big", signed=False)
|
||
payload_msg = payload[4:]
|
||
result['payload_size'] = payload_size
|
||
|
||
if payload_msg:
|
||
if message_compression == cls.GZIP:
|
||
payload_msg = gzip.decompress(payload_msg)
|
||
if serialization_method == cls.JSON:
|
||
payload_msg = json.loads(str(payload_msg, "utf-8"))
|
||
result['payload_msg'] = payload_msg
|
||
|
||
elif message_type == cls.SERVER_ERROR_RESPONSE:
|
||
code = int.from_bytes(payload[:4], "big", signed=False)
|
||
result['code'] = code
|
||
payload_size = int.from_bytes(payload[4:8], "big", signed=False)
|
||
result['payload_size'] = payload_size
|
||
|
||
return result
|
||
|
||
|
||
class AudioProcessor:
|
||
"""音频处理器"""
|
||
|
||
@staticmethod
|
||
def read_wav_file(file_path: str) -> tuple:
|
||
"""读取WAV文件，返回音频数据和参数"""
|
||
with wave.open(file_path, 'rb') as wf:
|
||
# 获取音频参数
|
||
channels = wf.getnchannels()
|
||
sampwidth = wf.getsampwidth()
|
||
framerate = wf.getframerate()
|
||
nframes = wf.getnframes()
|
||
|
||
# 读取音频数据
|
||
audio_data = wf.readframes(nframes)
|
||
|
||
return audio_data, {
|
||
'channels': channels,
|
||
'sampwidth': sampwidth,
|
||
'framerate': framerate,
|
||
'nframes': nframes
|
||
}
|
||
|
||
@staticmethod
|
||
def create_wav_file(audio_data: bytes, output_path: str,
|
||
sample_rate: int = 24000, channels: int = 1,
|
||
sampwidth: int = 2) -> None:
|
||
"""创建WAV文件，适配树莓派播放"""
|
||
with wave.open(output_path, 'wb') as wf:
|
||
wf.setnchannels(channels)
|
||
wf.setsampwidth(sampwidth)
|
||
wf.setframerate(sample_rate)
|
||
wf.writeframes(audio_data)
|
||
|
||
|
||
class DoubaoClient:
|
||
"""豆包客户端"""
|
||
|
||
def __init__(self, config: DoubaoConfig):
|
||
self.config = config
|
||
self.session_id = str(uuid.uuid4())
|
||
self.ws = None
|
||
self.log_id = ""
|
||
|
||
async def connect(self) -> None:
|
||
"""建立WebSocket连接"""
|
||
print(f"连接豆包服务器: {self.config.base_url}")
|
||
|
||
self.ws = await websockets.connect(
|
||
self.config.base_url,
|
||
additional_headers=self.config.get_headers(),
|
||
ping_interval=None
|
||
)
|
||
|
||
# 获取log_id
|
||
if hasattr(self.ws, 'response_headers'):
|
||
self.log_id = self.ws.response_headers.get("X-Tt-Logid")
|
||
elif hasattr(self.ws, 'headers'):
|
||
self.log_id = self.ws.headers.get("X-Tt-Logid")
|
||
|
||
print(f"连接成功, log_id: {self.log_id}")
|
||
|
||
# 发送StartConnection请求
|
||
await self._send_start_connection()
|
||
|
||
# 发送StartSession请求
|
||
await self._send_start_session()
|
||
|
||
async def _send_start_connection(self) -> None:
|
||
"""发送StartConnection请求"""
|
||
request = bytearray(DoubaoProtocol.generate_header())
|
||
request.extend(int(1).to_bytes(4, 'big'))
|
||
|
||
payload_bytes = b"{}"
|
||
payload_bytes = gzip.compress(payload_bytes)
|
||
request.extend(len(payload_bytes).to_bytes(4, 'big'))
|
||
request.extend(payload_bytes)
|
||
|
||
await self.ws.send(request)
|
||
response = await self.ws.recv()
|
||
parsed_response = DoubaoProtocol.parse_response(response)
|
||
print(f"StartConnection响应: {parsed_response}")
|
||
|
||
async def _send_start_session(self) -> None:
|
||
"""发送StartSession请求"""
|
||
session_config = {
|
||
"asr": {
|
||
"extra": {
|
||
"end_smooth_window_ms": 1500,
|
||
},
|
||
},
|
||
"tts": {
|
||
"speaker": "zh_female_vv_jupiter_bigtts",
|
||
"audio_config": {
|
||
"channel": 1,
|
||
"format": "pcm",
|
||
"sample_rate": 24000
|
||
},
|
||
},
|
||
"dialog": {
|
||
"bot_name": "豆包",
|
||
"system_role": "你使用活泼灵动的女声，性格开朗，热爱生活。",
|
||
"speaking_style": "你的说话风格简洁明了，语速适中，语调自然。",
|
||
"location": {"city": "北京"},
|
||
"extra": {
|
||
"strict_audit": False,
|
||
"audit_response": "支持客户自定义安全审核回复话术。",
|
||
"recv_timeout": 10,
|
||
"input_mod": "audio_file", # 使用音频文件模式
|
||
},
|
||
},
|
||
}
|
||
|
||
request = bytearray(DoubaoProtocol.generate_header())
|
||
request.extend(int(100).to_bytes(4, 'big'))
|
||
request.extend(len(self.session_id).to_bytes(4, 'big'))
|
||
request.extend(self.session_id.encode())
|
||
|
||
payload_bytes = json.dumps(session_config).encode()
|
||
payload_bytes = gzip.compress(payload_bytes)
|
||
request.extend(len(payload_bytes).to_bytes(4, 'big'))
|
||
request.extend(payload_bytes)
|
||
|
||
await self.ws.send(request)
|
||
response = await self.ws.recv()
|
||
parsed_response = DoubaoProtocol.parse_response(response)
|
||
print(f"StartSession响应: {parsed_response}")
|
||
|
||
async def send_audio_file(self, file_path: str) -> bytes:
|
||
"""发送音频文件并返回响应音频"""
|
||
print(f"处理音频文件: {file_path}")
|
||
|
||
# 读取音频文件
|
||
audio_data, audio_info = AudioProcessor.read_wav_file(file_path)
|
||
print(f"音频参数: {audio_info}")
|
||
|
||
# 计算分块大小（200ms）
|
||
chunk_size = int(audio_info['framerate'] * audio_info['channels'] *
|
||
audio_info['sampwidth'] * 0.2)
|
||
|
||
# 分块发送音频数据
|
||
total_chunks = (len(audio_data) + chunk_size - 1) // chunk_size
|
||
print(f"开始分块发送音频，共 {total_chunks} 块")
|
||
|
||
received_audio = b""
|
||
error_count = 0
|
||
|
||
for i in range(0, len(audio_data), chunk_size):
|
||
chunk = audio_data[i:i + chunk_size]
|
||
is_last = (i + chunk_size >= len(audio_data))
|
||
|
||
# 发送音频块
|
||
await self._send_audio_chunk(chunk, is_last)
|
||
|
||
# 接收响应
|
||
try:
|
||
response = await asyncio.wait_for(self.ws.recv(), timeout=2.0)
|
||
parsed_response = DoubaoProtocol.parse_response(response)
|
||
|
||
# 检查是否是错误响应
|
||
if 'code' in parsed_response and parsed_response['code'] != 0:
|
||
print(f"服务器返回错误: {parsed_response}")
|
||
error_count += 1
|
||
if error_count > 3:
|
||
raise Exception(f"服务器连续返回错误: {parsed_response}")
|
||
continue
|
||
|
||
# 处理音频响应
|
||
if (parsed_response.get('message_type') == 'SERVER_ACK' and
|
||
isinstance(parsed_response.get('payload_msg'), bytes)):
|
||
audio_chunk = parsed_response['payload_msg']
|
||
received_audio += audio_chunk
|
||
print(f"接收到音频数据块，大小: {len(audio_chunk)} 字节")
|
||
|
||
# 检查会话状态
|
||
event = parsed_response.get('event')
|
||
if event in [359, 152, 153]: # 这些事件表示会话相关状态
|
||
print(f"会话事件: {event}")
|
||
if event in [152, 153]: # 会话结束
|
||
print("检测到会话结束事件")
|
||
break
|
||
|
||
except asyncio.TimeoutError:
|
||
print("等待响应超时，继续发送")
|
||
|
||
# 模拟实时发送的延迟
|
||
await asyncio.sleep(0.05)
|
||
|
||
print("音频文件发送完成")
|
||
return received_audio
|
||
|
||
async def _send_audio_chunk(self, audio_data: bytes, is_last: bool = False) -> None:
|
||
"""发送音频块"""
|
||
request = bytearray(
|
||
DoubaoProtocol.generate_header(
|
||
message_type=DoubaoProtocol.CLIENT_AUDIO_ONLY_REQUEST,
|
||
message_type_specific_flags=DoubaoProtocol.NO_SEQUENCE,
|
||
serial_method=DoubaoProtocol.NO_SERIALIZATION, # 音频数据不需要序列化
|
||
compression_type=DoubaoProtocol.GZIP
|
||
)
|
||
)
|
||
request.extend(int(200).to_bytes(4, 'big'))
|
||
request.extend(len(self.session_id).to_bytes(4, 'big'))
|
||
request.extend(self.session_id.encode())
|
||
|
||
# 压缩音频数据
|
||
compressed_audio = gzip.compress(audio_data)
|
||
request.extend(len(compressed_audio).to_bytes(4, 'big')) # payload size(4 bytes)
|
||
request.extend(compressed_audio)
|
||
|
||
await self.ws.send(request)
|
||
|
||
async def close(self) -> None:
|
||
"""关闭连接"""
|
||
if self.ws:
|
||
try:
|
||
# 发送FinishSession
|
||
await self._send_finish_session()
|
||
|
||
# 发送FinishConnection
|
||
await self._send_finish_connection()
|
||
|
||
except Exception as e:
|
||
print(f"关闭会话时出错: {e}")
|
||
finally:
|
||
# 确保WebSocket连接关闭
|
||
try:
|
||
await self.ws.close()
|
||
except:
|
||
pass
|
||
print("连接已关闭")
|
||
|
||
async def _send_finish_session(self) -> None:
|
||
"""发送FinishSession请求"""
|
||
request = bytearray(DoubaoProtocol.generate_header())
|
||
request.extend(int(102).to_bytes(4, 'big'))
|
||
request.extend(len(self.session_id).to_bytes(4, 'big'))
|
||
request.extend(self.session_id.encode())
|
||
|
||
payload_bytes = b"{}"
|
||
payload_bytes = gzip.compress(payload_bytes)
|
||
request.extend(len(payload_bytes).to_bytes(4, 'big'))
|
||
request.extend(payload_bytes)
|
||
|
||
await self.ws.send(request)
|
||
|
||
async def _send_finish_connection(self) -> None:
|
||
"""发送FinishConnection请求"""
|
||
request = bytearray(DoubaoProtocol.generate_header())
|
||
request.extend(int(2).to_bytes(4, 'big'))
|
||
|
||
payload_bytes = b"{}"
|
||
payload_bytes = gzip.compress(payload_bytes)
|
||
request.extend(len(payload_bytes).to_bytes(4, 'big'))
|
||
request.extend(payload_bytes)
|
||
|
||
await self.ws.send(request)
|
||
response = await self.ws.recv()
|
||
parsed_response = DoubaoProtocol.parse_response(response)
|
||
print(f"FinishConnection响应: {parsed_response}")
|
||
|
||
|
||
class DoubaoProcessor:
|
||
"""豆包音频处理器"""
|
||
|
||
def __init__(self):
|
||
self.config = DoubaoConfig()
|
||
self.client = DoubaoClient(self.config)
|
||
|
||
async def process_audio_file(self, input_file: str, output_file: str = None) -> str:
|
||
"""处理音频文件
|
||
|
||
Args:
|
||
input_file: 输入音频文件路径
|
||
output_file: 输出音频文件路径，如果为None则自动生成
|
||
|
||
Returns:
|
||
输出音频文件路径
|
||
"""
|
||
if not os.path.exists(input_file):
|
||
raise FileNotFoundError(f"音频文件不存在: {input_file}")
|
||
|
||
# 生成输出文件名
|
||
if output_file is None:
|
||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||
output_file = f"doubao_response_{timestamp}.wav"
|
||
|
||
try:
|
||
# 连接豆包服务器
|
||
await self.client.connect()
|
||
|
||
# 等待一会确保会话建立
|
||
await asyncio.sleep(0.5)
|
||
|
||
# 发送音频文件并获取响应
|
||
received_audio = await self.client.send_audio_file(input_file)
|
||
|
||
if received_audio:
|
||
print(f"总共接收到音频数据: {len(received_audio)} 字节")
|
||
# 转换为WAV格式保存（适配树莓派播放）
|
||
AudioProcessor.create_wav_file(
|
||
received_audio,
|
||
output_file,
|
||
sample_rate=24000, # 豆包返回的音频采样率
|
||
channels=1,
|
||
sampwidth=2 # 16-bit
|
||
)
|
||
print(f"响应音频已保存到: {output_file}")
|
||
|
||
# 显示文件信息
|
||
file_size = os.path.getsize(output_file)
|
||
print(f"输出文件大小: {file_size} 字节")
|
||
|
||
else:
|
||
print("警告: 未接收到音频响应")
|
||
|
||
return output_file
|
||
|
||
except Exception as e:
|
||
print(f"处理音频文件时出错: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
raise
|
||
finally:
|
||
await self.client.close()
|
||
|
||
|
||
async def main():
|
||
"""测试函数"""
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description="豆包音频处理测试")
|
||
parser.add_argument("--input", type=str, required=True, help="输入音频文件路径")
|
||
parser.add_argument("--output", type=str, help="输出音频文件路径")
|
||
|
||
args = parser.parse_args()
|
||
|
||
processor = DoubaoProcessor()
|
||
try:
|
||
output_file = await processor.process_audio_file(args.input, args.output)
|
||
print(f"处理完成，输出文件: {output_file}")
|
||
except Exception as e:
|
||
print(f"处理失败: {e}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main()) |