File size: 2,220 Bytes
ac573f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | """Live-path test of server.py over a real WebSocket (real-time paced 32 ms frames).
python ws_test.py <ws_url> <user1.wav> [<barge.wav>] -> prints events; saves ws_test_agent.wav"""
import asyncio, json, sys, time, numpy as np, soundfile as sf, librosa, websockets
async def main(url, w1, w2=None):
def load(p):
y, sr = sf.read(p, dtype="float32"); y = y.mean(1) if y.ndim > 1 else y
return (librosa.resample(y, orig_sr=sr, target_sr=16000) * 32767).astype(np.int16)
u1 = load(w1); u2 = load(w2) if w2 else None
t0 = time.time(); ev = []; audio = []; first_audio = [None]; got_reply = [False]; got_reply = [False]
async with websockets.connect(url, max_size=2 ** 24) as ws:
async def rx():
async for m in ws:
t = round(time.time() - t0, 3)
if isinstance(m, bytes):
audio.append(np.frombuffer(m, np.int16)); first_audio[0] = first_audio[0] or (t if got_reply[0] else None); ev.append((t, "audio", len(m) // 2 / 24000))
else:
d = json.loads(m); got_reply[0] = got_reply[0] or d.get("type") == "reply"; ev.append((t, d.get("type"), {k: v for k, v in d.items() if k != "type"})); print(t, d, flush=True)
task = asyncio.create_task(rx())
async def send(y):
for i in range(0, len(y) - 511, 512):
await ws.send(y[i:i + 512].tobytes()); await asyncio.sleep(0.032)
z = np.zeros(512, np.int16)
await send(u1); print(round(time.time() - t0, 3), "user1 sent", flush=True)
for _ in range(int(6.0 / 0.032)):
await ws.send(z.tobytes()); await asyncio.sleep(0.032)
if u2 is not None and first_audio[0] and time.time() - t0 > first_audio[0] + 1.0: break
if u2 is not None:
print(round(time.time() - t0, 3), "barge-in user2 starts", flush=True); await send(u2)
for _ in range(int(8 / 0.032)):
await ws.send(z.tobytes()); await asyncio.sleep(0.032)
task.cancel()
if audio: sf.write("ws_test_agent.wav", np.concatenate(audio).astype(np.float32) / 32768, 24000)
print("WS_TEST_DONE", "first_audio_at", first_audio[0])
asyncio.run(main(*sys.argv[1:]))
|