bench.py - Benchmark Every DevPlace Model

#!/usr/bin/env python3
"""Benchmark every devplace model via the OpenAI-compatible gateway API."""

import argparse
import json
import logging
import os
import statistics
import sys
import time
from dataclasses import dataclass, field
from urllib.error import HTTPError, URLError
from urllib.parse import urlencode
from urllib.request import Request, urlopen

BASE_URL = "https://devplace.net"
API_KEY_ENV_VAR = "DEVPLACE_API_KEY"
LOG_FORMAT = "%(asctime)s [%(levelname)s] %(message)s"

logging.basicConfig(level=logging.DEBUG, format=LOG_FORMAT)
logger = logging.getLogger(__name__)


@dataclass
class ModelStats:
    """Collected statistics for a single model benchmark run."""
    source_model: str
    target_model: str | None = None
    provider: str = ""
    success: bool = False
    error: str = ""
    total_time_ms: float = 0.0
    upstream_latency_ms: float = 0.0
    gateway_overhead_ms: float = 0.0
    queue_wait_ms: float = 0.0
    connect_ms: float = 0.0
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    cache_hit_tokens: int = 0
    cache_miss_tokens: int = 0
    reasoning_tokens: int = 0
    context_window: int = 0
    context_utilization_pct: float = 0.0
    cost_usd: float = 0.0
    input_cost_usd: float = 0.0
    output_cost_usd: float = 0.0
    cost_native: bool = False
    tokens_per_second: float = 0.0
    response_text: str = ""


def get_api_key() -> str:
    """Read the API key from environment variable or .env file."""
    env_value = os.environ.get(API_KEY_ENV_VAR)
    if env_value:
        logger.info("API key found in environment variable %s", API_KEY_ENV_VAR)
        return env_value

    env_file = __import__("pathlib").Path(".env")
    if env_file.exists():
        for line in env_file.read_text().splitlines():
            line = line.strip()
            if line.startswith("#") or "=" not in line:
                continue
            key, _, value = line.partition("=")
            if key.strip() == API_KEY_ENV_VAR and value.strip():
                logger.info("API key found in .env file")
                return value.strip()

    logger.error("No API key found. Set %s env var or add to .env file.", API_KEY_ENV_VAR)
    sys.exit(1)


def list_models() -> list[dict]:
    """List all gateway models via the admin API."""
    url = f"{BASE_URL}/admin/gateway/models"
    headers = {"Accept": "application/json", "X-API-KEY": get_api_key()}
    req = Request(url, headers=headers, method="GET")
    try:
        with urlopen(req, timeout=30) as resp:
            raw = resp.read().decode("utf-8")
            data = json.loads(raw)
            return data.get("models", [])
    except Exception as exc:
        logger.error("Failed to list models: %s", exc)
        return []


def bench_model(model_name: str, prompt: str, max_tokens: int = 512) -> ModelStats:
    """Run a single benchmark against one model."""
    stats = ModelStats(source_model=model_name)
    url = f"{BASE_URL}/openai/v1/chat/completions"
    payload = {
        "model": model_name,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "stream": False,
    }
    headers = {
        "Authorization": f"Bearer {get_api_key()}",
        "Content-Type": "application/json",
        "X-App-Reference": "bench-py",
    }
    body = json.dumps(payload).encode("utf-8")
    req = Request(url, data=body, headers=headers, method="POST")

    start = time.monotonic()
    try:
        with urlopen(req, timeout=120) as resp:
            elapsed_ms = (time.monotonic() - start) * 1000
            raw = resp.read().decode("utf-8")
            data = json.loads(raw)
            stats.success = True
            stats.total_time_ms = elapsed_ms

            # Parse X-Gateway-* headers
            stats.target_model = resp.headers.get("X-Gateway-Model")
            stats.provider = resp.headers.get("X-Gateway-Backend", "")
            stats.prompt_tokens = int(resp.headers.get("X-Gateway-Prompt-Tokens", 0))
            stats.completion_tokens = int(resp.headers.get("X-Gateway-Completion-Tokens", 0))
            stats.total_tokens = int(resp.headers.get("X-Gateway-Total-Tokens", 0))
            stats.cache_hit_tokens = int(resp.headers.get("X-Gateway-Cache-Hit-Tokens", 0))
            stats.cache_miss_tokens = int(resp.headers.get("X-Gateway-Cache-Miss-Tokens", 0))
            stats.reasoning_tokens = int(resp.headers.get("X-Gateway-Reasoning-Tokens", 0))
            stats.context_window = int(resp.headers.get("X-Gateway-Context-Window", 0))
            stats.cost_usd = float(resp.headers.get("X-Gateway-Cost-USD", 0))
            stats.input_cost_usd = float(resp.headers.get("X-Gateway-Input-Cost-USD", 0))
            stats.output_cost_usd = float(resp.headers.get("X-Gateway-Output-Cost-USD", 0))
            stats.cost_native = resp.headers.get("X-Gateway-Cost-Native") == "1"
            stats.tokens_per_second = float(resp.headers.get("X-Gateway-Tokens-Per-Second", 0))
            stats.upstream_latency_ms = float(resp.headers.get("X-Gateway-Upstream-Latency-Ms", 0))
            stats.gateway_overhead_ms = float(resp.headers.get("X-Gateway-Gateway-Overhead-Ms", 0))
            stats.queue_wait_ms = float(resp.headers.get("X-Gateway-Queue-Wait-Ms", 0))
            stats.connect_ms = float(resp.headers.get("X-Gateway-Connect-Ms", 0))

            ctx_util = resp.headers.get("X-Gateway-Context-Utilization")
            if ctx_util:
                stats.context_utilization_pct = float(ctx_util) * 100

            # Extract response text
            choices = data.get("choices", [])
            if choices:
                msg = choices[0].get("message", {})
                stats.response_text = msg.get("content", "")[:200]

    except HTTPError as exc:
        stats.success = False
        stats.error = f"HTTP {exc.code}: {exc.read().decode('utf-8')[:200]}"
        stats.total_time_ms = (time.monotonic() - start) * 1000
    except URLError as exc:
        stats.success = False
        stats.error = f"URL Error: {exc.reason}"
        stats.total_time_ms = (time.monotonic() - start) * 1000
    except Exception as exc:
        stats.success = False
        stats.error = str(exc)[:200]
        stats.total_time_ms = (time.monotonic() - start) * 1000

    return stats


def format_table(results: list[tuple[dict, ModelStats]]) -> str:
    """Format benchmark results as a human-readable table."""
    lines: list[str] = []
    sep = "-" * 160

    lines.append(sep)
    lines.append(f"{'Model':<30} {'Provider':<12} {'Status':<8} {'Time(ms)':<10} {'TPS':<8} "
                 f"{'Prompt':<8} {'Output':<8} {'Cache%':<8} {'Cost($)':<10} {'Tokens/s':<10}")
    lines.append(sep)

    for model_info, stats in results:
        status = "OK" if stats.success else "FAIL"
        provider = stats.provider or model_info.get("provider", "")
        time_ms = f"{stats.total_time_ms:.0f}" if stats.success else "N/A"
        tps = f"{stats.tokens_per_second:.1f}" if stats.success else "N/A"
        prompt = str(stats.prompt_tokens) if stats.success else "-"
        output = str(stats.completion_tokens) if stats.success else "-"
        cache_pct = f"{(stats.cache_hit_tokens / stats.prompt_tokens * 100):.0f}" if stats.success and stats.prompt_tokens > 0 else "-"
        cost = f"{stats.cost_usd:.6f}" if stats.success else "-"

        lines.append(
            f"{stats.source_model:<30} {provider:<12} {status:<8} {time_ms:<10} {tps:<8} "
            f"{prompt:<8} {output:<8} {cache_pct:<8} {cost:<10} {tps:<10}"
        )

    lines.append(sep)
    return "\n".join(lines)


def format_detailed(results: list[tuple[dict, ModelStats]]) -> str:
    """Format detailed per-model statistics."""
    lines: list[str] = []
    lines.append("")
    lines.append("=" * 160)
    lines.append("DETAILED STATISTICS")
    lines.append("=" * 160)

    for model_info, stats in results:
        lines.append(f"\n--- {stats.source_model} ---")
        lines.append(f"  Target model:     {stats.target_model or 'N/A'}")
        lines.append(f"  Provider:         {stats.provider or model_info.get('provider', 'N/A')}")
        lines.append(f"  Status:           {'SUCCESS' if stats.success else 'FAILED'}")
        if not stats.success:
            lines.append(f"  Error:            {stats.error}")
            continue

        lines.append(f"  Total time:       {stats.total_time_ms:.0f} ms")
        lines.append(f"  Upstream latency: {stats.upstream_latency_ms:.0f} ms")
        lines.append(f"  Gateway overhead: {stats.gateway_overhead_ms:.0f} ms")
        lines.append(f"  Queue wait:       {stats.queue_wait_ms:.0f} ms")
        lines.append(f"  Connect time:     {stats.connect_ms:.0f} ms")
        lines.append(f"  Tokens/sec:       {stats.tokens_per_second:.1f}")
        lines.append(f"  Prompt tokens:    {stats.prompt_tokens}")
        lines.append(f"  Completion tokens:{stats.completion_tokens}")
        lines.append(f"  Total tokens:     {stats.total_tokens}")
        lines.append(f"  Cache hit:        {stats.cache_hit_tokens} ({stats.cache_hit_tokens/stats.prompt_tokens*100:.0f}%)" if stats.prompt_tokens > 0 else "  Cache hit:        N/A")
        lines.append(f"  Cache miss:       {stats.cache_miss_tokens}")
        lines.append(f"  Reasoning tokens: {stats.reasoning_tokens}")
        lines.append(f"  Context window:   {stats.context_window:,} tokens")
        lines.append(f"  Context util:     {stats.context_utilization_pct:.1f}%")
        lines.append(f"  Cost (USD):       ${stats.cost_usd:.6f}")
        lines.append(f"    Input cost:     ${stats.input_cost_usd:.6f}")
        lines.append(f"    Output cost:    ${stats.output_cost_usd:.6f}")
        lines.append(f"  Cost native:      {'Yes' if stats.cost_native else 'Computed'}")
        if stats.response_text:
            lines.append(f"  Response preview: {stats.response_text[:100]}...")

    return "\n".join(lines)


def format_summary(results: list[tuple[dict, ModelStats]]) -> str:
    """Format summary statistics across all models."""
    successful = [(mi, s) for mi, s in results if s.success]
    failed = [(mi, s) for mi, s in results if not s.success]

    lines: list[str] = []
    lines.append("")
    lines.append("=" * 80)
    lines.append("SUMMARY")
    lines.append("=" * 80)

    if not successful:
        lines.append("No models completed successfully.")
        return "\n".join(lines)

    times = [s.total_time_ms for _, s in successful]
    tpss = [s.tokens_per_second for _, s in successful]
    costs = [s.cost_usd for _, s in successful]
    outputs = [s.completion_tokens for _, s in successful]

    lines.append(f"  Models tested:    {len(results)}")
    lines.append(f"  Successful:       {len(successful)}")
    lines.append(f"  Failed:           {len(failed)}")
    lines.append("")
    lines.append(f"  {'Metric':<25} {'Min':<14} {'Max':<14} {'Mean':<14} {'Median':<14} {'StDev':<14}")
    lines.append(f"  {'-'*25} {'-'*14} {'-'*14} {'-'*14} {'-'*14} {'-'*14}")

    def fmt_stats(values, fmt="{:.0f}", unit=""):
        if len(values) < 2:
            return f"{fmt.format(values[0])}{unit}"
        mn, mx = min(values), max(values)
        mean_val = statistics.mean(values)
        med = statistics.median(values)
        stdev = statistics.stdev(values)
        return f"{fmt.format(mn)}{unit:<10} {fmt.format(mx)}{unit:<10} {fmt.format(mean_val)}{unit:<10} {fmt.format(med)}{unit:<10} {fmt.format(stdev)}{unit}"

    lines.append(f"  {'Total time (ms)':<25} {fmt_stats(times)}")
    lines.append(f"  {'Upstream lat (ms)':<25} {fmt_stats([s.upstream_latency_ms for _, s in successful])}")
    lines.append(f"  {'Tokens/sec':<25} {fmt_stats(tpss, '{:.1f}', '')}")
    lines.append(f"  {'Output tokens':<25} {fmt_stats(outputs, '{}', '')}")
    lines.append(f"  {'Cost (USD)':<25} {fmt_stats(costs, '${:.6f}', '')}")

    # Best performers
    best_tps = max(successful, key=lambda x: x[1].tokens_per_second)
    cheapest = min(successful, key=lambda x: x[1].cost_usd)
    fastest = min(successful, key=lambda x: x[1].total_time_ms)
    longest = max(successful, key=lambda x: x[1].completion_tokens)

    lines.append("")
    lines.append(f"  Fastest TPS:        {best_tps[1].source_model} ({best_tps[1].tokens_per_second:.1f} tok/s)")
    lines.append(f"  Cheapest:           {cheapest[1].source_model} (${cheapest[1].cost_usd:.6f})")
    lines.append(f"  Lowest latency:     {fastest[1].source_model} ({fastest[1].total_time_ms:.0f} ms)")
    lines.append(f"  Longest response:   {longest[1].source_model} ({longest[1].completion_tokens} tokens)")

    if failed:
        lines.append("")
        lines.append("  Failed models:")
        for _, s in failed:
            lines.append(f"    - {s.source_model}: {s.error}")

    return "\n".join(lines)


def main() -> int:
    parser = argparse.ArgumentParser(description="Benchmark all devplace models")
    parser.add_argument("--prompt", type=str, required=True, help="Prompt to send to each model")
    parser.add_argument("--max-tokens", type=int, default=512, help="Maximum output tokens per model")
    parser.add_argument("--api-key", type=str, default="", help="Override API key inline")
    parser.add_argument("--json", action="store_true", help="Output results as JSON only")
    args = parser.parse_args()

    if args.api_key:
        os.environ[API_KEY_ENV_VAR] = args.api_key
        logger.info("Using provided API key")

    models = list_models()
    if not models:
        logger.error("No models found on the gateway")
        return 1

    logger.info("Found %d models to benchmark", len(models))
    print(f"Benchmarking {len(models)} models with prompt: {args.prompt[:80]}...")
    print("")

    results: list[tuple[dict, ModelStats]] = []
    for i, model_info in enumerate(models, 1):
        source = model_info.get("source_model", "")
        print(f"[{i}/{len(models)}] Benchmarking {source}...", end=" ", flush=True)
        stats = bench_model(source, args.prompt, args.max_tokens)
        results.append((model_info, stats))
        status = "OK" if stats.success else f"FAIL ({stats.error[:50]})"
        if stats.success:
            print(f"{status} | {stats.completion_tokens} tokens | {stats.tokens_per_second:.1f} tok/s | ${stats.cost_usd:.6f}")
        else:
            print(status)

    if args.json:
        output = []
        for model_info, stats in results:
            entry = {
                "source_model": stats.source_model,
                "target_model": stats.target_model,
                "provider": stats.provider or model_info.get("provider", ""),
                "success": stats.success,
                "error": stats.error,
                "total_time_ms": round(stats.total_time_ms, 1),
                "upstream_latency_ms": round(stats.upstream_latency_ms, 1),
                "gateway_overhead_ms": round(stats.gateway_overhead_ms, 1),
                "queue_wait_ms": round(stats.queue_wait_ms, 1),
                "connect_ms": round(stats.connect_ms, 1),
                "prompt_tokens": stats.prompt_tokens,
                "completion_tokens": stats.completion_tokens,
                "total_tokens": stats.total_tokens,
                "cache_hit_tokens": stats.cache_hit_tokens,
                "cache_miss_tokens": stats.cache_miss_tokens,
                "reasoning_tokens": stats.reasoning_tokens,
                "context_window": stats.context_window,
                "context_utilization_pct": round(stats.context_utilization_pct, 1),
                "cost_usd": round(stats.cost_usd, 8),
                "input_cost_usd": round(stats.input_cost_usd, 8),
                "output_cost_usd": round(stats.output_cost_usd, 8),
                "cost_native": stats.cost_native,
                "tokens_per_second": round(stats.tokens_per_second, 2),
            }
            output.append(entry)
        print(json.dumps(output, indent=2))
    else:
        print()
        print(format_table(results))
        print(format_detailed(results))
        print(format_summary(results))

    return 0


if __name__ == "__main__":
    sys.exit(main())