ALchemt/llm-eval

1"""Run test suites against a set of LLM configs.2 3For each run config in configs.yaml and each prompt across the specified4suites, call the LLM and append the raw response as a jsonl line into5`runs/raw_<run_id>.jsonl`.6 7Usage:8    python -m src.runner                 # real LLM calls (needs HF_TOKEN)9    python -m src.runner --dry-run       # deterministic mock LLM (no token)10    python -m src.runner --runs baseline # limit to a single run id11"""12 13from __future__ import annotations14 15import argparse16import json17import os18import sys19import time20from dataclasses import asdict, dataclass21from pathlib import Path22 23import yaml24from dotenv import load_dotenv25 26ROOT = Path(__file__).resolve().parent.parent27SUITES_DIR = ROOT / "suites"28RUNS_DIR = ROOT / "runs"29CONFIGS_PATH = ROOT / "configs.yaml"30 31load_dotenv(ROOT / ".env")32 33 34@dataclass35class Sample:36    run_id: str37    suite: str38    prompt_id: str39    prompt: str40    expected: str41    rubric: str42    response: str43    latency_ms: int44    input_tokens: int45    output_tokens: int46    model: str47    temperature: float48    seed: int49    mock: bool50 51 52def load_configs() -> dict:53    with CONFIGS_PATH.open() as f:54        return yaml.safe_load(f)55 56 57def load_suite(suite_name: str) -> list[dict]:58    path = SUITES_DIR / f"{suite_name}.jsonl"59    with path.open() as f:60        return [json.loads(line) for line in f if line.strip()]61 62 63def discover_suites() -> list[str]:64    return sorted(p.stem for p in SUITES_DIR.glob("*.jsonl"))65 66 67def mock_call(prompt: str, system_prompt: str, temperature: float) -> tuple[str, int, int]:68    """Deterministic echo response for offline scaffold runs."""69    text = f"[MOCK T={temperature}] Echo of: {prompt[:80]}"70    return text, len(prompt.split()), len(text.split())71 72 73PROVIDER_ENDPOINTS = {74    "openai": {"base_url": None, "env_key": "OPENAI_API_KEY"},75    "openrouter": {"base_url": "https://openrouter.ai/api/v1", "env_key": "OPENROUTER_API_KEY"},76    "groq": {"base_url": "https://api.groq.com/openai/v1", "env_key": "GROQ_API_KEY"},77    "together": {"base_url": "https://api.together.xyz/v1", "env_key": "TOGETHER_API_KEY"},78}79 80 81def live_call(82    prompt: str,83    system_prompt: str,84    model: str,85    provider: str,86    temperature: float,87    max_tokens: int,88) -> tuple[str, int, int]:89    """OpenAI-compatible chat completion. Works against OpenAI, OpenRouter,90    Groq, Together, etc. — all expose the same Chat Completions schema."""91    from openai import OpenAI92 93    cfg = PROVIDER_ENDPOINTS.get(provider)94    if cfg is None:95        raise RuntimeError(f"Unknown provider '{provider}'. Add it to PROVIDER_ENDPOINTS.")96 97    api_key = os.getenv(cfg["env_key"])98    if not api_key:99        raise RuntimeError(f"{cfg['env_key']} not set. Copy .env.example to .env or pass --dry-run.")100 101    client = OpenAI(api_key=api_key, base_url=cfg["base_url"])102    resp = client.chat.completions.create(103        model=model,104        messages=[105            {"role": "system", "content": system_prompt},106            {"role": "user", "content": prompt},107        ],108        temperature=temperature,109        max_tokens=max_tokens,110    )111    text = resp.choices[0].message.content or ""112    usage = resp.usage113    return (114        text,115        getattr(usage, "prompt_tokens", 0) or 0,116        getattr(usage, "completion_tokens", 0) or 0,117    )118 119 120def run_one(121    run_cfg: dict,122    suite_name: str,123    prompts: list[dict],124    dry_run: bool,125    seed: int,126) -> list[Sample]:127    samples: list[Sample] = []128    for p in prompts:129        t0 = time.time()130        if dry_run:131            text, in_tok, out_tok = mock_call(132                p["prompt"], run_cfg["system_prompt"], run_cfg["temperature"]133            )134        else:135            text, in_tok, out_tok = live_call(136                p["prompt"],137                run_cfg["system_prompt"],138                run_cfg["model"],139                run_cfg["provider"],140                run_cfg["temperature"],141                run_cfg.get("max_tokens", 512),142            )143        latency_ms = int((time.time() - t0) * 1000)144        samples.append(145            Sample(146                run_id=run_cfg["id"],147                suite=suite_name,148                prompt_id=p["id"],149                prompt=p["prompt"],150                expected=p.get("expected", ""),151                rubric=p.get("rubric", "judge"),152                response=text,153                latency_ms=latency_ms,154                input_tokens=in_tok,155                output_tokens=out_tok,156                model=run_cfg["model"],157                temperature=run_cfg["temperature"],158                seed=seed,159                mock=dry_run,160            )161        )162    return samples163 164 165def write_raw(run_id: str, samples: list[Sample]) -> Path:166    RUNS_DIR.mkdir(exist_ok=True)167    out = RUNS_DIR / f"raw_{run_id}.jsonl"168    with out.open("w") as f:169        for s in samples:170            f.write(json.dumps(asdict(s)) + "\n")171    return out172 173 174def main() -> int:175    ap = argparse.ArgumentParser()176    ap.add_argument("--dry-run", action="store_true", help="use mock LLM, no HF calls")177    ap.add_argument("--runs", nargs="*", help="limit to these run ids (default: all)")178    ap.add_argument("--suites", nargs="*", help="limit to these suites (default: all)")179    ap.add_argument("--seed", type=int, default=42)180    args = ap.parse_args()181 182    cfg = load_configs()183    run_cfgs = cfg["runs"]184    if args.runs:185        run_cfgs = [r for r in run_cfgs if r["id"] in args.runs]186        if not run_cfgs:187            print(f"No runs match {args.runs}. Available: {[r['id'] for r in cfg['runs']]}")188            return 1189 190    suite_names = args.suites or discover_suites()191    if not suite_names:192        print("No suites found in suites/*.jsonl")193        return 1194 195    for run_cfg in run_cfgs:196        all_samples: list[Sample] = []197        for suite_name in suite_names:198            prompts = load_suite(suite_name)199            samples = run_one(run_cfg, suite_name, prompts, args.dry_run, args.seed)200            all_samples.extend(samples)201        out = write_raw(run_cfg["id"], all_samples)202        mode = "DRY" if args.dry_run else "LIVE"203        print(f"[{mode}] {run_cfg['id']}: wrote {len(all_samples)} samples to {out.relative_to(ROOT)}")204    return 0205 206 207if __name__ == "__main__":208    sys.exit(main())209