lean-bench/tests/test_pipeline.py

from __future__ import annotations

from pathlib import Path

from leanbench.config import Settings
from leanbench.db import Database
from leanbench.document import Document
from leanbench.llm import ScriptedClient
from leanbench.pipeline import ingest

from .conftest import detection_reply, formalize_reply

DETECT_PAGE_ONE = detection_reply(
    [
        {
            "kind": "statement",
            "label": "Definition 1.6.1",
            "summary": "divisibility",
            "start_quote": "Definition 1.6.1. Let a and b be integers.",
            "end_quote": "b = a k.",
            "continues_past_end": False,
        },
        {
            "kind": "statement",
            "label": "Theorem 1.6.3",
            "summary": "gcd exists",
            "start_quote": "Theorem 1.6.3. Let a and b be nonzero integers.",
            "end_quote": "the set of common",
            "continues_past_end": True,
        },
    ]
)

DETECT_BOTH_PAGES = detection_reply(
    [
        {
            "kind": "statement",
            "label": "Definition 1.6.1",
            "summary": "divisibility",
            "start_quote": "Definition 1.6.1. Let a and b be integers.",
            "end_quote": "b = a k.",
            "continues_past_end": False,
        },
        {
            "kind": "statement",
            "label": "Theorem 1.6.3",
            "summary": "gcd exists",
            "start_quote": "Theorem 1.6.3. Let a and b be nonzero integers.",
            "end_quote": "the greatest common divisor.",
            "continues_past_end": False,
        },
        {
            "kind": "problem",
            "label": "Exercise 1.6.4",
            "summary": "mutual divisibility",
            "start_quote": "Exercise 1.6.4. Show that if a divides b and b divides a",
            "end_quote": "a = b or a = -b.",
            "continues_past_end": False,
        },
    ]
)


def _settings(tmp_path: Path) -> Settings:
    settings = Settings.from_env(dotenv=None)
    settings.database = str(tmp_path / "leanbench.sqlite")
    settings.repair_attempts = 0
    return settings


def test_end_to_end_run_stores_indexed_items(tmp_path: Path, two_page_book, fake_checker) -> None:
    client = ScriptedClient(
        [DETECT_PAGE_ONE, DETECT_BOTH_PAGES]
        + [
            formalize_reply(name="divides_def",
                            lean_code="def divides (a b : Int) : Prop := ∃ k : Int, b = a * k"),
            formalize_reply(name="gcd_exists",
                            lean_code="theorem gcd_exists (a b : Int) : True := by\n  sorry"),
            formalize_reply(name="dvd_antisymm'",
                            lean_code="theorem dvd_antisymm' (a b : Int) : True := by\n  sorry"),
        ]
    )
    settings = _settings(tmp_path)
    database = Database(settings.database)
    result = ingest(
        two_page_book.path,
        settings=settings,
        client=client,
        database=database,
        checker=fake_checker,
        title="Test Book",
    )
    summary = result.summary()
    assert summary["detected"] == 3
    assert summary["problems"] == 1
    assert summary["multi_page"] == 1
    assert summary["stored"] == 3
    assert summary["verified"] == 3

    assert database.stats()["items"] == 3
    by_type = database.search(types=["Int"])
    assert len(by_type) == 3, "every item is reachable through its types"
    problems = database.search(kind="problem")
    assert len(problems) == 1 and problems[0]["start_page"] == 2
    database.close()


def test_reingesting_a_page_range_replaces_only_that_range(
    tmp_path: Path, two_page_book: Document, fake_checker
) -> None:
    settings = _settings(tmp_path)
    database = Database(settings.database)

    first = ScriptedClient([DETECT_PAGE_ONE, DETECT_BOTH_PAGES] + [formalize_reply()] * 3)
    ingest(two_page_book.path, settings=settings, client=first, database=database,
           checker=fake_checker)
    assert database.stats()["items"] == 3

    # redo page 2 only: the page-1 item survives, page-2 items are rebuilt
    second = ScriptedClient(
        [
            detection_reply(
                [
                    {
                        "kind": "problem",
                        "label": "Exercise 1.6.5",
                        "summary": "odd plus odd",
                        "start_quote": "Exercise 1.6.5. Prove that the sum of two odd",
                        "end_quote": "odd integers is even.",
                        "continues_past_end": False,
                    }
                ]
            ),
            formalize_reply(name="odd_add_odd"),
        ]
    )
    ingest(two_page_book.path, settings=settings, client=second, database=database,
           checker=fake_checker, first_page=2, last_page=2)

    pages = sorted(item["start_page"] for item in database.search(limit=100))
    assert pages == [1, 1, 2], "page 1 untouched, page 2 replaced by the new run"
    database.close()


def test_a_run_without_lean_stores_items_but_claims_nothing(
    tmp_path: Path, two_page_book: Document
) -> None:
    from leanbench.verify import NullChecker

    client = ScriptedClient([DETECT_PAGE_ONE, DETECT_BOTH_PAGES] + [formalize_reply()] * 3)
    settings = _settings(tmp_path)
    database = Database(settings.database)
    result = ingest(two_page_book.path, settings=settings, client=client, database=database,
                    checker=NullChecker())
    assert result.summary()["stored"] == 3
    assert result.summary()["verified"] == 0
    assert database.search(verified=True) == []
    database.close()