112 lines
5.5 KiB
Python
112 lines
5.5 KiB
Python
"""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}
|