Tolan/backend/stt_service.py

"""
STT服务模块 - 基于 DashScope OmniRealtimeConversation API
参考官方demo实现，提供流式语音识别功能
"""
import asyncio
import time
import logging
import traceback
from typing import Optional, Callable, Awaitable
import dashscope
from dashscope.audio.qwen_omni import OmniRealtimeConversation, OmniRealtimeCallback
from dashscope.audio.qwen_omni.omni_realtime import MultiModality, TranscriptionParams
from config import settings, get_dashscope_ws_url

logger = logging.getLogger(__name__)


class STTCallback(OmniRealtimeCallback):
    """STT识别回调处理器（参考官方demo实现）"""
    
    def __init__(
        self,
        request_id: str,
        on_text_update: Optional[Callable[[str, bool], Awaitable[None]]] = None,
        on_error: Optional[Callable[[str], Awaitable[None]]] = None,
        event_loop: Optional[asyncio.AbstractEventLoop] = None
    ):
        """
        初始化STT回调处理器
        
        Args:
            request_id: 请求ID（用于日志标识）
            on_text_update: 文本更新回调 (text, is_final) -> None
            on_error: 错误回调 (error_msg) -> None
            event_loop: 事件循环（用于线程安全的异步调用）
        """
        self.request_id = request_id
        self.on_text_update = on_text_update
        self.on_error = on_error
        self.event_loop = event_loop
        
        # 识别结果存储
        self.current_text = ""
        self.is_final = False
        self.final_event = asyncio.Event() if event_loop else None
        
        # 时间统计
        self.last_update_time = 0.0
        self.first_text_time = None  # 首个文本生成时间
        self.session_id: Optional[str] = None
        
        # 语音检测统计
        self.speech_detected = False  # 是否检测到语音
        self.speech_start_count = 0   # 语音开始次数
        self.speech_stop_count = 0    # 语音停止次数
        
        # 事件处理器映射（参考官方demo）
        self.handlers = {
            'session.created': self._handle_session_created,
            'conversation.item.input_audio_transcription.completed': self._handle_final_text,
            'conversation.item.input_audio_transcription.text': self._handle_stash_text,
            'input_audio_buffer.speech_started': self._handle_speech_start,
            'input_audio_buffer.speech_stopped': self._handle_speech_stop,
            'error': self._handle_error,
        }
    
    def on_open(self):
        """连接打开回调"""
        pass
    
    def on_close(self, code: int, msg: str):
        """连接关闭回调"""
        
        # 非正常关闭（1000是正常关闭码，None是主动关闭）
        if code is not None and code != 1000 and self.on_error:
            error_msg = f"STT连接异常关闭: {msg} (code: {code})"
            logger.error(
                f"[STT异常] request_id={self.request_id}, "
                f"错误类型=连接异常关闭, "
                f"关闭码={code}, "
                f"错误消息={msg}"
            )
            self._async_call(self.on_error(error_msg))
    
    def on_event(self, response: dict):
        """事件分发（参考官方demo）"""
        try:
            event_type = response.get('type')
            handler = self.handlers.get(event_type)
            
            # 日志记录（过滤掉高频事件）
            # if event_type not in ['input_audio_buffer.speech_started', 'input_audio_buffer.speech_stopped']:
            #     print(f"[{self.request_id}] 📡 STT事件: {event_type}")
            
            if handler:
                handler(response)
            else:
                # 记录未处理的事件（排除常见的已知事件）
                known_events = [
                    'session.created', 
                    'session.updated', 
                    'input_audio_buffer.committed',
                    'conversation.item.created'  # 会话项创建事件
                ]
                if event_type not in known_events:
                    pass
                    
        except Exception as e:
            logger.error(
                f"[STT异常] request_id={self.request_id}, "
                f"错误类型=事件处理异常, "
                f"异常类型={type(e).__name__}, "
                f"错误消息={str(e)}, "
                f"事件类型={response.get('type', 'unknown')}, "
                f"堆栈跟踪:\n{traceback.format_exc()}"
            )
    
    # ========== 事件处理器（参考官方demo） ==========
    
    def _handle_session_created(self, response: dict):
        """处理会话创建事件"""
        session_id = response.get('session', {}).get('id', '')
        self.session_id = session_id
    
    def _handle_final_text(self, response: dict):
        """处理最终识别结果（参考官方demo）"""
        text = response.get('transcript', '')
        
        # 更新状态
        self.current_text = text
        self.is_final = True
        self.last_update_time = time.time()
        
        # 触发最终结果事件
        if self.final_event and self.event_loop:
            try:
                self.event_loop.call_soon_threadsafe(self.final_event.set)
            except Exception as e:
                logger.error(
                    f"[STT异常] request_id={self.request_id}, "
                    f"错误类型=事件触发异常, "
                    f"异常类型={type(e).__name__}, "
                    f"错误消息={str(e)}, "
                    f"堆栈跟踪:\n{traceback.format_exc()}"
                )
        
        # 回调通知
        if self.on_text_update:
            self._async_call(self.on_text_update(text, True))
    
    def _handle_stash_text(self, response: dict):
        """处理中间识别结果（参考官方demo）"""
        # 合并 text（已确认部分）和 stash（正在识别部分）
        confirmed_text = response.get('text', '')
        stash_text = response.get('stash', '')
        full_text = confirmed_text + stash_text
        
        # 更新状态（不打印中间识别，减少日志）
        if full_text:
            self.current_text = full_text
            self.is_final = False
            current_time = time.time()
            self.last_update_time = current_time
            
            # 记录首个文本生成时间
            if self.first_text_time is None and full_text:
                self.first_text_time = current_time
            
            # 回调通知
            if self.on_text_update:
                self._async_call(self.on_text_update(full_text, False))
    
    def _handle_speech_start(self, response: dict):
        """处理语音开始事件（静默处理，减少日志）"""
        self.speech_detected = True
        self.speech_start_count += 1
    
    def _handle_speech_stop(self, response: dict):
        """处理语音停止事件（静默处理，减少日志）"""
        self.speech_stop_count += 1
    
    def _handle_error(self, response: dict):
        """处理错误事件"""
        error = response.get('error', {})
        error_msg = error.get('message', 'Unknown error')
        error_code = error.get('code', '')
        
        # 详细记录STT API错误
        logger.error(
            f"[STT异常] request_id={self.request_id}, "
            f"错误类型=STT API错误, "
            f"错误码={error_code}, "
            f"错误消息={error_msg}, "
            f"完整响应={response}"
        )
        
        # 回调通知
        if self.on_error:
            self._async_call(self.on_error(f"{error_msg} (code: {error_code})"))
    
    # ========== 辅助方法 ==========
    
    def _async_call(self, coro):
        """线程安全的异步调用"""
        if self.event_loop:
            try:
                asyncio.run_coroutine_threadsafe(coro, self.event_loop)
            except Exception as e:
                logger.error(
                    f"[STT异常] request_id={self.request_id}, "
                    f"错误类型=异步调用异常, "
                    f"异常类型={type(e).__name__}, "
                    f"错误消息={str(e)}, "
                    f"堆栈跟踪:\n{traceback.format_exc()}"
                )
    
    def get_text(self) -> str:
        """获取当前识别文本"""
        return self.current_text
    
    def is_final_result(self) -> bool:
        """是否为最终结果"""
        return self.is_final


class STTSession:
    """STT识别会话管理器（参考官方demo实现）"""
    
    def __init__(
        self,
        request_id: str,
        on_text_update: Optional[Callable[[str, bool], Awaitable[None]]] = None,
        on_error: Optional[Callable[[str], Awaitable[None]]] = None,
        event_loop: Optional[asyncio.AbstractEventLoop] = None
    ):
        """
        初始化STT会话
        
        Args:
            request_id: 请求ID
            on_text_update: 文本更新回调
            on_error: 错误回调
            event_loop: 事件循环
        """
        self.request_id = request_id
        self.callback = STTCallback(
            request_id=request_id,
            on_text_update=on_text_update,
            on_error=on_error,
            event_loop=event_loop
        )
        self.conversation: Optional[OmniRealtimeConversation] = None
        self.is_connected = False
        self.is_closed = False
    
    async def connect(self) -> bool:
        """
        连接STT服务并初始化会话
        
        Returns:
            是否连接成功
        """
        try:
            start_time = time.time()
            
            # 设置API Key
            dashscope.api_key = settings.active_dashscope_api_key
            
            # 创建会话对象（参考官方demo）
            def create_conversation():
                ws_url = get_dashscope_ws_url()
                # print(f"[{self.request_id}] STT WebSocket URL: {ws_url}")
                
                conversation = OmniRealtimeConversation(
                    model='qwen3-asr-flash-realtime',
                    url=ws_url,
                    callback=self.callback
                )
                # 注入自身引用（回调中可能需要）
                self.callback.conversation = conversation
                return conversation
            
            # 异步创建会话
            self.conversation = await asyncio.to_thread(create_conversation)
            
            # 连接WebSocket
            await asyncio.to_thread(self.conversation.connect)
            
            # 配置转录参数（参考官方demo）
            # ⚠️ 重要：language参数是必需的！
            # 官方支持: 'zh'(中文), 'en'(英文), 'ja'(日文), 'ko'(韩文)
            # 尝试支持: 'auto'(自动检测) - 需要测试是否有效
            stt_language = getattr(settings, 'STT_LANGUAGE', 'zh')
            
            transcription_params = TranscriptionParams(
                language=stt_language,
                sample_rate=16000,
                input_audio_format="pcm"
                # corpus_text=""  # 可选：辅助识别的语料
            )
            
            # 更新会话配置（禁用VAD，用户通过按钮控制录音）
            await asyncio.to_thread(
                self.conversation.update_session,
                output_modalities=[MultiModality.TEXT],
                enable_input_audio_transcription=True,
                transcription_params=transcription_params,
                enable_turn_detection=False
            )
            
            self.is_connected = True
            
            return True
            
        except Exception as e:
            logger.error(
                f"[STT异常] request_id={self.request_id}, "
                f"错误类型=STT连接失败, "
                f"异常类型={type(e).__name__}, "
                f"错误消息={str(e)}, "
                f"API Key={settings.active_dashscope_api_key[:10]}... (已脱敏), "
                f"使用国际版={settings.use_dashscope_intl}, "
                f"堆栈跟踪:\n{traceback.format_exc()}"
            )
            if self.callback.on_error:
                await self.callback.on_error(f"连接失败: {e}")
            return False
    
    async def send_audio(self, audio_data: bytes) -> bool:
        """
        发送音频数据（参考官方demo）
        
        Args:
            audio_data: PCM音频数据（原始bytes）
            
        Returns:
            是否发送成功
        """
        if not self.is_connected or not self.conversation:
            return False
        
        try:
            # Base64编码（API要求）
            import base64
            audio_b64 = base64.b64encode(audio_data).decode('ascii')
            
            # 发送音频帧（参考官方demo）
            await asyncio.to_thread(self.conversation.append_audio, audio_b64)
            return True
            
        except Exception as e:
            logger.error(
                f"[STT异常] request_id={self.request_id}, "
                f"错误类型=发送音频失败, "
                f"异常类型={type(e).__name__}, "
                f"错误消息={str(e)}, "
                f"音频数据长度={len(audio_data)} bytes, "
                f"连接状态=is_connected={self.is_connected}, "
                f"堆栈跟踪:\n{traceback.format_exc()}"
            )
            if self.callback.on_error:
                await self.callback.on_error(f"发送音频失败: {e}")
            return False
    
    async def commit_audio(self) -> bool:
        """
        提交音频以触发识别（禁用VAD时需要手动调用）
        
        Returns:
            是否提交成功
        """
        if not self.is_connected or not self.conversation:
            return False
        
        try:
            await asyncio.to_thread(self.conversation.commit)
            return True
        except Exception as e:
            logger.error(
                f"[STT异常] request_id={self.request_id}, "
                f"错误类型=提交音频失败, "
                f"异常类型={type(e).__name__}, "
                f"错误消息={str(e)}, "
                f"连接状态=is_connected={self.is_connected}, "
                f"堆栈跟踪:\n{traceback.format_exc()}"
            )
            return False
    
    async def wait_for_final_result(self, timeout: float = 3.0) -> Optional[str]:
        """
        等待最终识别结果
        
        Args:
            timeout: 超时时间（秒）
            
        Returns:
            识别文本，超时或失败返回None
        """
        if not self.callback.final_event:
            # print(f"[{self.request_id}] ⚠️  无法等待最终结果（未设置事件循环）")
            return self.callback.get_text()
        
        try:
            await asyncio.wait_for(self.callback.final_event.wait(), timeout=timeout)
            return self.callback.get_text()
            
        except asyncio.TimeoutError:
            # 超时后，检查是否有中间结果
            current_text = self.callback.get_text()
            logger.warning(
                f"[STT警告] request_id={self.request_id}, "
                f"错误类型=等待识别结果超时, "
                f"超时时间={timeout}秒, "
                f"当前识别文本={current_text}"
            )
            return current_text
        except Exception as e:
            logger.error(
                f"[STT异常] request_id={self.request_id}, "
                f"错误类型=等待识别结果异常, "
                f"异常类型={type(e).__name__}, "
                f"错误消息={str(e)}, "
                f"超时时间={timeout}秒, "
                f"堆栈跟踪:\n{traceback.format_exc()}"
            )
            return None
    
    async def close(self):
        """关闭STT会话（参考官方demo）"""
        if self.is_closed:
            return
        
        try:
            if self.conversation:
                # 异步关闭连接
                await asyncio.to_thread(self.conversation.close)
            
            self.is_connected = False
            self.is_closed = True
            
        except Exception as e:
            logger.error(
                f"[STT异常] request_id={self.request_id}, "
                f"错误类型=关闭会话异常, "
                f"异常类型={type(e).__name__}, "
                f"错误消息={str(e)}, "
                f"连接状态=is_connected={self.is_connected}, is_closed={self.is_closed}, "
                f"堆栈跟踪:\n{traceback.format_exc()}"
            )
    
    def get_current_text(self) -> str:
        """获取当前识别文本"""
        return self.callback.get_text()
    
    def get_last_update_time(self) -> float:
        """获取最后更新时间"""
        return self.callback.last_update_time
    
    def get_first_text_time(self) -> Optional[float]:
        """获取首个文本生成时间"""
        return self.callback.first_text_time


class STTService:
    """STT服务管理器（单例）"""
    
    _instance = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance
    
    def __init__(self):
        """初始化STT服务"""
        if not hasattr(self, '_initialized'):
            self._initialized = True
            
            # 设置API Key
            dashscope.api_key = settings.active_dashscope_api_key
            
            # 国际版endpoint配置
            if settings.use_dashscope_intl:
                from config import get_dashscope_domain
                dashscope_domain = get_dashscope_domain()
                dashscope.base_http_api_url = f"https://{dashscope_domain}/api/v1"
                dashscope.base_websocket_api_url = f"wss://{dashscope_domain}/api-ws/v1"
    
    async def create_session(
        self,
        request_id: str,
        on_text_update: Optional[Callable[[str, bool], Awaitable[None]]] = None,
        on_error: Optional[Callable[[str], Awaitable[None]]] = None,
        event_loop: Optional[asyncio.AbstractEventLoop] = None
    ) -> STTSession:
        """
        创建STT识别会话
        
        Args:
            request_id: 请求ID
            on_text_update: 文本更新回调
            on_error: 错误回调
            event_loop: 事件循环
            
        Returns:
            STT会话对象
        """
        session = STTSession(
            request_id=request_id,
            on_text_update=on_text_update,
            on_error=on_error,
            event_loop=event_loop
        )
        
        # 自动连接
        success = await session.connect()
        
        return session


# 全局STT服务实例
stt_service = STTService()