lean-bench/leanbench/srcsimplex.py

"""The source-side simplex: what the *prose* gives, assumes and claims.

One model call reads only the book passage -- never the Lean -- and lists its roles.
Keeping the Lean out of view is the point: the schema must come from the source, so
that comparing it with the Lean-side schema is a check and not an echo.
"""
from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Any

from .llm import ChatClient, parse_json

PROMPT_VERSION = "srcsimplex-v1"

INTRODUCED = ("given", "forall", "exists")

SYSTEM_PROMPT = """\
You read ONE passage from a mathematics book and describe its logical parameters.
Do not write Lean. Do not solve anything. Describe only what the passage itself says.

Return a JSON object:
{
  "roles": [
    {"kind": "object", "symbols": ["G"], "description": "finite group",
     "keywords": ["finite", "group"], "introduced": "given"},
    {"kind": "fact", "symbols": ["n"], "description": "n is positive",
     "keywords": ["positive"], "introduced": "given"},
    {"kind": "goal", "symbols": [], "description": "every element's order divides |G|",
     "keywords": ["order", "divides"], "introduced": "given"}
  ],
  "goal_has_exists": false
}

Rules:
- "object": a mathematical object the passage is about (a group G, an integer n, a
  function f). One entry per object, even if several share a description.
- "fact": a stated assumption or condition on those objects ("n > 0", "G is
  finite", "a and b are coprime"). Include conditions hidden in prose ("let G be a
  finite group" gives an object G AND a fact that G is finite).
- "goal": exactly one entry: the claim being asserted, the property being defined, or
  for an exercise the thing to be shown. For a definition, the goal is what is defined.
- "symbols": the letters/names used in the text, as written. Empty if none.
- "keywords": 1-5 lowercase words naming the concepts involved (e.g. "integer",
  "prime", "subgroup", "homomorphism", "finite"). Words from the text, not synonyms.
- "introduced" (objects only): "given" if the passage fixes it ("let n be..."),
  "forall" if it ranges over all of them ("for every n"), "exists" if the passage asserts
  one exists ("there is an n such that"). Use "given" when unsure.
- "goal_has_exists": true only if the claim itself asserts that something exists.
- At least two roles in total, and exactly one goal.
Reply with the JSON object only."""


@dataclass
class SourceRole:
    kind: str
    symbols: list[str] = field(default_factory=list)
    description: str = ""
    keywords: list[str] = field(default_factory=list)
    introduced: str = "given"
    alias: str = ""

    def as_dict(self) -> dict[str, Any]:
        return {
            "alias": self.alias, "kind": self.kind, "symbols": self.symbols,
            "description": self.description, "keywords": self.keywords,
            "introduced": self.introduced,
        }


@dataclass
class SourceSimplex:
    roles: list[SourceRole]
    goal_has_exists: bool = False
    notes: list[str] = field(default_factory=list)

    def by_kind(self, kind: str) -> list[SourceRole]:
        return [r for r in self.roles if r.kind == kind]

    @property
    def goal(self) -> SourceRole | None:
        goals = self.by_kind("goal")
        return goals[0] if goals else None

    def as_dict(self) -> dict[str, Any]:
        return {
            "roles": [r.as_dict() for r in self.roles],
            "goal_has_exists": self.goal_has_exists,
            "notes": self.notes,
        }


def _strings(value: Any) -> list[str]:
    if isinstance(value, str):
        value = [value]
    if not isinstance(value, list):
        return []
    return [str(v).strip() for v in value if str(v).strip()]


def parse_source_simplex(payload: Any) -> SourceSimplex:
    """Validate a model reply; raises ``ValueError`` when it is not a usable schema."""
    if not isinstance(payload, dict) or not isinstance(payload.get("roles"), list):
        raise ValueError("reply has no 'roles' list")
    roles: list[SourceRole] = []
    counters: dict[str, int] = {}
    notes: list[str] = []
    for item in payload["roles"]:
        if not isinstance(item, dict):
            continue
        kind = str(item.get("kind", "")).strip().lower()
        if kind not in ("object", "fact", "goal"):
            notes.append(f"dropped role with kind {kind!r}")
            continue
        introduced = str(item.get("introduced", "given")).strip().lower()
        counters[kind] = counters.get(kind, 0) + 1
        roles.append(SourceRole(
            kind=kind,
            symbols=_strings(item.get("symbols")),
            description=str(item.get("description", "")).strip(),
            keywords=[k.lower() for k in _strings(item.get("keywords"))],
            introduced=introduced if introduced in INTRODUCED else "given",
            alias=f"{kind}{counters[kind]}",
        ))
    goals = [r for r in roles if r.kind == "goal"]
    if len(goals) != 1:
        raise ValueError(f"expected exactly one goal role, got {len(goals)}")
    if len(roles) < 2:
        raise ValueError("a simplex needs at least two roles")
    return SourceSimplex(roles, bool(payload.get("goal_has_exists")), notes)


def passage_prompt(text: str, label: str = "") -> list[dict[str, str]]:
    header = f"Label: {label}\n\n" if label else ""
    body = re.sub(r"[ \t]+", " ", text).strip()
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"{header}Passage:\n<<<BEGIN>>>\n{body}\n<<<END>>>"},
    ]


def extract_source_simplex(
    client: ChatClient, text: str, label: str = "", *, attempts: int = 2
) -> SourceSimplex:
    """Ask the model for the passage's roles; retry once if the reply is unusable."""
    messages = passage_prompt(text, label)
    last: ValueError | None = None
    for _ in range(attempts):
        payload = client.complete_json(messages)  # type: ignore[attr-defined]
        try:
            return parse_source_simplex(payload)
        except ValueError as error:
            last = error
            messages = messages + [
                {"role": "assistant", "content": str(payload)[:3000]},
                {"role": "user", "content": f"Not usable: {error}. Reply with the JSON only."},
            ]
    raise ValueError(f"no usable source simplex: {last}")


__all__ = [
    "PROMPT_VERSION", "SYSTEM_PROMPT", "SourceRole", "SourceSimplex",
    "extract_source_simplex", "parse_json", "parse_source_simplex", "passage_prompt",
]