praxis/scripts/probe_ollama.py

#!/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())