lean-bench/leanbench/detect.py

"""Step 2: decide what on a page is worth turning into Lean, and exactly where it is.

Two kinds of hit:

``statement``  mathematics asserted as fact -- definition, theorem, proposition,
               lemma, corollary, a precise claim inside an example. We deliberately
               do *not* record which of those it is: all that matters is that it can
               become useful Lean 4.
``problem``    a question posed and left open to the reader -- an exercise, an
               unsolved problem. Also formalizable, but as a statement with no proof.

Each hit comes back as a span: (start page, start index) -> (end page, end index),
indices being offsets into that page's text. Several hits per page are normal.

If a hit runs off the bottom of the window, the model says so and the window grows
by one page and the page is re-read, repeatedly, so an item spread over three pages
is found whole rather than in pieces.
"""
from __future__ import annotations

import json
from collections.abc import Callable, Sequence
from dataclasses import asdict, dataclass, field
from typing import Any

from .anchors import resolve_span
from .document import Document, Page, PageWindow

KIND_STATEMENT = "statement"
KIND_PROBLEM = "problem"
KINDS = (KIND_STATEMENT, KIND_PROBLEM)

MAX_WINDOW_CHARS = 28000

SYSTEM_PROMPT = """\
You are an exacting extraction engine. You are given the raw text of one or more
consecutive pages of a mathematics book, exactly as extracted from the PDF (line
breaks, hyphenation and all).

Find every passage that could be turned into useful Lean 4 code. There are two kinds:

- "statement": mathematics asserted as fact -- a definition, theorem, proposition,
  lemma, corollary, or a precise claim made inside a remark or worked example. Do NOT
  tell me which of those it is; the only question is whether it can become useful
  Lean 4.
- "problem": a precise mathematical question left open to the reader -- an exercise,
  a problem, an unsolved question -- stated without its solution or proof here.

Ignore: running heads and page numbers, table of contents, index, bibliography,
motivational prose with no precise claim, figure captions, and instructions with no
mathematical content ("read section 2", "discuss with a classmate").

Rules:
- Quote verbatim. "start_quote" must be copied character-for-character from the text
  given to you and must begin at the very first character of the passage (including
  its printed label, e.g. "Theorem 1.6.3."). "end_quote" must likewise be copied
  verbatim and end at the very last character of the passage.
- start_quote: 40-90 characters. end_quote: 30-60 characters. Never paraphrase,
  never normalise whitespace, never fix hyphenation.
- List passages in order of appearance. Passages must not overlap.
- "continues_past_end": true only when the passage is cut off by the end of the text
  you were given -- that is, the book continues it on the following page.
- If nothing qualifies, return an empty list.

Reply with JSON only, in exactly this shape:
{"items": [{"kind": "statement" | "problem",
            "label": "Theorem 1.6.3" or "",
            "summary": "one short English line",
            "start_quote": "...",
            "end_quote": "...",
            "continues_past_end": false}]}
"""


@dataclass
class Detection:
    """One located item, ready for step 3."""

    kind: str
    label: str
    summary: str
    start_page: int
    start_index: int
    end_page: int
    end_index: int
    text: str
    truncated: bool = False
    match_method: str = "exact"
    match_score: float = 1.0
    pages_scanned: list[int] = field(default_factory=list)

    @property
    def spans_pages(self) -> int:
        return self.end_page - self.start_page + 1

    @property
    def locator(self) -> str:
        return f"p{self.start_page}:{self.start_index}-p{self.end_page}:{self.end_index}"

    def as_dict(self) -> dict[str, Any]:
        data = asdict(self)
        data["spans_pages"] = self.spans_pages
        data["locator"] = self.locator
        return data


def _window_text(window: PageWindow) -> str:
    text = window.text
    if len(text) <= MAX_WINDOW_CHARS:
        return text
    return text[:MAX_WINDOW_CHARS]


def detect_window(client: Any, window: PageWindow, *, temperature: float = 0.0) -> list[dict]:
    """Ask the model what is in this window. Returns raw (unresolved) item dicts."""
    pages = window.pages_covered()
    label = f"page {pages[0]}" if len(pages) == 1 else f"pages {pages[0]}-{pages[-1]}"
    user = (
        f"Text of {label} of the book follows between the markers.\n"
        f"<<<BEGIN {label.upper()}>>>\n{_window_text(window)}\n<<<END>>>"
    )
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user},
    ]
    if hasattr(client, "complete_json"):
        payload = client.complete_json(messages, temperature=temperature)
    else:  # any bare ChatClient
        from .llm import parse_json

        payload = parse_json(client.complete(messages, temperature=temperature))

    items = payload.get("items") if isinstance(payload, dict) else payload
    if not isinstance(items, list):
        return []
    clean: list[dict] = []
    for item in items:
        if not isinstance(item, dict):
            continue
        kind = str(item.get("kind", "")).strip().lower()
        if kind not in KINDS:
            kind = KIND_PROBLEM if "problem" in kind or "exercise" in kind else KIND_STATEMENT
        clean.append(
            {
                "kind": kind,
                "label": str(item.get("label") or "").strip(),
                "summary": str(item.get("summary") or "").strip(),
                "start_quote": str(item.get("start_quote") or ""),
                "end_quote": str(item.get("end_quote") or ""),
                "continues_past_end": bool(item.get("continues_past_end")),
            }
        )
    return clean


def resolve_items(
    window: PageWindow,
    items: Sequence[dict],
    *,
    still_truncated: bool = False,
) -> tuple[list[Detection], list[dict]]:
    """Turn quoted items into page/index spans. Returns (resolved, unresolved)."""
    resolved: list[Detection] = []
    unresolved: list[dict] = []
    text = window.text
    cursor = 0
    for item in items:
        span = resolve_span(text, item["start_quote"], item["end_quote"], search_from=cursor)
        if span is None or not span.resolved:
            unresolved.append({**item, "reason": "quotes not found in page text"})
            continue
        start_page, start_index = window.locate(span.start)
        end_page, end_index = window.locate(span.end)
        if (end_page, end_index) <= (start_page, start_index):
            unresolved.append({**item, "reason": "empty span"})
            continue
        cursor = max(cursor, span.start + 1)
        resolved.append(
            Detection(
                kind=item["kind"],
                label=item["label"],
                summary=item["summary"],
                start_page=start_page,
                start_index=start_index,
                end_page=end_page,
                end_index=end_index,
                text=text[span.start : span.end],
                truncated=bool(item.get("continues_past_end")) and still_truncated,
                match_method=span.method,
                match_score=round(span.score, 3),
                pages_scanned=window.pages_covered(),
            )
        )
    return resolved, unresolved


@dataclass
class DetectionReport:
    detections: list[Detection] = field(default_factory=list)
    unresolved: list[dict] = field(default_factory=list)
    windows_used: int = 0
    pages_scanned: int = 0
    llm_calls: int = 0

    def summary(self) -> str:
        statements = sum(1 for d in self.detections if d.kind == KIND_STATEMENT)
        problems = sum(1 for d in self.detections if d.kind == KIND_PROBLEM)
        multipage = sum(1 for d in self.detections if d.spans_pages > 1)
        return (
            f"{len(self.detections)} items "
            f"({statements} statements, {problems} problems, {multipage} spanning pages) "
            f"from {self.pages_scanned} pages in {self.llm_calls} model calls; "
            f"{len(self.unresolved)} unresolved"
        )


def detect_document(
    client: Any,
    document: Document,
    *,
    first_page: int | None = None,
    last_page: int | None = None,
    max_window_pages: int = 3,
    progress: Callable[[str], None] | None = None,
    on_window_error: str = "skip",
) -> DetectionReport:
    """Scan a document page by page, growing the window when an item is cut off."""
    pages: list[Page] = document.page_range(first_page, last_page)
    report = DetectionReport(pages_scanned=len(pages))
    position = 0

    while position < len(pages):
        size = 1
        window = PageWindow(pages[position : position + 1])
        raw: list[dict] = []
        needs_more = False
        while True:
            window = PageWindow(pages[position : position + size])
            try:
                raw = detect_window(client, window)
            except Exception as error:  # noqa: BLE001 - one bad page must not kill a book
                if on_window_error == "raise":
                    raise
                report.unresolved.append(
                    {"pages": window.pages_covered(), "reason": f"detector failed: {error}"}
                )
                raw = []
            report.llm_calls += 1
            report.windows_used += 1
            needs_more = any(item.get("continues_past_end") for item in raw)
            room_left = position + size < len(pages)
            if needs_more and size < max_window_pages and room_left:
                size += 1
                if progress:
                    progress(
                        f"  page {window.first_page}: item runs on, re-reading with "
                        f"{size} pages"
                    )
                continue
            break

        resolved, unresolved = resolve_items(window, raw, still_truncated=needs_more)
        report.detections.extend(resolved)
        report.unresolved.extend(unresolved)
        if progress:
            progress(
                f"  pages {window.first_page}-{window.last_page}: "
                f"{len(resolved)} item(s)" + (f", {len(unresolved)} unresolved" if unresolved else "")
            )
        position += size

    report.detections = deduplicate(report.detections)
    return report


def deduplicate(detections: Sequence[Detection]) -> list[Detection]:
    """Keep the widest span when two hits start at the same place."""
    best: dict[tuple[int, int], Detection] = {}
    for detection in detections:
        key = (detection.start_page, detection.start_index)
        current = best.get(key)
        if current is None or (detection.end_page, detection.end_index) > (
            current.end_page,
            current.end_index,
        ):
            best[key] = detection
    return sorted(best.values(), key=lambda d: (d.start_page, d.start_index))


def write_jsonl(path: str, detections: Sequence[Detection]) -> None:
    with open(path, "w", encoding="utf-8") as stream:
        for detection in detections:
            stream.write(json.dumps(detection.as_dict(), ensure_ascii=False) + "\n")