ASR-demo/realtime_asr_optimization_demo/auxiliary_service.py

"""独立 WebSocket Demo 使用的 VAD 和说话人辅助服务客户端。"""

from __future__ import annotations

import io
import wave
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from aiohttp import ClientSession, ClientTimeout, FormData


@dataclass(frozen=True)
class AuxiliaryServiceConfig:
    """辅助模型服务的 HTTP 连接配置。"""

    base_url: str = "http://127.0.0.1:8010"
    timeout_seconds: float = 45.0


def pcm16_to_wav(pcm_bytes: bytes, sample_rate: int = 16000) -> bytes:
    """将 Demo 内部的 16kHz 单声道 PCM16 封装成辅助服务可读取的 WAV。"""
    output = io.BytesIO()
    with wave.open(output, "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(sample_rate)
        wav_file.writeframes(pcm_bytes)
    return output.getvalue()


class AuxiliaryModelService:
    """调用独立辅助模型服务，不在 WebSocket 进程内加载 GPU 模型。"""

    def __init__(self, config: AuxiliaryServiceConfig) -> None:
        self.config = config
        self._session: ClientSession | None = None

    async def start(self) -> None:
        """创建可复用的 HTTP 会话，避免每个片段重复建立 TCP 连接。"""
        self._session = ClientSession(timeout=ClientTimeout(total=self.config.timeout_seconds))

    async def close(self) -> None:
        """关闭辅助服务 HTTP 会话。"""
        if self._session is not None:
            await self._session.close()
            self._session = None

    async def health(self) -> dict[str, Any]:
        """读取辅助服务健康状态，避免服务不可达时只能看到 ASR 的降级结果。"""
        if self._session is None:
            raise RuntimeError("auxiliary model service is not started")
        endpoint = self.config.base_url.rstrip("/") + "/health"
        async with self._session.get(endpoint) as response:
            body = await response.text()
            if response.status >= 400:
                raise RuntimeError(f"auxiliary health check failed ({response.status}): {body[:500]}")
            try:
                decoded = await response.json(content_type=None)
            except ValueError as exc:
                raise RuntimeError(f"auxiliary health check returned invalid JSON: {body[:500]}") from exc
        if not isinstance(decoded, dict):
            raise RuntimeError("auxiliary health check returned a non-object JSON value")
        return decoded

    async def resolve_speaker(
        self,
        pcm_bytes: bytes,
        session_id: str,
        start_time_ms: float,
        end_time_ms: float,
    ) -> dict[str, Any] | None:
        """提交一个已经由实时 VAD 完成的 turn，获取在线聚类结果。

        每次请求只包含当前 turn，不上传整段会话；辅助服务通过 session_id
        保存聚类中心，因此同一说话人在 A→B→A 场景下仍能保持同一标签。
        """
        if self._session is None:
            raise RuntimeError("auxiliary model service is not started")
        form = FormData()
        form.add_field("file", pcm16_to_wav(pcm_bytes), filename="turn.wav", content_type="audio/wav")
        form.add_field("session_id", session_id)
        form.add_field("start_time_ms", str(start_time_ms))
        form.add_field("end_time_ms", str(end_time_ms))
        endpoint = self.config.base_url.rstrip("/") + "/v1/speaker/resolve"
        async with self._session.post(endpoint, data=form) as response:
            body = await response.text()
            if response.status >= 400:
                raise RuntimeError(f"auxiliary speaker resolve failed ({response.status}): {body[:500]}")
            try:
                decoded = await response.json(content_type=None)
            except ValueError as exc:
                raise RuntimeError(f"auxiliary speaker resolve returned invalid JSON: {body[:500]}") from exc
        if not isinstance(decoded, dict):
            raise RuntimeError("auxiliary speaker resolve returned a non-object JSON value")
        if decoded.get("error"):
            raise RuntimeError(str(decoded["error"]))
        # 保留无标签响应里的具体原因；由组装器统一判断可信度，避免这里静默丢弃。
        return decoded

    async def reset_speaker_session(self, session_id: str) -> None:
        """通知辅助服务释放当前 WebSocket 对应的在线聚类状态。"""
        if self._session is None:
            return
        endpoint = self.config.base_url.rstrip("/") + "/v1/speaker/reset"
        try:
            async with self._session.post(endpoint, json={"session_id": session_id}) as response:
                await response.read()
        except Exception:
            # 清理失败不能影响已经完成的 ASR 结果，辅助服务会自行过期清理。
            return

    async def diarize(
        self,
        audio_bytes: bytes,
        source: str = "mic",
        file_name: str = "audio.wav",
    ) -> list[dict[str, Any]]:
        """提交完整会话音频，返回带毫秒时间范围和标签的聚类片段。

        麦克风、PCM 和 WAV 在 WebSocket 层已经能被识别为 16kHz PCM；
        MP3、M4A 等压缩文件必须保留原始容器，否则把压缩字节直接包装成
        PCM 会得到不可用的声纹输入。
        """
        if self._session is None:
            raise RuntimeError("auxiliary model service is not started")
        suffix = Path(file_name).suffix.lower()
        is_pcm = source == "mic" or suffix == ".pcm"
        if is_pcm:
            payload = pcm16_to_wav(audio_bytes)
            upload_name = "session.wav"
            content_type = "audio/wav"
        elif suffix == ".wav":
            payload = audio_bytes
            upload_name = "session.wav"
            content_type = "audio/wav"
        else:
            payload = audio_bytes
            upload_name = Path(file_name).name or "session.audio"
            content_type = {
                ".mp3": "audio/mpeg",
                ".m4a": "audio/mp4",
                ".ogg": "audio/ogg",
                ".opus": "audio/ogg",
            }.get(suffix, "application/octet-stream")
        form = FormData()
        form.add_field("file", payload, filename=upload_name, content_type=content_type)
        endpoint = self.config.base_url.rstrip("/") + "/v1/diarization"
        async with self._session.post(endpoint, data=form) as response:
            body = await response.text()
            if response.status >= 400:
                raise RuntimeError(f"auxiliary diarization failed ({response.status}): {body[:500]}")
            try:
                decoded = await response.json(content_type=None)
            except ValueError as exc:
                raise RuntimeError(f"auxiliary diarization returned invalid JSON: {body[:500]}") from exc
        raw_segments = decoded.get("segments", []) if isinstance(decoded, dict) else []
        if not isinstance(raw_segments, list):
            return []
        return [segment for segment in raw_segments if isinstance(segment, dict)]