lean-bench/leanbench/pipeline.py

"""The whole run, in order: load -> detect -> formalize -> verify -> store.

``ingest`` is what the CLI, the Docker entrypoint and CI all call. Re-ingesting a
book (or just some of its pages) deletes exactly the rows those pages produced last
time and writes the new ones, so a source can be redone without collateral damage.
"""
from __future__ import annotations

from collections.abc import Callable, Sequence
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from .config import Settings
from .db import Database, ItemRow, item_id
from .detect import DetectionReport, detect_document
from .document import Document
from .formalize import Formalization, formalize_all, verification_mode
from .llm import build_client
from .verify import LeanChecker, NullChecker


@dataclass
class IngestResult:
    document: Document
    report: DetectionReport
    formalizations: list[Formalization] = field(default_factory=list)
    stored: list[str] = field(default_factory=list)
    removed: int = 0

    def summary(self) -> dict[str, Any]:
        verified = sum(1 for f in self.formalizations if f.verified)
        proved = sum(1 for f in self.formalizations if f.proved)
        skipped = sum(1 for f in self.formalizations if f.skipped)
        return {
            "source": self.document.title,
            "source_hash": self.document.sha256,
            "pages": len(self.document),
            "detected": len(self.report.detections),
            "statements": sum(1 for d in self.report.detections if d.kind == "statement"),
            "problems": sum(1 for d in self.report.detections if d.kind == "problem"),
            "multi_page": sum(1 for d in self.report.detections if d.spans_pages > 1),
            "unresolved": len(self.report.unresolved),
            "formalized": len(self.formalizations) - skipped,
            "skipped": skipped,
            "verified": verified,
            "proved": proved,
            "stored": len(self.stored),
            "removed": self.removed,
        }


def build_checker(settings: Settings, *, enabled: bool = True):
    if not enabled:
        return NullChecker()
    checker = LeanChecker.from_settings(settings.lean)
    return checker if checker.available() else NullChecker()


def to_row(document: Document, formalization: Formalization) -> ItemRow:
    detection = formalization.detection
    identifier = item_id(
        document.sha256,
        detection.start_page,
        detection.start_index,
        detection.end_page,
        detection.end_index,
    )
    verify = formalization.verify
    return ItemRow(
        id=identifier,
        source_hash=document.sha256,
        kind=detection.kind,
        label=detection.label,
        summary=detection.summary,
        informal=formalization.informal,
        start_page=detection.start_page,
        start_index=detection.start_index,
        end_page=detection.end_page,
        end_index=detection.end_index,
        page_checksum=document.page(detection.start_page).checksum,
        source_text=detection.text,
        lean_name=formalization.name,
        lean_code=formalization.lean_code,
        imports=formalization.imports,
        has_proof=formalization.has_proof,
        verified=bool(verify and verify.ok),
        proved=bool(verify and verify.proved),
        verify_mode=verify.mode if verify else verification_mode(detection),
        diagnostics=(verify.diagnostics if verify else ""),
        confidence=formalization.confidence,
        match_method=detection.match_method,
        match_score=detection.match_score,
        truncated=detection.truncated,
        types=formalization.index_types,
    )


def store(database: Database, document: Document,
          formalizations: Sequence[Formalization]) -> list[str]:
    database.upsert_source(
        document.sha256,
        title=document.title,
        path=str(document.path),
        pages=len(document),
        extractor=document.extractor,
        meta=document.meta,
    )
    stored: list[str] = []
    for formalization in formalizations:
        if formalization.skipped or not formalization.lean_code:
            continue
        stored.append(database.upsert_item(to_row(document, formalization)))
    return stored


def ingest(
    source: str | Path,
    *,
    settings: Settings | None = None,
    client: Any = None,
    database: Database | None = None,
    checker: Any = None,
    first_page: int | None = None,
    last_page: int | None = None,
    title: str = "",
    replace: bool = True,
    store_results: bool = True,
    progress: Callable[[str], None] | None = None,
) -> IngestResult:
    settings = settings or Settings.from_env()
    client = client or build_client(settings.llm)
    checker = checker if checker is not None else build_checker(settings)
    owns_database = database is None
    database = database or Database(settings.database)

    say = progress or (lambda _message: None)
    document = Document.load(source, title=title)
    say(f"{document.title}: {len(document)} pages, sha256 {document.sha256[:12]}, "
        f"extractor {document.extractor}")

    result = IngestResult(document=document, report=DetectionReport())
    try:
        if store_results and replace:
            if first_page or last_page:
                result.removed = database.delete_pages(
                    document.sha256, first_page or 1, last_page or len(document)
                )
            else:
                result.removed = database.delete_source(document.sha256)
            if result.removed:
                say(f"removed {result.removed} item(s) from the previous run")

        say("detecting formalizable passages...")
        result.report = detect_document(
            client,
            document,
            first_page=first_page,
            last_page=last_page,
            max_window_pages=settings.max_window_pages,
            progress=progress,
        )
        say(result.report.summary())

        say("formalizing...")
        result.formalizations = formalize_all(
            client,
            result.report.detections,
            checker=checker,
            document_title=document.title,
            repair_attempts=settings.repair_attempts,
            progress=progress,
        )

        if store_results:
            result.stored = store(database, document, result.formalizations)
            say(f"stored {len(result.stored)} item(s) in {database.path}")
    finally:
        if owns_database:
            database.close()
    return result