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_ocr(settings: Settings, client: Any = None):
    """An OcrReader for this configuration, or None when OCR is switched off."""
    if settings.ocr.mode == "text":
        return None
    from .ocr import OcrCache, OcrReader, build_backend

    backend = build_backend(settings.ocr, client)
    cache = OcrCache(settings.ocr.cache_dir) if settings.ocr.cache_dir else None
    return OcrReader(backend=backend, cache=cache, dpi=settings.ocr.dpi)


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)
    reader = build_ocr(settings, client) if settings.ocr.mode != "text" else None
    if reader is not None:
        say(f"reading pages with {reader.backend.name} ({settings.ocr.mode} mode)")
    document = Document.load(
        source,
        title=title,
        mode=settings.ocr.mode,
        ocr=reader,
        threshold=settings.ocr.threshold,
        first_page=first_page,
        last_page=last_page,
        progress=progress,
    )
    say(f"{document.title}: {len(document)} pages, sha256 {document.sha256[:12]}, "
        f"extractor {document.extractor}")
    if reader is not None and reader.cache is not None:
        say(f"  OCR cache: {reader.cache.hits} hit(s), {reader.cache.misses} miss(es)")

    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