diff --git a/scripts/probe_ollama.py b/scripts/probe_ollama.py new file mode 100755 index 0000000..2f48ccd --- /dev/null +++ b/scripts/probe_ollama.py @@ -0,0 +1,226 @@ +#!/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()) \ No newline at end of file