Tolan/backend/websocket_handler.py

"""
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)