tinymind-native-8b-remote-handoff/bundle/evaluation/llm_stats_gateway.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.request import Request, urlopen

from evaluation.llm_stats_integration import _env_api_key


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


class LLMStatsGatewayClient:
    def __init__(
        self,
        api_key: str | None = None,
        *,
        base_url: str = "https://gateway.llm-stats.com/v1",
        timeout_s: float = 60.0,
    ):
        self.api_key = api_key or os.environ.get("LLM_STATS_GATEWAY_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[:5] + "***" + self.api_key[-4:]

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

    def chat_completion(
        self,
        *,
        model: str,
        messages: list[dict[str, str]],
        temperature: float = 0.2,
        max_tokens: int = 512,
        fetcher: GatewayFetcher | None = None,
    ) -> dict[str, Any]:
        payload = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        fn = fetcher or _default_gateway_fetcher
        return fn(f"{self.base_url}/chat/completions", self.headers(), payload, self.timeout_s)


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


def _choice_text(response: dict[str, Any]) -> str:
    try:
        return str(response["choices"][0]["message"]["content"])
    except (KeyError, IndexError, TypeError):
        return ""


def build_llm_stats_gateway_probe(
    out_dir: str | Path,
    *,
    api_key: str | None = None,
    model: str = "glm-5.1",
    prompt: str = "What is machine learning?",
    fetcher: GatewayFetcher | None = None,
) -> dict[str, Any]:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    client = LLMStatsGatewayClient(api_key=api_key)
    if not client.api_key:
        report = {
            "schema_version": "tinymind-llm-stats-gateway-probe-v1",
            "created_at": datetime.now(timezone.utc).isoformat(),
            "model": model,
            "api_key_redacted": "",
            "response_text": "",
            "raw_response": {},
            "claim_gate": {
                "gateway_probe_completed": False,
                "reason": "LLM_STATS_GATEWAY_API_KEY or LLM_STATS_API_KEY is not visible to this process.",
            },
        }
    else:
        response = client.chat_completion(
            model=model,
            messages=[{"role": "user", "content": prompt}],
            fetcher=fetcher,
        )
        report = {
            "schema_version": "tinymind-llm-stats-gateway-probe-v1",
            "created_at": datetime.now(timezone.utc).isoformat(),
            "model": model,
            "base_url": client.base_url,
            "api_key_redacted": client.redacted_key(),
            "prompt_sha256": __import__("hashlib").sha256(prompt.encode("utf-8")).hexdigest(),
            "response_text": _choice_text(response),
            "usage": response.get("usage", {}),
            "raw_response": response,
            "claim_gate": {
                "gateway_probe_completed": True,
                "tiny_model_rank_claim_allowed": False,
                "reason": "This probes an external gateway model, not TinyMind's own submitted benchmark result.",
            },
        }
    path = out / "llm_stats_gateway_probe.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