Tolan/backend/doubao_api/realtime_dialog_client.py

"""
豆包API实时对话客户端
基于官方demo实现
"""
import gzip
import json
import uuid
import asyncio
from typing import Dict, Any, Optional
import websockets
from loguru import logger

from . import protocol


class RealtimeDialogClient:
    """豆包API实时对话客户端"""
    
    def __init__(
        self, 
        api_key: str,
        app_id: str,
        base_url: str = "wss://openspeech.bytedance.com/api/v3/realtime/dialogue",
        output_audio_format: str = "ogg_opus",  # pcm 或 ogg_opus
        recv_timeout: int = 15,
        connect_timeout: int = 10
    ):
        self.api_key = api_key
        self.app_id = app_id
        self.base_url = base_url
        self.output_audio_format = output_audio_format
        self.recv_timeout = recv_timeout
        self.connect_timeout = connect_timeout
        
        # 生成唯一会话ID
        self.session_id = f"session_{uuid.uuid4().hex[:16]}"
        self.ws = None
        self.logid = ""
        self.connected = False
        
        # 配置请求头
        self.headers = {
            "X-Api-App-ID": app_id,
            "X-Api-Access-Key": api_key,
            "X-Api-Resource-Id": "volc.speech.dialog",
            "X-Api-App-Key": "PlgvMymc7f3tQnJ6",
            "X-Api-Connect-Id": str(uuid.uuid4()),
        }
    
    async def connect(self):
        """建立WebSocket连接并初始化会话"""
        try:
            logger.info(f"连接豆包API: {self.base_url}")
            logger.debug(f"会话ID: {self.session_id}")
            
            # 建立WebSocket连接
            self.ws = await asyncio.wait_for(
                websockets.connect(
                    self.base_url,
                    extra_headers=self.headers,
                    ping_interval=None
                ),
                timeout=self.connect_timeout
            )
            
            # 获取logid
            try:
                if hasattr(self.ws, 'response_headers'):
                    self.logid = self.ws.response_headers.get("X-Tt-Logid", "")
                else:
                    self.logid = ""
            except:
                self.logid = ""
            
            logger.info(f"豆包API连接成功，LogID: {self.logid}")
            
            # 发送StartConnection请求
            await self._send_start_connection()
            
            # 发送StartSession请求
            await self._send_start_session()
            
            self.connected = True
            logger.success(f"豆包会话已建立: {self.session_id}")
            
        except Exception as e:
            logger.error(f"连接豆包API失败: {e}")
            raise
    
    async def _send_start_connection(self):
        """发送StartConnection请求"""
        start_connection_request = bytearray(protocol.generate_header())
        start_connection_request.extend(int(1).to_bytes(4, 'big'))
        payload_bytes = str.encode("{}")
        payload_bytes = gzip.compress(payload_bytes)
        start_connection_request.extend((len(payload_bytes)).to_bytes(4, 'big'))
        start_connection_request.extend(payload_bytes)
        
        await self.ws.send(start_connection_request)
        response = await self.ws.recv()
        result = protocol.parse_response(response)
        logger.debug(f"StartConnection响应: {result}")
    
    async def _send_start_session(self):
        """发送StartSession请求"""
        # 构建会话配置
        session_config = {
            "asr": {
                "extra": {
                    "end_smooth_window_ms": 1200,  # 增加到1.2秒，避免中间停顿误判
                },
            },
            "tts": {
                "speaker": "zh_male_yunzhou_jupiter_bigtts",  # 使用官方音色
                "audio_config": {
                    "channel": 1,
                    "format": self.output_audio_format,  # pcm 或 ogg_opus
                    "sample_rate": 24000 if self.output_audio_format == "pcm" else 24000
                },
            },
            "dialog": {
                "bot_name": "豆包",
                "system_role": "你是一位英语口语陪练老师，帮助用户提高英语口语水平。请用简洁自然的英语回复，每次回复不超过2-3句话。",
                "speaking_style": "你的说话风格友好亲切，回复简短有力，鼓励用户多说英语。",
                "location": {
                    "city": "北京",
                },
                "extra": {
                    "strict_audit": False,
                    "recv_timeout": 120,  # 增加超时时间到120秒
                    "input_mod": "audio"  # 音频模式
                }
            }
        }
        
        payload_bytes = str.encode(json.dumps(session_config))
        payload_bytes = gzip.compress(payload_bytes)
        
        start_session_request = bytearray(protocol.generate_header())
        start_session_request.extend(int(100).to_bytes(4, 'big'))
        start_session_request.extend((len(self.session_id)).to_bytes(4, 'big'))
        start_session_request.extend(str.encode(self.session_id))
        start_session_request.extend((len(payload_bytes)).to_bytes(4, 'big'))
        start_session_request.extend(payload_bytes)
        
        await self.ws.send(start_session_request)
        response = await self.ws.recv()
        result = protocol.parse_response(response)
        logger.debug(f"StartSession响应: {result}")
    
    def is_connected(self) -> bool:
        """检查连接状态"""
        return self.connected and self.ws and self.ws.open
    
    async def send_audio(self, audio_data: bytes):
        """
        发送音频数据到豆包API
        audio_data: PCM音频数据（16kHz, 16bit, 单声道）
        """
        if not self.is_connected():
            raise ConnectionError("未连接到豆包API")
        
        try:
            # 构建音频请求（使用CLIENT_AUDIO_ONLY_REQUEST类型）
            task_request = bytearray(
                protocol.generate_header(
                    message_type=protocol.CLIENT_AUDIO_ONLY_REQUEST,
                    serial_method=protocol.NO_SERIALIZATION
                )
            )
            task_request.extend(int(200).to_bytes(4, 'big'))  # event: 200 (音频数据)
            task_request.extend((len(self.session_id)).to_bytes(4, 'big'))
            task_request.extend(str.encode(self.session_id))
            
            # 压缩音频数据
            payload_bytes = gzip.compress(audio_data)
            task_request.extend((len(payload_bytes)).to_bytes(4, 'big'))
            task_request.extend(payload_bytes)
            
            await self.ws.send(task_request)
            # 只记录原始大小，避免混淆
            logger.debug(f"发送音频到豆包: {len(audio_data)} bytes")
            
        except Exception as e:
            logger.error(f"发送音频失败: {e}")
            raise
    
    async def receive_message(self) -> Optional[Dict[str, Any]]:
        """
        接收豆包API返回的消息
        返回解析后的消息字典
        """
        if not self.is_connected():
            raise ConnectionError("未连接到豆包API")
        
        try:
            response = await asyncio.wait_for(self.ws.recv(), timeout=5.0)
            data = protocol.parse_response(response)
            return data
            
        except asyncio.TimeoutError:
            logger.debug("接收消息超时")
            return None
        except Exception as e:
            logger.error(f"接收消息失败: {e}")
            raise
    
    async def receive_audio(self) -> Optional[bytes]:
        """
        接收豆包API返回的音频数据
        返回值:
        - bytes: 音频数据
        - None: 无音频数据，继续接收
        - False: 音频流结束
        """
        try:
            data = await self.receive_message()
            
            if not data:
                return None
            
            # 检查事件类型
            event = data.get('event')
            payload_msg = data.get('payload_msg')
            
            # Event 352: TTS音频数据（豆包实际返回的事件类型）
            # Event 400: 其他音频数据
            if event in [352, 400] and isinstance(payload_msg, bytes):
                logger.debug(f"收到豆包音频数据 (Event {event}): {len(payload_msg)} bytes")
                return payload_msg
            
            # Event 350: TTS文本（AI回复的文本内容）
            elif event == 350 and isinstance(payload_msg, dict):
                tts_text = payload_msg.get('text', '')
                tts_type = payload_msg.get('tts_type', 'unknown')
                logger.info(f"📤 [TTS文本] Type: {tts_type}, Text: {tts_text}")
                return None
            
            # Event 351: TTS音频结束（真正的结束标志）
            # Event 401: 对话结束
            elif event in [351, 401]:
                logger.info(f"豆包音频流结束 (Event {event})")
                return False  # 返回False表示结束
            
            # Event 450: ASR中间结果（用户正在说话）
            # Event 451: ASR识别结果（实际事件类型）
            elif event in [450, 451] and isinstance(payload_msg, dict):
                # 从results数组中提取文本
                asr_text = ''
                results = payload_msg.get('results', [])
                if results and len(results) > 0:
                    asr_text = results[0].get('text', '')
                is_final = results[0].get('is_final', False) if results else False
                logger.info(f"📥 [ASR识别] {'[最终]' if is_final else '[中间]'} Text: {asr_text}")
                # 返回ASR文本给websocket_handler
                return ('asr_text', asr_text)
            
            # Event 459: ASR最终结果（用户说话结束）
            elif event == 459 and isinstance(payload_msg, dict):
                # 从results数组中提取文本
                asr_text = ''
                results = payload_msg.get('results', [])
                if results and len(results) > 0:
                    asr_text = results[0].get('text', '')
                logger.info(f"📥 [ASR最终结果] Text: {asr_text}")
                return None
            
            # Event 550: LLM流式文本输出
            elif event == 550 and isinstance(payload_msg, dict):
                content = payload_msg.get('content', '')
                logger.debug(f"💬 [LLM文本片段] {content}")
                # 返回一个特殊标记，让websocket_handler知道这是LLM文本
                return ('llm_text', content)
            
            # Event 359: 无内容
            elif event == 359:
                logger.debug("豆包返回无内容")
                return False  # 返回False表示结束
            
            # 其他事件（文本、状态等），继续等待
            else:
                if isinstance(payload_msg, dict):
                    # 检查是否是错误消息
                    if 'error' in payload_msg:
                        logger.error(f"❌ 豆包错误 (Event {event}): {payload_msg['error']}")
                    else:
                        logger.debug(f"收到豆包事件 {event}: {payload_msg}")
                else:
                    logger.debug(f"收到豆包事件: {event}")
                return None  # 继续接收
                
        except Exception as e:
            logger.error(f"接收音频失败: {e}")
            return None
    
    async def finish_session(self):
        """结束会话"""
        if not self.is_connected():
            return
        
        try:
            finish_session_request = bytearray(protocol.generate_header())
            finish_session_request.extend(int(102).to_bytes(4, 'big'))
            payload_bytes = str.encode("{}")
            payload_bytes = gzip.compress(payload_bytes)
            finish_session_request.extend((len(self.session_id)).to_bytes(4, 'big'))
            finish_session_request.extend(str.encode(self.session_id))
            finish_session_request.extend((len(payload_bytes)).to_bytes(4, 'big'))
            finish_session_request.extend(payload_bytes)
            
            await self.ws.send(finish_session_request)
            logger.debug("发送FinishSession请求")
            
        except Exception as e:
            logger.error(f"结束会话失败: {e}")
    
    async def finish_connection(self):
        """结束连接"""
        if not self.is_connected():
            return
        
        try:
            finish_connection_request = bytearray(protocol.generate_header())
            finish_connection_request.extend(int(2).to_bytes(4, 'big'))
            payload_bytes = str.encode("{}")
            payload_bytes = gzip.compress(payload_bytes)
            finish_connection_request.extend((len(payload_bytes)).to_bytes(4, 'big'))
            finish_connection_request.extend(payload_bytes)
            
            await self.ws.send(finish_connection_request)
            response = await self.ws.recv()
            result = protocol.parse_response(response)
            logger.debug(f"FinishConnection响应: {result}")
            
        except Exception as e:
            logger.error(f"结束连接失败: {e}")
    
    async def close(self):
        """关闭连接"""
        try:
            if self.is_connected():
                # 先结束会话
                await self.finish_session()
                # 再结束连接
                await self.finish_connection()
                # 关闭WebSocket
                await self.ws.close()
                
            self.connected = False
            logger.info(f"豆包API连接已关闭: {self.session_id}")
            
        except Exception as e:
            logger.error(f"关闭连接失败: {e}")