Files
ai-companion/client/test_client.py
T

364 lines
14 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
"""最小 aiortc 客户端(实验用)。
**为什么需要它**:Pipecat 的官方客户端 SDK 覆盖浏览器 / 移动端 / C++ / ESP32,
唯独没有适配"跑 Linux 的 x86 盒子接 USB 麦克风阵列"——也就是 N5095 的形态。
所以客户端必须自己写,这个文件是雏形。
它同时是验证工具:不依赖浏览器就能测服务端。用法:
# 发一段 wav,收回应答音频
python client/test_client.py --server http://<服务器>:7860 \
--wav p0/stt_tts_bench/wavs/stt_1.wav --out /tmp/reply.wav
# 只做连接测试,不发音频
python client/test_client.py --server http://<服务器>:7860 --connect-only
**它还不是什么**:不做 AEC、不做播客参考、不做打断(那些要真实声卡和硬件),
只用文件音频验证信令与媒体通路。
"""
from __future__ import annotations
import argparse
import asyncio
import fractions
import json
import logging
import sys
import time
import urllib.request
import wave
from pathlib import Path
import numpy as np
from aiortc import (
RTCConfiguration,
RTCPeerConnection,
RTCSessionDescription,
)
from aiortc.mediastreams import AudioStreamTrack, MediaStreamError
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("client")
class SilentTrack(AudioStreamTrack):
"""静音音轨。只为触发 ICE 候选收集,不发真实音频。"""
kind = "audio"
def __init__(self, sample_rate: int = 16000):
super().__init__()
self.sample_rate = sample_rate
self._n = 0
async def recv(self):
import fractions as _fr
from av import AudioFrame as _AF
await asyncio.sleep(0.02)
pcm = np.zeros(self.sample_rate // 50, dtype=np.int16)
f = _AF(format="s16", layout="mono", samples=len(pcm))
f.planes[0].update(pcm.tobytes())
f.sample_rate = self.sample_rate
f.pts = self._n
f.time_base = _fr.Fraction(1, self.sample_rate)
self._n += len(pcm)
return f
class FileAudioTrack(AudioStreamTrack):
"""把 wav 文件当作麦克风音源,按实时速率吐出来。
实时速率很重要——如果一次性灌完,服务端的 VAD 会把它当成
一段超长的连续语音,测出来的延迟也不真实。
"""
kind = "audio"
def __init__(self, wav_path: Path, *, sample_rate: int = 16000,
tail_silence_s: float = 1.2):
super().__init__()
with wave.open(str(wav_path)) as w:
assert w.getnchannels() == 1, "需要单声道"
self.src_rate = w.getframerate()
pcm = np.frombuffer(w.readframes(w.getnframes()), dtype=np.int16)
self.samples = pcm.astype(np.float32) / 32768.0
self.sample_rate = sample_rate
if self.src_rate != sample_rate:
n = int(len(self.samples) * sample_rate / self.src_rate)
self.samples = np.interp(np.linspace(0, len(self.samples) - 1, n),
np.arange(len(self.samples)), self.samples)
self._source_pos = 0
self._timestamp_samples = 0
self._tail_remaining = int(tail_silence_s * sample_rate)
self.source_finished_at: float | None = None
self._t0: float | None = None
self._frame_samples = sample_rate // 100 # 10ms
async def recv(self):
"""按 10ms 一块、以真实时间节奏产出。"""
if self._t0 is None:
self._t0 = time.perf_counter()
self._source_pos = 0
frame_pts = self._timestamp_samples
if self._source_pos >= len(self.samples):
if self.source_finished_at is None:
self.source_finished_at = time.perf_counter()
audio = np.zeros(self._frame_samples, dtype=np.float32)
self._tail_remaining = max(0, self._tail_remaining - self._frame_samples)
else:
audio = self.samples[self._source_pos:self._source_pos + self._frame_samples]
if len(audio) < self._frame_samples:
audio = np.pad(audio, (0, self._frame_samples - len(audio)))
self._source_pos += self._frame_samples
# 即使源文件已经结束,静音帧也必须持续推进时钟和 PTS。
self._timestamp_samples += self._frame_samples
# 实时节奏
target = self._t0 + self._timestamp_samples / self.sample_rate
delay = target - time.perf_counter()
if delay > 0:
await asyncio.sleep(delay)
pcm = (np.clip(audio, -1, 1) * 32767).astype(np.int16)
from av import AudioFrame
f = AudioFrame(format="s16", layout="mono", samples=len(pcm))
f.planes[0].update(pcm.tobytes())
f.sample_rate = self.sample_rate
f.pts = int(frame_pts * 90000 / self.sample_rate)
f.time_base = fractions.Fraction(1, 90000)
return f
@property
def finished(self) -> bool:
return self._source_pos >= len(self.samples) and self._tail_remaining <= 0
class ReplyRecorder:
"""收服务端音频,写 wav,并记录首音延迟。"""
def __init__(self, out_path: Path, sample_rate: int = 24000,
speech_rms_threshold: float = 80.0):
self.out_path = out_path
self.sample_rate = sample_rate
self.chunks: list[np.ndarray] = []
self.first_non_silent_at: float | None = None
self.speech_rms_threshold = speech_rms_threshold
def feed(self, frame) -> None:
arr = frame.to_ndarray()
channels = len(frame.layout.channels)
if frame.format.is_planar:
raw = arr.astype(np.float32).mean(axis=0)
else:
raw = arr.astype(np.float32).reshape(-1, channels).mean(axis=1)
# WebRTC/Opus 通常以 48kHz 交付,即使服务端原始输出是 24kHz。
# WAV 必须记录接收帧的实际采样率,否则时长会被算错。
if not self.chunks:
self.sample_rate = frame.sample_rate
elif frame.sample_rate != self.sample_rate:
raise ValueError(
f"回复音轨采样率发生变化: {self.sample_rate} -> {frame.sample_rate}")
raw = raw.reshape(-1)
rms = float(np.sqrt(np.mean(raw ** 2))) if raw.size else 0.0
if self.first_non_silent_at is None and rms >= self.speech_rms_threshold:
self.first_non_silent_at = time.perf_counter()
self.chunks.append(raw / 32768.0)
def save(self) -> float:
if not self.chunks:
return 0.0
audio = np.concatenate(self.chunks)
pcm = (np.clip(audio, -1, 1) * 32767).astype(np.int16)
self.out_path.parent.mkdir(parents=True, exist_ok=True)
with wave.open(str(self.out_path), "wb") as w:
w.setnchannels(1)
w.setsampwidth(2)
w.setframerate(self.sample_rate)
w.writeframes(pcm.tobytes())
return len(audio) / self.sample_rate
@property
def has_speech(self) -> bool:
return self.first_non_silent_at is not None
def first_audio_ms_from(self, started_at: float | None) -> float | None:
if started_at is None or self.first_non_silent_at is None:
return None
return (self.first_non_silent_at - started_at) * 1000
def post_offer(server: str, sdp: str, type_: str) -> dict:
"""把 offer 发给服务端,取回 answer。"""
payload = json.dumps({"sdp": sdp, "type": type_, "pc_id": None,
"restart_pc": False, "request_data": None}).encode()
req = urllib.request.Request(f"{server.rstrip('/')}/api/offer", data=payload,
headers={"Content-Type": "application/json"},
method="POST")
with urllib.request.urlopen(req, timeout=30) as r:
return json.loads(r.read().decode())
async def run(server: str, wav: Path | None, out: Path, *, connect_only: bool,
timeout_s: float) -> int:
pc = RTCPeerConnection(RTCConfiguration(iceServers=[]))
recorder = ReplyRecorder(out)
track = None
got_audio = asyncio.Event()
@pc.on("connectionstatechange")
async def _state():
log.info("连接状态: %s", pc.connectionState)
if pc.connectionState in ("failed", "closed"):
got_audio.set()
@pc.on("track")
def _track(t):
log.info("收到服务端音轨: %s", t.kind)
async def _pump():
try:
while True:
frame = await t.recv()
recorder.feed(frame)
except (MediaStreamError, asyncio.CancelledError):
pass
finally:
got_audio.set()
asyncio.ensure_future(_pump())
if wav and not connect_only:
track = FileAudioTrack(wav)
pc.addTrack(track)
log.info("音源: %s(%d 采样点)", wav.name, len(track.samples))
else:
# --connect-only 也要加一条(静音)音轨。
# **不能省**:aiortc 的 offer 里没有媒体时完全不收集 ICE 候选,
# 现象是"信令成功、SDP 交换正常、但 ICE 永远停在 new"。
pc.addTrack(SilentTrack())
log.info("连接测试模式:已加静音音轨(仅为触发 ICE 收集)")
# **必须建 data channel 并定期发 "ping"。**
# Pipecat 的 SmallWebRTC 用 data channel 上的消息来判断"客户端还活着":
# if isinstance(message, str) and message.startswith("ping"):
# self._last_received_time = time.time()
# 而 connection 的 connected 事件、进而 on_client_connected 回调、
# 进而 StartFrame(音频推流的开关)都依赖它。
# 不发 ping 的现象是:ICE 连通、音轨协商成功、但**服务端一帧音频都不推**、
# 助手永远不说话,而且日志里只有一句 "Data channel not established within 10s"。
channel = pc.createDataChannel("pipecat")
dc_state = {"ok": False}
@channel.on("open")
def _dc_open():
dc_state["ok"] = True
log.info("data channel 已打开,开始发心跳")
async def _keepalive():
while pc.connectionState not in ("closed", "failed"):
if dc_state["ok"]:
try:
channel.send("ping")
except Exception: # noqa: BLE001
pass
await asyncio.sleep(1.0)
asyncio.ensure_future(_keepalive())
offer = await pc.createOffer()
await pc.setLocalDescription(offer)
# 等 ICE 收集完成再发 offer——**这步不能省**。
# 不等的话 SDP 里没有任何候选地址,服务端拿到的是一份"空地址"的 offer,
# 信令看着成功,但 ICE 永远连不上(状态停在 new)。
# 这叫非 trickle ICE:把候选都塞进 SDP 一次发完。
# 真实产品会用 trickle(边收集边发候选),但那样需要额外的信令端点。
for _ in range(100):
if pc.iceGatheringState == "complete":
break
await asyncio.sleep(0.1)
n_cand = pc.localDescription.sdp.count("candidate:")
log.info("ICE 收集: %s,SDP 含 %d 个候选", pc.iceGatheringState, n_cand)
if n_cand == 0:
log.warning("SDP 里没有候选地址,ICE 很可能连不上")
log.info("本地 SDP 已生成,发给服务端…")
t_send = time.perf_counter()
answer = post_offer(server, pc.localDescription.sdp, pc.localDescription.type)
if answer.get("type") not in ("answer", None):
log.error("服务端返回异常: %s", answer)
return 1
if answer.get("sdp"):
await pc.setRemoteDescription(RTCSessionDescription(
sdp=answer["sdp"], type=answer["type"]))
log.info("远端 SDP 已设置(往返 %.0f ms)", (time.perf_counter() - t_send) * 1000)
else:
log.info("服务端未返回 SDP(可能是 pc_id 复用): %s", answer)
# 等 ICE 连通
for _ in range(100):
if pc.connectionState in ("connected", "completed"):
break
if pc.connectionState in ("failed", "closed"):
log.error("连接失败")
await pc.close()
return 1
await asyncio.sleep(0.1)
log.info("✅ ICE 已连通(%s)", pc.connectionState)
print(f" ✅ 连接建立,状态 {pc.connectionState}")
if track:
# 等音频放完 + 留出响应时间
t0 = time.perf_counter()
while not track.finished and time.perf_counter() - t0 < timeout_s:
await asyncio.sleep(0.2)
log.info("音频已发完,等待回复…")
# 等首个非静音回复;SmallWebRTC 空闲时也会持续发送静音。
t_wait = time.perf_counter()
while time.perf_counter() - t_wait < timeout_s:
if recorder.has_speech:
break
await asyncio.sleep(0.1)
# 再收一会儿,让回复说完
await asyncio.sleep(6)
dur = recorder.save()
await pc.close()
got_audio.set()
print()
print("=== 结果 ===")
print(f" 回复: {'✅ 收到非静音语音' if recorder.has_speech else '❌ 只有静音'}")
if dur:
print(f" 录制音轨: {dur:.2f} 秒 → {out}")
fa = recorder.first_audio_ms_from(track.source_finished_at if track else None)
if fa is not None:
print(f" 首音延迟(源音频结束→首个非静音输出): {fa:.0f} ms")
return 0 if (connect_only or recorder.has_speech) else 1
def main() -> int:
p = argparse.ArgumentParser()
p.add_argument("--server", default="http://cadadmin-server.taild0f91d.ts.net:7860",
help="服务端地址。默认用 tailscale 域名——LAN 地址会被重新编址"
"(2026-09-15 实测从 192.168.0.x 变成 172.16.1.x)")
p.add_argument("--wav", type=Path)
p.add_argument("--out", type=Path, default=Path("/tmp/client_reply.wav"))
p.add_argument("--connect-only", action="store_true")
p.add_argument("--timeout", type=float, default=30.0)
a = p.parse_args()
if not a.connect_only and not a.wav:
print("需要 --wav 或 --connect-only", file=sys.stderr)
return 2
return asyncio.run(run(a.server, a.wav, a.out,
connect_only=a.connect_only, timeout_s=a.timeout))
if __name__ == "__main__":
sys.exit(main())