first tests

scripts/ws_client_test.py

@@ -0,0 +1,42 @@
+"""Drive the server's WebSocket like the browser does: stream a wav in real time in 2048-sample chunks, 3 turns."""
+
+import asyncio
+import json
+import sys
+
+import numpy as np
+import soundfile as sf
+import websockets
+
+
+async def main(path, url="ws://127.0.0.1:8000/ws"):
+    audio, rate = sf.read(path, dtype="float32")
+    async with websockets.connect(url, max_size=None) as ws:
+        print("<", json.loads(await ws.recv())["type"])
+        for turn in range(3):  # later turns exercise the cached history
+            await ws.send(json.dumps({"type": "start", "rate": rate}))
+            for i in range(0, len(audio), 2048):
+                await ws.send(audio[i:i + 2048].tobytes())
+                await asyncio.sleep(2048 / rate)  # real-time pacing, like a mic
+            await ws.send(json.dumps({"type": "stop"}))
+            pcm_bytes = 0
+            while True:
+                msg = await ws.recv()
+                if isinstance(msg, bytes):
+                    pcm_bytes += len(msg)
+                    continue
+                m = json.loads(msg)
+                if m["type"] in ("token", "audio_chunk", "log"):
+                    continue
+                if m["type"] == "done":
+                    keys = ("first_text_ms", "first_frame_ms", "first_audio_ms", "flush_ms", "prefill_calls",
+                            "prefill_busy_ms", "context_tokens", "context_tokens_after", "rtf", "gen_ms", "codec_decode_ms", "audio_out_s")
+                    print("< done", m["text"][:80], {k: m["stats"].get(k) for k in keys})
+                elif m["type"] != "prompt":
+                    print("<", m["type"], {k: v for k, v in m.items() if k != "type"})
+                if m["type"] in ("done", "error"):
+                    break
+            print(f"  received {pcm_bytes / 4 / 24000:.2f}s of PCM")
+
+
+asyncio.run(main(sys.argv[1]))