Implement provider-neutral Lean reasoning and evaluation foundation

lean_agent/core.py

@@ -0,0 +1,111 @@
+"""Small provider-neutral proof repair loop. Dataset statements are trusted input."""
+import json
+import os
+import re
+import subprocess
+import tempfile
+import time
+import urllib.request
+from dataclasses import dataclass, asdict
+from pathlib import Path
+from typing import Protocol
+
+class Provider(Protocol):
+    def complete(self, messages: list[dict[str, str]]) -> str: ...
+
+class HTTPProvider:
+    """OpenAI-compatible chat completions, including local serving endpoints."""
+    def __init__(self, model, base_url, api_key_env="LLM_API_KEY", timeout=120):
+        self.model, self.base_url = model, base_url.rstrip('/')
+        self.api_key_env, self.timeout = api_key_env, timeout
+
+    def complete(self, messages):
+        headers = {"Content-Type": "application/json"}
+        key = os.environ.get(self.api_key_env)
+        if key:
+            headers["Authorization"] = "Bearer " + key
+        request = urllib.request.Request(self.base_url + '/chat/completions',
+            data=json.dumps({"model": self.model, "messages": messages}).encode(), headers=headers)
+        with urllib.request.urlopen(request, timeout=self.timeout) as response:
+            content = json.load(response)["choices"][0]["message"]["content"]
+        if not isinstance(content, str):
+            raise ValueError("Provider returned non-text content")
+        return content
+
+class CommandProvider:
+    """Any model: executable receives {messages: [...]} on stdin, returns proof text."""
+    def __init__(self, argv, timeout=120):
+        if not argv: raise ValueError("Empty provider command")
+        self.argv, self.timeout = argv, timeout
+
+    def complete(self, messages):
+        proc = subprocess.run(self.argv, input=json.dumps({"messages": messages}),
+                              text=True, capture_output=True, timeout=self.timeout)
+        if proc.returncode: raise RuntimeError("Provider command failed")
+        return proc.stdout
+
+@dataclass
+class Check:
+    success: bool
+    diagnostics: str
+
+class LeanVerifier:
+    def __init__(self, command=None, cwd=None, timeout=30):
+        self.command = command or ["lean"]
+        self.cwd = str(Path(cwd).resolve()) if cwd else None
+        self.timeout = timeout
+
+    def check(self, statement, proof, prelude=""):
+        # Only a tactic body is accepted. Fail closed on holes and command injection.
+        if re.search(r'\b(sorry|admit|axiom|unsafe|implemented_by|run_tac|elab|macro|syntax)\b|--|/\*|/\-|\-/|#', proof):
+            return Check(False, "Forbidden proof token (holes, metaprogramming, comments or commands)")
+        if not proof.strip(): return Check(False, "Empty proof")
+        source = prelude + '\nset_option maxHeartbeats 200000\n'
+        source += 'theorem agent_target : ' + statement + ' := by\n'
+        source += '\n'.join('  ' + line for line in proof.splitlines())
+        source += '\n#print axioms agent_target\n'
+        with tempfile.TemporaryDirectory(prefix="lean-check-", dir=self.cwd) as directory:
+            path = Path(directory) / 'Candidate.lean'
+            path.write_text(source)
+            try:
+                result = subprocess.run([*self.command, str(path)], cwd=self.cwd,
+                    capture_output=True, text=True, timeout=self.timeout)
+            except subprocess.TimeoutExpired:
+                return Check(False, "Lean timeout")
+            diagnostics = result.stdout + result.stderr
+        # Reject warnings and require the axiom report; allow only Lean's standard axioms.
+        reports = re.findall(r"depends on axioms:\s*\[([^\]]*)\]", diagnostics)
+        clean = "does not depend on any axioms" in diagnostics
+        if reports:
+            clean = all(set(x.strip() for x in report.split(',') if x.strip()) <=
+                        {"propext", "Classical.choice", "Quot.sound"} for report in reports)
+        return Check(result.returncode == 0 and clean and 'warning:' not in diagnostics,
+                     diagnostics[-16000:])
+
+def extract(text):
+    match = re.search(r'```(?:lean4?|text)?\s*\n(.*?)```', text, re.S)
+    proof = (match.group(1) if match else text).strip()
+    if proof == 'by': return ''
+    if proof.startswith('by\n') or proof.startswith('by '): proof = proof[2:].strip()
+    return proof
+
+def solve(provider, verifier, task, max_attempts=3):
+    if max_attempts < 1: raise ValueError("max_attempts must be positive")
+    start = time.monotonic()
+    messages = [{"role": "system", "content": "Prove the Lean 4 theorem. Return ONLY the tactic body after by. No theorem declaration, comments, sorry, admit, metaprogramming or extra commands."},
+                {"role": "user", "content": task.get('prelude', '') + '\n⊢ ' + task['statement']}]
+    attempts = []
+    for _ in range(max_attempts):
+        try:
+            proof = extract(provider.complete(messages))
+            check = verifier.check(task['statement'], proof, task.get('prelude', ''))
+        except Exception as exc:
+            # Do not serialize exception bodies: providers may include credentials.
+            return {"id": task['id'], "success": False, "error": type(exc).__name__,
+                    "attempts": attempts, "seconds": time.monotonic() - start}
+        attempts.append({"proof": proof, **asdict(check)})
+        if check.success: break
+        messages.extend([{"role": "assistant", "content": proof},
+                         {"role": "user", "content": "Lean rejected this proof. Repair it:\n" + check.diagnostics}])
+    return {"id": task['id'], "success": attempts[-1]['success'],
+            "attempts": attempts, "seconds": time.monotonic() - start}