ai4math

ai4math

Tools for AI-assisted mathematics. The headline feature: ask for a proof, get one a theorem prover has actually checked.

A language model asked for a hard proof will produce something fluent and often wrong, and it will sound equally confident either way. So don't take its word for it. ai4math sends the proof to Lean, and when Lean rejects it, sends the compiler's errors back to the model to try again — up to three attempts by default.

The verdict comes from Lean's kernel, not from the model. A proof that does not compile is not a proof, and a proof that leans on sorry is rejected too — that one compiles cleanly while proving nothing.

The demo

pip install ai4math

export ANTHROPIC_AUTH_TOKEN=...   # Claude Code must be installed and authenticated
python -m ai4math.demo --example

Real output from that command (the natural-language section abridged; the Lean proof and the verdict are verbatim):

Problem: Let n be a positive integer. Prove that n^3 - n is divisible by 6 for every integer n.

Asking sonnet for a solution and a Lean proof (up to 3 attempt(s))...

[attempt 1] Lean accepted the proof.

=== Solution (natural language) ===
Factor: n^3 - n = n(n-1)(n+1), the product of three consecutive integers.
- Divisibility by 2: among n-1 and n, one is even, so the product is even.
- Divisibility by 3: among any three consecutive integers, exactly one is a
  multiple of 3, since residues mod 3 cycle through 0, 1, 2.
Since 2 and 3 are coprime and both divide n(n-1)(n+1), their product 6 divides n^3-n.

=== Lean proof (verified) ===
import Mathlib

theorem n_cube_sub_n_dvd_six (n : ℤ) : (6 : ℤ) ∣ n ^ 3 - n := by
  have key : ∀ x : ZMod 6, x ^ 3 - x = 0 := by decide
  have h : ((n ^ 3 - n : ℤ) : ZMod 6) = 0 := by
    push_cast
    exact key (n : ZMod 6)
  exact (ZMod.intCast_zmod_eq_zero_iff_dvd _ _).mp h

=== Result: proved and verified by Lean after 1 attempt ===

Your own problem, and a look at each attempt:

python -m ai4math.demo --show-proof \
  "Let a, b, c be positive reals with a + b + c = 3. Prove that ab + bc + ca <= 3."

Exit status is 0 for a verified proof, 1 if none was found in the attempt budget, 2 for a setup problem such as a missing toolchain.

When the first attempt fails

This is the part that earns its keep. Lean's diagnostics are specific enough to act on, so they go straight back to the model. A real run, proving Gauss's summation formula by induction:

$ python -m ai4math.demo "Prove that for every natural number n, the sum 1 + 2 + ... + n equals n*(n+1)/2."

[attempt 1] rejected: 1 Lean error(s).
    | 11:4: error: omega could not prove the goal:
    | a possible counterexample may satisfy the constraints
    |   -1 ≤ 2*b - d ≤ 0
    |   a - b + c ≤ -2
    | where
    |  a := ↑(k * (k + 1)) / 2
    |  b := ↑((k + 1) * (k + 1 + 1)) / 2
[attempt 2] Lean accepted the proof.

=== Result: proved and verified by Lean after 2 attempts ===

The first attempt reached for omega on a goal mixing natural-number division with multiplication, which it cannot discharge. Told exactly that — with the counterexample constraints and the offending terms — the model established the divisibility facts first and then closed the goal.

That is the difference between this and sampling repeatedly: the goal state travels with the error, so the second attempt is informed rather than another guess. It is the loop behind Draft-Sketch-Prove (Jiang et al. 2023) and Baldur (First et al. 2023).

Setup

The demo needs two external tools. Neither is a Python dependency.

Claude Code, authenticated:

export ANTHROPIC_AUTH_TOKEN=...     # or ANTHROPIC_API_KEY

Lean 4, via elan:

curl -sSfL https://elan.lean-lang.org/elan-init.sh | sh -s -- -y
export PATH="$HOME/.elan/bin:$PATH"

Anything past core Lean needs Mathlib, which lives in a Lake project:

lake +leanprover/lean4:v4.19.0 init mathdemo math
cd mathdemo && lake update && lake exe cache get     # downloads a prebuilt cache
export AI4MATH_LEAN_PROJECT="$PWD"                   # or pass --project

Without it you still get a working demo, limited to what core Lean can prove.

From Python

from ai4math.formal import prove_with_feedback, LeanCliBackend
from ai4math.models import ClaudeCLI

outcome = prove_with_feedback(
    "Prove that the square of an odd integer is odd.",
    ClaudeCLI(model="sonnet"),
    LeanCliBackend(project_dir="mathdemo"),
    max_attempts=3,
)

if outcome:                      # truthy only if Lean accepted a proof
    print(outcome.proof)
print(outcome.summary())         # 'proved and verified by Lean after 2 attempts'

ClaudeCLI is only a convenience. Any Callable[[str], str] works — an SDK client, a local vLLM server, a stub in a test — so the library needs no LLM dependency:

prove_with_feedback(problem, lambda prompt: my_model.generate(prompt), backend)

Every attempt is kept in outcome.rounds, each with the response, the extracted Lean, and the compiler's verdict.

The rest of the library

Verified proof search is the demo, not the whole package. ai4math also covers the informal side — the plumbing that every AI-for-mathematics project rewrites, where the answer is a string and grading it is deceptively hard (0.5, 1/2 and \frac{1}{2} are the same number; (0,1) as an interval is not (0,1) as a point):

from ai4math import extract_answer, verify

answer = extract_answer(r"...so the answer is \boxed{\frac{5}{6}}.")
bool(verify(r"\frac{5}{6}", answer))            # True
bool(verify(r"\frac{5}{6}", "0.8333333333"))    # True  — numeric tolerance
bool(verify("[0,1]", "(0,1)"))                  # False — different objects

Plus pass@k and majority-vote metrics, a guarded SymPy tool surface for tool-using models, prompt templates paired with the parsers that match them, and dataset loaders. The core needs only SymPy.

→ Full guide for all of it, with the reasoning behind the grading rules.

Development

git clone https://github.com/wangyu9/ai4math
cd ai4math
pip install -e ".[dev]"
pytest              # tests requiring Lean skip cleanly without it
ruff check src tests
mypy src

License

Apache-2.0. See LICENSE.