asr_latency_harness.py

"""Realtime ASR latency harness — Socket.IO `audio_stream` path.

Shared so both sides can measure the same thing the same way. Run it against your own audio
and send back the JSON from --out; we'll compare directly against our numbers.

WHAT IT MEASURES

A call is a series of short turns, not one long stream, so this benchmarks per turn: one
transcription session per utterance over a single persistent WebSocket, audio paced at 1x
realtime (as it would arrive from a live caller).

  ttfp_ms    first audio frame of the turn -> first non-empty partial transcript
  final_ms   last audio frame of the turn  -> final transcript      <-- turn-taking latency
  rtt_ms     TCP connect time to the endpoint (your network path, measured 5x)

`final_ms` is the number that determines how quickly an agent can start responding.
`ttfp_ms` is heavily influenced by where speech starts inside the clip — a turn that opens
with 1.5s of silence or line noise will report a correspondingly larger ttfp. It is not a
like-for-like number unless both sides use identical audio.

REQUIREMENTS

  pip install humain-sautech certifi

USAGE

  export SAUTECH_API_KEY=<your key>

  # your own telephony audio (any sample rate / channel count, 16-bit WAV — it is
  # converted to the 16kHz mono PCM16 the API expects)
  python3 asr_latency_harness.py --audio call_sample.wav --language ar --out results.json

  # match our published run exactly
  python3 asr_latency_harness.py --audio sample.wav --turn-seconds 5 --max-turns 8 --frame-ms 20

  # sweep framing to check whether your frame size matters on your path
  python3 asr_latency_harness.py --audio sample.wav --frame-ms 20,50,100

PROTOCOL NOTE (for cross-checking your client)

Each `audio_stream` emit is a binary frame:

    [ transcription_id: 16 bytes ][ flags: 1 byte ][ language: 1 byte ][ PCM16 audio ]

    flags bit 0 (0x01)  is_start  — set on the FIRST frame of an utterance
    flags bit 1 (0x02)  is_final  — set on the LAST frame of an utterance

Reuse one transcription_id for all frames of one utterance, and use a new one per utterance.
Frames of any size are accepted. Results arrive on `transcription_result`; treat `is_final`
as end-of-utterance.
"""

from __future__ import annotations

import argparse
import array
import asyncio
import contextlib
import json
import os
import pathlib
import platform
import socket
import statistics
import sys
import time
import wave
from datetime import datetime, timezone

import certifi

# python.org builds cannot find the system CA store; point OpenSSL at certifi's bundle.
os.environ.setdefault("SSL_CERT_FILE", certifi.where())

from sautech.stt import Language, RealtimeClient  # noqa: E402

BASE = "https://sautech.humain.com"
API_PATH = "/realtime/socket.io/"
TARGET_SR = 16000
BYTES_PER_SEC = TARGET_SR * 2  # mono PCM16

LANGS = {"ar": Language.Ar, "en": Language.En, "codeswitch": Language.ArEn}


# --------------------------------------------------------------------------- audio

FFMPEG_HINT = (
    "convert it first, e.g.:\n"
    "    ffmpeg -i {name} -ar 16000 -ac 1 -c:a pcm_s16le converted.wav"
)


class AudioError(Exception):
    """Unreadable or unsupported input audio — reported without a traceback."""


def load_pcm16_16k(path: pathlib.Path) -> tuple[bytes, str]:
    """Read a 16-bit PCM WAV as 16kHz mono PCM16. Returns (pcm, description of conversion).

    Only linear-PCM WAV is supported. Compressed WAV (mu-law/A-law, common on telephony
    recordings) and non-WAV containers (mp3/aac/mp4) are rejected with a conversion hint
    rather than a stack trace.
    """
    try:
        with contextlib.closing(wave.open(str(path))) as w:
            sr, ch, width = w.getframerate(), w.getnchannels(), w.getsampwidth()
            raw = w.readframes(w.getnframes())
    except wave.Error as exc:
        raise AudioError(
            f"{path.name}: not a linear-PCM WAV file ({exc}). "
            f"mu-law/A-law and mp3/aac/mp4 are not read directly — "
            + FFMPEG_HINT.format(name=path.name)
        ) from exc
    except OSError as exc:
        raise AudioError(f"{path.name}: cannot read ({exc})") from exc

    if width != 2:
        raise AudioError(
            f"{path.name}: need 16-bit PCM, got {width * 8}-bit — "
            + FFMPEG_HINT.format(name=path.name)
        )

    samples = array.array("h")
    samples.frombytes(raw)
    steps = [f"{sr}Hz/{ch}ch"]

    if ch > 1:
        mono = array.array("h", [0] * (len(samples) // ch))
        for i in range(len(mono)):
            mono[i] = sum(samples[i * ch:(i + 1) * ch]) // ch
        samples = mono
        steps.append("downmixed to mono")

    if sr != TARGET_SR:
        ratio = TARGET_SR / sr
        out = array.array("h", [0] * int(len(samples) * ratio))
        for i in range(len(out)):
            src = i / ratio
            lo = int(src)
            hi = min(lo + 1, len(samples) - 1)
            frac = src - lo
            out[i] = int(samples[lo] * (1 - frac) + samples[hi] * frac)
        samples = out
        steps.append(f"resampled to {TARGET_SR}Hz")

    return samples.tobytes(), ", ".join(steps)


def speech_onset_s(pcm: bytes, threshold: float = 300.0) -> float | None:
    """Approximate first moment of speech, so ttfp can be read in context."""
    samples = array.array("h")
    samples.frombytes(pcm)
    step = int(TARGET_SR * 0.1)
    for i in range(0, len(samples) - step, step):
        seg = samples[i:i + step]
        rms = (sum(x * x for x in seg) / len(seg)) ** 0.5
        if rms > threshold:
            return round(i / TARGET_SR, 2)
    return None


# --------------------------------------------------------------------------- network

def measure_rtt_ms(host: str, port: int = 443, samples: int = 5) -> float | None:
    """Median TCP connect time — one round trip, no TLS or application work."""
    times = []
    for _ in range(samples):
        t0 = time.perf_counter()
        try:
            with socket.create_connection((host, port), timeout=5):
                times.append((time.perf_counter() - t0) * 1000)
        except OSError:
            continue
    return round(statistics.median(times), 1) if times else None


# --------------------------------------------------------------------------- bench

def pct(values: list[float], p: float) -> float | None:
    if not values:
        return None
    ordered = sorted(values)
    idx = min(len(ordered) - 1, int(round((p / 100) * (len(ordered) - 1))))
    return ordered[idx]


async def run_turn(client: RealtimeClient, pcm: bytes, lang: str,
                   frame_bytes: int, settle: float) -> dict:
    frames = [pcm[i:i + frame_bytes] for i in range(0, len(pcm), frame_bytes)]
    frames = [f for f in frames if len(f) == frame_bytes]
    if not frames:
        return {"error": "turn shorter than one frame", "events": 0,
                "ttfp_ms": None, "final_ms": None, "text": ""}
    frame_dur = frame_bytes / BYTES_PER_SEC

    st = {"t_first": None, "t_last": None, "ttfp": None, "final_at": None,
          "events": 0, "texts": []}

    def on_response(model) -> None:
        now = time.perf_counter()
        st["events"] += 1
        text = (getattr(model, "transcription", None) or getattr(model, "text", None) or "")
        if text.strip():
            st["texts"].append(text)
            if st["ttfp"] is None and st["t_first"] is not None:
                st["ttfp"] = (now - st["t_first"]) * 1000
        if getattr(model, "is_final", False) and st["final_at"] is None:
            st["final_at"] = now

    rec: dict = {"audio_s": round(len(frames) * frame_dur, 2), "frames": len(frames),
                 "ttfp_ms": None, "final_ms": None, "events": 0, "text": "", "error": None}
    try:
        stream = await client.start_stream(
            language=LANGS[lang], on_response=on_response, on_error=lambda e: None
        )
        t0 = time.perf_counter()
        st["t_first"] = t0
        for i, frame in enumerate(frames):
            delay = (t0 + i * frame_dur) - time.perf_counter()
            if delay > 0:
                await asyncio.sleep(delay)
            is_last = i == len(frames) - 1
            if is_last:
                st["t_last"] = time.perf_counter()
            await stream.send(frame, is_last=is_last)

        deadline = time.perf_counter() + settle
        while time.perf_counter() < deadline and st["final_at"] is None:
            await asyncio.sleep(0.02)
        with contextlib.suppress(Exception):
            await stream.close(timeout=1.0)
    except Exception as exc:  # noqa: BLE001 - diagnostic harness
        rec["error"] = f"{type(exc).__name__}: {exc}"

    rec["events"] = st["events"]
    if st["ttfp"] is not None:
        rec["ttfp_ms"] = round(st["ttfp"], 1)
    if st["final_at"] and st["t_last"]:
        rec["final_ms"] = round((st["final_at"] - st["t_last"]) * 1000, 1)
    if st["texts"]:
        rec["text"] = st["texts"][-1]
    return rec


async def bench(path: pathlib.Path, pcm: bytes, conversion: str,
                lang: str, frame_ms: float, args) -> dict:
    frame_bytes = int(frame_ms / 1000 * BYTES_PER_SEC)
    turn_bytes = int(args.turn_seconds * BYTES_PER_SEC)
    segments = [pcm[i:i + turn_bytes] for i in range(0, len(pcm), turn_bytes)]
    segments = [s for s in segments if len(s) >= turn_bytes // 2][:args.max_turns]

    out: dict = {
        "audio_file": path.name,
        "audio_source": conversion,
        "audio_seconds": round(len(pcm) / BYTES_PER_SEC, 1),
        "speech_onset_s": speech_onset_s(pcm[:turn_bytes]),
        "language": lang,
        "frame_ms": frame_ms,
        "frame_bytes": frame_bytes,
        "turn_seconds": args.turn_seconds,
        "turns": [],
    }

    client = RealtimeClient(api_url=BASE, api_key=os.environ["SAUTECH_API_KEY"],
                            api_path=API_PATH)
    try:
        for seg in segments:
            out["turns"].append(await run_turn(client, seg, lang, frame_bytes, args.settle))
            await asyncio.sleep(0.3)
    finally:
        with contextlib.suppress(Exception):
            await client.disconnect()

    ok = [t for t in out["turns"] if not t["error"]]
    ttfp = [t["ttfp_ms"] for t in ok if t["ttfp_ms"] is not None]
    final = [t["final_ms"] for t in ok if t["final_ms"] is not None]
    out["summary"] = {
        "turns": len(out["turns"]),
        "errored": len(out["turns"]) - len(ok),
        "silent": len([t for t in ok if t["events"] == 0]),
        "ttfp_p50": pct(ttfp, 50), "ttfp_p95": pct(ttfp, 95),
        "final_p50": pct(final, 50), "final_p95": pct(final, 95),
        "final_max": max(final) if final else None,
    }
    return out


def f(v, width=10, dec=0):
    return f"{'-':>{width}}" if v is None else f"{v:>{width}.{dec}f}"


async def main() -> int:
    ap = argparse.ArgumentParser(description="Realtime ASR latency harness (Socket.IO path)")
    ap.add_argument("--audio", action="append", required=True,
                    help="16-bit WAV; repeatable. Any sample rate / channel count.")
    ap.add_argument("--language", default="ar", choices=sorted(LANGS),
                    help="ar | en | codeswitch (default ar)")
    ap.add_argument("--frame-ms", default="20",
                    help="frame size(s) in ms, comma-separated (default 20)")
    ap.add_argument("--turn-seconds", type=float, default=5.0,
                    help="utterance length per turn (default 5)")
    ap.add_argument("--max-turns", type=int, default=8, help="cap turns per file (default 8)")
    ap.add_argument("--settle", type=float, default=8.0,
                    help="max seconds to wait for the final event (default 8)")
    ap.add_argument("--out", default=None, help="write raw results as JSON")
    args = ap.parse_args()

    if not os.environ.get("SAUTECH_API_KEY"):
        sys.exit("set SAUTECH_API_KEY")

    try:
        frame_sizes = [float(x) for x in args.frame_ms.split(",")]
    except ValueError:
        sys.exit(f"--frame-ms: expected comma-separated numbers, got {args.frame_ms!r}")
    if any(f * BYTES_PER_SEC / 1000 < 2 for f in frame_sizes):
        sys.exit("--frame-ms: each frame must cover at least one 16-bit sample")

    # Pre-flight: read and convert every file before printing a table or hitting the network,
    # so a bad input fails immediately with a clear message.
    paths = [pathlib.Path(a) for a in args.audio]
    loaded: list[tuple[pathlib.Path, bytes, str]] = []
    for p in paths:
        if not p.exists():
            sys.exit(f"not found: {p}")
        try:
            pcm, conversion = load_pcm16_16k(p)
        except AudioError as exc:
            sys.exit(str(exc))
        if len(pcm) < BYTES_PER_SEC // 2:
            sys.exit(f"{p.name}: only {len(pcm) / BYTES_PER_SEC:.2f}s of audio — "
                     f"need at least 0.5s")
        loaded.append((p, pcm, conversion))

    host = BASE.split("://", 1)[1]
    rtt = measure_rtt_ms(host)

    print(f"endpoint   {BASE}{API_PATH}")
    print(f"network    TCP connect RTT to {host}: "
          f"{f'{rtt:.1f}ms' if rtt else 'unavailable'}")
    print(f"client     python {platform.python_version()} on {platform.system()}")
    print(f"mode       turns of {args.turn_seconds}s, max {args.max_turns} per file, "
          f"paced at 1x realtime\n")

    header = (f"{'audio':<24}{'frame':>8}{'lang':>11}{'turns':>6}{'ttfp_p50':>10}"
              f"{'final_p50':>11}{'final_p95':>11}{'final_max':>11}{'silent':>8}{'err':>5}")
    print(header)
    print("-" * len(header))

    results = []
    for path, pcm, conversion in loaded:
        for frame_ms in frame_sizes:
            out = await bench(path, pcm, conversion, args.language, frame_ms, args)
            results.append(out)
            s = out["summary"]
            print(f"{path.name[:23]:<24}{str(int(frame_ms)) + 'ms':>8}{args.language:>11}"
                  f"{s['turns']:>6}{f(s['ttfp_p50'], 10)}{f(s['final_p50'], 11)}"
                  f"{f(s['final_p95'], 11)}{f(s['final_max'], 11)}"
                  f"{s['silent']:>8}{s['errored']:>5}")
            if out["speech_onset_s"] is not None:
                print(f"{'':<24}  └ speech starts ~{out['speech_onset_s']}s into the first "
                      f"turn (ttfp includes this)")

    all_final = [t["final_ms"] for r in results for t in r["turns"]
                 if not t["error"] and t["final_ms"] is not None]
    if all_final:
        print(f"\nacross {len(all_final)} turns — final: p50 {pct(all_final,50):.0f}ms  "
              f"p95 {pct(all_final,95):.0f}ms  max {max(all_final):.0f}ms")
    print("\nfinal = last audio frame -> final transcript (turn-taking latency)")

    if args.out:
        pathlib.Path(args.out).write_text(json.dumps({
            "generated_utc": datetime.now(timezone.utc).isoformat(),
            "endpoint": BASE + API_PATH,
            "tcp_rtt_ms": rtt,
            "client": {"python": platform.python_version(), "os": platform.system(),
                       "release": platform.release()},
            "args": {k: v for k, v in vars(args).items()},
            "results": results,
        }, indent=2, ensure_ascii=False))
        print(f"raw -> {args.out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(asyncio.run(main()))