diff --git a/scripts/probe_cartesia.py b/scripts/probe_cartesia.py new file mode 100755 index 0000000..a2e0ba1 --- /dev/null +++ b/scripts/probe_cartesia.py @@ -0,0 +1,166 @@ +#!/usr/bin/env python3 +"""R2 probe — Cartesia Sonic TTS first-audio-byte latency. + +Per PLAN.md SLICE-01 TASK-01-03: WebSocket to Cartesia Sonic, send a sample text +chunk, measure first-audio-byte latency over 20 iterations; log min/median/p95. + +Exit code 0 in all cases: + - If CARTESIA_API_KEY is missing, print KEY_MISSING and exit 0. + - If present, run the live probe and print a latency table. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import statistics +import sys +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +try: + from dotenv import load_dotenv + + load_dotenv() +except ImportError: # pragma: no cover + pass + + +def _banner(msg: str) -> None: + print("\n" + "=" * 72) + print(msg) + print("=" * 72 + "\n") + + +def _require_key() -> str | None: + key = os.environ.get("CARTESIA_API_KEY", "").strip() + if not key: + _banner( + "KEY_MISSING — CARTESIA_API_KEY not set.\n" + " Cannot run live Cartesia probe. Probe infrastructure is built\n" + " and ready; live measurements are pending API key provisioning.\n" + " Set CARTESIA_API_KEY in .env (see .env.example) and re-run." + ) + return None + return key + + +SAMPLE_TEXT = ( + "Hi, I received my order yesterday and the item is cracked. " + "I want my money back." +) + +CARTESIA_WS_URL = "wss://api.cartesia.ai/tts/websocket" +DEFAULT_VOICE_ID = "a3536a36-1d18-4efb-a95a-7c44b7b5e384" + + +async def _probe_once(api_key: str, voice_id: str, model_id: str) -> float | None: + """Open Cartesia WS, request TTS, return ms-to-first-audio-byte.""" + import websockets + + headers = [("x-api-key", api_key), ("cartesia-version", "2024-06-10")] + t0 = time.perf_counter() + first_audio_ms: float | None = None + + try: + async with websockets.connect( + CARTESIA_WS_URL, additional_headers=headers, open_timeout=10 + ) as ws: + req = { + "model_id": model_id, + "transcript": SAMPLE_TEXT, + "voice": {"id": voice_id}, + "output_format": { + "container": "raw", + "encoding": "pcm_s16le", + "sample_rate": 24000, + }, + "stream": True, + } + await ws.send(json.dumps(req)) + # Read frames until we get the first audio chunk. + while True: + msg = await asyncio.wait_for(ws.recv(), timeout=10) + if isinstance(msg, (bytes, bytearray)): + first_audio_ms = (time.perf_counter() - t0) * 1000.0 + break + # JSON control messages (e.g. done) — ignore until audio. + if isinstance(msg, str): + data = json.loads(msg) + if data.get("type") == "done": + break + except Exception as exc: # pragma: no cover - network/auth errors + print(f" [probe] Cartesia connection failed: {exc}") + return None + + return first_audio_ms + + +async def run_live(api_key: str, iterations: int, voice_id: str, model_id: str) -> list[float]: + samples: list[float] = [] + print(f" Running {iterations} Cartesia Sonic iterations (voice={voice_id})...") + for i in range(iterations): + ms = await _probe_once(api_key, voice_id, model_id) + if ms is not None: + samples.append(ms) + print(f" [{i + 1:2d}/{iterations}] first-audio: {ms:6.1f} ms") + else: + print(f" [{i + 1:2d}/{iterations}] no audio received (skipped)") + await asyncio.sleep(0.3) + return samples + + +def _summarize(samples: list[float], label: str) -> dict: + if not samples: + print(f"\n {label}: no samples collected.\n") + return {"label": label, "n": 0} + s = sorted(samples) + p95 = s[int(0.95 * (len(s) - 1))] + row = { + "label": label, + "n": len(s), + "min_ms": round(min(s), 1), + "median_ms": round(statistics.median(s), 1), + "p95_ms": round(p95, 1), + "mean_ms": round(statistics.mean(s), 1), + } + print( + f" {label}: n={row['n']} min={row['min_ms']:.1f} " + f"median={row['median_ms']:.1f} p95={row['p95_ms']:.1f} " + f"mean={row['mean_ms']:.1f} (ms)" + ) + return row + + +async def amain() -> int: + parser = argparse.ArgumentParser(description="R2 Cartesia Sonic latency probe") + parser.add_argument("--iterations", type=int, default=20) + parser.add_argument("--voice-id", default=os.environ.get("CARTESIA_VOICE_ID", DEFAULT_VOICE_ID)) + parser.add_argument("--model-id", default="sonic-2") + parser.add_argument("--out", default=None) + args = parser.parse_args() + + _banner("R2 PROBE — Cartesia Sonic first-audio-byte latency") + api_key = _require_key() + if api_key is None: + return 0 + + samples = await run_live(api_key, args.iterations, args.voice_id, args.model_id) + summary = _summarize(samples, "cartesia_sonic_first_audio") + print() + if args.out: + Path(args.out).write_text(json.dumps(summary, indent=2)) + print(f" Wrote {args.out}") + return 0 + + +def main() -> int: + return asyncio.run(amain()) + + +if __name__ == "__main__": + sys.exit(main()) \ No newline at end of file diff --git a/scripts/probe_deepgram.py b/scripts/probe_deepgram.py index 8ad8683..4d73c1f 100755 --- a/scripts/probe_deepgram.py +++ b/scripts/probe_deepgram.py @@ -78,66 +78,58 @@ def _synth_pcm(duration_s: float = 2.0, sample_rate: int = 16000) -> bytes: return bytes(frames) +DEEPGRAM_WS_URL = "wss://api.deepgram.com/v1/listen" + + async def _probe_once(api_key: str, model: str, pcm: bytes, sample_rate: int) -> float | None: - """Open a Deepgram streaming WebSocket, send PCM, return ms-to-first-partial.""" - from deepgram import ( - DeepgramClient, - LiveTranscriptionEvents, - LiveOptions, - ) - - client = DeepgramClient(api_key) - latency_holder: dict[str, float | None] = {"first_partial_ms": None} - - def _on_message(_result, *args, **kwargs): # deepgram callback signature - if latency_holder["first_partial_ms"] is not None: - return # only the first partial matters - # deepgram returns an object with .type or a dict-like; handle both. - rtype = getattr(_result, "type", None) or ( - _result.get("type") if isinstance(_result, dict) else None - ) - if rtype == "Results": - t = time.perf_counter() - latency_holder["first_partial_ms"] = (t - latency_holder["t0"]) * 1000.0 - - def _on_open(*args, **kwargs): - latency_holder["t0"] = time.perf_counter() - - dg_connection = client.listen.asyncwebsocket.v1 - dg_connection.on(LiveTranscriptionEvents.Open, _on_open) - dg_connection.on(LiveTranscriptionEvents.Transcript, _on_message) - - options = LiveOptions( - model=model, - language="en", - encoding="linear16", - channels=1, - sample_rate=sample_rate, - interim_results=True, - endpointing=300, + """Open a Deepgram streaming WebSocket, send PCM, return ms-to-first-partial. + + Uses the raw Deepgram streaming WebSocket API (not the SDK) so the probe is + independent of SDK version churn and measures the actual network path. + """ + import websockets + + params = ( + f"?model={model}&language=en&encoding=linear16&channels=1" + f"&sample_rate={sample_rate}&interim_results=true&endpointing=300" ) + headers = [("Authorization", f"Token {api_key}")] + t0 = time.perf_counter() + first_partial_ms: float | None = None try: - await dg_connection.start(options) + async with websockets.connect( + DEEPGRAM_WS_URL + params, additional_headers=headers, open_timeout=10 + ) as ws: + # Send in small chunks to mimic real streaming. + chunk = 3200 # 100ms of 16kHz mono 16-bit + for i in range(0, len(pcm), chunk): + await ws.send(pcm[i : i + chunk]) + await asyncio.sleep(0.02) + # Wait for the first transcript message. + try: + while True: + msg = await asyncio.wait_for(ws.recv(), timeout=5) + if isinstance(msg, str): + data = json.loads(msg) + if data.get("type") == "Results": + channel = data.get("channel", {}) + alts = channel.get("alternatives", []) + if alts and alts[0].get("transcript", "").strip(): + first_partial_ms = (time.perf_counter() - t0) * 1000.0 + break + except asyncio.TimeoutError: + pass + # Signal close. + try: + await ws.send(json.dumps({"type": "CloseStream"})) + except Exception: + pass except Exception as exc: # pragma: no cover - network/auth errors print(f" [probe] Deepgram connection failed: {exc}") return None - try: - # Send in small chunks to mimic real streaming. - chunk = 3200 # 100ms of 16kHz mono 16-bit - for i in range(0, len(pcm), chunk): - await dg_connection.send(pcm[i : i + chunk]) - await asyncio.sleep(0.02) - # Wait a bit for the final partial to arrive. - await asyncio.sleep(1.5) - finally: - try: - await dg_connection.finish() - except Exception: - pass - - return latency_holder.get("first_partial_ms") + return first_partial_ms async def run_live(api_key: str, iterations: int, model: str) -> list[float]: