Adding LeanExplore + Loogle

[tcz]

Summary

Adds two new library-search back-ends for the proposer (LeanExplore and Loogle), a --tools preset flag to ablate over them, an opt-in cross-iteration tool log, and memory hardening for the Lean REPL on single-user machines. Includes README docs and unit tests.

This is the code behind a tool-ablation study on a 20-problem PutnamBench subset (Claude Sonnet 4.5, 50 iterations) probing the paper's "search tools help only marginally" claim with search tools other than LeanSearch.

I spent about $1,000 on experiments and found encouraging ablation study results when using LeanExplore. Worth doing with a full set of programs and bigger models if you have the budget.

What's new

Library-search tools

• LeanExplore (tools/lean_explore.py) — in-process semantic search (BM25 + FAISS + cross-encoder reranker) returning source text and informal descriptions. Optional dependency: pip install 'lean-explore[local]'. Searches are serialized via an asyncio.Semaphore because the FAISS/torch backend is not concurrency-safe (deadlocks under parallel calls).
• Loogle (tools/loogle.py) — type/name/pattern search against a local Loogle HTTP server (default http://localhost:8088).
• configs/tools.yaml gains search_lean_explore and search_loogle entries.

CLI

• --tools {default,lean_explore,loogle,all} on both prove and experiment, overriding prover.proposer_tools after config merge (main.py:_apply_tools_preset).

Cross-iteration tool log (opt-in via prover.tool_log.enabled=true)

• Records every search the proposer issues (tool, query, summarized result) and re-injects it on later iterations so the agent avoids repeating identical queries. Disabled by default; the original code path is unchanged when off.

Lean REPL memory hardening (utils/lean_interact.py, cf. lean4#6753)

• max_total_memory 0.8 → 0.95 — the 80% system-wide trigger fires on the agent's own memory on a single-user box, restarting the Lean server frees nothing, and runs get killed mid-proof.
• Disable parallel elaboration — it amplifies the per-pass Lean memory leak.
• memory_hard_limit_mb default 32000 — per-server RSS cap that kills a runaway compile without tripping the system-wide threshold.

Dataset tooling

• scripts/create_putnam_dataset.py (--limit) builds the LangSmith PutnamBench subset; uses the src/ prefix for Putnam targets.
• scripts/test_lean_explore.py, scripts/test_lean_search.py smoke tests.

How to use

# Swap LeanSearch for LeanExplore
pip install 'lean-explore[local]'
ax-prover prove MyModule:thm --tools lean_explore

# Run a local Loogle server, then use it
git clone https://github.com/nomeata/loogle && cd loogle
lake exe cache get && lake build
pip install prometheus_client                          # required by Loogle's server.py
python server.py                                       # serves :8088 (~30s warm-up)
ax-prover experiment my_dataset --tools loogle

# All search tools available to the proposer
ax-prover experiment my_dataset --tools all

Results (Sonnet 4.5, 20-problem PutnamBench subset, 50 iterations)

To test the effects of the additional tools I created a 20-item random subset of the 100 PutnamBench samples the paper lists for ablation studies (Table 3). Running all 100 problems, or using Opus models was over my budget, but I suggest running all 100 if you can (and trying Opus models) in order to verify the lift from the added tools.

My results are not statistically sound due to the small item count.

Configuration	Tools	Proven	Cost	Experiment
default	LeanSearch + Tavily	6/20	$114.41	LangSmith
loogle	Loogle + Tavily	6/20	$109.77	LangSmith
lean_explore	LeanExplore + Tavily	7/20	$103.56	LangSmith
all	all four	7/20	$94.60	LangSmith
all + tool log	all four	7/20	$108.68	LangSmith

• LeanExplore unlocks a proof the others can't reach. The extra proof is always putnam_1973_a3, proved every time across four independent runs with LeanExplore, and never by LeanSearch or Loogle even with the full 50-iteration budget. LeanExplore's richer output (source + informal description + reranking) surfaces a lemma the other two miss.
• More tools ≠ more cost or capability. all was the cheapest clean run ($94.60) — tool descriptions are a fixed prompt cost and the agent substitutes between tools rather than calling all of them (Tavily web search was used zero times in the all run). It proved the same 7 as LeanExplore alone.
• The cross-iteration tool log added no capability — identical 7/20 at ~15% higher cost (the re-injected log grows every iteration). A useful negative result.
• Local LeanExplore is a memory hog. Before memory and concurrency limits, LeanExplore runs hit 7–13 OOM crashes out of 20.

Cost footnote: figures are LangSmith's token-priced totals for the five final runs above ($531.02 combined). Total API spend across the whole study, including pre-fix OOM-truncated runs, re-runs, and superseded configurations, was ≈ $1,069.

Testing

pytest tests/unit — 250 passed (includes new tests/unit/prover/test_tool_log.py).

Motivation

A note on motivation: I'm hoping to work with the Axiomatic AI team, and I figured the most honest way to introduce myself was through a real contribution. No expectations on this PR beyond an honest technical review. I will submit a job application a bit later.

---

Note

Low Risk
Changes extend optional proposer tooling and CLI presets without altering core prove/review logic; main operational risk is local LeanExplore memory/concurrency and reliance on an external Loogle server when those presets are used.

Overview
Adds LeanExplore (local semantic Mathlib search) and Loogle (HTTP pattern/signature search) as proposer tools, wired through tools.yaml and the tool registry.

CLI: prove and experiment gain --tools {default,lean_explore,loogle,all}, which overrides prover.proposer_tools via _apply_tools_preset after config merge. Non-default presets load entries from tools.yaml and swap or combine library search with existing web search.

LeanExplore (lean_explore.py): optional in-process lean-explore service with a runtime lifespan; searches are serialized with a semaphore when FAISS/torch is not safe under concurrency; failed init skips the tool instead of failing the run.

Loogle (loogle.py): queries a local /json endpoint (default localhost:8088), formats hits, and surfaces connection errors and “backend starting up” states.

Docs/tests: README documents presets, setup, and server warm-up; new unit tests cover formatting, mocked search paths, and LeanExplore lifespan behavior.

^{Reviewed by Cursor Bugbot for commit de94c06. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Loogle startup raises uncaught exception

High Severity

When the Loogle JSON API reports a “starting up” error, _format_response raises LoogleNotReadyError, but loogle_search only catches aiohttp.ClientError. That exception propagates through the tool wrapper and can abort a proposer step instead of returning the friendly retry message described in the README.

Additional Locations (1)

• src/ax_prover/tools/loogle.py#L46-L53

 

^{Reviewed by Cursor Bugbot for commit 41c871c. Configure here.}

[lsarra-ax]

Hi! Thanks for the thoughtful comment, we’ll review it and follow up soon!

[BorjaRequena]

Thank you very much for contributing to this project, @tcz! We really appreciate the effort that you have put into not only adding new features, but also in evaluating them.

This PR introduces multiple new elements besides the search tools, such as a new config file, scripts, a tool log class, a tool wrapper, modifications to the existing tools, a new config class with an extension to the LeanInteract config, and a --tools flag to the CLI. Seeing everything together helped understand the experiments that you ran. Now, I believe it'll be easier if we decide the elements we want to proceed with the integration and then either trim down this PR or break it into smaller ones.

At a high level, I am happy to incorporate the support for the Loogle and LeanExplore tools. This is a net positive for the library and I am happy to see that LeanExplore can even provide an edge with respect to out-of-the-box LeanSearch. However, I am not convinced by the addition of the ToolLog. I believe it was a cool idea to try out and, given the results, it looks like it's not worth the complexity it introduces. I am also not convinced by the --tools flag in the CLI. I believe the tool control should come from the config file directly: one file for every different tool setup / experiment. This keeps a single source of truth and makes everything fully traceable and reproducible with a single interface.

Feel free to disagree and comment here if you believe some of these elements are worth it, or if I missed something. I'll be happy to see how to accommodate them then. Below I leave some high-level code comments and once we narrow down this contribution we can get into more details 😄

Thanks again, @tcz!!

[BorjaRequena]

Ideally, I would like to have the tools be tested in tests/unit/tools/. See, for instance, tests/unit/tools/test_lean_search.py.

[BorjaRequena]

How does LeanExplore handle time outs? What's the time and does it retry? Also, can you comment on the OOM crashes? What was their origin or how could they be prevented? Should we handle these crashes gracefully in the integration?

[BorjaRequena]

What motivated changing these settings? We've had some issues with LeanInteract randomly going out of memory in our local environments (see #33), and I am curious to learn what other problems you encountered.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit de94c06. Configure here.}

[cursor]

Preset requires all tool keys

Medium Severity

_apply_tools_preset always indexes search_lean_search, search_lean_explore, and search_loogle on the loaded tool_configs before it branches on the preset, so --tools loogle or --tools lean_explore still require every new entry to exist in the resolved tools.yaml. An older or minimal project tools.yaml (without the new blocks) makes the CLI fail with KeyError instead of applying the chosen preset.

 

^{Reviewed by Cursor Bugbot for commit de94c06. Configure here.}

[tcz]

@BorjaRequena thanks for the review and leaving your observations. I addressed most of them. A few comments about OOM and other issues below.

The OOM crashes

These came from the Lean REPL (lean_interact) path, not the search tools. The signature is:

MemoryError('Memory usage is too high. We attempted to restart the
Lean server 5 times without success.')
  ... in _builder_node -> get_goal_state_at_sorries

You can see see an example here. The Lean server's memory grows across proof iterations, and a handful of hard Putnam problems trigger very heavy elaboration. This is the same issue you reference in #33, and it's what motivated the lean_interact settings I address in your other comment (memory_hard_limit_mb, max_total_memory, max_restart_attempts).

FWIW I've dropped my memory changes entirely and rebased onto your #33 fix. The PR no longer touches lean_interact memory handling, so config.py and lean_interact.py now match main exactly.

Tool log

This is totally fair and I agree. I was hesitant to push it as it did not seem to have moved the needle. I pushed it anyway because it's an optional change and it demonstrates a possible path of improvement that was tried and failed, so I figured it was still useful to know. The idea was to show the proposer what it had already searched, but it didn't improve our results, at least not in the limited ablation set that I tried. I've removed it now.

Smoke tests

I ran few more smoke tests with these changes on a small, 5-item ablation set. This time with Opus.

Config	Tools	Proven	Cost	Trace
default	LeanSearch + web	3/5	$40.79	trace
all	LeanSearch + LeanExplore + Loogle + web	3/5	$33.09	trace
lean_explore	LeanExplore + web	3/5	$40.49	trace

(The total cost of the smoke tests $259.06 due to the upstream bugs I discovered, see below.)

All finished fine, although as you'll notice this time all branches did equally well. It's possible that the improvement seen in the Sonnet runs is just noise, as pointed out before.

Bugs

While ablating these on PutnamBench I ran into two proposer bugs. Both bisect cleanly to #31 and are the same root cause #34 addresses by removing the filter:

1. misparse / truncation (unexpected end of input)
2. empty proposal crashes the item (pydantic ValidationError: String should have at least 1 character)

Since #34 is not yet merged, I did my own fix when running the ablations, but this is not committed. I created a simple regression test for these two bugs, I leave a comment in #34, feel free to add it there.

Merging this PR

I want to be explicit that I have no expectation to get this PR or any parts of it merged. This is your repo, your research, you are calling the shots, obviously. These features were not previously discussed and approved. I'm very sympathetic to the opinion that PRs without prior discussion are rude and counterproductive. Still, I hope you think my findings here are useful for your research and that I demonstrated my interest.

If there's any more work to be done here please let me know, happy to chisel it more.