x
This commit is contained in:
Home
@@ -0,0 +1,319 @@
|
||||
from fastapi import WebSocket, WebSocketDisconnect
|
||||
from typing import Dict, List, Optional, Callable,Any
|
||||
from dataclasses import dataclass, field
|
||||
import asyncio
|
||||
import uuid
|
||||
import logging
|
||||
from audio_ai_chat.config.logger import logger
|
||||
from audio_ai_chat.core.asr.asr_manager import ASRManager
|
||||
from audio_ai_chat.core.connection import ConnectionManager, ConnectionContext
|
||||
from audio_ai_chat.codec.ProtocolCodec import ProtocolCodec, MessageType
|
||||
from audio_ai_chat.utils.exceptions import ServiceCallError
|
||||
from audio_ai_chat.core.llm.dify.dify import LLMConversation, llm_client
|
||||
from audio_ai_chat.core.tts.tts_client import TTSManager
|
||||
from functools import partial
|
||||
|
||||
# 全局WebSocket连接管理器（单例模式，确保全局统一）
|
||||
class WebSocketConnectionManager:
|
||||
_instance: Optional["WebSocketConnectionManager"] = None
|
||||
|
||||
def __init__(self):
|
||||
self.connection_manager = None
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance.active_connections: List[WebSocket] = []
|
||||
cls._instance.client_context_map: Dict[str, ConnectionContext] = {}
|
||||
cls._instance.connection_manager: Optional[ConnectionManager] = None
|
||||
cls._instance.consume_wakeup = asyncio.Event()
|
||||
return cls._instance
|
||||
|
||||
async def initialize(self):
|
||||
"""初始化：获取ConnectionManager全局单例"""
|
||||
self.connection_manager = await ConnectionManager.get_instance()
|
||||
logger.info("WebSocketConnectionManager 初始化成功")
|
||||
|
||||
async def connect(self, client_id: str, websocket: WebSocket) -> ConnectionContext:
|
||||
"""建立连接+身份校验"""
|
||||
await websocket.accept()
|
||||
self.active_connections.append(websocket)
|
||||
logger.info(
|
||||
f"连接 {client_id} 已接受，等待身份信息（5秒超时），当前连接数: {len(self.active_connections)}"
|
||||
)
|
||||
|
||||
# 超时接收身份包
|
||||
try:
|
||||
ping_packet = await asyncio.wait_for(websocket.receive_bytes(), timeout=5.0)
|
||||
except asyncio.TimeoutError:
|
||||
error_msg = f"连接 {client_id} 身份校验超时"
|
||||
logger.warning(error_msg)
|
||||
error_packet = ProtocolCodec.pack(
|
||||
MessageType.ERROR, {"code": 1008, "message": "身份校验超时，请重试"}
|
||||
)
|
||||
await websocket.send_bytes(error_packet)
|
||||
raise TimeoutError(error_msg)
|
||||
|
||||
# 解包并验证包类型
|
||||
msg_type, _, identity_data = ProtocolCodec.unpack(ping_packet)
|
||||
if msg_type != MessageType.IDENTITY:
|
||||
error_msg = f"连接 {client_id} 首个包类型错误"
|
||||
logger.error(error_msg)
|
||||
error_packet = ProtocolCodec.pack(
|
||||
MessageType.ERROR, {"code": 4001, "message": "非法请求：首个包必须是身份校验包"}
|
||||
)
|
||||
await websocket.send_bytes(error_packet)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# 校验身份信息
|
||||
user_id = identity_data.get("user_id")
|
||||
token = identity_data.get("token")
|
||||
name = identity_data.get("name") or f"用户{user_id}"
|
||||
if not all([user_id, token]):
|
||||
error_msg = f"连接 {client_id} 身份信息不完整"
|
||||
logger.error(error_msg)
|
||||
error_packet = ProtocolCodec.pack(
|
||||
MessageType.ERROR, {"code": 4003, "message": "身份信息不完整：必须包含user_id和token"}
|
||||
)
|
||||
await websocket.send_bytes(error_packet)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# 创建/获取连接上下文
|
||||
context = await self.connection_manager.create_or_reconnect_context(
|
||||
new_client_id=client_id, user_id=user_id
|
||||
)
|
||||
context.set_user_info(token, user_id, name)
|
||||
self.client_context_map[client_id] = context
|
||||
|
||||
# 响应身份校验成功
|
||||
success_packet = ProtocolCodec.pack(
|
||||
MessageType.IDENTITY,
|
||||
{
|
||||
"code": 200,
|
||||
"message": "身份校验成功，连接已就绪",
|
||||
"data": {"client_id": client_id, "user_id": user_id, "name": name}
|
||||
}
|
||||
)
|
||||
await websocket.send_bytes(success_packet)
|
||||
logger.info(f"用户 {user_id}（{name}）身份校验通过（client_id: {client_id}）")
|
||||
return context
|
||||
|
||||
|
||||
def _create_asr_callback(self, context: ConnectionContext) -> Callable[[dict], None]:
|
||||
"""
|
||||
闭包：为当前连接创建专属的ASR回调函数
|
||||
回调内部持有ConnectionContext引用，直接操作其消息队列
|
||||
"""
|
||||
async def asr_result_callback(result: dict):
|
||||
"""专属回调：将ASR结果打包后插入当前连接的消息队列"""
|
||||
try:
|
||||
print('result', 'result', result)
|
||||
# 处理错误结果
|
||||
if result.get("error"):
|
||||
logger.error(f"ASR错误（client_id: {context.client_id}）：{result['error']}")
|
||||
# 打包错误消息
|
||||
error_packet = ProtocolCodec.pack(
|
||||
MessageType.ERROR,
|
||||
{"code": 5001, "message": f"ASR服务错误：{result['error']}"}
|
||||
)
|
||||
context.message_queue.put_nowait(error_packet)
|
||||
else:
|
||||
# 3. 调用Dify流式接口（传入ASR元数据，用于存储到对话历史）
|
||||
print('context', context.user_id)
|
||||
final_asr_text = result.get("text", "")
|
||||
# 2. 调用大模型（异步）
|
||||
if final_asr_text:
|
||||
asyncio.create_task(
|
||||
self.call_llm_and_send(
|
||||
context=context,
|
||||
query=final_asr_text,
|
||||
conversation=context.llm_session
|
||||
)
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"ASR结果入队（client_id: {context.client_id}）："
|
||||
f"文本={result['text']}，最终结果={result['is_final']}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"ASR回调处理失败（client_id: {context.client_id}）：{str(e)}")
|
||||
|
||||
return asr_result_callback
|
||||
|
||||
# ====================== 调用大模型 ======================
|
||||
# ====================== 大模型流式回调 ======================
|
||||
@staticmethod
|
||||
async def llm_stream_callback(context, chunk: str, conversation_id: str, is_finished: bool):
|
||||
"""大模型流式回调（纯异步，无阻塞）"""
|
||||
if not chunk:
|
||||
return
|
||||
print('大模型流式回调', chunk)
|
||||
req_id = await context.tts_client.synthesize(chunk)
|
||||
print('req_id', req_id)
|
||||
|
||||
async def call_llm_and_send(self,context ,query: str, conversation: LLMConversation):
|
||||
"""调用大模型，流式结果转发前端 + TTS"""
|
||||
if not query:
|
||||
return
|
||||
logger.info(f"调用大模型 - 用户(): {query}")
|
||||
|
||||
try:
|
||||
|
||||
stream_callback = partial(WebSocketConnectionManager.llm_stream_callback, context)
|
||||
|
||||
conv_id, full_reply = await llm_client.send_message(
|
||||
query=query,
|
||||
conversation=conversation,
|
||||
stream_callback=stream_callback, # 传递绑定后的回调
|
||||
response_mode="streaming"
|
||||
)
|
||||
logger.info(f"大模型回复完成 - 会话ID: {conv_id}, 完整回复: {full_reply}")
|
||||
except Exception as e:
|
||||
logger.error(f"大模型调用失败: {str(e)}")
|
||||
|
||||
|
||||
|
||||
|
||||
async def handle_connection(self, websocket: WebSocket):
|
||||
"""处理单个WebSocket连接的完整生命周期"""
|
||||
client_id = str(id(websocket)) # 生成唯一连接ID
|
||||
logger.info(f"新WebSocket连接：client_id={client_id}")
|
||||
context: Optional[ConnectionContext] = None
|
||||
asr_conn = None
|
||||
llm_conn = None # 新增：LLM连接变量
|
||||
try:
|
||||
# 1. 建立连接并获取上下文
|
||||
context = await self.connect(client_id, websocket)
|
||||
if not context:
|
||||
logger.error(f"连接 {client_id} 上下文创建失败")
|
||||
return
|
||||
|
||||
# 2. 获取ASR连接和专属回调
|
||||
asr_client = ASRManager.get_instance()
|
||||
asr_conn = await asr_client.get_connection()
|
||||
if not asr_conn:
|
||||
raise ServiceCallError("获取ASR连接失败")
|
||||
# 创建当前连接的专属ASR回调（闭包绑定context）
|
||||
asr_callback = self._create_asr_callback(context)
|
||||
|
||||
# 3. 启动ASR通信任务（传入专属回调）
|
||||
communication_task = asyncio.create_task(
|
||||
asr_client.start_communication(asr_conn, asr_callback)
|
||||
)
|
||||
|
||||
context.llm_session = LLMConversation(
|
||||
user_id=context.user_id,
|
||||
scene_description="语音识别对话场景"
|
||||
)
|
||||
|
||||
context.tts_client = TTSManager()
|
||||
def handle_tts_result(context, req_id: str, result: Dict[str, Any]):
|
||||
"""处理TTS结果回调"""
|
||||
status = result.get("status")
|
||||
print('处理TTS结果回调')
|
||||
if status == "completed":
|
||||
audio_data = result.get("audio_data")
|
||||
pack_data = ProtocolCodec.pack(MessageType.AUDIO_DATA, audio_data)
|
||||
context.message_queue.put_nowait(pack_data)
|
||||
print('插入', len(audio_data))
|
||||
|
||||
stream_callback = partial(handle_tts_result, context)
|
||||
context.tts_client.set_result_callback(stream_callback)
|
||||
# 3. 设置是否播放（可选，默认True）
|
||||
context.tts_client.set_playback_enabled(False) # 设置为False则不播放
|
||||
|
||||
# 4. 初始化连接
|
||||
await context.tts_client.initialize()
|
||||
|
||||
# 4. 定义前端数据接收任务
|
||||
async def recv_frontend_data():
|
||||
while True:
|
||||
try:
|
||||
raw_bytes = await websocket.receive_bytes()
|
||||
msg_type, sequence, data = ProtocolCodec.unpack(raw_bytes)
|
||||
|
||||
if msg_type == MessageType.AUDIO_DATA:
|
||||
# 推送音频数据到ASR
|
||||
success = await asr_client.push_audio(asr_conn, data)
|
||||
if not success:
|
||||
logger.warning(f"连接 {client_id} 音频推送失败（队列满/连接失效）")
|
||||
elif msg_type == MessageType.CONTROL:
|
||||
# 处理控制指令（如暂停/继续ASR）
|
||||
logger.info(f"连接 {client_id} 收到控制指令：{data}")
|
||||
if data.get("action") == "stop_asr":
|
||||
asr_conn.stop_event.set()
|
||||
else:
|
||||
logger.warning(f"连接 {client_id} 收到未知消息类型：{msg_type.value}")
|
||||
except WebSocketDisconnect:
|
||||
logger.info(f"前端 {client_id} 主动断开连接")
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"连接 {client_id} 接收前端数据失败：{str(e)}")
|
||||
break
|
||||
|
||||
# 5. 定义ASR结果发送任务（从上下文队列取数据）
|
||||
async def send_asr_result():
|
||||
while True:
|
||||
try:
|
||||
# 从当前连接的消息队列获取ASR结果（超时0.05秒避免阻塞）
|
||||
result_packet = await asyncio.wait_for(
|
||||
context.message_queue.get(), timeout=0.05
|
||||
)
|
||||
print('发送', result_packet)
|
||||
await websocket.send_bytes(result_packet)
|
||||
except asyncio.TimeoutError:
|
||||
continue # 无数据时继续等待
|
||||
except Exception as e:
|
||||
logger.error(f"连接 {client_id} 发送ASR结果失败：{str(e)}")
|
||||
break
|
||||
|
||||
# 6. 启动任务并等待完成
|
||||
task_send = asyncio.create_task(send_asr_result())
|
||||
task_recv = asyncio.create_task(recv_frontend_data())
|
||||
|
||||
|
||||
done, pending = await asyncio.wait(
|
||||
[task_recv, task_send, communication_task],
|
||||
return_when=asyncio.FIRST_COMPLETED
|
||||
)
|
||||
|
||||
# 取消未完成的任务
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
await asyncio.gather(*pending, return_exceptions=True)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"连接 {client_id} 处理异常：{str(e)}")
|
||||
# 异常时发送错误消息给前端
|
||||
if websocket.state == "CONNECTED":
|
||||
error_packet = ProtocolCodec.pack(
|
||||
MessageType.ERROR, {"code": 5000, "message": f"服务异常：{str(e)}"}
|
||||
)
|
||||
await websocket.send_bytes(error_packet)
|
||||
finally:
|
||||
# 7. 资源清理（关键）
|
||||
# 停止ASR通信任务
|
||||
# if communication_task and not communication_task.done():
|
||||
# communication_task.cancel()
|
||||
# try:
|
||||
# await communication_task
|
||||
# except Exception as e:
|
||||
# logger.warning(f"连接 {client_id} ASR任务取消异常：{str(e)}")
|
||||
|
||||
# 关闭ASR连接
|
||||
# if asr_conn:
|
||||
# await asr_client.close_connection(asr_conn)
|
||||
|
||||
# 关闭WebSocket连接
|
||||
if websocket.state == "CONNECTED":
|
||||
await websocket.close(code=1008, reason="连接终止")
|
||||
|
||||
# 移除连接和上下文
|
||||
if websocket in self.active_connections:
|
||||
self.active_connections.remove(websocket)
|
||||
if client_id in self.client_context_map:
|
||||
del self.client_context_map[client_id]
|
||||
|
||||
logger.info(
|
||||
f"连接 {client_id} 资源清理完成，当前连接数: {len(self.active_connections)}"
|
||||
)
|
||||
Reference in New Issue
Block a user