src/endopath/textsource.py

"""Pluggable text source: how the ingest path obtains a report's text.

The ingest path selects a report's text through a `TextSource`, never a
hardcoded reader (docs/prd.md section 7, "Two points plug out"). Three
implementations ship:

* `CorpusOcrTextSource` returns the TCGA corpus recognition row for a case.
* `PdfTextLayerTextSource` reads the report's own embedded PDF text layer (#45),
  for a folder source that carries no corpus OCR row.
* `VisionLlmTextSource` reads the rendered pages through a vision model, the MVP
  path for a report with no usable recognized text.

`VisionLlmTextSource` delegates the model call to a `PageReader`, a
provider-agnostic callable. No call site here names a specific provider: the
endpoint selects a source by id and calls `text_for`. A local, on-premise model
for privacy plugs in by implementing `PageReader` and registering it with
`set_default_page_reader`, changing nothing at the call site.
"""

from __future__ import annotations

import dataclasses
from pathlib import Path
from typing import Callable, Optional, Protocol, runtime_checkable

# Stable ids for each source, so a call site (the Intake endpoint) and a stored
# job record name a source by identity rather than by provider. The MVP default
# reads the rendered pages with a vision model.
CORPUS_OCR = "corpus_ocr"
PDF_TEXT_LAYER = "pdf_text_layer"
VISION_LLM = "vision_llm"

DEFAULT_SOURCE_ID = VISION_LLM

SOURCE_LABELS: dict[str, str] = {
    CORPUS_OCR: "TCGA corpus OCR",
    PDF_TEXT_LAYER: "Report PDF text layer",
    VISION_LLM: "Vision model read of rendered pages",
}


@dataclasses.dataclass(frozen=True)
class TextRequest:
    """One case's inputs, enough for any source to resolve its text without
    knowing which source will run. `corpus_text` is the recognition row when the
    case came from the prepared corpus; `page_images` are the rendered pages a
    PDF-text-layer or vision source reads; `pdf_path` is the original scan."""

    case_barcode: str
    patient_filename: Optional[str] = None
    corpus_text: Optional[str] = None
    page_images: tuple[Path, ...] = ()
    pdf_path: Optional[Path] = None


@runtime_checkable
class PageReader(Protocol):
    """Turns a report's rendered pages into text. The seam a local model plugs
    into: an on-premise reader implements this one method and never surfaces its
    provider to the call site."""

    def read(self, page_images: list[Path]) -> str: ...


@runtime_checkable
class TextSource(Protocol):
    """Resolves the text the extraction pass reads for one case. Returns None
    when this source cannot supply text for the case, so the caller records a
    miss rather than a fabricated value (CLAUDE.md invariant 2)."""

    id: str
    label: str

    def text_for(self, request: TextRequest) -> Optional[str]: ...


class CorpusOcrTextSource:
    """The prepared TCGA corpus recognition text, carried on the request. The
    default text source for a case ingested from the prepared corpus."""

    id = CORPUS_OCR
    label = SOURCE_LABELS[CORPUS_OCR]

    def text_for(self, request: TextRequest) -> Optional[str]:
        text = request.corpus_text
        if text is None:
            return None
        text = str(text).strip()
        return text or None


class PdfTextLayerTextSource:
    """The report's own embedded PDF text layer (#45). A folder source carries
    no corpus OCR row, so this reads the text a born-digital or recognition-
    embedded PDF already holds. Returns None when the scan has no text layer,
    which is the common case for an image-only scan and the reason the vision
    source exists."""

    id = PDF_TEXT_LAYER
    label = SOURCE_LABELS[PDF_TEXT_LAYER]

    def text_for(self, request: TextRequest) -> Optional[str]:
        if request.pdf_path is None or not Path(request.pdf_path).exists():
            return None
        import fitz  # PyMuPDF, already an ingestion dependency

        doc = fitz.open(request.pdf_path)
        try:
            pages = [doc.load_page(i).get_text().strip() for i in range(doc.page_count)]
        finally:
            doc.close()
        joined = "\n\n".join(page for page in pages if page).strip()
        return joined or None


class VisionLlmTextSource:
    """A vision model read of the rendered pages, the MVP path for a report with
    no usable recognized text. The model call goes through a `PageReader`, so
    this class names no provider. With no reader configured it returns None
    rather than raising, keeping ingestion robust until a reader (hosted or
    local) is wired."""

    id = VISION_LLM
    label = SOURCE_LABELS[VISION_LLM]

    def __init__(self, page_reader: Optional[PageReader] = None) -> None:
        self._page_reader = page_reader

    def text_for(self, request: TextRequest) -> Optional[str]:
        reader = self._page_reader or _default_page_reader
        if reader is None or not request.page_images:
            return None
        text = reader.read(list(request.page_images))
        text = str(text).strip() if text else ""
        return text or None


# The default reader a `VisionLlmTextSource` uses when none is passed. Left unset
# so the MVP degrades to "no text from this source" rather than crashing when no
# model is configured; an operator wires one (hosted or local) once, here.
_default_page_reader: Optional[PageReader] = None


def set_default_page_reader(reader: Optional[PageReader]) -> None:
    """Register the reader every `VisionLlmTextSource` uses by default. The one
    place a provider is named: a local on-premise model registers here and every
    call site keeps working unchanged."""
    global _default_page_reader
    _default_page_reader = reader


def select(source_id: str, *, page_reader: Optional[PageReader] = None) -> TextSource:
    """The text source for an id. `page_reader` overrides the vision default for
    this selection, which is how a caller injects a specific reader (a local
    model, or a test fake) without touching the registry."""
    if source_id == CORPUS_OCR:
        return CorpusOcrTextSource()
    if source_id == PDF_TEXT_LAYER:
        return PdfTextLayerTextSource()
    if source_id == VISION_LLM:
        return VisionLlmTextSource(page_reader=page_reader)
    raise ValueError(f"unknown text source id: {source_id!r} (known: {sorted(SOURCE_LABELS)})")


def available_sources() -> list[dict[str, str]]:
    """The selectable sources, id + label, in the order the Intake screen offers
    them: corpus OCR, PDF text layer, then the vision MVP default last."""
    order = (CORPUS_OCR, PDF_TEXT_LAYER, VISION_LLM)
    return [{"id": sid, "label": SOURCE_LABELS[sid]} for sid in order]


# A `PageReader` built from a plain callable, for wiring a function (or a test
# fake) without a class. `set_default_page_reader(callable_reader(fn))`.
def callable_reader(fn: Callable[[list[Path]], str]) -> PageReader:
    class _CallableReader:
        def read(self, page_images: list[Path]) -> str:
            return fn(page_images)

    return _CallableReader()