lean-bench/textbook/search_database.py

"""Rank premises from the database against the current proof goals.

This is the retrieval step of the intended lean-agent loop: rather than asking the
LLM to prove a target in one shot, the agent keeps a list of open goals (the proof
state, as vscode-lean4 displays it), and at each step searches the premise database
for entries that might apply to the goal in front of it. The ranked results are what
gets shown to the LLM alongside the goal, so the next line of Lean it writes can
cite a real premise by name instead of inventing a lemma.

It reads the database JSONL directly and derives each premise's declaration name,
identifiers and proof status on the fly, sharing those functions with
build_database.py. There is no precomputed index: at this scale deriving them costs
microseconds, and a second generated copy of the data would only create something
that can silently go stale against the source. A precomputed index earns its place
when it holds something genuinely expensive, i.e. embeddings.

Scoring here is deliberately simple and dependency-free: overlap of identifiers and
symbols between the goal and each premise, weighted so that a rare identifier
matters more than a common one (`IsGCDOfList` appearing in both is a much stronger
signal than `Int`), plus a small bonus for an already-proven premise over one merely
asserted as an axiom from the book. It is a ranking heuristic, not a decision
procedure: it narrows 20-odd premises to the few worth showing, and cannot tell you
whether a premise actually applies -- only Lean can. Swapping in embeddings later
means replacing score_premise() and nothing else.
"""
import argparse
import json
import math
import re
from collections import Counter
from pathlib import Path

# Qualified names (Nat.Prime) first so they aren't split, then bare identifiers, then
# the type and operator symbols that carry signal. Arithmetic operators are included
# deliberately: a goal that is pure algebra ("a * (b + c) = a * b + a * c") has no
# identifiers at all, and without these it would tokenize to nothing and retrieve
# nothing. They appear in nearly every premise, so inverse-document-frequency
# weighting already discounts them to almost zero -- they only decide ties.
TOKEN_RE = re.compile(
    r'[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)+'
    r'|[A-Za-z_][A-Za-z0-9_]*'
    r'|[∣∀∃↔→∧∨¬≤≥≠∈ℤℕℚℝ×∑∏]'
    r'|[+*=/-]')
STOPWORDS = {'let', 'in', 'fun', 'if', 'then', 'else', 'true', 'false', 'Prop', 'Type', 'this'}


try:
    from .build_database import decl_name, resolve_proof, statement_identifiers
except ImportError:
    from build_database import decl_name, resolve_proof, statement_identifiers


def load_database(path, results=None):
    """Database entries in the shape search() ranks, derived from the source file."""
    entries = [json.loads(line) for line in path.read_text(encoding='utf-8').splitlines() if line.strip()]
    records = []
    for entry in entries:
        proof, status = resolve_proof(entry, (results or {}).get(entry['id']))
        records.append({
            "name": decl_name(entry['id']),
            "id": entry['id'],
            "source_label": entry.get('source_label'),
            "kind": entry.get('kind'),
            "status": status,
            "proof": proof or entry.get('proof'),
            "statement": entry['statement'],
            "natural_language": entry.get('natural_language', '').split('[fix note:')[0].strip(),
            "identifiers": statement_identifiers(entry['statement']),
            "requires": entry.get('requires', []),
        })
    return records


SYMBOLS = set('∣∀∃↔→∧∨¬≤≥≠∈ℤℕℚℝ×∑∏+*=/-')


def tokenize(text):
    """Identifiers and meaningful symbols. Single-letter variable names are dropped
    (matching `a` against `a` says nothing), but single-character symbols are kept."""
    return [t for t in TOKEN_RE.findall(text or '')
            if t not in STOPWORDS and (len(t) > 1 or t in SYMBOLS)]


def document_frequencies(records):
    """How many premises each token appears in, for inverse-document-frequency weighting."""
    frequencies = Counter()
    for record in records:
        for token in set(tokenize(record['statement']) + record.get('identifiers', [])
                          + record.get('requires', [])):
            frequencies[token] += 1
    return frequencies


def score_premise(goal_tokens, record, frequencies, total):
    premise_tokens = set(tokenize(record['statement']) + record.get('identifiers', [])
                          + record.get('requires', []))
    score = 0.0
    for token in set(goal_tokens) & premise_tokens:
        # Rare tokens are far more informative than ones every premise mentions.
        score += math.log(1 + total / (1 + frequencies.get(token, 0)))
    if record.get('status') == 'proved':
        score *= 1.1  # prefer a real theorem over an axiom asserted from the book
    return score


def search(records, goal, limit=5):
    """Premises ranked against one goal, best first. Zero-score entries are dropped."""
    frequencies = document_frequencies(records)
    total = max(len(records), 1)
    goal_tokens = tokenize(goal)
    scored = []
    for record in records:
        score = score_premise(goal_tokens, record, frequencies, total)
        if score > 0:
            scored.append((score, record))
    scored.sort(key=lambda pair: (-pair[0], pair[1]['name']))
    return [{**record, "score": round(score, 3)} for score, record in scored[:limit]]


def format_for_prompt(hits):
    """The premise list as it would be shown to the LLM next to the current goal."""
    if not hits:
        return '(no database premise looks relevant to this goal)'
    lines = []
    for hit in hits:
        status = 'proven' if hit['status'] == 'proved' else 'asserted from the book (axiom)'
        lines.append(f"{hit['name']}  [{hit['source_label']}, {status}, score {hit['score']}]\n"
                     f"  {hit['statement']}")
    return '\n'.join(lines)


def main():
    parser = argparse.ArgumentParser(description='Rank database premises against a proof goal')
    parser.add_argument('database', type=Path, help='database JSONL, e.g. database/goodman-1.6.jsonl')
    parser.add_argument('goal', help='the current goal, as Lean prints it')
    parser.add_argument('--limit', type=int, default=5)
    parser.add_argument('--json', action='store_true', help='emit raw JSON instead of prompt text')
    args = parser.parse_args()

    hits = search(load_database(args.database), args.goal, args.limit)
    print(json.dumps(hits, ensure_ascii=False, indent=2) if args.json else format_for_prompt(hits))


if __name__ == '__main__':
    main()