diff --git a/audio_ai_chat/audio_ai_chat/core/websocket_handler_buck2.py b/audio_ai_chat/audio_ai_chat/core/websocket_handler_buck2.py new file mode 100644 index 0000000..c913812 --- /dev/null +++ b/audio_ai_chat/audio_ai_chat/core/websocket_handler_buck2.py @@ -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)}" + )