tinymind-native-colab-handoff/bundle/evaluation/llm_stats_integration.py

from __future__ import annotations

from datetime import datetime, timezone
import json
import os
from pathlib import Path
from typing import Any, Callable
from urllib.parse import urlencode
from urllib.request import Request, urlopen


Fetcher = Callable[[str, dict[str, str], float], dict[str, Any]]


class LLMStatsClient:
    def __init__(
        self,
        api_key: str | None = None,
        *,
        base_url: str = "https://api.llm-stats.com/stats/v1",
        timeout_s: float = 20.0,
    ):
        self.api_key = api_key or _env_api_key()
        self.base_url = base_url.rstrip("/")
        self.timeout_s = timeout_s

    def redacted_key(self) -> str:
        if not self.api_key:
            return ""
        if len(self.api_key) <= 8:
            return self.api_key[:2] + "***"
        return self.api_key[:3] + "***" + self.api_key[-4:]

    def headers(self) -> dict[str, str]:
        if not self.api_key:
            raise RuntimeError("LLM_STATS_API_KEY is not set")
        return {
            "Authorization": f"Bearer {self.api_key}",
            "Accept": "application/json",
            "User-Agent": "TinyMind-Eval/1.0",
        }

    def get(self, path: str, params: dict[str, Any] | None = None, *, fetcher: Fetcher | None = None) -> dict[str, Any]:
        query = "?" + urlencode({k: v for k, v in (params or {}).items() if v is not None}) if params else ""
        url = f"{self.base_url}/{path.lstrip('/')}{query}"
        fn = fetcher or _default_fetcher
        return fn(url, self.headers(), self.timeout_s)


def _env_api_key() -> str:
    key = os.environ.get("LLM_STATS_API_KEY", "")
    if key:
        return key
    # Windows child processes only inherit process env set before launch. This
    # fallback lets Codex see user/machine env vars without logging the secret.
    if os.name == "nt":
        try:
            import subprocess

            script = "[Environment]::GetEnvironmentVariable('LLM_STATS_API_KEY','User'); [Environment]::GetEnvironmentVariable('LLM_STATS_API_KEY','Machine')"
            out = subprocess.check_output(["powershell", "-NoProfile", "-Command", script], text=True, stderr=subprocess.DEVNULL, timeout=5)
            for line in out.splitlines():
                if line.strip():
                    return line.strip()
        except Exception:
            return ""
    return ""


def _default_fetcher(url: str, headers: dict[str, str], timeout: float) -> dict[str, Any]:
    req = Request(url, headers=headers, method="GET")
    with urlopen(req, timeout=timeout) as resp:
        payload = resp.read().decode("utf-8", errors="replace")
    return json.loads(payload)


def build_llm_stats_report(
    out_dir: str | Path,
    *,
    api_key: str | None = None,
    category: str = "coding",
    model_ids: list[str] | None = None,
    fetcher: Fetcher | None = None,
) -> dict[str, Any]:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    client = LLMStatsClient(api_key=api_key)
    if not client.api_key:
        report = {
            "schema_version": "tinymind-llm-stats-report-v1",
            "created_at": datetime.now(timezone.utc).isoformat(),
            "provider": "llm-stats",
            "api_key_redacted": "",
            "category": category,
            "model_ids": model_ids or [],
            "rankings": {},
            "benchmarks": {},
            "scores": {},
            "claim_gate": {
                "external_stats_imported": False,
                "tiny_model_rank_claim_allowed": False,
                "reason": "LLM_STATS_API_KEY is not visible to this process. Set it in the Codex process/User environment and rerun.",
            },
        }
        path = out / "llm_stats_report.json"
        report["json_path"] = str(path)
        path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
        return report
    rankings = client.get("rankings", {"category": category}, fetcher=fetcher)
    benchmarks = client.get("benchmarks", fetcher=fetcher)
    scores = client.get("scores", {"models": ",".join(model_ids)} if model_ids else None, fetcher=fetcher)
    report = {
        "schema_version": "tinymind-llm-stats-report-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "provider": "llm-stats",
        "api_key_redacted": client.redacted_key(),
        "category": category,
        "model_ids": model_ids or [],
        "endpoints": {
            "rankings": f"{client.base_url}/rankings",
            "benchmarks": f"{client.base_url}/benchmarks",
            "scores": f"{client.base_url}/scores",
        },
        "rankings": rankings,
        "benchmarks": benchmarks,
        "scores": scores,
        "claim_gate": {
            "external_stats_imported": True,
            "tiny_model_rank_claim_allowed": False,
            "reason": "This imports external leaderboard/stat data. TinyMind rank claims require submitting TinyMind itself and matching returned official model id/results.",
        },
    }
    path = out / "llm_stats_report.json"
    report["json_path"] = str(path)
    path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    return report