praxis/scripts/probe_deepgram.py

#!/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("<h", sample)
    return bytes(frames)


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,
    )

    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())