Skip to content

feat: local Lean project search + CSLib search tools#13

Draft
KrystianNowakowski wants to merge 40 commits into
mainfrom
local-lean-search-tool
Draft

feat: local Lean project search + CSLib search tools#13
KrystianNowakowski wants to merge 40 commits into
mainfrom
local-lean-search-tool

Conversation

@KrystianNowakowski

@KrystianNowakowski KrystianNowakowski commented Jun 3, 2026

Copy link
Copy Markdown
Collaborator

What

Adds two project-local search tools the remote Mathlib LeanSearch doesn't cover, plus the declaration-parsing robustness they need:

  • search_lean_local — searches the project's own .lean files for declarations by name (namespace-qualified, modifier/attribute-aware), with a body-search fallback for identifiers that aren't declaration names (e.g. where/let rec helpers, constructors).
  • search_cslib — searches the vendored CSLib dependency under .lake/packages/cslib, which local search excludes and Mathlib search doesn't index.
  • Shared lean_parsing.py improvements: modifier/attribute prefixes (private, @[simp], …), namespace-qualified names, and per-simple-name occurrence disambiguation.
  • Proposer-prompt routing: pick the search tool by the file's imports — Mathlib.*search_lean_search_tool, Cslib.*search_cslib_tool, other/local imports → search_lean_local_tool.

Code-review fixes included

This branch also carries the search/parsing fixes from the experiment-branch review:

  • DECL_NAME_END accepts universe-polymorphic names (theorem foo.{u}).
  • extract_function_from_content keeps a leading open … in (and other … in) command prefix instead of dropping it.
  • Comment-stripping vs raw-content occurrence misalignment fixed (comment_spans/non_comment_matches): a declaration keyword inside a comment no longer shifts occurrence indices.
  • _format_results always shows at least the first match (truncated with a marker) instead of returning zero bodies when the only match exceeds max_chars.
  • Local/CSLib search reads each .lean file once per query (single walk shared with the body fallback) instead of twice on a name-miss.
  • TemporaryProposal resolves the target's namespace-qualified occurrence (resolve_target_occurrence) instead of always extracting occurrence 0.

Tests

Full unit suite passes on the branch; ruff clean. Includes dedicated tests for both tools (isolation, namespaces, duplicate names, modifiers, body fallback, edge cases) and for each fix above.

🤖 Generated with Claude Code


Note

Medium Risk
Changes the default prover tool set and proposer context (verbatim definition injection) plus shared Lean parsing used when applying proposals; broad surface area but covered by extensive unit tests and bounded caching behavior.

Overview
Adds local Lean search (search_lean_local) over the project’s own .lean files and CSLib search (search_cslib) under .lake/packages/cslib, both built on a shared _search_root pipeline: exact name match, then body-identifier fallback (e.g. where helpers), then fuzzy name suggestions when still empty.

Prover integration: default config registers both tools; create_tool can pass base_folder into scoped searchers. The proposer is told to pick search by imports (Mathlib vs CSLib vs local). WS1 caching records what local search returned, intersects with identifiers used in the latest proposal, stores append-only used_definitions on graph state, and injects a bounded <local-definitions> block so the agent need not re-query the same defs each iteration.

Parsing/build fixes supporting search and apply: declaration prefixes (private, @[simp], …), qualified names and occurrence disambiguation, comment-aware matching, open … in prefixes in extracted blocks, universe .{u} name boundaries, single file read per query, and resolve_target_occurrence when applying proposals.

Also adds Claude Opus 4.7/4.8 LLM configs and ablation YAMLs, pins langchain-anthropic via a git rev in pyproject.toml, plus design/implementation docs for the search workstreams.

Reviewed by Cursor Bugbot for commit b9e7989. Bugbot is set up for automated code reviews on this repo. Configure here.

KrystianNowakowski and others added 14 commits June 3, 2026 16:17
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Implements LocalLeanSearcher with cached root resolution (walk-up then walk-down), full-file scan excluding .lake/, match extraction via extract_function_from_content, and result formatting with max_results/max_chars caps and overflow-by-name listing. All edge cases covered: empty query, no lakefile, multiple projects, dot-lake exclusion.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add Iterator annotation, _compute_root docstring, consolidate test
imports, add max_chars-budget overflow test, clarify walk-down test.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Also removes the temporary module-level ruff noqa now that the
registration factory consumes the remaining imports.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Update register_tool docstring for optional base_folder param, clarify
_declaration_line comment, and test alternate lakefile markers.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Logs "LocalLeanSearch: Found N matches for 'X'" / "No results" at INFO
and a debug invocation line, mirroring the lean_search tool so tool
usage is visible in the prover's terminal output.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
A multi-word query now matches declaration names containing every
whitespace-separated token (e.g. 'Treap insert' -> Treap.insert),
instead of requiring the literal multi-word substring (which never
matched a name). Single-word queries are unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add the search_lean_local tool_config to the bundled tools.yaml and
include it in the default proposer_tools, so the local library search
is on by default rather than requiring an opt-in config overlay.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The tool is now defined in the bundled configs/tools.yaml; the
top-level experiment config no longer needs its own copy.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Comment thread scripts/resume_ai4math_experiment.py Outdated
Comment thread src/ax_prover/tools/local_lean_search.py Outdated
`\b` after a declaration name treated `.` as a boundary, so locating
"Treap" could match an earlier "Treap.insert" and return the wrong
block/line. Replace the trailing `\b` in extract_function_from_content
and _declaration_line with a shared DECL_NAME_END lookahead that
requires the name token to actually end (whitespace or a delimiter).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Comment thread src/ax_prover/tools/local_lean_search.py
KrystianNowakowski and others added 4 commits June 4, 2026 11:26
Repos that vendor the same file under multiple directories (e.g.
Challenges/ and Challenges_Remainder/) made search_lean_local return the
same declaration several times, wasting context. Group matches by
(name, block): identical declarations collapse to one entry listing all
locations ("-- path:line (also: …)"), while same-named declarations with
different bodies stay separate. Result counts and caps now count unique
declarations (CoverageIntervalDefs: 6 -> 3).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
CI uses ruff 0.15.15, whose formatter joins this implicit string
concatenation onto one line. Match it so `ruff format --check` passes.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat: claude 4.7

* chore: update langchain-anthropic

* fix: use ax leansearch

* feat: claude 4.8

* fix: install langchain-anthropic from main
…es (#16)

* Local search: handle modifiers/attributes and namespace-qualified names

The local search tool missed many declarations because the line-based parser
keyed off the first whitespace-delimited word. For `noncomputable def foo`,
`@[simp] def foo`, `private def foo`, etc. that word wasn't a keyword, so the
declaration was dropped. All four "missing" Tier-1 identifiers (dijkstra_rec,
relax_neighbors, isAncestor, weightSum) are `noncomputable def`. It also parsed
`def insert` inside `namespace BinaryTree` as plain `insert`, so the query
"BinaryTree insert" couldn't match.

Shared parser (lean_parsing.py): match an optional inline attribute/modifier
prefix before the keyword, longest-first so `noncomputable def` beats `def`.
Applied consistently to the lister, block extractor (start + end), and
find_declaration_at_line. This also fixes a latent bug where a `noncomputable
def` target was invisible to the builder (false MissingTargetTheorem).

Local search tool (local_lean_search.py): track namespace/section scope and
match queries against the fully-qualified name, while keeping the simple source
name for block extraction and line lookup.

Verified end-to-end against the AI4Math challenges: dijkstra_rec,
relax_neighbors, weightSum, isAncestor, and "BinaryTree insert" now all resolve.

Deferred: `where` / `let rec` nested helpers (e.g. query_aux).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Fix wrong block/line for duplicate simple names across namespaces

Code review (cursor bot): when one file declared the same simple name in two
namespaces (e.g. `def insert` in `namespace A` and `namespace B`), extraction
and line lookup used re.search (first match), so the later namespace's entry
showed the first declaration's body and line.

_matching_declaration_names now returns a 0-based occurrence index per simple
name (counted across all declarations, aligning with re.finditer order).
extract_function_from_content and _declaration_line take an `occurrence` arg
(default 0, backward compatible) and select the N-th match.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Comment thread src/ax_prover/utils/lean_parsing.py
Comment thread src/ax_prover/tools/local_lean_search.py
Comment thread src/ax_prover/tools/cslib_search.py
KrystianNowakowski and others added 2 commits June 5, 2026 14:51
Tell the proposer to pick the search tool by which library the file
imports: Mathlib.* -> search_lean_search_tool, Cslib.* ->
search_cslib_tool, other/local imports -> search_lean_local_tool.
Replaces the generic "Reference Mathlib lemmas" nudge, which gave no
guidance for the CSLib and local-project search tools now available.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Comment thread src/ax_prover/tools/local_lean_search.py
KrystianNowakowski and others added 6 commits June 5, 2026 15:45
The name-boundary assertion excluded `.`, so a universe-polymorphic
declaration like `theorem foo.{u} ...` was unparseable and
extract_function_from_content returned None (breaking TemporaryProposal
for such targets). Allow a literal `.{` universe binder to follow the
name while still rejecting a qualified-name continuation like `foo.bar`.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
A declaration can be preceded by contiguous `in`-terminated commands that
bind to it (e.g. `open Nat in`, `set_option ... in`). extract_function_
from_content dropped them, so the applied/compiled block lost the binding
-> "unknown identifier". Walk backward over contiguous preceding lines
that end in a standalone `in` keyword and include them, stopping at a
blank line or any non-prefix line so ordinary decls/comments are never
absorbed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
list_all_declarations_in_lean_code enumerates declarations over
comment-stripped text, but extract_function_from_content and
_declaration_line matched re.finditer over RAW content. A declaration
keyword at line-start inside a comment shifted the occurrence index, so
local search returned the commented-out block and wrong line.

Add comment_spans()/non_comment_matches() helpers (adapting the
strip_comments state machine to record comment ranges) and use them in
both matchers so occurrence counting agrees everywhere, while the
returned block is still sliced from raw content.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
When a single matched declaration's block was larger than config.max_chars,
the budget gate pushed it to overflow, leaving `shown` empty: the user got a
header plus an "Additional matches (not shown)" note but no source at all.

Now the first/most-relevant match is always shown, truncating its block to the
remaining budget with a clear truncation marker when it would not fit.
Genuine overflow of later entries is unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
_search_root previously called _collect_matches twice (name pass, then body
fallback), each independently walking _iter_lean_files and read_text-ing every
.lean file. On a name-miss over a large tree this doubled disk I/O and parsing.

Walk the tree once into a cached list of (relative_path, content); both the name
pass and the (only-if-needed) body fallback run against the cached contents.
Behavior is preserved: body-fallback results are used only when the name pass
found zero matches across all files, and log messages / formatting are identical.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
TemporaryProposal.__enter__ extracted the proposed declaration at
occurrence 0, so when proposal.code contained multiple top-level decls
sharing the target's simple name (e.g. A.insert and B.insert), the wrong
block was applied and committed. It also could not resolve a
namespace-qualified location name (e.g. B.insert) against source written
with the simple name (def insert inside namespace B).

Add resolve_target_occurrence() in lean_parsing.py (namespace-aware,
reusing list_all_declarations_in_lean_code) and use it in build.py to
select the correct (simple_name, occurrence) before extraction, falling
back to occurrence 0 when resolution is ambiguous or finds nothing so the
common single-declaration case is unchanged.

The resolution helper lives in lean_parsing (already imported by both
utils.build and tools.local_lean_search) to avoid a utils -> tools
circular/layering dependency.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Comment thread src/ax_prover/utils/build.py
@KrystianNowakowski KrystianNowakowski changed the title Local lean search tool feat: local Lean project search + CSLib search tools Jun 5, 2026
KrystianNowakowski and others added 2 commits June 8, 2026 11:23
* docs: spec for local-search used-definition caching + fuzzy fallback

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* docs: WS1 + WS2 implementation plans for local-search caching & fuzzy fallback

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* refactor: _search_root returns (text, decls); callers unpack

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* feat: pure helpers for run-scoped used-definition cache

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test: cover None target and empty pool in accumulate_used_definitions

* feat: LocalLeanSearcher records returned declarations for the run

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test: cover first-seen-wins on repeat local search

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* feat: add used_definitions field to ProverAgentState

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* feat: local-definitions prompt template and proposer nudge

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix: make local-definitions nudge conditional on the section being present

* feat: wire used-definition cache through memory and proposer nodes

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test: behavioral check that memory node preserves experience and adds used_definitions

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat: fuzzy scoring primitives for local search

* feat: fuzzy declaration matchers; blend whole+token scores for tie-breaking ranking

* fix: average token scores so single-token matches keep recall while ties still rank

* feat: fuzzy header rendering in _format_results

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* refactor: collapse _format_results header branch; assert exact fuzzy header

* feat: fuzzy fallback tier in local search (exact -> body -> fuzzy)

* test: end-to-end fuzzy fallback through LocalLeanSearcher.search

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* docs: align WS2 spec with shipped tier order (exact->body->fuzzy) and header text

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes and found 2 potential issues.

Fix All in Cursor

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

Reviewed by Cursor Bugbot for commit a33b203. Configure here.

Comment thread src/ax_prover/tools/local_lean_search.py
Comment thread src/ax_prover/tools/local_lean_search.py
- #1: char cap skips (continue) an over-budget definition instead of break, so
  smaller later definitions are still cached
- #2: ambiguous bare names (min/map/insert/...) cached only on a qualified
  reference, avoiding false positives when the proof used a Mathlib/core identifier
- #4: injected <local-definitions> block is bounded to a fraction of the model
  input budget, dropping whole trailing entries (verbatim source never chopped)
- #7: warn when the local-search tool is present but its searcher can't be
  recovered (used-definition caching silently disabled)
- #8: normalize the fuzzy query once per search instead of per declaration

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@lsarra-ax
lsarra-ax marked this pull request as draft June 8, 2026 12:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants