#!/usr/bin/env python3
"""Test which FreeLLMAPI model is best for Lean proofs."""
import urllib.request
import json

API = "http://127.0.0.1:3001/v1"
KEY = "freellmapi-ef5610fa456735f2bcb6205faf03f8725fe41265f129cd56"

models = ["codestral", "deepseek-v4-flash", "deepseek-v4-pro", "phi-4-reasoning", "hermes-3-405b", "nous-coder"]

prompt = "theorem add_comm (a b : Nat) : a + b = b + a := by\n  sorry\n\nFill the sorry. Output ONLY: rfl or the proof term."

for model in models:
    body = json.dumps({
        "model": model,
        "messages": [
            {"role": "system", "content": "You are a Lean 4 proof engineer. Output raw Lean code only."},
            {"role": "user", "content": prompt},
        ],
        "temperature": 0.1,
        "max_tokens": 200,
    }).encode()
    req = urllib.request.Request(
        f"{API}/chat/completions", data=body,
        headers={"Content-Type": "application/json",
                 "Authorization": f"Bearer {KEY}"},
    )
    try:
        resp = urllib.request.urlopen(req, timeout=60)
        data = json.loads(resp.read())
        content = data["choices"][0]["message"]["content"][:100]
        routed = data.get("_routed_via", {})
        print(f"  {model:25s} → {content}")
    except Exception as e:
        print(f"  {model:25s} → ERROR: {str(e)[:60]}")