246 lines
8.3 KiB
Python
246 lines
8.3 KiB
Python
# asr_client.py
|
||
import asyncio
|
||
import json
|
||
import websockets
|
||
from typing import Optional, List, Dict, Callable, Awaitable
|
||
from dataclasses import dataclass, field
|
||
|
||
# 音频参数
|
||
AUDIO_PARAMS = {
|
||
"sample_rate": 16000,
|
||
"channels": 1,
|
||
"sample_width": 2,
|
||
"frame_size": 1024
|
||
}
|
||
|
||
# ASR 服务配置
|
||
ASR_CONFIG = {
|
||
"host": "10.10.10.202",
|
||
"port": 10096,
|
||
"mode": "2pass",
|
||
"chunk_size": [5, 10, 5],
|
||
"chunk_interval": 10,
|
||
"use_itn": 1,
|
||
"hotwords": "",
|
||
"reconnect_max_times": 3,
|
||
"pool_size": 5,
|
||
"audio_queue_size": 10000
|
||
}
|
||
|
||
# 定义回调函数类型（异步函数，接收 ASR 结果字典）
|
||
ASRResultCallback = Callable[[Dict], Awaitable[None]]
|
||
|
||
|
||
@dataclass
|
||
class ASRConnection:
|
||
"""ASR 连接对象（内置音频队列）"""
|
||
ws: Optional[websockets.WebSocketClientProtocol] = None
|
||
is_busy: bool = False
|
||
is_alive: bool = False
|
||
reconnect_count: int = 0
|
||
audio_queue: asyncio.Queue = field(default_factory=lambda: asyncio.Queue(maxsize=ASR_CONFIG["audio_queue_size"]))
|
||
stop_event: asyncio.Event = field(default_factory=asyncio.Event)
|
||
|
||
|
||
# 全局连接池
|
||
asr_connection_pool: List[ASRConnection] = []
|
||
pool_lock = asyncio.Lock()
|
||
|
||
|
||
async def create_asr_connection() -> Optional[ASRConnection]:
|
||
"""创建单个 ASR 连接并初始化"""
|
||
asr_conn = ASRConnection()
|
||
asr_uri = f"ws://{ASR_CONFIG['host']}:{ASR_CONFIG['port']}"
|
||
|
||
try:
|
||
ws = await websockets.connect(
|
||
asr_uri,
|
||
subprotocols=["binary"],
|
||
ping_interval=None,
|
||
open_timeout=5
|
||
)
|
||
asr_conn.ws = ws
|
||
asr_conn.is_alive = True
|
||
|
||
# 发送初始化配置
|
||
init_msg = json.dumps({
|
||
"mode": ASR_CONFIG["mode"],
|
||
"chunk_size": ASR_CONFIG["chunk_size"],
|
||
"chunk_interval": ASR_CONFIG["chunk_interval"],
|
||
"wav_name": "pool_connection",
|
||
"is_speaking": True,
|
||
"hotwords": ASR_CONFIG["hotwords"],
|
||
"itn": bool(ASR_CONFIG["use_itn"]),
|
||
"audio_fs": AUDIO_PARAMS["sample_rate"]
|
||
})
|
||
await ws.send(init_msg)
|
||
print("ASR 连接初始化成功")
|
||
return asr_conn
|
||
|
||
except Exception as e:
|
||
print(f"创建 ASR 连接失败：{e}")
|
||
asr_conn.is_alive = False
|
||
return asr_conn
|
||
|
||
|
||
async def init_asr_pool():
|
||
"""初始化 ASR 连接池"""
|
||
global asr_connection_pool
|
||
print(f"开始初始化 ASR 连接池，大小：{ASR_CONFIG['pool_size']}")
|
||
|
||
tasks = [create_asr_connection() for _ in range(ASR_CONFIG["pool_size"])]
|
||
connections = await asyncio.gather(*tasks)
|
||
asr_connection_pool = [conn for conn in connections if conn.is_alive]
|
||
print(f"ASR 连接池初始化完成，有效连接数：{len(asr_connection_pool)}")
|
||
|
||
|
||
async def get_idle_asr_connection() -> Optional[ASRConnection]:
|
||
"""从连接池获取空闲连接"""
|
||
async with pool_lock:
|
||
idle_conns = [
|
||
conn for conn in asr_connection_pool
|
||
if not conn.is_busy and conn.is_alive
|
||
]
|
||
if idle_conns:
|
||
conn = idle_conns[0]
|
||
conn.is_busy = True
|
||
conn.stop_event.clear()
|
||
return conn
|
||
|
||
if len(asr_connection_pool) < ASR_CONFIG["pool_size"]:
|
||
new_conn = await create_asr_connection()
|
||
if new_conn.is_alive:
|
||
new_conn.is_busy = True
|
||
asr_connection_pool.append(new_conn)
|
||
return new_conn
|
||
|
||
print("ASR 连接池无空闲连接")
|
||
return None
|
||
|
||
|
||
async def push_audio_data(asr_conn: ASRConnection, audio_data: bytes) -> bool:
|
||
"""插入音频数据到 ASR 内置队列"""
|
||
if not asr_conn or not asr_conn.is_alive or asr_conn.stop_event.is_set():
|
||
return False
|
||
try:
|
||
asr_conn.audio_queue.put_nowait(audio_data)
|
||
return True
|
||
except asyncio.QueueFull:
|
||
print("ASR 音频队列已满，丢弃当前音频帧")
|
||
return False
|
||
|
||
|
||
async def release_asr_connection(conn: ASRConnection):
|
||
"""释放 ASR 连接"""
|
||
async with pool_lock:
|
||
conn.is_busy = False
|
||
conn.stop_event.set()
|
||
# 清空队列
|
||
while not conn.audio_queue.empty():
|
||
try:
|
||
conn.audio_queue.get_nowait()
|
||
except asyncio.QueueEmpty:
|
||
break
|
||
# 重连逻辑
|
||
if not conn.is_alive and conn.reconnect_count < ASR_CONFIG["reconnect_max_times"]:
|
||
print(f"尝试重连 ASR 连接（次数：{conn.reconnect_count + 1}）")
|
||
new_conn = await create_asr_connection()
|
||
if new_conn.is_alive:
|
||
idx = asr_connection_pool.index(conn)
|
||
asr_connection_pool[idx] = new_conn
|
||
else:
|
||
conn.reconnect_count += 1
|
||
elif conn.reconnect_count >= ASR_CONFIG["reconnect_max_times"]:
|
||
asr_connection_pool.remove(conn)
|
||
print("ASR 连接重连次数耗尽，已移除")
|
||
|
||
|
||
async def handle_asr_communication(
|
||
asr_conn: ASRConnection,
|
||
result_callback: ASRResultCallback # 替换为回调函数
|
||
):
|
||
"""
|
||
处理 ASR 通信（结果通过回调函数返回）
|
||
:param asr_conn: ASR 连接对象
|
||
:param result_callback: 异步回调函数，接收 ASR 结果字典
|
||
"""
|
||
if not asr_conn or not asr_conn.ws:
|
||
# 错误结果通过回调返回
|
||
await result_callback({"error": "无可用 ASR 连接", "text": ""})
|
||
return
|
||
|
||
# 发送音频到 ASR 服务
|
||
async def send_audio():
|
||
while not asr_conn.stop_event.is_set() and asr_conn.is_alive:
|
||
try:
|
||
pcm_data = await asyncio.wait_for(asr_conn.audio_queue.get(), timeout=1.0)
|
||
if pcm_data and asr_conn.is_alive:
|
||
await asr_conn.ws.send(pcm_data)
|
||
await asyncio.sleep(0.005)
|
||
except asyncio.TimeoutError:
|
||
continue
|
||
except Exception as e:
|
||
print(f"发送音频到 ASR 失败：{e}")
|
||
asr_conn.is_alive = False
|
||
await result_callback({"error": f"音频发送失败：{str(e)}", "text": ""})
|
||
asr_conn.stop_event.set()
|
||
break
|
||
|
||
# 接收 ASR 结果并调用回调
|
||
async def recv_result():
|
||
while not asr_conn.stop_event.is_set() and asr_conn.is_alive:
|
||
try:
|
||
asr_result = await asr_conn.ws.recv()
|
||
result_json = json.loads(asr_result)
|
||
print(result_json.get("text", ""))
|
||
if result_json.get("timestamp", "") == '':
|
||
continue
|
||
result = {
|
||
"text": result_json.get("text", ""),
|
||
"mode": result_json.get("mode", ""),
|
||
"timestamp": result_json.get("timestamp", ""),
|
||
"is_final": result_json.get("is_final", False),
|
||
"error": ""
|
||
}
|
||
# 调用回调函数，传递结果
|
||
await result_callback(result)
|
||
except websockets.exceptions.ConnectionClosed:
|
||
print("ASR 连接已关闭")
|
||
asr_conn.is_alive = False
|
||
await result_callback({"error": "ASR 连接断开", "text": ""})
|
||
asr_conn.stop_event.set()
|
||
break
|
||
except Exception as e:
|
||
print(f"接收 ASR 结果失败：{e}")
|
||
asr_conn.is_alive = False
|
||
await result_callback({"error": f"接收结果失败：{str(e)}", "text": ""})
|
||
asr_conn.stop_event.set()
|
||
break
|
||
|
||
try:
|
||
send_task = asyncio.create_task(send_audio())
|
||
recv_task = asyncio.create_task(recv_result())
|
||
await asyncio.gather(send_task, recv_task)
|
||
finally:
|
||
send_task.cancel()
|
||
recv_task.cancel()
|
||
try:
|
||
await send_task
|
||
await recv_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
await release_asr_connection(asr_conn)
|
||
|
||
|
||
async def close_asr_pool():
|
||
"""关闭所有 ASR 连接"""
|
||
async with pool_lock:
|
||
for conn in asr_connection_pool:
|
||
conn.stop_event.set()
|
||
# if conn.ws and not conn.ws.closed:
|
||
# try:
|
||
# await conn.ws.close()
|
||
# print("ASR 连接已关闭")
|
||
# except Exception as e:
|
||
# print(f"关闭 ASR 连接失败：{e}")
|
||
asr_connection_pool.clear() |