lean-bench/tests/test_search_database.py

import unittest

import json
import tempfile
from pathlib import Path

from textbook.search_database import format_for_prompt, load_database, search, tokenize

INDEX = [
    {"name": "goodman_1_6_20", "id": "goodman-1-6-20", "source_label": "Corollary 1.6.20",
     "kind": "corollary", "status": "asserted",
     "statement": "\u2200 (p : \u2115) (as : List \u2124), Nat.Prime p \u2192 (p:\u2124) \u2223 List.prod as \u2192 \u2203 a \u2208 as, (p:\u2124) \u2223 a",
     "natural_language": "", "identifiers": ["Nat.Prime", "List.prod"], "requires": ["Nat.Prime"]},
    {"name": "goodman_1_6_24", "id": "goodman-1-6-24", "source_label": "Proposition 1.6.24",
     "kind": "proposition", "status": "asserted",
     "statement": "\u2200 (as : List \u2124), \u2203 d, IsGCDOfList d as",
     "natural_language": "", "identifiers": ["IsGCDOfList"], "requires": []},
    {"name": "goodman_1_6_1", "id": "goodman-1-6-1", "source_label": "Proposition 1.6.1",
     "kind": "proposition", "status": "asserted",
     "statement": "\u2200 (a : \u2124), a + 0 = a", "natural_language": "",
     "identifiers": [], "requires": []},
]


SOURCE_ENTRIES = [
    {"id": "goodman-1-6-20", "source_label": "Corollary 1.6.20", "kind": "corollary",
     "statement": "\u2200 (p : \u2115), Nat.Prime p \u2192 (p:\u2124) \u2223 List.prod as",
     "prelude": "import Mathlib", "natural_language": "x [fix note: dropped]",
     "requires": ["Nat.Prime"]},
    {"id": "goodman-1-6-24", "source_label": "Proposition 1.6.24", "kind": "proposition",
     "statement": "\u2200 (as : List \u2124), \u2203 d, IsGCDOfList d as",
     "prelude": "import Mathlib", "natural_language": "", "requires": [],
     "proof": "exact foo", "proof_checked": True},
    {"id": "goodman-1-6-1", "source_label": "Proposition 1.6.1", "kind": "proposition",
     "statement": "\u2200 (a : \u2124), a + 0 = a", "prelude": "import Mathlib",
     "natural_language": "", "requires": [], "proof": "ring", "proof_checked": False},
]


class LoadDatabaseTests(unittest.TestCase):
    """load_database derives name/status/identifiers from the source file, so there is
    no separate index that could go stale against it."""

    def _load(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / 'db.jsonl'
            path.write_text('\n'.join(json.dumps(e) for e in SOURCE_ENTRIES) + '\n',
                            encoding='utf-8')
            return load_database(path)

    def test_derives_lean_name_and_strips_fix_notes(self):
        records = {r['id']: r for r in self._load()}
        self.assertEqual(records['goodman-1-6-20']['name'], 'goodman_1_6_20')
        self.assertNotIn('fix note', records['goodman-1-6-20']['natural_language'])

    def test_derives_proof_status_from_proof_checked(self):
        records = {r['id']: r for r in self._load()}
        self.assertEqual(records['goodman-1-6-24']['status'], 'proved')
        self.assertEqual(records['goodman-1-6-1']['status'], 'proof_pending')
        self.assertEqual(records['goodman-1-6-20']['status'], 'asserted')

    def test_search_works_directly_on_loaded_source(self):
        hits = search(self._load(), '\u22a2 IsGCDOfList d as', limit=2)
        self.assertEqual(hits[0]['name'], 'goodman_1_6_24')


class SearchDatabaseTests(unittest.TestCase):
    def test_tokenize_keeps_qualified_names_and_drops_stopwords(self):
        tokens = tokenize('let x := Nat.Prime p')
        self.assertIn('Nat.Prime', tokens)
        self.assertNotIn('let', tokens)

    def test_tokenize_keeps_symbols_but_drops_single_letter_variables(self):
        # A pure-algebra goal has no identifiers at all, so the symbols have to
        # survive or it retrieves nothing; matching 'a' against 'a' says nothing.
        tokens = tokenize('\u2200 (a b : \u2124), a * (b + c) = 0')
        self.assertIn('\u2124', tokens)
        self.assertIn('*', tokens)
        self.assertIn('+', tokens)
        self.assertNotIn('a', tokens)

    def test_exact_premise_ranks_first(self):
        hits = search(INDEX, '\u22a2 IsGCDOfList d as', limit=3)
        self.assertEqual(hits[0]['name'], 'goodman_1_6_24')

    def test_rare_identifier_outweighs_common_one(self):
        # Both premises mention List; only one mentions Nat.Prime, so the goal about
        # primes must prefer it rather than tying on the common token.
        hits = search(INDEX, '\u22a2 Nat.Prime p \u2192 (p:\u2124) \u2223 List.prod as', limit=3)
        self.assertEqual(hits[0]['name'], 'goodman_1_6_20')

    def test_irrelevant_premises_are_dropped_not_ranked_last(self):
        hits = search(INDEX, '\u22a2 IsGCDOfList d as', limit=5)
        self.assertNotIn('goodman_1_6_1', [h['name'] for h in hits])

    def test_proven_premise_is_preferred_over_an_asserted_one(self):
        index = [dict(INDEX[1]), dict(INDEX[1], name='proved_twin', id='proved-twin',
                                       status='proved', source_label='Proposition 1.6.24b')]
        hits = search(index, '\u22a2 IsGCDOfList d as', limit=2)
        self.assertEqual(hits[0]['name'], 'proved_twin')

    def test_limit_is_respected_and_empty_result_is_explained(self):
        self.assertEqual(len(search(INDEX, '\u22a2 List.prod as', limit=1)), 1)
        self.assertIn('no database premise', format_for_prompt([]))

    def test_format_for_prompt_marks_axioms_distinctly(self):
        text = format_for_prompt(search(INDEX, '\u22a2 IsGCDOfList d as', limit=1))
        self.assertIn('asserted from the book (axiom)', text)
        self.assertIn('goodman_1_6_24', text)


if __name__ == '__main__':
    unittest.main()