chat-ai-api/app/core/asr_client.py

"""
火山引擎大模型流式语音识别 (ASR) 客户端
协议地址: wss://openspeech.bytedance.com/api/v3/sauc/bigmodel
文档参考: https://www.volcengine.com/docs/6561/1354869
"""

import asyncio
import gzip
import json
import uuid
import subprocess
import tempfile
import os

import websockets

from ..config.config import Config

# ============ 协议常量 ============
PROTOCOL_VERSION = 0b0001
DEFAULT_HEADER_SIZE = 0b0001

# Message Type
FULL_CLIENT_REQUEST = 0b0001
AUDIO_ONLY_REQUEST = 0b0010
FULL_SERVER_RESPONSE = 0b1001
SERVER_ACK = 0b1011
SERVER_ERROR_RESPONSE = 0b1111

# Message Type Specific Flags
NO_SEQUENCE = 0b0000
POS_SEQUENCE = 0b0001
NEG_SEQUENCE = 0b0010
NEG_WITH_SEQUENCE = 0b0011

# Message Serialization
NO_SERIALIZATION = 0b0000
JSON_SERIALIZATION = 0b0001

# Message Compression
NO_COMPRESSION = 0b0000
GZIP_COMPRESSION = 0b0001

# ASR 服务地址
ASR_WS_URL = "wss://openspeech.bytedance.com/api/v3/sauc/bigmodel"


def _generate_header(
    message_type=FULL_CLIENT_REQUEST,
    message_type_specific_flags=NO_SEQUENCE,
    serial_method=JSON_SERIALIZATION,
    compression_type=GZIP_COMPRESSION,
    reserved_data=0x00,
) -> bytearray:
    """
    生成协议头 (4 bytes):
    - protocol_version (4 bits) + header_size (4 bits)
    - message_type (4 bits) + message_type_specific_flags (4 bits)
    - serialization_method (4 bits) + message_compression (4 bits)
    - reserved (8 bits)
    """
    header = bytearray()
    header_size = 1
    header.append((PROTOCOL_VERSION << 4) | header_size)
    header.append((message_type << 4) | message_type_specific_flags)
    header.append((serial_method << 4) | compression_type)
    header.append(reserved_data)
    return header


def _generate_before_payload(sequence: int) -> bytearray:
    """添加序列号信息 (4 bytes, big-endian, signed)"""
    before_payload = bytearray()
    before_payload.extend(sequence.to_bytes(4, 'big', signed=True))
    return before_payload


def _parse_response(res: bytes) -> dict:
    """解析服务器响应帧"""
    protocol_version = res[0] >> 4
    header_size = res[0] & 0x0f
    message_type = res[1] >> 4
    message_type_specific_flags = res[1] & 0x0f
    serialization_method = res[2] >> 4
    message_compression = res[2] & 0x0f
    reserved = res[3]
    header_extensions = res[4:header_size * 4]
    payload = res[header_size * 4:]

    result = {
        'is_last_package': False,
        'message_type': message_type,
    }
    payload_msg = None
    payload_size = 0

    if message_type_specific_flags & 0x01:
        seq = int.from_bytes(payload[:4], "big", signed=True)
        result['payload_sequence'] = seq
        payload = payload[4:]

    if message_type_specific_flags & 0x02:
        result['is_last_package'] = True

    if message_type == FULL_SERVER_RESPONSE:
        payload_size = int.from_bytes(payload[:4], "big", signed=True)
        payload_msg = payload[4:]
    elif message_type == SERVER_ACK:
        seq = int.from_bytes(payload[:4], "big", signed=True)
        result['seq'] = seq
        if len(payload) >= 8:
            payload_size = int.from_bytes(payload[4:8], "big", signed=False)
            payload_msg = payload[8:]
    elif message_type == SERVER_ERROR_RESPONSE:
        code = int.from_bytes(payload[:4], "big", signed=False)
        result['code'] = code
        payload_size = int.from_bytes(payload[4:8], "big", signed=False)
        payload_msg = payload[8:]

    if payload_msg is None:
        return result

    if message_compression == GZIP_COMPRESSION:
        payload_msg = gzip.decompress(payload_msg)
    if serialization_method == JSON_SERIALIZATION:
        payload_msg = json.loads(str(payload_msg, "utf-8"))
    elif serialization_method != NO_SERIALIZATION:
        payload_msg = str(payload_msg, "utf-8")

    result['payload_msg'] = payload_msg
    result['payload_size'] = payload_size
    return result


def _convert_webm_to_pcm(webm_data: bytes) -> bytes:
    """
    使用 ffmpeg 将 WebM/Opus 音频转为 PCM 16kHz 16-bit 单声道。
    返回原始 PCM 字节。
    """
    inp_path = None
    out_path = None
    try:
        with tempfile.NamedTemporaryFile(suffix='.webm', delete=False) as inp:
            inp.write(webm_data)
            inp_path = inp.name
        out_path = inp_path.replace('.webm', '.pcm')

        subprocess.run(
            [
                'ffmpeg', '-y', '-i', inp_path,
                '-f', 's16le', '-ar', '16000', '-ac', '1',
                out_path
            ],
            check=True,
            capture_output=True,
        )
        with open(out_path, 'rb') as f:
            return f.read()
    finally:
        if inp_path and os.path.exists(inp_path):
            os.unlink(inp_path)
        if out_path and os.path.exists(out_path):
            os.unlink(out_path)


async def transcribe_audio(audio_data: bytes) -> str:
    """
    将音频二进制数据 (WebM/Opus) 通过火山引擎大模型 ASR 转写为文本。

    流程:
    1. ffmpeg 转码 WebM → PCM 16kHz 16bit mono
    2. WebSocket 连接 ASR 服务
    3. 发送 FULL_CLIENT_REQUEST（音频参数）
    4. 分片发送 AUDIO_ONLY_REQUEST（音频数据）
    5. 发送结束帧
    6. 接收并拼接识别结果
    """
    config = Config()

    if not config.VOLC_ASR_API_KEY:
        raise ValueError("ASR 配置缺失: 请设置 VOLC_ASR_API_KEY 环境变量")

    # 1. 转码为 PCM
    pcm_data = await asyncio.get_event_loop().run_in_executor(
        None, _convert_webm_to_pcm, audio_data
    )

    if not pcm_data:
        raise ValueError("音频转码失败: PCM 数据为空")

    # 2. 建立 WebSocket 连接
    connect_id = str(uuid.uuid4())
    headers = {
        "X-Api-Key": config.VOLC_ASR_API_KEY,
        "X-Api-Resource-Id": config.VOLC_ASR_RESOURCE_ID,
        "X-Api-Connect-Id": connect_id,
    }

    # 3. 构建初始请求参数
    request_params = {
        "user": {
            "uid": connect_id,
        },
        "audio": {
            "format": "pcm",
            "rate": 16000,
            "bits": 16,
            "channel": 1,
            "codec": "raw",
        },
        "request": {
            "model_name": "bigmodel",
            "enable_itn": True,
            "result_type": "full",
        }
    }

    final_text = ""

    async with websockets.connect(
        ASR_WS_URL,
        extra_headers=headers,
        max_size=10 * 1024 * 1024,
    ) as ws:
        # 发送 FULL_CLIENT_REQUEST
        payload_bytes = gzip.compress(json.dumps(request_params).encode('utf-8'))
        full_request = bytearray(_generate_header(
            message_type=FULL_CLIENT_REQUEST,
            message_type_specific_flags=POS_SEQUENCE,
        ))
        full_request.extend(_generate_before_payload(sequence=1))
        full_request.extend(len(payload_bytes).to_bytes(4, 'big'))
        full_request.extend(payload_bytes)
        await ws.send(full_request)

        # 接收 ACK
        res = await ws.recv()
        ack_result = _parse_response(res)
        if ack_result.get('message_type') == SERVER_ERROR_RESPONSE:
            raise RuntimeError(f"ASR 连接错误: {ack_result.get('payload_msg', ack_result)}")

        # 4. 分片发送音频数据 (每片 3200 bytes ≈ 100ms @16kHz 16bit mono)
        chunk_size = 3200
        seq = 2
        total_chunks = (len(pcm_data) + chunk_size - 1) // chunk_size

        for i in range(0, len(pcm_data), chunk_size):
            chunk = pcm_data[i:i + chunk_size]
            is_last = (i + chunk_size >= len(pcm_data))

            compressed_chunk = gzip.compress(chunk)

            if is_last:
                # 最后一帧用 NEG_WITH_SEQUENCE 标志
                audio_request = bytearray(_generate_header(
                    message_type=AUDIO_ONLY_REQUEST,
                    message_type_specific_flags=NEG_WITH_SEQUENCE,
                    serial_method=NO_SERIALIZATION,
                    compression_type=GZIP_COMPRESSION,
                ))
                audio_request.extend(_generate_before_payload(sequence=-seq))
            else:
                audio_request = bytearray(_generate_header(
                    message_type=AUDIO_ONLY_REQUEST,
                    message_type_specific_flags=POS_SEQUENCE,
                    serial_method=NO_SERIALIZATION,
                    compression_type=GZIP_COMPRESSION,
                ))
                audio_request.extend(_generate_before_payload(sequence=seq))

            audio_request.extend(len(compressed_chunk).to_bytes(4, 'big'))
            audio_request.extend(compressed_chunk)
            await ws.send(audio_request)

            # 接收中间结果
            res = await ws.recv()
            result = _parse_response(res)

            if result.get('message_type') == SERVER_ERROR_RESPONSE:
                raise RuntimeError(f"ASR 识别错误: {result.get('payload_msg', result)}")

            # 提取文本结果
            if 'payload_msg' in result and isinstance(result['payload_msg'], dict):
                text = result['payload_msg'].get('result', {}).get('text', '')
                if text:
                    final_text = text  # 每次全量更新（result_type=full）

            seq += 1

            # 如果是最后一帧且收到结束标记
            if result.get('is_last_package'):
                break

            # 控制发送速率，避免过快
            if not is_last:
                await asyncio.sleep(0.02)

        # 5. 如果还没收到最终结果，继续接收
        if not final_text or not _parse_response(res).get('is_last_package', False):
            try:
                while True:
                    res = await asyncio.wait_for(ws.recv(), timeout=5.0)
                    result = _parse_response(res)
                    if 'payload_msg' in result and isinstance(result['payload_msg'], dict):
                        text = result['payload_msg'].get('result', {}).get('text', '')
                        if text:
                            final_text = text
                    if result.get('is_last_package'):
                        break
            except (asyncio.TimeoutError, websockets.exceptions.ConnectionClosed):
                pass

    return final_text.strip()