lean-bench/leanbench/formalize.py

"""Step 3: every detected span becomes Lean 4 code, then Lean gets the last word.

A ``statement`` item becomes whatever declaration fits -- ``def``, ``structure``,
``theorem`` -- with a proof only if the source actually gave one. A ``problem`` item
is an open question, so it becomes a statement with ``sorry`` for the proof and only
the statement's legality is checked.

When Lean rejects the code, the diagnostics go back to the model and it gets another
go: a compiler error is the cheapest, most honest feedback available.
"""
from __future__ import annotations

import re
from collections.abc import Callable
from dataclasses import asdict, dataclass, field
from typing import Any

from .detect import KIND_PROBLEM, Detection
from .llm import extract_code
from .typeindex import index_types
from .verify import MODE_FULL, MODE_STATEMENT, VerifyResult, contains_sorry

SYSTEM_PROMPT = """\
You translate mathematics from books into Lean 4 (Mathlib). You write Lean 4 with
Mathlib naming and syntax -- never Lean 3, never `begin ... end`.

You are given one passage, extracted from a PDF, so expect broken line breaks,
hyphenation and mangled symbols; reconstruct what the mathematics must have been.

Produce one self-contained top-level declaration that captures the passage:
- a definition in the passage -> `def`, `abbrev`, `structure`, `class` or `instance`;
- a theorem / proposition / lemma / corollary -> `theorem`;
- an exercise or open problem -> `theorem` whose statement is the thing to be shown.

Rules:
- Prefer existing Mathlib definitions over re-defining them. Never invent Mathlib
  lemma or definition names: if you are unsure a name exists, restate the maths with
  primitives you are sure of.
- State the mathematics faithfully. A statement that is trivially true, or that
  quantifies over the wrong thing, is a failure even if it compiles.
- Give the declaration a descriptive snake_case name.
- Include a proof ONLY when the passage itself contains the proof and you are
  confident in it; otherwise end the declaration with `sorry`.
- If the passage is not formalizable as a single declaration, say so with
  "skip": true and a short reason.

Reply with JSON only:
{"skip": false,
 "reason": "",
 "imports": ["Mathlib"],
 "name": "snake_case_name",
 "lean_code": "theorem snake_case_name (n : \\u2115) : ... := by\\n  sorry",
 "has_proof": false,
 "informal": "one-line English restatement",
 "index_types": ["Int", "Nat"],
 "confidence": 0.0}

"index_types" lists the Lean types this item is about: for a theorem, the types of
its inputs and hypotheses; for a definition, the type of the thing being defined.
"""

REPAIR_PROMPT = """\
Lean rejected that code. Fix it and reply with the same JSON shape.

Do not paper over the error with `sorry` in place of a statement, and do not change
what the passage says in order to make it compile. If a Mathlib name does not exist,
restate the mathematics with primitives that do.

Lean diagnostics:
{diagnostics}
"""


@dataclass
class Formalization:
    """The Lean rendering of one detected item, plus what Lean thought of it."""

    detection: Detection
    lean_code: str = ""
    imports: list[str] = field(default_factory=lambda: ["Mathlib"])
    name: str = ""
    informal: str = ""
    has_proof: bool = False
    confidence: float = 0.0
    skipped: bool = False
    skip_reason: str = ""
    attempts: int = 0
    verify: VerifyResult | None = None
    index_types: list[tuple[str, str]] = field(default_factory=list)

    @property
    def verified(self) -> bool:
        return bool(self.verify and self.verify.ok)

    @property
    def proved(self) -> bool:
        return bool(self.verify and self.verify.proved)

    def as_dict(self) -> dict[str, Any]:
        return {
            "detection": self.detection.as_dict(),
            "lean_code": self.lean_code,
            "imports": self.imports,
            "name": self.name,
            "informal": self.informal,
            "has_proof": self.has_proof,
            "confidence": self.confidence,
            "skipped": self.skipped,
            "skip_reason": self.skip_reason,
            "attempts": self.attempts,
            "verify": self.verify.as_dict() if self.verify else None,
            "index_types": [{"name": n, "role": r} for n, r in self.index_types],
        }


def verification_mode(detection: Detection) -> str:
    """Open problems carry no proof, so only their statement can be checked."""
    return MODE_STATEMENT if detection.kind == KIND_PROBLEM else MODE_FULL


def _user_prompt(detection: Detection, document_title: str = "") -> str:
    where = (
        f"pages {detection.start_page}-{detection.end_page}"
        if detection.end_page != detection.start_page
        else f"page {detection.start_page}"
    )
    kind_note = (
        "This is an exercise / open problem: state it, do not solve it. End with `sorry`."
        if detection.kind == KIND_PROBLEM
        else "This is asserted as mathematics in the book."
    )
    header = f"Source: {document_title or 'book'}, {where}"
    if detection.label:
        header += f", labelled {detection.label!r}"
    return (
        f"{header}\n{kind_note}\n\nPassage:\n<<<BEGIN>>>\n{detection.text}\n<<<END>>>"
    )


def _clean_code(payload: dict) -> str:
    code = payload.get("lean_code") or payload.get("code") or ""
    if not isinstance(code, str):
        return ""
    code = extract_code(code, "lean")
    # models occasionally keep the import lines inside the declaration body
    lines = [line for line in code.splitlines() if not re.match(r"^\s*import\s", line)]
    return "\n".join(lines).strip()


def _imports(payload: dict) -> list[str]:
    raw = payload.get("imports")
    if isinstance(raw, str):
        raw = [raw]
    if not isinstance(raw, list):
        raw = []
    modules = [str(m).strip() for m in raw if str(m).strip()]
    return modules or ["Mathlib"]


def formalize(
    client: Any,
    detection: Detection,
    *,
    checker: Any = None,
    document_title: str = "",
    repair_attempts: int = 1,
) -> Formalization:
    """One detected span -> one Lean declaration, verified as far as it can be."""
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": _user_prompt(detection, document_title)},
    ]
    result = Formalization(detection=detection)
    mode = verification_mode(detection)

    for attempt in range(repair_attempts + 1):
        payload = client.complete_json(messages)
        if not isinstance(payload, dict):
            payload = {}
        result.attempts = attempt + 1

        if payload.get("skip"):
            result.skipped = True
            result.skip_reason = str(payload.get("reason") or "model declined")
            return result

        code = _clean_code(payload)
        if not code:
            result.skipped = True
            result.skip_reason = "model returned no Lean code"
            return result

        result.lean_code = code
        result.imports = _imports(payload)
        result.name = str(payload.get("name") or "").strip()
        result.informal = str(payload.get("informal") or "").strip()
        result.has_proof = bool(payload.get("has_proof")) and not contains_sorry(code)
        try:
            result.confidence = float(payload.get("confidence") or 0.0)
        except (TypeError, ValueError):
            result.confidence = 0.0
        extra = payload.get("index_types")
        result.index_types = index_types(code, extra if isinstance(extra, list) else [])

        if checker is None:
            return result
        result.verify = checker.check(code, mode=mode, imports=result.imports)
        if result.verify.ok or result.verify.skipped or attempt >= repair_attempts:
            return result

        messages = messages + [
            {"role": "assistant", "content": _echo(payload, code)},
            {
                "role": "user",
                "content": REPAIR_PROMPT.format(
                    diagnostics=(result.verify.diagnostics or "")[-3000:]
                ),
            },
        ]
    return result


def _echo(payload: dict, code: str) -> str:
    import json

    return json.dumps({**payload, "lean_code": code}, ensure_ascii=False)[:6000]


def formalize_all(
    client: Any,
    detections: list[Detection],
    *,
    checker: Any = None,
    document_title: str = "",
    repair_attempts: int = 1,
    progress: Callable[[str], None] | None = None,
    on_error: str = "skip",
) -> list[Formalization]:
    results: list[Formalization] = []
    for position, detection in enumerate(detections, 1):
        try:
            formalization = formalize(
                client,
                detection,
                checker=checker,
                document_title=document_title,
                repair_attempts=repair_attempts,
            )
        except Exception as error:  # noqa: BLE001 - one bad item must not kill a book
            if on_error == "raise":
                raise
            formalization = Formalization(
                detection=detection, skipped=True, skip_reason=f"formalizer failed: {error}"
            )
        results.append(formalization)
        if progress:
            status = (
                "skipped"
                if formalization.skipped
                else "proved"
                if formalization.proved
                else "verified"
                if formalization.verified
                else "unverified"
            )
            progress(
                f"  [{position}/{len(detections)}] {detection.label or detection.kind} "
                f"@ {detection.locator}: {status}"
            )
    return results


def as_dicts(formalizations: list[Formalization]) -> list[dict]:
    return [f.as_dict() for f in formalizations]


__all__ = [
    "Formalization",
    "formalize",
    "formalize_all",
    "verification_mode",
    "as_dicts",
    "asdict",
]