plugins/audio/core.py

import base64
import logging
import os
import tempfile
from pathlib import Path

logger = logging.getLogger(__name__)

# 支持的音频格式后缀
SUPPORTED_AUDIO_SUFFIXES = {".wav", ".mp3", ".ogg", ".flac", ".m4a", ".aac", ".wma", ".opus", ".webm"}


class AudioTranscriptionService:
    """基于 transformers ASR pipeline 的音频转录服务。"""

    def __init__(self):
        self._pipeline = None
        self.model_name = os.getenv("ASR_MODEL", "openai/whisper-tiny")
        self._load_error = None  # 记录模型加载错误

    def get_pipeline(self):
        if self._pipeline is None:
            if self._load_error:
                raise RuntimeError(f"模型加载失败（已缓存错误）: {self._load_error}")
            logger.info(f"正在加载音频转录模型: {self.model_name}（首次加载约 10-30 秒）")
            from transformers import pipeline

            try:
                self._pipeline = pipeline(
                    "automatic-speech-recognition",
                    model=self.model_name,
                )
                logger.info("音频转录模型加载完成")
            except Exception as e:
                self._load_error = str(e)
                logger.error(f"音频转录模型加载失败: {e}")
                raise
        return self._pipeline

    def transcribe_file(self, audio_path: str, language: str | None = None) -> dict:
        """转录本地音频文件。"""
        try:
            path = Path(audio_path)
            if not path.exists():
                return {
                    "success": False,
                    "text": "",
                    "segments": [],
                    "error": f"音频文件不存在: {audio_path}",
                }

            # 验证文件格式
            suffix = path.suffix.lower()
            if suffix not in SUPPORTED_AUDIO_SUFFIXES:
                logger.warning(f"音频格式 {suffix} 可能不被支持，继续尝试转录")

            # 检查文件大小（警告大文件）
            file_size_mb = path.stat().st_size / (1024 * 1024)
            if file_size_mb > 50:
                logger.warning(f"音频文件较大 ({file_size_mb:.1f}MB)，转录可能需要较长时间")

            logger.info(f"开始转录: {path.name} ({file_size_mb:.1f}MB)")

            options = {"return_timestamps": True}
            if language:
                options["generate_kwargs"] = {"language": language}

            result = self.get_pipeline()(str(path), **options)
            text = result.get("text", "") if isinstance(result, dict) else str(result)

            logger.info(f"转录完成: {len(text)} 字符")

            return {
                "success": True,
                "text": text.strip(),
                "segments": result.get("chunks", []) if isinstance(result, dict) else [],
                "error": None,
            }
        except Exception as e:
            logger.error(f"音频转录失败: {e}", exc_info=True)
            return {
                "success": False,
                "text": "",
                "segments": [],
                "error": f"转录失败: {str(e)}",
            }

    def transcribe_base64(
        self,
        audio_base64: str,
        suffix: str = ".wav",
        language: str | None = None,
    ) -> dict:
        """转录 Base64 音频。"""
        try:
            audio_bytes = base64.b64decode(audio_base64)
        except Exception as e:
            return {
                "success": False,
                "text": "",
                "segments": [],
                "error": f"Base64 解码失败: {e}",
            }

        # 确保 suffix 以点开头
        if suffix and not suffix.startswith("."):
            suffix = f".{suffix}"

        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
            tmp_file.write(audio_bytes)
            tmp_path = tmp_file.name

        try:
            return self.transcribe_file(tmp_path, language)
        finally:
            Path(tmp_path).unlink(missing_ok=True)


_transcription_service: AudioTranscriptionService | None = None


def get_transcription_service() -> AudioTranscriptionService:
    global _transcription_service
    if _transcription_service is None:
        _transcription_service = AudioTranscriptionService()
    return _transcription_service