src/endopath/reocr.py

"""Re-OCR the rendered page images to a cleaner ground-truth report text (#212).

The prepared TCGA corpus OCR (`data/raw/TCGA_Reports.csv`) is pervasively damaged: 90% of reports carry
some damage and half are heavily damaged, with tables shredded into period-separated fragments
(`benchmark.ocr_damage_density`). Extraction reads that text, so a damaged transcription produces damaged
drafts and evidence spans that will not match the page.

This module transcribes the page images directly with a vision model, preserving layout and tables, to
produce a report text that both the extraction pass and the evidence check can rely on. The model is
reached through the app's provider seam (`endopath.llm.get_client`), so the same call runs against the
hosted Anthropic API or a local, on-premise model with no change here.

`ClaudePageReader` implements the `textsource.PageReader` protocol, so it plugs into the existing
`VisionLlmTextSource`: registering it with `textsource.set_default_page_reader(ClaudePageReader())` makes
the ingest path read pages through it. `scripts/reocr_pages.py` uses it to re-OCR a selected set of cases
and persist the result as `reports.report_text` (`storage.set_report_text`).

Cost guard: every page transcription is cached by the page image's sha256 checksum
(`PageTranscriptCache`), so re-running the script never re-bills a page whose pixels are unchanged, and a
`--limit`/`--worst`/`--case` selection in the script bounds how many pages are sent. No page is skipped
silently: the reader reports per-run token counts and cache hits.
"""

from __future__ import annotations

import base64
import dataclasses
import hashlib
import json
import os
from pathlib import Path
from typing import Optional

# Consistent with every other LLM call site in the app (llm_extraction, induction, resolution, the chat
# co-pilots): the codebase standardizes on this model, and the vision extract in resolution.py already
# reads page images with it. Override with REOCR_MODEL to trade cost for capability on a bulk run;
# claude-opus-4-8 is the higher-capability, higher-cost alternative reachable through the same seam.
REOCR_MODEL = (os.environ.get("REOCR_MODEL") or "").strip() or "claude-sonnet-5"

# A single page of a pathology report transcribes to at most a few thousand tokens; this ceiling leaves
# ample headroom and stays well under the streaming threshold, so a plain non-streaming call is safe.
DEFAULT_MAX_TOKENS = 8000

DEFAULT_CACHE_PATH = Path("data") / "interim" / "reocr_cache.json"

_IMAGE_MEDIA_TYPES = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg"}

# The pages of one report are joined with a blank line, matching the corpus's own paragraph spacing, so
# the concatenated text reads as one continuous document. No synthetic page markers are injected: a
# marker is text the model could later quote, and invariant 2 verifies an evidence quote as a verbatim
# substring of exactly this stored text.
PAGE_SEPARATOR = "\n\n"

TRANSCRIPTION_PROMPT = """\
You are transcribing a scanned pathology report page into plain text for a clinical dataset. Reproduce \
the page exactly as printed.

Rules:
- Transcribe every word, number, and punctuation mark verbatim. Do not paraphrase, summarize, correct, \
translate, re-order, or add commentary of your own.
- Preserve the reading order and the line breaks of the page.
- When the page contains a table, render it as a Markdown table: one row per printed row, columns in \
their printed order, so the rows stay readable instead of collapsing into a list of fragments.
- Keep a field label and its value on the same line, as printed ("Label: value").
- For a checkbox or a selectable option, write [x] for a filled or checked mark and [ ] for an empty one.
- If a word is genuinely illegible, write [illegible] in its place. Do not guess at it.
- Output only the transcription. No preamble, no explanation, no code fences.
"""


def page_checksum(path: Path | str) -> str:
    """sha256 of a page image's bytes: the cache key, so an unchanged page is never re-transcribed and a
    re-rendered page (different pixels) is."""
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


class PageTranscriptCache:
    """A JSON-backed page-transcript cache keyed by `checksum` and validated against the model that
    produced it, so switching REOCR_MODEL re-transcribes rather than serving another model's text. The
    file is written atomically. `dirty` lets a caller batch a single save after a run instead of writing
    per page."""

    def __init__(self, path: Path | str = DEFAULT_CACHE_PATH) -> None:
        self.path = Path(path)
        self.dirty = False
        self._data: dict[str, dict] = {}
        if self.path.exists():
            try:
                self._data = json.loads(self.path.read_text(encoding="utf-8"))
            except (json.JSONDecodeError, OSError):
                self._data = {}

    def get(self, checksum: str, model: str) -> Optional[str]:
        entry = self._data.get(checksum)
        if entry is None or entry.get("model") != model:
            return None
        return entry.get("text")

    def put(self, checksum: str, model: str, text: str) -> None:
        self._data[checksum] = {"model": model, "text": text, "chars": len(text)}
        self.dirty = True

    def save(self) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        tmp = self.path.with_suffix(self.path.suffix + ".tmp")
        tmp.write_text(json.dumps(self._data, ensure_ascii=False, indent=0), encoding="utf-8")
        tmp.rename(self.path)
        self.dirty = False

    def __len__(self) -> int:
        return len(self._data)


@dataclasses.dataclass
class PageTranscript:
    """One page's transcription and the tokens it cost. A cache hit reports zero tokens; the caller sums
    these across a run for a real cost figure rather than an estimate."""

    text: str
    checksum: str
    cached: bool
    input_tokens: int = 0
    output_tokens: int = 0


def _image_block(path: Path) -> dict:
    media_type = _IMAGE_MEDIA_TYPES.get(path.suffix.lower())
    if media_type is None:
        raise ValueError(f"unsupported image type for re-OCR: {path}")
    return {
        "type": "image",
        "source": {
            "type": "base64",
            "media_type": media_type,
            "data": base64.standard_b64encode(path.read_bytes()).decode("ascii"),
        },
    }


def _strip_code_fences(text: str) -> str:
    """Drop a wrapping ``` fence if the model added one despite the instruction, so the stored text is
    the transcription and not a fenced block. A fence that is not on the first/last line is left alone
    (it could be part of the page content)."""
    lines = text.strip().splitlines()
    if len(lines) >= 2 and lines[0].startswith("```") and lines[-1].strip() == "```":
        return "\n".join(lines[1:-1]).strip()
    return text.strip()


def transcribe_page(
    path: Path | str,
    *,
    client=None,
    model: str = REOCR_MODEL,
    cache: Optional[PageTranscriptCache] = None,
    max_tokens: int = DEFAULT_MAX_TOKENS,
) -> PageTranscript:
    """Transcribe one page image, returning its text and token cost. Uses `cache` when the page's
    checksum + model already produced text, so a re-run pays nothing for an unchanged page. The model is
    reached through the provider seam (`endopath.llm.get_client`), so a local model swaps in with no
    change here."""
    path = Path(path)
    checksum = page_checksum(path)
    if cache is not None:
        hit = cache.get(checksum, model)
        if hit is not None:
            return PageTranscript(text=hit, checksum=checksum, cached=True)

    if client is None:
        from endopath import llm

        client = llm.get_client()

    response = client.messages.create(
        model=model,
        max_tokens=max_tokens,
        messages=[
            {
                "role": "user",
                "content": [_image_block(path), {"type": "text", "text": TRANSCRIPTION_PROMPT}],
            }
        ],
    )
    if getattr(response, "stop_reason", None) == "max_tokens":
        raise RuntimeError(
            f"re-OCR of {path} hit max_tokens before the page finished; raise max_tokens (currently "
            f"{max_tokens})."
        )
    text = "".join(block.text for block in response.content if getattr(block, "type", None) == "text")
    text = _strip_code_fences(text)

    usage = getattr(response, "usage", None)
    input_tokens = getattr(usage, "input_tokens", None) or 0
    output_tokens = getattr(usage, "output_tokens", None) or 0

    if cache is not None:
        cache.put(checksum, model, text)

    return PageTranscript(
        text=text,
        checksum=checksum,
        cached=False,
        input_tokens=input_tokens,
        output_tokens=output_tokens,
    )


@dataclasses.dataclass
class RunStats:
    pages: int = 0
    cached_pages: int = 0
    input_tokens: int = 0
    output_tokens: int = 0

    @property
    def billed_pages(self) -> int:
        return self.pages - self.cached_pages


class ClaudePageReader:
    """A `textsource.PageReader` that transcribes a report's rendered pages with a vision model and joins
    them into one report text. Registering it with `textsource.set_default_page_reader(...)` makes the
    ingest path read pages through it; `scripts/reocr_pages.py` uses it directly. Token cost and cache
    hits from the most recent `read` are on `last_run` for a caller that reports cost."""

    def __init__(
        self,
        *,
        client=None,
        model: str = REOCR_MODEL,
        cache: Optional[PageTranscriptCache] = None,
        max_tokens: int = DEFAULT_MAX_TOKENS,
        page_separator: str = PAGE_SEPARATOR,
    ) -> None:
        self._client = client
        self._model = model
        self._cache = cache
        self._max_tokens = max_tokens
        self._page_separator = page_separator
        self.last_run = RunStats()

    def read(self, page_images: list[Path]) -> str:
        self.last_run = RunStats()
        parts: list[str] = []
        for path in page_images:
            transcript = transcribe_page(
                path,
                client=self._client,
                model=self._model,
                cache=self._cache,
                max_tokens=self._max_tokens,
            )
            self.last_run.pages += 1
            if transcript.cached:
                self.last_run.cached_pages += 1
            self.last_run.input_tokens += transcript.input_tokens
            self.last_run.output_tokens += transcript.output_tokens
            if transcript.text:
                parts.append(transcript.text)
        if self._cache is not None and self._cache.dirty:
            self._cache.save()
        return self._page_separator.join(parts)