""" WebSocket处理器 - 新版本（STT+LLM+TTS流水线，支持流式LLM） """ import json import base64 import asyncio from typing import Dict, Optional from fastapi import WebSocket from models import ( AudioChunkMessage, AudioEndMessage, CancelMessage, TextResultMessage, EmotionMessage, AudioResponseMessage, ErrorMessage, PingMessage, PongMessage, ) from ai_service import ai_service class WebSocketHandler: """WebSocket连接处理器""" def __init__(self, websocket: WebSocket): self.websocket = websocket self.audio_buffers: Dict[str, bytearray] = {} # requestId -> audio_buffer self.cancelled_requests: set = set() # 已取消的请求ID self.recognition_sessions: Dict[str, any] = {} # requestId -> Recognition对象 self.recognition_results: Dict[str, str] = {} # requestId -> 识别结果文本 self.recognition_failed: set = set() # STT识别失败的请求ID self.event_loop: Optional[asyncio.AbstractEventLoop] = None # 保存事件循环引用 self.request_start_times: Dict[str, float] = {} # 请求开始时间 async def handle_connection(self): """处理WebSocket连接""" await self.websocket.accept() print("客户端已连接") # 保存事件循环引用，供回调线程使用 self.event_loop = asyncio.get_running_loop() try: while True: # 接收消息 data = await self.websocket.receive_text() message_data = json.loads(data) # 处理消息 await self.handle_message(message_data) except Exception as e: print(f"WebSocket错误: {e}") finally: print("客户端已断开") async def handle_message(self, message_data: dict): """处理接收到的消息""" message_type = message_data.get("type") request_id = message_data.get("requestId") # 检查是否为已取消的请求 if request_id in self.cancelled_requests: return if message_type == "audio_chunk": await self.handle_audio_chunk(AudioChunkMessage(**message_data)) elif message_type == "audio_end": await self.handle_audio_end(AudioEndMessage(**message_data)) elif message_type == "cancel": await self.handle_cancel(CancelMessage(**message_data)) elif message_type == "ping": await self.handle_ping(PingMessage(**message_data)) async def handle_audio_chunk(self, message: AudioChunkMessage): """处理音频数据块 - 流式STT""" request_id = message.requestId # 检查是否取消或STT已失败 if request_id in self.cancelled_requests or request_id in self.recognition_failed: return # Base64解码音频数据 audio_data = base64.b64decode(message.data) chunk_size = len(audio_data) # 初始化buffer（用于备份，以防STT失败） if request_id not in self.audio_buffers: self.audio_buffers[request_id] = bytearray() import time self.request_start_times[request_id] = time.time() print(f"\n[{request_id}] 🎤 开始接收音频流 [时间: {time.strftime('%H:%M:%S.%f')[:-3]}]") print(f"[{request_id}] 音频参数: {message.sampleRate}Hz, {message.channels}ch, {message.bitDepth}bit") self.audio_buffers[request_id].extend(audio_data) total_size = len(self.audio_buffers[request_id]) # 计算已录制时长（16000 sample_rate, 16bit = 2 bytes per sample, mono） duration_seconds = total_size / (16000 * 2) # 打印音频块信息 print(f"[{request_id}] 📦 接收chunk#{message.chunkIndex}: " f"大小={chunk_size}字节, 累计={total_size}字节 ({duration_seconds:.2f}秒)") # 第一个chunk：启动流式STT识别 if request_id not in self.recognition_sessions: print(f"[{request_id}] 🚀 首个chunk，启动STT流式识别...") success = await self.start_streaming_stt(request_id) if not success: # STT启动失败，标记为失败状态 print(f"[{request_id}] ❌ STT启动失败") self.recognition_failed.add(request_id) return # 发送音频帧到STT服务 recognition = self.recognition_sessions.get(request_id) if recognition: try: # 在线程池中发送音频帧 import time send_start = time.time() await asyncio.to_thread(recognition.send_audio_frame, audio_data) send_duration = (time.time() - send_start) * 1000 print(f"[{request_id}] ➡️ 发送至STT: {chunk_size}字节 (耗时: {send_duration:.1f}ms)") except Exception as e: # 发送失败，标记为失败状态，避免后续继续发送 if request_id not in self.recognition_failed: print(f"[{request_id}] ❌ 发送音频帧错误，停止发送: {e}") self.recognition_failed.add(request_id) async def handle_audio_end(self, message: AudioEndMessage): """处理音频结束 - 使用流式LLM优化""" request_id = message.requestId import time # 统计音频流信息 total_size = len(self.audio_buffers.get(request_id, bytearray())) duration_seconds = total_size / (16000 * 2) total_chunks = message.totalChunks print(f"\n[{request_id}] 🎤 音频录制结束 [时间: {time.strftime('%H:%M:%S.%f')[:-3]}]") print(f"[{request_id}] 📊 音频流统计: 总大小={total_size}字节, 时长={duration_seconds:.2f}秒, 共{total_chunks}个chunk") # 检查是否取消 if request_id in self.cancelled_requests: print(f"[{request_id}] 请求已取消") self.cancelled_requests.remove(request_id) await self.cleanup_request(request_id) return try: # 1. STT: 等待完整识别结果（保证准确性） stt_start = time.time() await self.stop_streaming_stt(request_id) # 获取STT识别结果 text = self.recognition_results.get(request_id, "") if not text: print(f"[{request_id}] STT未识别到文本，使用默认") text = "Hello, how can I help you?" # 发送最终识别结果 await self.send_text_result(request_id, text, True) stt_end = time.time() stt_duration = stt_end - self.request_start_times.get(request_id, stt_start) print(f"[{request_id}] STT识别结果: {text}") print(f"[{request_id}] ⏱️ STT耗时: {stt_duration:.2f}秒") if request_id in self.cancelled_requests: return # 2. LLM: 生成回复（使用非流式，保证音频连贯性） print(f"[{request_id}] 步骤 2/3: LLM 生成回复") llm_start = time.time() # 调用非流式 LLM，获取完整回复和情感 response_text, emotion = await ai_service.generate_response(text) llm_end = time.time() llm_duration = llm_end - llm_start print(f"[{request_id}] ⏱️ LLM耗时: {llm_duration:.2f}秒") if request_id in self.cancelled_requests: return # 发送情感标签 await self.send_emotion(request_id, emotion) # 3. TTS: 一次性生成完整音频（保证音色一致、无停顿） print(f"[{request_id}] 步骤 3/3: TTS 生成语音") tts_start = time.time() chunk_index = 0 async for audio_chunk in ai_service.text_to_speech_stream(response_text, emotion): if request_id in self.cancelled_requests: return # Base64编码 audio_base64 = base64.b64encode(audio_chunk).decode() # 发送音频数据 await self.send_audio_response( request_id, audio_base64, chunk_index, False, ) chunk_index += 1 tts_end = time.time() tts_duration = tts_end - tts_start print(f"[{request_id}] ⏱️ TTS耗时: {tts_duration:.2f}秒 (共{chunk_index}个数据块)") # 发送TTS结束标记 await self.send_audio_response(request_id, "", chunk_index, True) # 计算总耗时 total_end = time.time() total_duration = total_end - self.request_start_times.get(request_id, stt_start) print(f"[{request_id}] ⏱️ 总耗时: {total_duration:.2f}秒 (STT: {stt_duration:.2f}s + LLM: {llm_duration:.2f}s + TTS: {tts_duration:.2f}s)") print(f"[{request_id}] ✅ 请求处理完成\n") except Exception as e: print(f"[{request_id}] 处理错误: {e}") import traceback traceback.print_exc() await self.send_error(request_id, "NETWORK_ERROR", str(e)) finally: # 清理资源 await self.cleanup_request(request_id) async def start_streaming_stt(self, request_id: str) -> bool: """启动流式STT识别，返回是否成功""" try: import time stt_init_start = time.time() print(f"[{request_id}] 🔊 启动流式STT识别...") # 发送"正在听..."的中间状态 await self.send_text_result(request_id, "Listening...", False) from dashscope.audio.asr import Recognition, RecognitionCallback, RecognitionResult # 定义回调处理类 handler = self req_id = request_id class StreamCallback(RecognitionCallback): def on_open(self): print(f"[{req_id}] 🔗 STT连接已打开") def on_close(self): print(f"[{req_id}] 🔌 STT连接已关闭") # 标记为失败，停止后续音频帧发送 handler.recognition_failed.add(req_id) def on_error(self, error): print(f"[{req_id}] ❌ STT错误: {error}") # 标记为失败 handler.recognition_failed.add(req_id) def on_event(self, result: RecognitionResult): # 获取识别的句子 sentence = result.get_sentence() if sentence and 'text' in sentence: text = sentence['text'] # 判断是否为最终结果（end_time 可能为 None） end_time = sentence.get('end_time') is_final = end_time is not None and end_time > 0 status = "✅最终" if is_final else "⏳中间" print(f"[{req_id}] 💬 STT识别({status}): {text}") # 保存最新的识别结果 handler.recognition_results[req_id] = text # 从回调线程安全地发送到事件循环 if handler.event_loop: try: asyncio.run_coroutine_threadsafe( handler.send_text_result(req_id, text, False), handler.event_loop ) except Exception as e: print(f"[{req_id}] ❌ 发送中间结果失败: {e}") # 创建识别对象 def create_recognition(): return Recognition( model='paraformer-realtime-v2', format='pcm', sample_rate=16000, callback=StreamCallback() ) recognition = await asyncio.to_thread(create_recognition) # 启动识别 await asyncio.to_thread(recognition.start) # 保存识别会话 self.recognition_sessions[request_id] = recognition stt_init_duration = (time.time() - stt_init_start) * 1000 print(f"[{request_id}] ✅ STT识别会话已启动 (耗时: {stt_init_duration:.1f}ms)") # 短暂延迟，确保连接稳定 await asyncio.sleep(0.1) # 检查是否已经失败 if request_id in self.recognition_failed: print(f"[{request_id}] ❌ STT连接启动后立即失败") return False return True except Exception as e: print(f"[{request_id}] 启动STT识别失败: {e}") import traceback traceback.print_exc() return False async def stop_streaming_stt(self, request_id: str): """停止流式STT识别""" try: recognition = self.recognition_sessions.get(request_id) if recognition: print(f"[{request_id}] 🛑 停止STT识别...") import time stop_start = time.time() await asyncio.to_thread(recognition.stop) stop_duration = (time.time() - stop_start) * 1000 print(f"[{request_id}] ✅ STT识别已停止 (耗时: {stop_duration:.1f}ms)") except Exception as e: print(f"[{request_id}] ❌ 停止STT识别错误: {e}") async def cleanup_request(self, request_id: str): """清理请求相关资源""" # 清理audio buffer if request_id in self.audio_buffers: del self.audio_buffers[request_id] # 清理识别会话 if request_id in self.recognition_sessions: del self.recognition_sessions[request_id] # 清理识别结果 if request_id in self.recognition_results: del self.recognition_results[request_id] # 清理取消标记 if request_id in self.cancelled_requests: self.cancelled_requests.remove(request_id) # 清理失败标记 if request_id in self.recognition_failed: self.recognition_failed.remove(request_id) # 清理时间记录 if request_id in self.request_start_times: del self.request_start_times[request_id] print(f"[{request_id}] 资源已清理") async def handle_cancel(self, message: CancelMessage): """处理取消请求""" request_id = message.requestId print(f"[{request_id}] 收到取消请求") self.cancelled_requests.add(request_id) # 停止STT识别 await self.stop_streaming_stt(request_id) # 清理资源 await self.cleanup_request(request_id) async def handle_ping(self, message: PingMessage): """处理Ping""" await self.send_pong(message.requestId) async def send_text_result(self, request_id: str, text: str, is_final: bool): """发送文本结果""" message = TextResultMessage( type="text_result", requestId=request_id, timestamp=self._get_timestamp(), text=text, isFinal=is_final, ) await self.websocket.send_text(message.model_dump_json()) async def send_emotion(self, request_id: str, emotion: str): """发送情感标签""" message = EmotionMessage( type="emotion", requestId=request_id, timestamp=self._get_timestamp(), emotion=emotion, ) await self.websocket.send_text(message.model_dump_json()) async def send_audio_response( self, request_id: str, data: str, chunk_index: int, is_last: bool ): """发送音频响应""" message = AudioResponseMessage( type="audio_response", requestId=request_id, timestamp=self._get_timestamp(), data=data, chunkIndex=chunk_index, isLast=is_last, ) await self.websocket.send_text(message.model_dump_json()) async def send_error(self, request_id: str, code: str, error_message: str): """发送错误消息""" message = ErrorMessage( type="error", requestId=request_id, timestamp=self._get_timestamp(), code=code, message=error_message, ) await self.websocket.send_text(message.model_dump_json()) async def send_pong(self, request_id: str): """发送Pong""" message = PongMessage( type="pong", requestId=request_id, timestamp=self._get_timestamp(), ) await self.websocket.send_text(message.model_dump_json()) def _get_timestamp(self) -> int: """获取当前时间戳（毫秒）""" import time return int(time.time() * 1000)