diff --git a/scripts/probe_deepgram.py b/scripts/probe_deepgram.py new file mode 100755 index 0000000..8ad8683 --- /dev/null +++ b/scripts/probe_deepgram.py @@ -0,0 +1,207 @@ +#!/usr/bin/env python3 +"""R1 probe — Deepgram Nova-3 streaming ASR first-partial-transcript latency. + +Per PLAN.md SLICE-01 TASK-01-02: measure first-partial-transcript latency from a +sample audio file (synthesized PCM) over 20 iterations; log min/median/p95. + +Exit code 0 in all cases: + - If DEEPGRAM_API_KEY is missing, print a clear KEY_MISSING banner and exit 0 + (the probe infrastructure is the deliverable; live numbers come when keys + are provisioned). + - If the key is present, run the live probe and print a latency table. + +Usage: + python scripts/probe_deepgram.py [--iterations N] [--model nova-3] +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import statistics +import sys +import time +from pathlib import Path + +# Make the project importable when run from the repo root. +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +try: + from dotenv import load_dotenv + + load_dotenv() +except ImportError: # pragma: no cover - dotenv is a declared dep + pass + + +def _banner(msg: str) -> None: + print("\n" + "=" * 72) + print(msg) + print("=" * 72 + "\n") + + +def _require_key() -> str | None: + """Return the Deepgram API key or None (with a printed banner if missing).""" + key = os.environ.get("DEEPGRAM_API_KEY", "").strip() + if not key: + _banner( + "KEY_MISSING — DEEPGRAM_API_KEY not set.\n" + " Cannot run live Deepgram probe. Probe infrastructure is built\n" + " and ready; live measurements are pending API key provisioning.\n" + " Set DEEPGRAM_API_KEY in .env (see .env.example) and re-run." + ) + return None + return key + + +def _synth_pcm(duration_s: float = 2.0, sample_rate: int = 16000) -> bytes: + """Synthesize a short mono 16-bit PCM buffer (silence + a low tone). + + Deepgram needs real audio frames; we generate a recognizable signal so the + streaming endpoint returns a partial. The exact transcript content is not + the point — the *latency to first partial* is. + """ + import math + import struct + + n = int(duration_s * sample_rate) + frames = bytearray() + for i in range(n): + # 220 Hz tone for the first 1.5s, then silence — a clearly voiced segment. + if i < int(1.5 * sample_rate): + sample = int(16000 * math.sin(2 * math.pi * 220 * i / sample_rate)) + else: + sample = 0 + frames += struct.pack(" 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, + ) + + try: + await dg_connection.start(options) + 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") + + +async def run_live(api_key: str, iterations: int, model: str) -> list[float]: + sample_rate = 16000 + pcm = _synth_pcm(duration_s=2.0, sample_rate=sample_rate) + samples: list[float] = [] + print(f" Running {iterations} Deepgram Nova-3 iterations (model={model})...") + for i in range(iterations): + ms = await _probe_once(api_key, model, pcm, sample_rate) + if ms is not None: + samples.append(ms) + print(f" [{i + 1:2d}/{iterations}] first-partial: {ms:6.1f} ms") + else: + print(f" [{i + 1:2d}/{iterations}] no partial 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="R1 Deepgram Nova-3 latency probe") + parser.add_argument("--iterations", type=int, default=20) + parser.add_argument("--model", default=os.environ.get("DEEPGRAM_MODEL", "nova-3")) + parser.add_argument("--out", default=None, help="optional JSON results path") + args = parser.parse_args() + + _banner("R1 PROBE — Deepgram Nova-3 first-partial-transcript latency") + api_key = _require_key() + if api_key is None: + return 0 + + samples = await run_live(api_key, args.iterations, args.model) + summary = _summarize(samples, "deepgram_nova3_first_partial") + 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