lean-bench/tests/test_ocr.py

from __future__ import annotations

from pathlib import Path

import pytest

from leanbench.document import Document, OcrUnavailable
from leanbench.ocr import (
    OcrCache,
    OcrReader,
    VisionLLMBackend,
    looks_broken,
    text_quality,
)

# Real extraction of page 90 of the project's test book: the maths font carries a
# custom encoding, so every operator arrives as a stray accent or ligature glyph.
BROKEN = """Proposition 1.11.7 (Chinese remainder theorem). Let a and b be relatively
prime natural numbers. There is an isomorphism of rings
Zab Š Za ˚ Zb;
defined by Œxab 7! .Œxa; Œxb/.
"""

# What the page actually says, as the vision model reads it.
CLEAN = """Proposition 1.11.7 (Chinese remainder theorem). Let $a$ and $b$ be relatively
prime natural numbers. There is an isomorphism of rings
$$\\mathbb{Z}_{ab} \\cong \\mathbb{Z}_a \\oplus \\mathbb{Z}_b,$$
defined by $[x]_{ab} \\mapsto ([x]_a, [x]_b)$.
"""

# The other corruption mode, from page 300 of the same book: the maths is made of
# ordinary ASCII, so a glyph-based check scores it a perfect 1.0.
OPERATORS_AS_LETTERS = """6.2.18. Let I and J be two ideals in a ring R. Show that
I C J D fa C b W a 2 I and b 2 J g
is an ideal in R.
6.2.19. Let I and J be two ideals in a ring R. Show that
IJ D fa1b1 C a2b2 C    C asbs W s  1; ai 2 I; bi 2 J g
is an ideal in R, and IJ  I \\ J .
"""

PROSE = """The chapter closes with a discussion of why the theorem matters, and a
remark on its history, with no displayed formulas at all.
"""


class FakeBackend:
    """Returns a canned transcription and counts how often it was asked."""

    name = "fake:v1"

    def __init__(self, text: str = CLEAN) -> None:
        self.text = text
        self.calls = 0

    def transcribe(self, image: bytes) -> str:
        self.calls += 1
        return self.text


class ExplodingBackend:
    name = "boom:v1"

    def transcribe(self, image: bytes) -> str:
        raise RuntimeError("model refused")


class FlakyBackend:
    """Fails on chosen pages only -- the ordinary case of one unreadable page."""

    name = "flaky:v1"

    def __init__(self, fail_on: set[int]) -> None:
        self.fail_on = fail_on
        self.seen = 0

    def transcribe(self, image: bytes) -> str:
        self.seen += 1
        if self.seen in self.fail_on:
            raise RuntimeError(f"page {self.seen} refused")
        return CLEAN


# -- the quality heuristic -------------------------------------------------


def test_mangled_maths_scores_far_below_clean_text() -> None:
    assert text_quality(BROKEN) < 0.75
    assert text_quality(CLEAN) > 0.9
    assert text_quality(PROSE) > 0.9


def test_operators_disguised_as_letters_are_caught() -> None:
    """The failure mode that a glyph-only check scores 1.0 and waves through."""
    assert text_quality(OPERATORS_AS_LETTERS) < 0.75
    assert looks_broken(OPERATORS_AS_LETTERS)


def test_ordinary_prose_with_numerals_survives() -> None:
    assert not looks_broken(
        "In chapter 2 we saw that the group has order 2 and index 2 in G.")


def test_looks_broken_separates_the_two() -> None:
    assert looks_broken(BROKEN)
    assert not looks_broken(CLEAN)
    assert not looks_broken(PROSE)


def test_empty_text_is_the_worst_case() -> None:
    assert text_quality("") == 0.0
    assert text_quality("   \n  ") == 0.0


def test_a_single_accented_word_is_not_damning() -> None:
    # Real prose does contain the odd accent; only a pile of them is evidence.
    assert not looks_broken("Consider the Poincare´ conjecture and its history.")


# -- the cache -------------------------------------------------------------


def test_cache_round_trips(tmp_path: Path) -> None:
    cache = OcrCache(tmp_path)
    assert cache.get("hash", 3, 200, "fake:v1") is None
    cache.put("hash", 3, 200, "fake:v1", "hello")
    assert cache.get("hash", 3, 200, "fake:v1") == "hello"
    assert cache.hits == 1 and cache.misses == 1


def test_cache_keys_on_dpi_and_backend(tmp_path: Path) -> None:
    cache = OcrCache(tmp_path)
    cache.put("hash", 3, 200, "fake:v1", "at 200")
    assert cache.get("hash", 3, 400, "fake:v1") is None, "different dpi is a different read"
    assert cache.get("hash", 3, 200, "other:v1") is None, "different backend too"


def test_cache_can_be_switched_off(tmp_path: Path) -> None:
    cache = OcrCache(tmp_path, enabled=False)
    cache.put("hash", 1, 200, "fake:v1", "x")
    assert cache.get("hash", 1, 200, "fake:v1") is None


# -- the reader ------------------------------------------------------------


def test_reader_pays_for_a_page_once(tmp_path: Path, monkeypatch) -> None:
    monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png")
    backend = FakeBackend()
    reader = OcrReader(backend=backend, cache=OcrCache(tmp_path))
    first = reader.page_text("book.pdf", "hash", 7)
    second = reader.page_text("book.pdf", "hash", 7)
    assert first == second
    assert backend.calls == 1, "the second read must come from cache"


def test_reader_records_failures_instead_of_raising(tmp_path: Path, monkeypatch) -> None:
    monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png")
    reader = OcrReader(backend=ExplodingBackend(), cache=OcrCache(tmp_path))
    assert reader.page_text("book.pdf", "hash", 4) is None
    assert "model refused" in reader.failures[4]


def test_reader_strips_a_markdown_fence(tmp_path: Path, monkeypatch) -> None:
    monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png")
    reader = OcrReader(backend=FakeBackend("```\nTheorem 1.\n```"), cache=None)
    assert reader.page_text("book.pdf", "hash", 1) == "Theorem 1."


# -- the vision backend's request shape ------------------------------------


def test_vision_backend_sends_an_image_part() -> None:
    captured: dict = {}

    class Client:
        def complete(self, messages, **kwargs):
            captured["messages"] = messages
            captured["kwargs"] = kwargs
            return "  transcribed  "

    backend = VisionLLMBackend(client=Client(), model="qwen-vl-max")
    assert backend.transcribe(b"\x89PNG") == "transcribed"
    parts = captured["messages"][0]["content"]
    assert parts[0]["type"] == "text"
    assert parts[1]["type"] == "image_url"
    assert parts[1]["image_url"]["url"].startswith("data:image/png;base64,")
    assert captured["kwargs"]["model"] == "qwen-vl-max"
    assert backend.name == "vision:qwen-vl-max"


# -- Document integration --------------------------------------------------


@pytest.fixture()
def mixed_book(tmp_path: Path) -> Path:
    """Page 1 is fine, page 2 is font-mangled."""
    path = tmp_path / "mixed.txt"
    path.write_text(PROSE + "\f" + BROKEN, encoding="utf-8")
    return path


def test_auto_mode_only_pays_for_the_broken_page(mixed_book: Path, monkeypatch) -> None:
    monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png")
    backend = FakeBackend()
    reader = OcrReader(backend=backend, cache=None)
    document = Document.load(mixed_book, mode="auto", ocr=reader, threshold=0.75)

    assert backend.calls == 1, "the clean page must not be sent to the model"
    assert document.meta["ocr"]["pages_replaced"] == [2]
    assert document.page(1).text.startswith("The chapter closes")
    assert "\\cong" in document.page(2).text
    assert "fake:v1" in document.extractor


def test_ocr_mode_reads_every_page(mixed_book: Path, monkeypatch) -> None:
    monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png")
    backend = FakeBackend()
    Document.load(mixed_book, mode="ocr", ocr=reader_for(backend))
    assert backend.calls == 2


def reader_for(backend) -> OcrReader:
    return OcrReader(backend=backend, cache=None)


def test_page_range_bounds_what_gets_read(mixed_book: Path, monkeypatch) -> None:
    monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png")
    backend = FakeBackend()
    document = Document.load(mixed_book, mode="ocr", ocr=reader_for(backend),
                             first_page=2, last_page=2)
    assert backend.calls == 1, "a 25-page trial of a 600-page book costs 25 calls"
    assert document.meta["ocr"]["pages_replaced"] == [2]


def test_text_mode_never_touches_the_model(mixed_book: Path) -> None:
    backend = FakeBackend()
    document = Document.load(mixed_book, mode="text", ocr=reader_for(backend))
    assert backend.calls == 0
    assert "Š" in document.page(2).text, "embedded text is left exactly as found"
    assert "ocr" not in document.meta


def test_a_failed_page_keeps_its_embedded_text(mixed_book: Path, monkeypatch) -> None:
    """One bad page falls back; the rest of the book is unaffected."""
    monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png")
    reader = OcrReader(backend=FlakyBackend(fail_on={2}), cache=None)
    document = Document.load(mixed_book, mode="ocr", ocr=reader)
    assert document.meta["ocr"]["pages_failed"] == [2]
    assert document.meta["ocr"]["pages_replaced"] == [1]
    assert "Š" in document.page(2).text, "we fall back, we do not lose the page"


def test_total_ocr_failure_is_loud(mixed_book: Path, monkeypatch) -> None:
    """Every page failing is a broken configuration, not a run to complete.

    Silently falling back here is the expensive failure: the fallback text is the
    mangled text OCR was called in to replace, so the run costs a full detection
    pass and stores nothing.
    """
    monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png")
    reader = OcrReader(backend=ExplodingBackend(), cache=None)
    with pytest.raises(OcrUnavailable) as caught:
        Document.load(mixed_book, mode="ocr", ocr=reader)
    assert "model refused" in str(caught.value), "the real cause must reach the operator"
    assert "boom:v1" in str(caught.value)


def test_auto_mode_with_nothing_to_do_does_not_raise(tmp_path: Path) -> None:
    """No page needed OCR, so nothing failed -- that is success, not an outage."""
    path = tmp_path / "clean.txt"
    path.write_text(PROSE, encoding="utf-8")
    document = Document.load(path, mode="auto", ocr=OcrReader(backend=ExplodingBackend()))
    assert document.meta["ocr"]["pages_replaced"] == []
    assert document.meta["ocr"]["pages_failed"] == []


def test_offsets_stay_usable_after_ocr(mixed_book: Path, monkeypatch) -> None:
    """The span contract must hold whichever extractor produced the text."""
    monkeypatch.setattr("leanbench.ocr.render_page", lambda *a, **k: b"png")
    document = Document.load(mixed_book, mode="ocr", ocr=reader_for(FakeBackend()))
    page = document.page(2)
    start = page.text.index("Proposition")
    end = page.text.index("rings") + len("rings")
    assert document.slice(2, start, 2, end).startswith("Proposition 1.11.7")
    # The extractor is part of the offset contract, so it is recorded on the source.
    assert "fake:v1" in document.extractor