rtpbench 0.3.0__py3-none-any.whl

rtpbench/detect.py

@@ -0,0 +1,145 @@
+"""Speech detectors for received audio: [start, end) sample spans of detected speech.
+
+  gate       |sample| >= 1000 on the 16-bit scale (the calibration's rule); gaps under 20 ms merged
+  rms        every 1 ms, the RMS of the 10 ms around it, mean removed, above 0.01 of full scale:
+             the onset rule of Coval's TTS time-to-first-audio benchmark (coval-ai/benchmarks,
+             runner/src/coval_bench/metrics/ttfa.py), applied to both edges of every span
+  webrtcvad  WebRTC VAD, aggressiveness 3, 10 ms frames
+  silero     Silero VAD 6.2.3 (ONNX) at 8 kHz, 32 ms windows; speech from p >= 0.5 until p < 0.35
+  silero16   the same on the audio upsampled to 16 kHz. Fed a whole call, the 8 kHz mode answers
+             speech that follows earlier speech and a long silence tens of ms late, sometimes not
+             at all; the 16 kHz mode on the same audio does not.
+
+All but gate run raw: no minimum durations and no padding, so their bias shows. They need the
+`vad` extra (webrtcvad-wheels, onnxruntime, numpy); the Silero model is fetched once from
+its PyPI wheel, checked against a pinned hash and cached.
+"""
+
+import hashlib
+import io
+import os
+import urllib.request
+import zipfile
+from array import array
+from collections.abc import Callable, Iterable
+from pathlib import Path
+from typing import Any
+
+from rtpbench.voice import GATE, RATE
+
+DETECTORS = ("gate", "rms", "webrtcvad", "silero", "silero16")
+GATE_GAP = 160
+RMS_FRAME, RMS_HOP, RMS_LEVEL = 80, 8, 0.01
+WEBRTC_MODE = 3
+WEBRTC_FRAME = 80
+SILERO_WINDOW = 256
+SILERO_ON, SILERO_OFF = 0.5, 0.35
+SILERO_WHEEL = (
+    "https://files.pythonhosted.org/packages/84/ef/"
+    "9099037ed6f180ea33220178df4107112c0ce2bf5fb4d6f6ab19db2844ed/"
+    "silero_vad-6.2.3-py3-none-any.whl",
+    "7b7f5436cfcb02fae583a05b512ea96467fd449fe54cb49a5e4f06c51a1e43b8",
+)
+
+
+def _spans(flags: Iterable[bool], size: int, gap: int = 0) -> list[tuple[int, int]]:
+    """Runs of True in per-unit flags as sample spans, merging gaps of at most `gap` samples."""
+    out: list[tuple[int, int]] = []
+    for i, on in enumerate(flags):
+        if not on:
+            continue
+        start, end = i * size, (i + 1) * size
+        if out and start - out[-1][1] <= gap:
+            out[-1] = (out[-1][0], end)
+        else:
+            out.append((start, end))
+    return out
+
+
+def gate(pcm: list[int]) -> list[tuple[int, int]]:
+    return _spans((abs(s) >= GATE for s in pcm), 1, GATE_GAP)
+
+
+def rms(pcm: list[int]) -> list[tuple[int, int]]:
+    import numpy as np
+
+    if not pcm:
+        return []  # edge padding needs a sample
+    x = np.pad(np.asarray(pcm, dtype=np.float64) / 32768, RMS_FRAME // 2, mode="edge")
+    s1, s2 = (np.concatenate([[0], np.cumsum(v)]) for v in (x, x * x))
+    at = np.arange(0, len(x) - RMS_FRAME + 1, RMS_HOP)
+    mean = (s1[at + RMS_FRAME] - s1[at]) / RMS_FRAME
+    power = (s2[at + RMS_FRAME] - s2[at]) / RMS_FRAME - mean * mean
+    return _spans((power > RMS_LEVEL**2).tolist(), RMS_HOP)
+
+
+def webrtc(pcm: list[int]) -> list[tuple[int, int]]:
+    import webrtcvad
+
+    vad = webrtcvad.Vad(WEBRTC_MODE)
+    n = len(pcm) // WEBRTC_FRAME
+    raw = array("h", pcm).tobytes()  # native order; webrtcvad wants little-endian (x86-64, arm64)
+    step = 2 * WEBRTC_FRAME
+    return _spans(
+        (vad.is_speech(raw[i * step : (i + 1) * step], RATE) for i in range(n)), WEBRTC_FRAME
+    )
+
+
+def _silero_model() -> bytes:
+    url, sha = SILERO_WHEEL
+    cache = Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache") / "rtpbench"
+    path = cache / f"silero_vad-{sha[:12]}.onnx"
+    if not path.exists():
+        with urllib.request.urlopen(url, timeout=60) as r:
+            wheel = r.read()
+        if hashlib.sha256(wheel).hexdigest() != sha:
+            raise RuntimeError(f"{url}: unexpected sha256")
+        with zipfile.ZipFile(io.BytesIO(wheel)) as z:
+            model = z.read("silero_vad/data/silero_vad.onnx")
+        cache.mkdir(parents=True, exist_ok=True)
+        path.write_bytes(model)
+    return path.read_bytes()
+
+
+_session: Any = None
+
+
+def _silero(pcm: list[int], rate: int) -> list[tuple[int, int]]:
+    """As silero_vad.OnnxWrapper: 32 ms windows, each with 4 ms of context. At 16 kHz the input
+    is `pcm` upsampled 2x (zero-stuffing, then a 63-tap windowed-sinc low-pass at 4 kHz)."""
+    import numpy as np
+    import onnxruntime
+
+    global _session
+    if _session is None:
+        opts = onnxruntime.SessionOptions()
+        opts.inter_op_num_threads = opts.intra_op_num_threads = 1
+        _session = onnxruntime.InferenceSession(_silero_model(), sess_options=opts)
+    x = np.asarray(pcm, dtype=np.float32) / 32768
+    if rate == 16000:
+        up = np.zeros(2 * len(x), dtype=np.float32)
+        up[::2] = x
+        n = np.arange(63) - 31
+        x = np.convolve(up, (np.sinc(n / 2) * np.hamming(63)).astype(np.float32), "same")
+    win, ctx = SILERO_WINDOW * rate // RATE, 32 * rate // RATE
+    state = np.zeros((2, 1, 128), dtype=np.float32)
+    context = np.zeros((1, ctx), dtype=np.float32)
+    sr = np.array(rate, dtype=np.int64)
+    flags, speaking = [], False
+    for i in range(len(x) // win):
+        chunk = np.concatenate([context, x[None, i * win : (i + 1) * win]], 1)
+        out, state = _session.run(None, {"input": chunk, "state": state, "sr": sr})
+        context = chunk[:, -ctx:]
+        p = float(out[0, 0])
+        speaking = p >= SILERO_ON or (speaking and p >= SILERO_OFF)
+        flags.append(speaking)
+    return _spans(flags, SILERO_WINDOW)
+
+
+DETECT: dict[str, Callable[[list[int]], list[tuple[int, int]]]] = {
+    "gate": gate,
+    "rms": rms,
+    "webrtcvad": webrtc,
+    "silero": lambda pcm: _silero(pcm, RATE),
+    "silero16": lambda pcm: _silero(pcm, 16000),
+}

rtpbench/g711.py

@@ -0,0 +1,82 @@
+"""G.711 µ-law (PT 0) and A-law (PT 8), after the Sun reference implementation."""
+
+import re
+from collections.abc import Iterable
+
+PCMU, PCMA = 0, 8
+NAMES = {PCMU: "PCMU", PCMA: "PCMA"}
+SILENCE = {PCMU: 0xFF, PCMA: 0xD5}
+
+_SEG_END = (0x3F, 0x7F, 0xFF, 0x1FF, 0x3FF, 0x7FF, 0xFFF, 0x1FFF)
+
+
+def _segment(value: int, ends: Iterable[int]) -> int:
+    for seg, end in enumerate(ends):
+        if value <= end:
+            return seg
+    return 8
+
+
+def ulaw_encode(sample: int) -> int:
+    value = sample >> 2
+    mask = 0xFF
+    if value < 0:
+        value, mask = -value, 0x7F
+    value = min(value, 8159) + 0x21
+    seg = _segment(value, _SEG_END)
+    if seg >= 8:
+        return 0x7F ^ mask
+    return ((seg << 4) | ((value >> (seg + 1)) & 0xF)) ^ mask
+
+
+def ulaw_decode(code: int) -> int:
+    code = ~code & 0xFF
+    t = (((code & 0xF) << 3) + 0x84) << ((code & 0x70) >> 4)
+    return 0x84 - t if code & 0x80 else t - 0x84
+
+
+def alaw_encode(sample: int) -> int:
+    value = sample >> 3
+    mask = 0xD5
+    if value < 0:
+        value, mask = -value - 1, 0x55
+    seg = _segment(value, (e >> 1 for e in _SEG_END))
+    if seg >= 8:
+        return 0x7F ^ mask
+    shift = 1 if seg < 2 else seg
+    return ((seg << 4) | ((value >> shift) & 0xF)) ^ mask
+
+
+def alaw_decode(code: int) -> int:
+    code ^= 0x55
+    t = (code & 0xF) << 4
+    seg = (code & 0x70) >> 4
+    if seg == 0:
+        t += 8
+    else:
+        t = (t + 0x108) << (seg - 1)
+    return t if code & 0x80 else -t
+
+
+_ENCODE = {PCMU: ulaw_encode, PCMA: alaw_encode}
+_DECODE = {PCMU: ulaw_decode, PCMA: alaw_decode}
+
+
+def encode(pt: int, samples: Iterable[int]) -> bytes:
+    enc = _ENCODE[pt]
+    return bytes(enc(s) for s in samples)
+
+
+def decode(pt: int, payload: bytes) -> list[int]:
+    dec = _DECODE[pt]
+    return [dec(b) for b in payload]
+
+
+def gate(pt: int, threshold: int) -> re.Pattern[bytes]:
+    """A regex matching one encoded sample whose magnitude reaches `threshold`.
+
+    `search` on a payload gives the first active sample; on the reversed payload, the last.
+    """
+    dec = _DECODE[pt]
+    codes = bytes(b for b in range(256) if abs(dec(b)) >= threshold)
+    return re.compile(b"[" + re.escape(codes) + b"]")

rtpbench/mock/compose.yaml

@@ -0,0 +1,18 @@
+# The mock engine (src/rtpbench/realtime.py) where a target expects its AI engine, merged over
+# ../compose.yaml and a target's compose file: `rtpbench target NAME --mock` points the target's
+# engine variables at it. The target loses its egress network: on the mock it needs none, and
+# cannot reach a paid engine. It logs to RTPBENCH_OUT/realtime.jsonl.
+services:
+  target:
+    networks: !override [agentnet]
+
+  realtime:
+    image: ${RTPBENCH_IMAGE:?}
+    init: true
+    networks: [agentnet]
+    volumes: ["${RTPBENCH_OUT:?}:/out"]
+    command: python -m rtpbench realtime --out /out
+    healthcheck:
+      test: test -e /tmp/ready
+      interval: 1s
+      retries: 30

rtpbench/pcap.py

@@ -0,0 +1,106 @@
+"""Classic pcap reader (as written by tcpdump -w): IPv4 UDP datagrams and TCP segments with
+data, with ns timestamps."""
+
+import socket
+import struct
+from dataclasses import dataclass
+from pathlib import Path
+
+_MAGIC = {
+    b"\xd4\xc3\xb2\xa1": ("<", 1000),
+    b"\x4d\x3c\xb2\xa1": ("<", 1),
+    b"\xa1\xb2\xc3\xd4": (">", 1000),
+    b"\xa1\xb2\x3c\x4d": (">", 1),
+}
+_ETHERNET, _RAW, _SLL, _SLL2 = 1, 101, 113, 276
+
+
+@dataclass(frozen=True, slots=True)
+class Datagram:
+    t_ns: int
+    src: str
+    sport: int
+    dst: str
+    dport: int
+    payload: bytes
+    tcp: bool = False
+    seq: int = 0  # TCP sequence number of the payload's first byte
+
+
+def _ip_offset(link: int, frame: bytes) -> int | None:
+    if link == _ETHERNET:
+        off, ethertype = 14, struct.unpack_from("!H", frame, 12)[0]
+        while ethertype == 0x8100:
+            ethertype = struct.unpack_from("!H", frame, off + 2)[0]
+            off += 4
+        return off if ethertype == 0x0800 else None
+    if link == _SLL:
+        return 16 if struct.unpack_from("!H", frame, 14)[0] == 0x0800 else None
+    if link == _SLL2:
+        return 20 if struct.unpack_from("!H", frame, 0)[0] == 0x0800 else None
+    return 0 if link == _RAW else None
+
+
+def read(path: Path, partial: bool = False) -> list[Datagram]:
+    """Its TCP segments that carry data and its UDP datagrams; with `partial`, also a datagram
+    the snaplen cut short, with what the capture kept of it."""
+    data = path.read_bytes()
+    if len(data) < 24 or data[:4] not in _MAGIC:
+        raise ValueError(f"{path}: not a classic pcap file")
+    order, scale = _MAGIC[data[:4]]
+    link = struct.unpack_from(order + "I", data, 20)[0] & 0x0FFFFFFF
+    record = struct.Struct(order + "IIII")
+    out, pos = [], 24
+    while pos + 16 <= len(data):
+        sec, frac, incl, _ = record.unpack_from(data, pos)
+        frame = data[pos + 16 : pos + 16 + incl]
+        pos += 16 + incl
+        try:
+            ip = _ip_offset(link, frame)
+        except struct.error:  # shorter than its link header says
+            continue
+        if ip is None or len(frame) < ip + 28 or frame[ip + 9] not in (6, 17):
+            continue
+        if struct.unpack_from("!H", frame, ip + 6)[0] & 0x3FFF:  # fragment
+            continue
+        t = sec * 1_000_000_000 + frac * scale
+        src, dst = (
+            socket.inet_ntoa(frame[ip + 12 : ip + 16]),
+            socket.inet_ntoa(frame[ip + 16 : ip + 20]),
+        )
+        l4 = ip + 4 * (frame[ip] & 0x0F)
+        if frame[ip + 9] == 6:  # TCP: segments that carry data
+            end = ip + struct.unpack_from("!H", frame, ip + 2)[0]
+            if len(frame) < l4 + 20 or len(frame) < end:
+                continue
+            sport, dport, seq = struct.unpack_from("!HHI", frame, l4)
+            body = frame[l4 + 4 * (frame[l4 + 12] >> 4) : end]
+            if body:
+                out.append(Datagram(t, src, sport, dst, dport, body, True, seq))
+            continue
+        if len(frame) < l4 + 8:
+            continue
+        sport, dport, length = struct.unpack_from("!HHH", frame, l4)
+        if length < 8 or (len(frame) < l4 + length and not partial):  # cut by the snaplen
+            continue
+        out.append(Datagram(t, src, sport, dst, dport, frame[l4 + 8 : l4 + length]))
+    return out
+
+
+def streams(grams: list[Datagram]) -> list[Datagram]:
+    """The TCP segments of `grams` that add bytes to their flow's stream, in capture order,
+    each cut to the bytes it adds (a retransmission adds none)."""
+    out: list[Datagram] = []
+    nxt: dict[tuple[str, int, str, int], int] = {}  # next sequence number per flow
+    for g in grams:
+        if not g.tcp:
+            continue
+        flow = (g.src, g.sport, g.dst, g.dport)
+        seen = (nxt.get(flow, g.seq) - g.seq) & 0xFFFFFFFF
+        if seen >= 1 << 31:  # a gap: bytes lost before the capture saw them
+            seen = 0
+        if seen >= len(g.payload):
+            continue
+        nxt[flow] = (g.seq + len(g.payload)) & 0xFFFFFFFF
+        out.append(Datagram(g.t_ns, *flow, g.payload[seen:], True, g.seq + seen))
+    return out

rtpbench/progress.py

@@ -0,0 +1,184 @@
+"""What a long run is doing and how far along it is: a live status bar on a terminal, a line per
+event anywhere else (CI logs, redirects, pipes). Host side only: the captures that time the calls
+are taken in the containers, and the bar redraws twice a second at most, from its own thread."""
+
+import re
+import signal
+import sys
+import time
+from types import TracebackType
+
+from rich import progress as rp
+from rich.console import Console
+from rich.markup import escape
+from rich.text import Text
+
+QUIET_S = 30.0  # without a terminal, a turn gets a line only after this long without one
+_TURN = re.compile(r"turn (\d+)/(\d+)")  # what the caller prints as each turn starts
+
+
+def eta(elapsed: float, done: float, total: int) -> float | None:
+    """Seconds left for `total` units at the rate of the `done` (fractional) done in `elapsed`
+    seconds; None before any is done."""
+    return elapsed * (total - done) / done if done > 0 else None
+
+
+def hms(seconds: float) -> str:
+    s = round(seconds)
+    return f"{s // 3600}:{s // 60 % 60:02d}:{s % 60:02d}"
+
+
+class _Stopped(SystemExit):
+    """A stop signal, told apart from any other SystemExit the run raises."""
+
+
+def _stop(signum: int, _frame: object) -> None:
+    """SIGINT and SIGTERM (systemd stop, a cancelled CI job, kill) unwind the run the same way,
+    so its `finally` blocks tear the stack down and hand the results back. A second signal
+    meanwhile would cut that short and leave containers and root-owned files behind."""
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+    signal.signal(signal.SIGTERM, signal.SIG_IGN)
+    raise _Stopped(128 + signum)
+
+
+class Progress:
+    """Units of work (calls) out of `plan`'s total: `unit` starts one, `turn` takes the caller's
+    output while it runs, `done` ends it; `phase` names what happens outside units. As a context
+    manager it draws the bar (on a terminal), says where a run failed, and turns SIGINT and
+    SIGTERM into a clean exit that runs the run's teardown."""
+
+    def __init__(self) -> None:
+        self.t0 = time.monotonic()
+        self.total = self.k = 0  # units planned, started
+        self.done_units = 0.0  # with the current one's turns as a fraction
+        self.label = ""  # the current unit or phase
+        self.turns: tuple[int, int] | None = None  # the current unit's turn, of how many
+        self.unit_t = self.first_t = self.said_t = 0.0
+        self.end_t: float | None = None  # when the run should end, at the rate so far
+        console = Console(highlight=False)
+        # rich's own test honours TERM=dumb and TTY_INTERACTIVE; FORCE_COLOR into a file must
+        # still not redraw
+        self.bar = None
+        if console.is_interactive and sys.stdout.isatty():
+            self.bar = rp.Progress(
+                rp.SpinnerColumn(),
+                rp.TextColumn("{task.description}"),
+                rp.BarColumn(),
+                rp.MofNCompleteColumn(),
+                rp.TimeElapsedColumn(),
+                _Left(self),
+                console=console,
+                transient=True,
+                refresh_per_second=2,
+            )
+            self.task = self.bar.add_task("", total=None)
+
+    def __enter__(self) -> Progress:
+        self.was = [signal.signal(s, _stop) for s in (signal.SIGINT, signal.SIGTERM)]
+        if self.bar:
+            self.bar.start()
+        return self
+
+    def __exit__(
+        self, exc: type[BaseException] | None, _v: BaseException | None, _t: TracebackType | None
+    ) -> None:
+        took = hms(time.monotonic() - self.t0)
+        try:
+            if exc:
+                how = "interrupted" if issubclass(exc, KeyboardInterrupt | _Stopped) else "failed"
+                self._event(False, f"{how} while {self.label or 'starting'}", f"after {took}")
+            else:
+                self._event(True, "finished", f"in {took}")
+            if self.bar:
+                self.bar.stop()
+        finally:
+            for s, handler in zip((signal.SIGINT, signal.SIGTERM), self.was, strict=True):
+                signal.signal(s, handler)
+
+    def plan(self, total: int) -> None:
+        self.total = total
+        if self.bar:
+            self.bar.update(self.task, total=total)
+
+    def phase(self, text: str) -> None:
+        self.label, self.turns = text, None
+        self._say(text)
+        self._draw()
+
+    def unit(self, label: str) -> None:
+        self.k += 1
+        self.label, self.turns = label, None
+        self.unit_t = time.monotonic()
+        self.first_t = self.first_t or self.unit_t  # the rate leaves out building and starting
+        self._say(f"{self._of()} {label}: calling")
+        self._draw()
+
+    def turn(self, line: str) -> None:
+        """A line of the caller's output; those that say a turn started move the bar."""
+        if not (m := _TURN.fullmatch(line.strip())):
+            return
+        k, n = self.turns = int(m[1]), int(m[2])
+        self._estimate(self.k - 1 + (k - 1) / n)
+        if time.monotonic() - self.said_t >= QUIET_S:
+            left = self.left()
+            self._say(
+                f"{self._of()} {self.label}: turn {k} of {n}" + (f" ({left})" if left else "")
+            )
+        self._draw()
+
+    def done(self, text: str = "", ok: bool | None = None) -> None:
+        """The unit ended with `text`; `ok`: whether that is what it should be (None: no
+        judgement)."""
+        text = text or (f"{self.turns[1]} turns" if self.turns else "done")
+        self._estimate(self.k)
+        if self.bar:
+            self._event(ok, f"{self.label}  {text}", hms(time.monotonic() - self.unit_t))
+        else:
+            left = self.left()
+            line = f"{self._of()} {self.label}: {text}" + (f" ({left})" if left else "")
+            self._say(line if ok is not False else f"FAILED {line}")
+        self._draw()
+
+    def left(self) -> str:
+        if self.end_t is None or self.done_units >= self.total:
+            return ""
+        return f"about {hms(max(0.0, self.end_t - time.monotonic()))} left"
+
+    def _of(self) -> str:
+        return f"[{self.k}/{self.total}]"
+
+    def _estimate(self, done: float) -> None:
+        now = time.monotonic()
+        self.done_units = done
+        left = eta(now - self.first_t, done, self.total)
+        self.end_t = None if left is None else now + left  # counts down between estimates
+
+    def _say(self, line: str) -> None:
+        """A line of the log, without a terminal (on one the bar shows the same)."""
+        if not self.bar:
+            print(f"{hms(time.monotonic() - self.t0)} {line}", flush=True)
+            self.said_t = time.monotonic()
+
+    def _event(self, ok: bool | None, text: str, note: str) -> None:
+        """A line that stays: above the bar on a terminal, else in the log."""
+        if not self.bar:
+            self._say(f"{text} {note}")
+            return
+        mark = {True: "[green]✓[/]", False: "[red]✗[/]", None: "[dim]·[/]"}[ok]
+        self.bar.console.print(f"{mark} {escape(text)}  [dim]{escape(note)}[/]")
+
+    def _draw(self) -> None:
+        if self.bar:
+            turn = f" · turn {self.turns[0]} of {self.turns[1]}" if self.turns else ""
+            self.bar.update(self.task, description=self.label + turn, completed=self.done_units)
+
+
+class _Left(rp.ProgressColumn):
+    """Time left, counting down between the estimates units and turns make."""
+
+    def __init__(self, progress: Progress) -> None:
+        super().__init__()
+        self.progress = progress
+
+    def render(self, task: rp.Task) -> Text:
+        return Text(self.progress.left(), style="progress.remaining")