#!/usr/bin/env python3 """R3 probe — Ollama Cloud direct-API time-to-first-token (TTFT). Per PLAN.md SLICE-01 TASK-01-04: direct API call to https://ollama.com/api/chat with OLLAMA_API_KEY bearer, model gemma4:cloud, stream=True, measure TTFT over 20 iterations; also probe deepseek-v4-flash:cloud no-think mode TTFT. Log min/median/p95 + any throttle events (R5). Exit code 0 in all cases: - If OLLAMA_API_KEY is missing, print KEY_MISSING and exit 0. - If present, run the live probe for both models and print TTFT tables. R6 note: this probe also confirms the Ollama Cloud direct API is callable with a bearer token (R6). If it returns 401/403, that is recorded as a throttle/auth event, not a crash. """ 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 _require_key() -> str | None: key = os.environ.get("OLLAMA_API_KEY", "").strip() if not key: _banner( "KEY_MISSING — OLLAMA_API_KEY not set.\n" " Cannot run live Ollama Cloud probe. Probe infrastructure is built\n" " and ready; live measurements are pending API key provisioning.\n" " Set OLLAMA_API_KEY in .env (see .env.example) and re-run." ) return None return key CHAT_URL = os.environ.get("OLLAMA_CHAT_URL", "https://ollama.com/api/chat") ROLEPLAY_MODEL = os.environ.get("OLLAMA_ROLEPLAY_MODEL", "gemma4:cloud") DEBRIEF_MODEL = os.environ.get("OLLAMA_DEBRIEF_MODEL", "deepseek-v4-flash:cloud") # A short role-play prompt that should produce a fast first token. ROLEPLAY_MESSAGES = [ { "role": "system", "content": ( "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." ), }, {"role": "user", "content": "Hi, I want to help you with your order. What happened?"}, ] # Debrief prompt — no_think mode for latency (D-020). DEBRIEF_MESSAGES = [ { "role": "system", "content": ( "You are a coaching mentor. Produce a concise (3-bullet) debrief " "about the learner's customer-service performance. " "Do not reason step-by-step; respond directly." ), }, { "role": "user", "content": "The learner said: 'I'm sorry you're upset. I can offer a refund.'", }, ] async def _probe_once( api_key: str, model: str, messages: list[dict], no_think: bool ) -> tuple[float | None, str | None]: """Call Ollama Cloud /api/chat streaming, return (ttft_ms, error_or_none).""" import httpx headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} body: dict = {"model": model, "messages": messages, "stream": True} if no_think: # Ollama no-think mode for deepseek-v4-flash:cloud (D-020). body["think"] = False t0 = time.perf_counter() ttft_ms: float | None = None 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, 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 msg = chunk.get("message", {}) content = msg.get("content", "") if content and ttft_ms is None: ttft_ms = (time.perf_counter() - t0) * 1000.0 break except Exception as exc: # pragma: no cover - network errors return None, f"connection error: {exc}" return ttft_ms, None async def run_model( api_key: str, model: str, messages: list[dict], iterations: int, label: str, no_think: bool ) -> tuple[list[float], list[str]]: samples: list[float] = [] errors: list[str] = [] print(f" Running {iterations} iterations for {label} (model={model}, no_think={no_think})...") for i in range(iterations): ms, err = await _probe_once(api_key, model, messages, no_think) if ms is not None: samples.append(ms) print(f" [{i + 1:2d}/{iterations}] TTFT: {ms:6.1f} ms") else: errors.append(err or "unknown") print(f" [{i + 1:2d}/{iterations}] error: {err}") await asyncio.sleep(0.5) return samples, errors 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="R3 Ollama Cloud TTFT probe") parser.add_argument("--iterations", type=int, default=20) parser.add_argument("--out", default=None) args = parser.parse_args() _banner("R3 PROBE — Ollama Cloud direct-API time-to-first-token") api_key = _require_key() if api_key is None: return 0 # R6: confirms the direct API + bearer works for the role-play model. rp_samples, rp_errors = await run_model( api_key, ROLEPLAY_MODEL, ROLEPLAY_MESSAGES, args.iterations, "ollama_gemma4_cloud_ttft", no_think=False, ) rp_summary = _summarize(rp_samples, "ollama_gemma4_cloud_ttft") print() # Debrief model with no_think (D-020). db_samples, db_errors = await run_model( api_key, DEBRIEF_MODEL, DEBRIEF_MESSAGES, args.iterations, "ollama_deepseek_v4_flash_nothink_ttft", no_think=True, ) db_summary = _summarize(db_samples, "ollama_deepseek_v4_flash_nothink_ttft") # R5: log throttle events (any error could indicate throttling/auth). all_errors = rp_errors + db_errors if all_errors: print(f"\n R5 — {len(all_errors)} error/throttle event(s) recorded:") for e in all_errors[:10]: print(f" - {e}") else: print("\n R5 — no throttle/auth events recorded.") print() if args.out: result = { "roleplay": rp_summary, "debrief": db_summary, "errors": all_errors, } 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())