198 lines
9.0 KiB
Python
198 lines
9.0 KiB
Python
"""独立 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
|
||
# 旧版辅助服务没有窗口声纹端点；探测到一次 404/405 后不再重复请求。
|
||
self.speaker_embedding_unsupported = False
|
||
|
||
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,
|
||
speaker_verified: bool = True,
|
||
) -> dict[str, Any] | None:
|
||
"""提交一个已经由实时 VAD 完成的 turn，获取在线聚类结果。
|
||
|
||
每次请求只包含当前 turn，不上传整段会话；辅助服务通过 session_id
|
||
保存聚类中心，因此同一说话人在 A→B→A 场景下仍能保持同一标签。
|
||
speaker_verified=False 表示换人强制切段产生的边界段，辅助服务
|
||
只匹配标签、不用它更新簇质心。
|
||
"""
|
||
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))
|
||
form.add_field("speaker_verified", "1" if speaker_verified else "0")
|
||
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 speaker_embedding(self, pcm_bytes: bytes) -> list[float] | None:
|
||
"""为一个活跃 turn 的短窗口提取归一化声纹，绝不读取或更新聚类状态。
|
||
|
||
任何失败都返回 None 而不是抛异常：窗比对只是切段辅助，绝不能
|
||
把辅助服务的抖动传导成 ASR 阻塞。404/405 视为服务版本过旧，置位
|
||
speaker_embedding_unsupported 供 WebSocket 侧停用该功能。"""
|
||
if self._session is None or self.speaker_embedding_unsupported:
|
||
return None
|
||
form = FormData()
|
||
form.add_field("file", pcm16_to_wav(pcm_bytes), filename="window.wav", content_type="audio/wav")
|
||
endpoint = self.config.base_url.rstrip("/") + "/v1/speaker/embedding"
|
||
try:
|
||
# 窗比对嵌在音频帧循环里，必须用远短于会话级 45s 的超时兜底。
|
||
async with self._session.post(endpoint, data=form, timeout=ClientTimeout(total=8.0)) as response:
|
||
if response.status in {404, 405}:
|
||
self.speaker_embedding_unsupported = True
|
||
return None
|
||
if response.status >= 400:
|
||
return None
|
||
decoded = await response.json(content_type=None)
|
||
except Exception:
|
||
return None
|
||
embedding = decoded.get("embedding") if isinstance(decoded, dict) else None
|
||
if not isinstance(embedding, list) or not embedding:
|
||
return None
|
||
try:
|
||
return [float(value) for value in embedding]
|
||
except (TypeError, ValueError):
|
||
return None
|
||
|
||
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)]
|