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.

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

TOKEN_RE = re.compile(r'[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)+|[A-Za-z_][A-Za-z0-9_]*|[∣∀∃↔→∧∨¬≤≥≠∈]')
STOPWORDS = {'let', 'in', 'fun', 'if', 'then', 'else', 'true', 'false', 'Prop', 'Type', 'this'}


def load_index(path):
    return [json.loads(line) for line in path.read_text(encoding='utf-8').splitlines() if line.strip()]


def tokenize(text):
    return [t for t in TOKEN_RE.findall(text or '') if t not in STOPWORDS and len(t) > 1]


def document_frequencies(index):
    """How many premises each token appears in, for inverse-document-frequency weighting."""
    frequencies = Counter()
    for record in index:
        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(index, goal, limit=5):
    """Premises ranked against one goal, best first. Zero-score entries are dropped."""
    frequencies = document_frequencies(index)
    total = max(len(index), 1)
    goal_tokens = tokenize(goal)
    scored = []
    for record in index:
        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('index', type=Path, help='search index from build_database.py')
    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_index(args.index), 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()