301 lines
12 KiB
Python
301 lines
12 KiB
Python
# frontend_ws.py
|
||
from fastapi import WebSocket, WebSocketDisconnect
|
||
from session_manager import create_session, close_session
|
||
|
||
# 导入大模型模块
|
||
from llm_client import call_llm, LLMConversation, llm_client
|
||
import numpy as np
|
||
from typing import Optional, Dict, Callable, Awaitable, List
|
||
from dataclasses import dataclass, field
|
||
import sys
|
||
import json
|
||
import asyncio
|
||
|
||
import logging
|
||
logger = logging.getLogger(__name__)
|
||
from ws_message_manager import (
|
||
ws_queue_manager,
|
||
ClientMsgType,
|
||
ServerMsgType,
|
||
)
|
||
|
||
from asr_client import (
|
||
init_asr_pool,
|
||
get_idle_asr_connection,
|
||
handle_asr_communication,
|
||
push_audio_data, # 导入音频插入接口
|
||
close_asr_pool
|
||
)
|
||
from tts_client import CosyVoiceTTSSocketClient
|
||
active_connections = []
|
||
user_llm_conversations: dict[str, LLMConversation] = {}
|
||
|
||
async def frontend_websocket_handler(websocket: WebSocket):
|
||
"""前端WebSocket入口处理器（保留Session，ASR/TTS无状态）"""
|
||
# 1. 握手+创建会话（核心：Session 管理连接上下文）
|
||
await websocket.accept()
|
||
active_connections.append(websocket)
|
||
conn_id = id(websocket)
|
||
user_id = f"user_{id(websocket)}" # 实际场景替换为真实用户ID
|
||
session = await create_session(websocket) # 创建会话（无用户逻辑）
|
||
logger.info(f"连接 {conn_id} 建立成功")
|
||
# await ws_queue_manager.create_queue(websocket)
|
||
tts_session_id = user_id
|
||
result_queue = asyncio.Queue(maxsize=100)
|
||
asr_task = None
|
||
asr_conn = None
|
||
|
||
# 初始化用户LLM会话
|
||
if user_id not in user_llm_conversations:
|
||
user_llm_conversations[user_id] = LLMConversation(
|
||
user_id=user_id,
|
||
scene_description="语音识别对话场景" # 自定义大模型场景
|
||
)
|
||
llm_conversation = user_llm_conversations[user_id]
|
||
|
||
# 1. 创建客户端
|
||
client = CosyVoiceTTSSocketClient(
|
||
ws_url="ws://10.10.10.202:50000/ws/tts",
|
||
max_queue_size=500 # 最大队列长度
|
||
)
|
||
|
||
# 2. 注册回调（按需自定义）
|
||
client.on_audio_chunk = lambda req_id, chunk: print(f"收到[{req_id[:8]}]音频块，长度: {len(chunk)}")
|
||
|
||
# 3. 建立连接（自动启动队列消费）
|
||
await client.connect()
|
||
|
||
# ---------------------- TTS 音频回调（转发前端） ----------------------
|
||
async def tts_audio_callback(audio_chunk, sample_rate, is_finished):
|
||
"""将音频块转发给前端的逻辑"""
|
||
print(f"音频块长度：{len(audio_chunk)}，采样率：{sample_rate}，是否结束：{is_finished}")
|
||
#todo
|
||
|
||
# ---------------------- 定义 ASR 结果回调函数 ----------------------
|
||
async def asr_result_callback(result: dict):
|
||
"""
|
||
ASR 结果回调：同时转发给前端 + 推送大模型
|
||
:param result: ASR 识别结果字典
|
||
"""
|
||
try:
|
||
print('结果回调：同时转发给前端 + 推送大模型', result)
|
||
final_asr_text = result["text"]
|
||
print(f"ASR最终识别结果: {final_asr_text}")
|
||
# 调用大模型（异步，不阻塞）
|
||
asyncio.create_task(
|
||
call_llm_and_send(
|
||
query=final_asr_text,
|
||
conversation=llm_conversation
|
||
)
|
||
)
|
||
# # 1. 转发给前端（确保连接未断开）
|
||
# if not websocket.client_state.disconnected:
|
||
# await websocket.send_json(result)
|
||
#
|
||
# # 2. 推送到大模型（过滤错误/空结果）
|
||
# if not result.get("error") and result.get("text"):
|
||
# # 异步推送，不阻塞回调
|
||
# asyncio.create_task(send_to_llm(result))
|
||
|
||
except Exception as e:
|
||
print(f"回调函数执行失败：{e}")
|
||
|
||
# -------------------------------------------------------------------
|
||
# ---------------------- 大模型流式回调（转发前端） ----------------------
|
||
async def llm_stream_callback(chunk: str, conversation_id: str, is_finished: bool):
|
||
req_id1 = client.add_tts_request(tts_text=chunk)
|
||
print('req_id1', req_id1)
|
||
"""
|
||
大模型流式回调：将实时回复转发给前端
|
||
"""
|
||
# if websocket.client_state.disconnected:
|
||
# return
|
||
# try:
|
||
# print('大模型流式回调', chunk)
|
||
# await websocket.send_json({
|
||
# "type": "llm_result",
|
||
# "data": {
|
||
# "chunk": chunk,
|
||
# "conversation_id": conversation_id,
|
||
# "is_finished": is_finished,
|
||
# "query": final_asr_text # 关联ASR查询文本
|
||
# }
|
||
# })
|
||
# except Exception as e:
|
||
# print(f"LLM回调执行失败: {e}")
|
||
|
||
|
||
# if websocket.client_state.disconnected:
|
||
# return
|
||
|
||
try:
|
||
# 1. 转发大模型文本给前端
|
||
# await websocket.send_json({
|
||
# "type": "llm_result",
|
||
# "data": {
|
||
# "chunk": chunk,
|
||
# "conversation_id": conversation_id,
|
||
# "is_finished": is_finished,
|
||
# "user_id": user_id
|
||
# }
|
||
# })
|
||
# print('大模型返回的', chunk)
|
||
# await tts_client.send_text_chunk(tts_session_id, chunk)
|
||
pass
|
||
|
||
# 2. 流式文本发送给 TTS（非错误文本 + 非结束标识）
|
||
# if chunk and not chunk.startswith("[错误]") and not is_finished:
|
||
# await tts_client.send_text_chunk(tts_session_id, chunk)
|
||
# # 合成结束标识
|
||
# elif is_finished:
|
||
# # 发送空文本块触发 TTS 结束
|
||
# await tts_client.send_text_chunk(tts_session_id, "")
|
||
|
||
except Exception as e:
|
||
print(f"LLM回调联动TTS失败({user_id}): {e}")
|
||
|
||
# ---------------------- 调用大模型并转发结果 ----------------------
|
||
async def call_llm_and_send(query: str, conversation: LLMConversation):
|
||
"""调用大模型，流式结果通过WebSocket返回前端"""
|
||
if not query:
|
||
return
|
||
print(f"调用大模型 - 用户({user_id}): {query}")
|
||
# 调用大模型（流式）
|
||
conv_id, full_reply = await llm_client.send_message(
|
||
query=query,
|
||
conversation=conversation,
|
||
stream_callback=llm_stream_callback,
|
||
response_mode="streaming"
|
||
)
|
||
print(f"大模型回复完成 - 会话ID: {conv_id}, 完整回复: {full_reply}")
|
||
# -------------------------------------------------------------------
|
||
try:
|
||
# 1. 获取 ASR 连接
|
||
asr_conn = await get_idle_asr_connection()
|
||
if not asr_conn:
|
||
await websocket.send_json({"error": "ASR 服务暂时不可用", "text": ""})
|
||
return
|
||
|
||
# 启动 ASR 协程（传递回调函数，移除结果队列）
|
||
asr_task = asyncio.create_task(
|
||
handle_asr_communication(asr_conn, asr_result_callback)
|
||
)
|
||
|
||
|
||
|
||
# 3. 并行处理：接收前端数据 + 发送 ASR 结果
|
||
async def recv_frontend_data():
|
||
"""接收前端所有数据，仅将音频数据插入 ASR 队列"""
|
||
while not asr_conn.stop_event.is_set():
|
||
|
||
try:
|
||
raw_bytes = await websocket.receive_bytes()
|
||
success = await push_audio_data(asr_conn, raw_bytes)
|
||
if not success:
|
||
print("音频数据插入失败（队列满/连接失效）")
|
||
# # 接收前端数据（自动区分类型）
|
||
# data_type = websocket.receive()
|
||
# if data_type.type == "bytes":
|
||
# # 二进制数据 = 音频数据，调用 ASR 接口插入
|
||
# raw_bytes = await websocket.receive_bytes()
|
||
# # 调用 ASR 模块的插入接口，无需管理队列
|
||
# success = await push_audio_data(asr_conn, raw_bytes)
|
||
# if not success:
|
||
# print("音频数据插入失败（队列满/连接失效）")
|
||
# elif data_type.type == "json":
|
||
# # JSON 数据 = 控制指令
|
||
# control_data = await websocket.receive_json()
|
||
# print(f"收到前端控制指令：{control_data}")
|
||
# if control_data.get("action") == "stop":
|
||
# asr_conn.stop_event.set() # 触发 ASR 停止
|
||
# await websocket.send_json({"text": "", "error": "用户主动停止"})
|
||
# break
|
||
# elif data_type.type == "text":
|
||
# # 文本数据 = 其他指令
|
||
# text_data = await websocket.receive_text()
|
||
# print(f"收到前端文本数据：{text_data}")
|
||
except WebSocketDisconnect:
|
||
print("前端主动断开连接")
|
||
asr_conn.stop_event.set()
|
||
break
|
||
except Exception as e:
|
||
print(f"接收前端数据失败：{e}")
|
||
asr_conn.stop_event.set()
|
||
await websocket.send_json({"error": f"接收数据失败：{str(e)}", "text": ""})
|
||
break
|
||
|
||
async def send_asr_result():
|
||
"""从结果队列取数据，发送给前端"""
|
||
while not asr_conn.stop_event.is_set():
|
||
try:
|
||
result = await asyncio.wait_for(result_queue.get(), timeout=1.0)
|
||
# if not websocket.client_state.disconnected:
|
||
print('result', result)
|
||
await websocket.send_json(result)
|
||
except asyncio.TimeoutError:
|
||
continue
|
||
except Exception as e:
|
||
print(f"发送 ASR 结果失败：{e}")
|
||
asr_conn.stop_event.set()
|
||
break
|
||
|
||
# 4. 并发执行
|
||
await asyncio.gather(recv_frontend_data(), send_asr_result())
|
||
|
||
except Exception as e:
|
||
print(f"WebSocket 处理异常：{e}")
|
||
if asr_conn:
|
||
asr_conn.stop_event.set()
|
||
await websocket.send_json({"error": str(e), "text": ""})
|
||
finally:
|
||
# 清理资源
|
||
if asr_conn:
|
||
asr_conn.stop_event.set()
|
||
# 取消 ASR 任务
|
||
if asr_task and not asr_task.done():
|
||
asr_task.cancel()
|
||
try:
|
||
await asr_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
# 移除活跃连接
|
||
if websocket in active_connections:
|
||
active_connections.remove(websocket)
|
||
print(f"客户端连接已清理，当前连接数：{len(active_connections)}")
|
||
# 关闭前端连接
|
||
# if not websocket.client_state.disconnected:
|
||
# await websocket.close()
|
||
|
||
# try:
|
||
#
|
||
# # 4. 持续接收前端消息（二进制，兼容音频/文本/指令）
|
||
# while not session.is_closed:
|
||
# raw_bytes = await websocket.receive_bytes()
|
||
# # print('raw_bytes', raw_bytes)
|
||
# # todo 前段暂时传入的全部都是语音消息
|
||
# if not audio_queue.full():
|
||
# audio_queue.put_nowait(raw_bytes)
|
||
# # await forward_frontend_to_asr(session, raw_bytes)
|
||
# # 解析前端消息类型（复用之前的编解码器）
|
||
# # msg_type, data = WsMsgCodec.decode_client_msg(raw_bytes)
|
||
# # # 分发处理（不同类型消息走不同逻辑）
|
||
# # if msg_type == ClientMsgType.AUDIO:
|
||
# # # 音频直接转发给ASR（Session 传递ASR客户端，无用户标识）
|
||
# # if session.asr_client and not session.is_closed:
|
||
# # await send_audio_to_asr(session.asr_client, data)
|
||
# # elif msg_type == ClientMsgType.TEXT:
|
||
# # # 文本直接走LLM（无需ASR）
|
||
# # await forward_text_to_llm(session, data)
|
||
# # elif msg_type == ClientMsgType.ACTION:
|
||
# # # 操作指令处理
|
||
# # await handle_frontend_action(session, data)
|
||
#
|
||
# except WebSocketDisconnect:
|
||
# logger.info(f"会话 [{session.session_id}] 前端主动断开")
|
||
# except Exception as e:
|
||
# logger.error(f"会话 [{session.session_id}] 异常: {str(e)}", exc_info=True)
|
||
# finally:
|
||
# # 5. 清理会话（一键关闭所有资源，核心价值）
|
||
# await close_session(session)
|
||
# await websocket.close()
|