x
This commit is contained in:
Home
@@ -0,0 +1,246 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user