runner.py

"""
Benchmark test execution engine.

Handles:
  - Official mode: runs all tests against the official API and records baseline results
  - Test mode: runs all tests against the relay service and compares with the baseline
"""

import json
import os
import re
import sys
import time
from dataclasses import dataclass, field, asdict
from datetime import datetime, timezone
from typing import Any, Optional
try:
    from tqdm import tqdm
except ModuleNotFoundError:
    def tqdm(iterable, **kwargs):
        return iterable

try:
    from llm_benchmark.api_client import APIClient, create_client
    from llm_benchmark.tests import (
        ALL_TESTS,
        TOKENIZER_TESTS,
        PROGRAMMING_TESTS,
        get_test_by_id,
    )
except ModuleNotFoundError:
    from api_client import APIClient, create_client
    from tests import (
        ALL_TESTS,
        TOKENIZER_TESTS,
        PROGRAMMING_TESTS,
        get_test_by_id,
    )


@dataclass
class TestResult:
    """Result of a single test."""

    test_id: str
    test_name: str
    category: str
    passed: bool
    score: float
    details: str
    elapsed_seconds: float = 0.0
    token_count: Optional[int] = None
    content: Optional[str] = None
    baseline: Optional[dict] = None
    metadata: Optional[dict] = None

    def to_dict(self) -> dict:
        d = asdict(self)
        if self.content and len(self.content) > 500:
            d["content"] = self.content[:500] + "..."
        return d


@dataclass
class BenchmarkRun:
    """Complete benchmark run result."""

    mode: str  # 'official' or 'test'
    api_type: str
    model: str
    base_url: Optional[str]
    timestamp: str
    results: list = field(default_factory=list)
    total_tests: int = 0
    passed_tests: int = 0
    overall_score: float = 0.0
    category_scores: dict = field(default_factory=dict)
    temperature: float = 0.0

    def calculate_summary(self):
        """Calculate summary statistics from results."""
        self.total_tests = len(self.results)
        self.passed_tests = sum(1 for r in self.results if r.passed)
        self.overall_score = (
            sum(r.score for r in self.results) / len(self.results)
            if self.results
            else 0.0
        )

        # Per-category scores
        categories = {}
        for r in self.results:
            if r.category not in categories:
                categories[r.category] = []
            categories[r.category].append(r.score)
        self.category_scores = {
            cat: sum(scores) / len(scores) for cat, scores in categories.items()
        }

    def to_dict(self) -> dict:
        return {
            "mode": self.mode,
            "api_type": self.api_type,
            "model": self.model,
            "base_url": self.base_url,
            "timestamp": self.timestamp,
            "temperature": self.temperature,
            "total_tests": self.total_tests,
            "passed_tests": self.passed_tests,
            "overall_score": round(self.overall_score, 4),
            "category_scores": {
                k: round(v, 4) for k, v in self.category_scores.items()
            },
            "results": [r.to_dict() for r in self.results],
        }


class BenchmarkRunner:
    """Executes benchmark tests against an API."""

    def __init__(
        self,
        api_type: str,
        api_key: str,
        model: str,
        base_url: Optional[str] = None,
        temperature: float = 0.0,
        top_p: float = 1.0,
        max_tokens: int = 4096,
        timeout: float = 120.0,
    ):
        self.api_type = api_type
        self.model = model
        self.base_url = base_url
        self.temperature = temperature
        self.top_p = top_p
        self.max_tokens = max_tokens
        self.timeout = timeout

        self.client = create_client(
            api_type=api_type,
            api_key=api_key,
            model=model,
            base_url=base_url,
            temperature=temperature,
            top_p=top_p,
            max_tokens=max_tokens,
            timeout=timeout,
        )

        # Programing test cache (code results from official run)
        self._programming_cache = {}

    def close(self):
        self.client.close()

    def run_test(self, test: dict, baseline_data: Optional[dict] = None) -> TestResult:
        """Execute a single test case."""
        test_id = test["id"]
        test_name = test["name"]
        category = test["category"]

        # Get baseline for this test if available
        test_baseline = None
        if baseline_data and test_id in baseline_data:
            test_baseline = baseline_data[test_id]

        start_time = time.time()

        try:
            # Preference-fingerprint tests need repeated stochastic samples, not
            # a single temperature=0 completion.
            if test.get("sample_count"):
                samples = []
                usage_totals = {
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": 0,
                }
                sample_count = int(test["sample_count"])

                for _ in range(sample_count):
                    response = self.client.chat(
                        test["messages"],
                        temperature=test.get("temperature", self.temperature),
                        top_p=test.get("top_p", self.top_p),
                        max_tokens=test.get("max_tokens", self.max_tokens),
                    )
                    samples.append(response.content)
                    if response.usage.prompt_tokens is not None:
                        usage_totals["prompt_tokens"] += response.usage.prompt_tokens
                    if response.usage.completion_tokens is not None:
                        usage_totals["completion_tokens"] += response.usage.completion_tokens
                    if response.usage.total_tokens is not None:
                        usage_totals["total_tokens"] += response.usage.total_tokens

                usage_dict = {
                    **usage_totals,
                    "samples": samples,
                    "sample_count": sample_count,
                }
                result = test["evaluator"]("", usage_dict, test_baseline)
                metadata = result.get("metadata", {})
                content = json.dumps(metadata.get("distribution", {}), ensure_ascii=False)

                return TestResult(
                    test_id=test_id,
                    test_name=test_name,
                    category=category,
                    passed=result["passed"],
                    score=result["score"],
                    details=result["details"],
                    elapsed_seconds=time.time() - start_time,
                    token_count=usage_totals["prompt_tokens"],
                    content=content,
                    metadata=metadata,
                )

            # For tokenizer tests, we need the token count from the API response
            if test.get("tokenizer_probe"):
                response = self.client.chat(test["messages"])
                token_count = response.usage.prompt_tokens
                content = response.content

                result = test["evaluator"](
                    content,
                    {"prompt_tokens": token_count},
                    test_baseline,
                )

                return TestResult(
                    test_id=test_id,
                    test_name=test_name,
                    category=category,
                    passed=result["passed"],
                    score=result["score"],
                    details=result["details"],
                    elapsed_seconds=time.time() - start_time,
                    token_count=token_count,
                    content=content[:200] if content else "",
                )

            # All other tests
            response = self.client.chat(test["messages"])
            content = response.content
            usage_dict = {
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens,
            }

            evaluator_baseline = None if test.get("llm_judge") and test_baseline else test_baseline

            # For programming tests, check cache first
            if category == "programming":
                if test_id in self._programming_cache:
                    result = self._programming_cache[test_id]
                else:
                    result = test["evaluator"](content, usage_dict, evaluator_baseline)
                    self._programming_cache[test_id] = result
            else:
                result = test["evaluator"](content, usage_dict, evaluator_baseline)

            return TestResult(
                test_id=test_id,
                test_name=test_name,
                category=category,
                passed=result["passed"],
                score=result["score"],
                details=result["details"],
                elapsed_seconds=time.time() - start_time,
                token_count=response.usage.prompt_tokens,
                content=content,
                metadata=result.get("metadata"),
            )

        except Exception as e:
            return TestResult(
                test_id=test_id,
                test_name=test_name,
                category=category,
                passed=False,
                score=0.0,
                details=f"Error: {type(e).__name__}: {str(e)[:300]}",
                elapsed_seconds=time.time() - start_time,
                content="",
            )

    def _extract_json_array(self, text: str) -> list[dict]:
        """Extract a JSON array from a judge response."""
        match = re.search(r"```(?:json)?\s*([\s\S]*?)```", text)
        if match:
            text = match.group(1)
        array_match = re.search(r"(\[[\s\S]*\])", text)
        if array_match:
            text = array_match.group(1)
        parsed = json.loads(text.strip())
        if not isinstance(parsed, list):
            raise ValueError("Judge response was not a JSON array")
        return [item for item in parsed if isinstance(item, dict)]

    def _judge_prompt(self, category: str, items: list[dict]) -> list[dict]:
        """Ask the model under test to judge a batch of baseline/relay comparisons."""
        payload = [
            {
                "test_id": item["test_id"],
                "test_name": item["test_name"],
                "official_baseline": item["official"],
                "provider_output": item["actual"],
            }
            for item in items
        ]
        instructions = (
            "You are judging whether a routed model output is consistent with an official "
            "baseline output for an LLM authenticity benchmark. Judge semantic consistency, "
            "not exact wording. For identity tests, accept equivalent organization/model facts "
            "and reject different model-family claims. For knowledge-boundary tests, accept "
            "answers that name the same entity/fact even with wording differences. For alignment "
            "tests, compare policy behavior: refusal vs conditional safe help vs procedural help, "
            "and the safety boundary being applied.\n\n"
            "Return ONLY a JSON array. Each item must be: "
            "{\"test_id\": string, \"consistent\": boolean, \"score\": number between 0 and 1, "
            "\"reason\": short string}. Use score 1 for clearly consistent, 0.7 for mostly "
            "consistent, 0.4 for ambiguous/partial, 0 for inconsistent."
        )
        response = self.client.chat(
            [
                {"role": "system", "content": instructions},
                {
                    "role": "user",
                    "content": (
                        f"Judge this {category} batch:\n"
                        + json.dumps(payload, ensure_ascii=False, indent=2)
                    ),
                },
            ],
            temperature=0.0,
            top_p=1.0,
            max_tokens=min(self.max_tokens, 2048),
        )
        return self._extract_json_array(response.content)

    def apply_llm_judging(self, results: list[TestResult], baseline_data: Optional[dict], tests: Optional[list] = None) -> None:
        """Batch-judge dynamic-text tests using the model under test.

        Deterministic tests keep their original evaluator scores. Tests opt into
        this path with ``llm_judge: True`` in their test definition.
        """
        if not baseline_data:
            return

        test_defs = {t["id"]: t for t in (tests if tests is not None else ALL_TESTS)}
        batches: dict[str, list[dict]] = {}
        result_map = {r.test_id: r for r in results}

        for result in results:
            test_def = test_defs.get(result.test_id)
            baseline = baseline_data.get(result.test_id) if baseline_data else None
            if not test_def or not test_def.get("llm_judge") or not baseline:
                continue
            batches.setdefault(result.category, []).append(
                {
                    "test_id": result.test_id,
                    "test_name": result.test_name,
                    "official": baseline.get("content", ""),
                    "actual": result.content or "",
                }
            )

        for category, items in batches.items():
            try:
                decisions = self._judge_prompt(category, items)
            except Exception as exc:
                for item in items:
                    result = result_map[item["test_id"]]
                    result.passed = False
                    result.score = 0.0
                    result.details = f"LLM judge failed: {type(exc).__name__}: {str(exc)[:220]}"
                    metadata = result.metadata or {}
                    metadata["llm_judge_error"] = str(exc)[:500]
                    result.metadata = metadata
                continue

            decision_map = {str(d.get("test_id")): d for d in decisions}
            for item in items:
                result = result_map[item["test_id"]]
                decision = decision_map.get(item["test_id"])
                if not decision:
                    result.passed = False
                    result.score = 0.0
                    result.details = "LLM judge omitted this test"
                    continue

                try:
                    score = float(decision.get("score", 0.0))
                except (TypeError, ValueError):
                    score = 0.0
                score = max(0.0, min(1.0, score))
                consistent = bool(decision.get("consistent")) or score >= 0.65
                reason = str(decision.get("reason", ""))[:300]

                result.passed = consistent and score >= 0.65
                result.score = score
                result.details = f"LLM judge: {reason}" if reason else "LLM judge completed"
                metadata = result.metadata or {}
                metadata["llm_judge"] = {
                    "consistent": consistent,
                    "score": score,
                    "reason": reason,
                }
                result.metadata = metadata
    def run_all(
        self,
        baseline_file: Optional[str] = None,
        max_retries: int = 2,
        tests: Optional[list] = None,
    ) -> BenchmarkRun:
        """Run all benchmark tests.

        Args:
            baseline_file: Path to baseline JSON (for test mode).
            max_retries: Number of retries for failed API calls.
            tests: Optional subset of tests to run. Defaults to ALL_TESTS.
        """
        # Load baseline data if provided
        baseline_data = None
        if baseline_file and os.path.exists(baseline_file):
            with open(baseline_file, "r", encoding="utf-8") as f:
                raw = json.load(f)
                # The baseline is stored in the "results" dict keyed by test_id
                if "results" in raw:
                    baseline_data = {}
                    for r in raw["results"]:
                        tid = r.get("test_id")
                        if tid:
                            baseline_data[tid] = r
                else:
                    baseline_data = raw
            print(f"  Loaded baseline from: {baseline_file}")

        test_list = tests if tests is not None else ALL_TESTS
        mode = "test" if baseline_data else "official"
        mode_label = "OFFICIAL (recording baseline)" if mode == "official" else "TEST (comparing)"

        print(f"\n{'='*60}")
        print(f"  LLM Relay Authentication Benchmark")
        print(f"  Mode: {mode_label}")
        print(f"  API: {self.api_type.upper()} | Model: {self.model}")
        if self.base_url:
            print(f"  Endpoint: {self.base_url}")
        else:
            print(f"  Endpoint: (default for {self.api_type})")
        print(f"{'='*60}\n")

        results = []

        for test in tqdm(test_list, desc="Running tests", unit="test"):
            test_id = test["id"]
            test_name = test["name"]

            for attempt in range(1 + max_retries):
                tr = self.run_test(test, baseline_data)
                if tr.score > 0 or attempt >= max_retries:
                    break
                if attempt < max_retries:
                    time.sleep(1)

            results.append(tr)

            # Short status line
            status = "PASS" if tr.passed else "FAIL"
            score_pct = f"{tr.score:.0%}"
            print(f"  [{status}] [{score_pct}] {test_id}: {test_name} ({tr.details[:80]})")

        if baseline_data:
            print("  Running LLM judge for dynamic-text comparisons...")
            self.apply_llm_judging(results, baseline_data, test_list)

        run = BenchmarkRun(
            mode=mode,
            api_type=self.api_type,
            model=self.model,
            base_url=self.base_url,
            timestamp=datetime.now(timezone.utc).isoformat(),
            results=results,
            temperature=self.temperature,
        )
        run.calculate_summary()

        return run