238 lines
12 KiB
Python
238 lines
12 KiB
Python
# 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}") |