# audio_ai_chat/websocket/manager.py from fastapi import WebSocket from typing import Dict, Optional from datetime import datetime import base64 from audio_ai_chat.asr.base import ASRBase, ASRResultCallback from audio_ai_chat.asr.asr_manager import ASRManager from audio_ai_chat.websocket.connection_context import ConnectionManager, ConnectionContext # 导入全局单例类 from audio_ai_chat.config.logger import logger class WebSocketConnectionManager: """全局唯一的WebSocket连接处理器（管理WebSocket连接生命周期）""" def __init__(self): self.asr_conn_map: Dict[str, Optional[object]] = {} # key=client_id，value=ASR连接 # 不实例化新的ConnectionManager，而是使用全局单例 self.connection_manager: Optional[ConnectionManager] = None async def initialize(self): """初始化：获取ConnectionManager全局单例（在FastAPI启动时调用）""" self.connection_manager = await ConnectionManager.get_instance() logger.info("WebSocketConnectionManager 初始化成功（绑定全局ConnectionManager）") async def handle_connection(self, websocket: WebSocket): """处理单个WebSocket连接的完整生命周期""" # 校验ConnectionManager是否初始化 if not self.connection_manager: await websocket.accept() await websocket.send_text("服务未初始化完成，请稍后重试") await websocket.close() logger.error("WebSocketConnectionManager 未初始化，拒绝连接") return # 1. 接受连接，生成client_id（用字符串类型，避免int溢出） await websocket.accept() client_id = str(id(websocket)) # client_id为字符串，与ConnectionManager的key类型一致 logger.info(f"新WebSocket连接：client_id={client_id}") try: # 2. 创建连接上下文（通过全局ConnectionManager） context = await self.connection_manager.create_connection(client_id=client_id) if not context: await websocket.send_text("连接上下文创建失败") await websocket.close() return # 3. 获取ASR实例 asr_client = ASRManager.get_instance() if not asr_client or not ASRManager.is_available(): await websocket.send_json({ "type": "error", "message": "ASR服务未初始化，无法提供转写服务", "timestamp": datetime.utcnow().isoformat() + "Z" }) await self.connection_manager.remove_connection(client_id=client_id) await websocket.close() return # 4. 获取ASR连接 asr_conn = await asr_client.get_connection() if not asr_conn: await websocket.send_json({ "type": "error", "message": "ASR无空闲连接，连接失败", "timestamp": datetime.utcnow().isoformat() + "Z" }) await self.connection_manager.remove_connection(client_id=client_id) await websocket.close() return self.asr_conn_map[client_id] = asr_conn # 5. 定义ASR结果回调（绑定当前上下文） async def asr_callback(result: Dict[str, Any]): if not context.is_active: logger.warning(f"连接已关闭，忽略ASR结果：client_id={client_id}") return # 处理ASR结果并存入上下文 context.add_asr_result(result) # 推送给前端 if result.get("error"): await websocket.send_json({ "type": "asr_error", "message": result["error"], "timestamp": datetime.utcnow().isoformat() + "Z" }) else: await websocket.send_json({ "type": "asr_progress" if not result["is_final"] else "asr_final", "text": result["text"], "is_final": result["is_final"], "timestamp": result.get("timestamp", datetime.utcnow().isoformat() + "Z") }) # 6. 启动ASR通信 asr_task = asyncio.create_task( asr_client.start_communication(conn=asr_conn, callback=asr_callback) ) # 7. 循环接收前端数据 while context.is_active: try: # 假设前端发送JSON格式数据（区分音频/文本/用户信息） data = await websocket.receive_json() data_type = data.get("type") # 处理用户信息（登录后发送） if data_type == "user_info": try: token = data.get("token") user_id = data.get("user_id") name = data.get("name", "匿名用户") context.set_user_info(token=token, user_id=user_id, name=name) await websocket.send_json({ "type": "info", "message": "用户信息设置成功", "timestamp": datetime.utcnow().isoformat() + "Z" }) except Exception as e: err_msg = f"用户信息设置失败：{str(e)}" context.add_system_message(err_msg) await websocket.send_json({ "type": "error", "message": err_msg, "timestamp": datetime.utcnow().isoformat() + "Z" }) # 处理Base64编码的音频数据 elif data_type == "audio_data": audio_base64 = data.get("audio_data") if not audio_base64: continue try: audio_data = base64.b64decode(audio_base64) success = await asr_client.push_audio(asr_conn, audio_data) if not success: await websocket.send_json({ "type": "warning", "message": "ASR音频队列已满，部分数据丢失", "timestamp": datetime.utcnow().isoformat() + "Z" }) except Exception as e: err_msg = f"音频解码失败：{str(e)}" logger.error(f"client_id={client_id}，{err_msg}") await websocket.send_json({ "type": "error", "message": err_msg, "timestamp": datetime.utcnow().isoformat() + "Z" }) # 处理纯文本输入 elif data_type == "text_input": text = data.get("text", "").strip() if text: context.add_chat_history({ "role": "user", "content": text, "source": "text", "asr_metadata": None }) await websocket.send_json({ "type": "info", "message": f"已接收文本：{text}", "timestamp": datetime.utcnow().isoformat() + "Z" }) # 处理大模型请求 elif data_type == "request_llm": if context.is_processing: await websocket.send_json({ "type": "warning", "message": "正在处理上一个请求，请稍后再试", "timestamp": datetime.utcnow().isoformat() + "Z" }) continue # 获取对话历史 chat_history = context.get_chat_history(limit=20) logger.debug(f"请求大模型：client_id={client_id}，历史条数={len(chat_history)}") # 模拟大模型调用（实际替换为真实LLM调用） context.is_processing = True try: # llm_response = await context.llm_session.generate(chat_history=chat_history) llm_response = f"模拟大模型回复：已收到你的{len(chat_history)}条对话历史" context.add_llm_result(llm_response) await websocket.send_json({ "type": "llm_response", "text": llm_response, "timestamp": datetime.utcnow().isoformat() + "Z" }) except Exception as e: err_msg = f"大模型调用失败：{str(e)}" context.add_system_message(err_msg) await websocket.send_json({ "type": "error", "message": err_msg, "timestamp": datetime.utcnow().isoformat() + "Z" }) finally: context.is_processing = False # 未知数据类型 else: err_msg = f"未知数据类型：{data_type}" await websocket.send_json({ "type": "error", "message": err_msg, "timestamp": datetime.utcnow().isoformat() + "Z" }) except Exception as e: # 捕获前端发送数据异常（如断开连接） logger.error(f"接收前端数据异常：client_id={client_id}，error={str(e)}") break except Exception as e: # 其他异常 err_msg = f"连接处理异常：{str(e)}" logger.error(f"client_id={client_id}，{err_msg}") await websocket.send_json({ "type": "error", "message": err_msg, "timestamp": datetime.utcnow().isoformat() + "Z" }) finally: # 8. 资源清理 # 取消ASR任务 asr_task.cancel() try: await asr_task except asyncio.CancelledError: pass # 释放ASR连接 if client_id in self.asr_conn_map: asr_conn = self.asr_conn_map.pop(client_id) await asr_client.release_connection(asr_conn) # 移除连接上下文 await self.connection_manager.remove_connection(client_id=client_id) # 关闭WebSocket await websocket.close() logger.info(f"WebSocket连接关闭：client_id={client_id}")