diff --git a/scripts/probe_e2e.py b/scripts/probe_e2e.py new file mode 100755 index 0000000..f60cac5 --- /dev/null +++ b/scripts/probe_e2e.py @@ -0,0 +1,348 @@ +#!/usr/bin/env python3 +"""R4 probe — integrated three-hop end-to-end latency. + +Per PLAN.md SLICE-01 TASK-01-05: feed a sample ASR transcript → Ollama +gemma4:cloud streaming → Cartesia TTS streaming; measure end-to-end +(transcript-in → first-audio-out). Run 10 iterations. Also measure the same +path with Piper self-hosted (if Piper can be stood up locally; otherwise note +as pending and pre-stage in SLICE-02). + +Exit code 0 in all cases: + - If OLLAMA_API_KEY or CARTESIA_API_KEY is missing, print KEY_MISSING and + exit 0 (the probe infrastructure is the deliverable). + - If present, run the live integrated probe and print e2e latency. + +The Piper leg is invoked only if PRAXIS_TTS=piper is set AND a Piper voice +model is available; otherwise it is documented as pre-staged (R4 mitigation). +""" + +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 _missing(keys: list[str]) -> None: + _banner( + "KEY_MISSING — " + ", ".join(keys) + " not set.\n" + " Cannot run live integrated e2e probe. Probe infrastructure is built\n" + " and ready; live measurements are pending API key provisioning.\n" + " Set the missing key(s) in .env (see .env.example) and re-run." + ) + + +CHAT_URL = os.environ.get("OLLAMA_CHAT_URL", "https://ollama.com/api/chat") +CARTESIA_WS_URL = "wss://api.cartesia.ai/tts/websocket" +ROLEPLAY_MODEL = os.environ.get("OLLAMA_ROLEPLAY_MODEL", "gemma4:cloud") +DEFAULT_VOICE_ID = "a3536a36-1d18-4efb-a95a-7c44b7b5e384" + +# The "transcript-in" — a realistic ASR final transcript from the learner. +SAMPLE_TRANSCRIPT = "Hi, I want to help you with your order. What happened?" +SYSTEM_PROMPT = ( + "You are Jordan, a customer who received a damaged product. " + "You are frustrated but not abusive. Stay in character. Keep responses " + "to 1-2 sentences." +) + + +async def _ollama_first_token(api_key: str) -> tuple[str | None, float | None, str | None]: + """Stream Ollama gemma4:cloud, return (full_text, ttft_ms, error).""" + import httpx + + headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} + body = { + "model": ROLEPLAY_MODEL, + "messages": [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": SAMPLE_TRANSCRIPT}, + ], + "stream": True, + } + t0 = time.perf_counter() + ttft_ms: float | None = None + chunks: list[str] = [] + + try: + async with httpx.AsyncClient(timeout=30.0) as client: + async with client.stream("POST", CHAT_URL, headers=headers, json=body) as resp: + if resp.status_code != 200: + text = await resp.aread() + return None, None, f"HTTP {resp.status_code}: {text[:200]!r}" + async for line in resp.aiter_lines(): + if not line: + continue + try: + chunk = json.loads(line) + except json.JSONDecodeError: + continue + content = chunk.get("message", {}).get("content", "") + if content: + if ttft_ms is None: + ttft_ms = (time.perf_counter() - t0) * 1000.0 + chunks.append(content) + except Exception as exc: # pragma: no cover + return None, None, f"connection error: {exc}" + + return "".join(chunks), ttft_ms, None + + +async def _cartesia_first_audio( + api_key: str, text: str, voice_id: str, model_id: str +) -> tuple[float | None, str | None]: + """Send text to Cartesia WS, return (first_audio_ms_from_t0, error).""" + 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": text, + "voice": {"id": voice_id}, + "output_format": { + "container": "raw", + "encoding": "pcm_s16le", + "sample_rate": 24000, + }, + "stream": True, + } + await ws.send(json.dumps(req)) + 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 + if isinstance(msg, str): + data = json.loads(msg) + if data.get("type") == "done": + break + except Exception as exc: # pragma: no cover + return None, f"cartesia error: {exc}" + + return first_audio_ms, None + + +async def _piper_first_audio(text: str) -> tuple[float | None, str | None]: + """Synthesize via Piper self-hosted, return (first_audio_ms, error). + + Piper pre-staging note (R4 mitigation): Piper is pre-staged as the + production v0.1 TTS fallback per ARCHITECTURE.md. The voice model must be + downloaded separately (see docs/latency-report.md). If not available, + returns an error string that the caller documents as pending. + """ + try: + from piper import PiperVoice # type: ignore + except ImportError: + return None, "piper-tts not installed (pre-staged for SLICE-02)" + + model_path = os.environ.get("PIPER_VOICE_MODEL", "") + if not model_path or not Path(model_path).exists(): + return None, "PIPER_VOICE_MODEL not set or file missing (pre-staged for SLICE-02)" + + import io + + t0 = time.perf_counter() + try: + voice = PiperVoice.load(model_path) + wav_bytes = io.BytesIO() + for chunk in voice.synthesize(text, chunk_size=200): + wav_bytes.write(chunk.audio_int16_bytes) + first_audio_ms = (time.perf_counter() - t0) * 1000.0 + return first_audio_ms, None + except Exception as exc: # pragma: no cover + return None, f"piper error: {exc}" + + +async def _e2e_once_cartesia(ollama_key: str, cartesia_key: str, voice_id: str, model_id: str) -> dict: + """Run the integrated ASR-transcript → Ollama → Cartesia path once.""" + t_start = time.perf_counter() + text, ttft_ms, llm_err = await _ollama_first_token(ollama_key) + if llm_err or not text: + return {"ok": False, "error": llm_err or "empty LLM output", "ttft_ms": None} + tts_ms, tts_err = await _cartesia_first_audio(cartesia_key, text, voice_id, model_id) + if tts_err or tts_ms is None: + return {"ok": False, "error": tts_err or "no TTS audio", "ttft_ms": ttft_ms} + e2e_ms = (time.perf_counter() - t_start) * 1000.0 + return { + "ok": True, + "ttft_ms": ttft_ms, + "tts_first_audio_ms": tts_ms, + "e2e_ms": e2e_ms, + "llm_text": text[:80], + } + + +async def _e2e_once_piper(ollama_key: str) -> dict: + """Run the integrated ASR-transcript → Ollama → Piper path once.""" + t_start = time.perf_counter() + text, ttft_ms, llm_err = await _ollama_first_token(ollama_key) + if llm_err or not text: + return {"ok": False, "error": llm_err or "empty LLM output", "ttft_ms": None} + tts_ms, tts_err = await _piper_first_audio(text) + if tts_err or tts_ms is None: + return {"ok": False, "error": tts_err or "no TTS audio", "ttft_ms": ttft_ms, "piper_pending": True} + e2e_ms = (time.perf_counter() - t_start) * 1000.0 + return { + "ok": True, + "ttft_ms": ttft_ms, + "tts_first_audio_ms": tts_ms, + "e2e_ms": e2e_ms, + "llm_text": text[:80], + } + + +def _summarize(samples: list[float], label: str) -> dict: + if not samples: + print(f" {label}: no samples collected.") + 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="R4 integrated e2e latency probe") + parser.add_argument("--iterations", type=int, default=10) + parser.add_argument("--voice-id", default=os.environ.get("CARTESIA_VOICE_ID", DEFAULT_VOICE_ID)) + parser.add_argument("--cartesia-model", default="sonic-2") + parser.add_argument("--out", default=None) + parser.add_argument("--piper", action="store_true", help="also run the Piper leg") + args = parser.parse_args() + + _banner("R4 PROBE — integrated three-hop e2e (transcript → Ollama → TTS)") + ollama_key = os.environ.get("OLLAMA_API_KEY", "").strip() + cartesia_key = os.environ.get("CARTESIA_API_KEY", "").strip() + + missing = [] + if not ollama_key: + missing.append("OLLAMA_API_KEY") + if not cartesia_key: + missing.append("CARTESIA_API_KEY") + if missing: + _missing(missing) + return 0 + + # ── Cartesia leg ──────────────────────────────────────────────────────── + print(f"\n Cartesia leg — {args.iterations} iterations:") + e2e_samples: list[float] = [] + ttft_samples: list[float] = [] + tts_samples: list[float] = [] + for i in range(args.iterations): + r = await _e2e_once_cartesia(ollama_key, cartesia_key, args.voice_id, args.cartesia_model) + if r.get("ok"): + e2e_samples.append(r["e2e_ms"]) + ttft_samples.append(r["ttft_ms"]) + tts_samples.append(r["tts_first_audio_ms"]) + print(f" [{i + 1:2d}/{args.iterations}] e2e={r['e2e_ms']:.1f}ms " + f"(llm_ttft={r['ttft_ms']:.1f}, tts={r['tts_first_audio_ms']:.1f})") + else: + print(f" [{i + 1:2d}/{args.iterations}] error: {r.get('error')}") + await asyncio.sleep(0.5) + + print() + e2e_summary = _summarize(e2e_samples, "e2e_cartesia") + ttft_summary = _summarize(ttft_samples, "e2e_cartesia_llm_ttft") + tts_summary = _summarize(tts_samples, "e2e_cartesia_tts_first_audio") + + # ── Piper leg (optional / pre-staged) ─────────────────────────────────── + piper_summary: dict = {} + if args.piper: + print(f"\n Piper leg — {args.iterations} iterations:") + p_e2e: list[float] = [] + p_ttft: list[float] = [] + p_tts: list[float] = [] + for i in range(args.iterations): + r = await _e2e_once_piper(ollama_key) + if r.get("ok"): + p_e2e.append(r["e2e_ms"]) + p_ttft.append(r["ttft_ms"]) + p_tts.append(r["tts_first_audio_ms"]) + print(f" [{i + 1:2d}/{args.iterations}] e2e={r['e2e_ms']:.1f}ms") + elif r.get("piper_pending"): + print(f" [{i + 1:2d}/{args.iterations}] Piper pre-staged (pending voice model) — skipping") + break + else: + print(f" [{i + 1:2d}/{args.iterations}] error: {r.get('error')}") + await asyncio.sleep(0.5) + print() + piper_summary = _summarize(p_e2e, "e2e_piper") + else: + print("\n Piper leg not requested (--piper). Piper is pre-staged as the R4 " + "mitigation per ARCHITECTURE.md; live Piper measurement pending " + "voice-model provisioning (see docs/latency-report.md).") + + # ── Budget comparison ─────────────────────────────────────────────────── + budget = 600.0 + print(f"\n Latency budget: {budget:.0f}ms") + if e2e_samples: + med = statistics.median(e2e_samples) + over = med > budget + print(f" Cartesia median e2e: {med:.1f}ms — {'OVER' if over else 'WITHIN'} budget " + f"(delta {med - budget:+.1f}ms)") + if piper_summary.get("n"): + # type: ignore + med = piper_summary.get("median_ms") + if med: + over = med > budget + print(f" Piper median e2e: {med:.1f}ms — {'OVER' if over else 'WITHIN'} budget " + f"(delta {med - budget:+.1f}ms)") + + print() + if args.out: + result = { + "e2e_cartesia": e2e_summary, + "e2e_cartesia_llm_ttft": ttft_summary, + "e2e_cartesia_tts_first_audio": tts_summary, + "e2e_piper": piper_summary, + "budget_ms": budget, + } + Path(args.out).write_text(json.dumps(result, 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