src/endopath/llm_extraction.py

"""LLM-based extraction for the CAP/ICCR checklist fields that require
interpreting narrative language (docs/prd.md section 7: "A language model
handles fields that require interpreting narrative language: grade, margin
status, lymphovascular invasion.").

The extraction tool schema is built from the generated field specs
(`endopath.fields`, the single source of truth `GET /api/fields` serves), never
a hardcoded field list. `build_checklist_tool` derives the tool's parameters
from a set of value schemas, so changing the confirmed mapping's field set (by
regenerating `fields.py`) or handing a different dictionary's value schemas to
`build_checklist_tool` changes what the model extracts with no edit here (#96,
#99).

Uses a forced tool choice as the schema-constrained decoding mechanism section 7
calls for ("should be used regardless of which model is chosen, to cut malformed
output"): the model cannot return anything except a structurally valid call to
the checklist tool. The model is reached through an `ExtractionModel` seam, so a
local, on-premise model swaps in for privacy without touching the extraction
call site (#99). The hosted default is `AnthropicExtractionModel`.

Every field comes back at NEEDS_REVIEW, never auto-confirmed: a person confirms
every value before it counts as data (docs/prd.md section 8.3).

Hallucination guard: each extracted value must come with an evidence_quote
the model claims is verbatim from the report. That claim is checked here --
if the quote doesn't actually appear in the source text, the value is kept
(a human still reviews it) but its confidence is forced down and the
evidence is marked unverified, rather than silently trusting an
unverifiable citation.
"""

from __future__ import annotations

from typing import Any, Optional, Protocol, runtime_checkable

import anthropic

from endopath.fields import ENUM_FIELDS, FIELD_VALUE_SCHEMAS
from endopath.schema import (
    ChecklistField,
    Case,
    EvidenceSpan,
    FieldStatus,
)

MODEL = "claude-sonnet-5"

# Both derive from endopath.fields.FIELDS, the single source of truth (#37). The
# enums are registry-driven (#38): a value set changed in the registry flows into
# the tool's allowed values and the coercion below with no edit here. Kept under
# their original private names so existing call sites and tests in this module
# (and resolution.py) read unchanged.
_ENUM_FIELDS = ENUM_FIELDS
_FIELD_VALUE_SCHEMAS = FIELD_VALUE_SCHEMAS

TOOL_NAME = "extract_endometrial_checklist"
TOOL_DESCRIPTION = (
    "Record extracted CAP/ICCR endometrial carcinoma checklist fields from a "
    "pathology report."
)


def _field_schema(value_schema: dict) -> dict:
    return {
        "type": "object",
        "properties": {
            "value": {"anyOf": [value_schema, {"type": "null"}]},
            "confidence": {"type": "number", "minimum": 0, "maximum": 1},
            "evidence_quote": {
                "anyOf": [{"type": "string"}, {"type": "null"}],
                "description": "Exact, verbatim substring copied from the report text "
                "that justifies this value. Null if value is null.",
            },
        },
        "required": ["value", "confidence", "evidence_quote"],
    }


def build_checklist_tool(
    field_value_schemas: Optional[dict[str, dict]] = None,
    *,
    name: str = TOOL_NAME,
    description: str = TOOL_DESCRIPTION,
) -> dict:
    """The extraction tool, derived from a set of field value schemas.

    `field_value_schemas` maps a field name to its JSON Schema value type, the
    same shape `endopath.fields.FIELD_VALUE_SCHEMAS` exposes and `GET /api/fields`
    serves. It defaults to the generated checklist set, so a field added or an
    enum changed in the specs flows into the tool with no edit here. Passing a
    different mapping (a different dictionary's fields) builds a different tool,
    which is how a dictionary change alters the extracted field set (#99).
    """
    schemas = field_value_schemas if field_value_schemas is not None else _FIELD_VALUE_SCHEMAS
    return {
        "name": name,
        "description": description,
        "input_schema": {
            "type": "object",
            "properties": {fname: _field_schema(vs) for fname, vs in schemas.items()},
            "required": list(schemas.keys()),
        },
    }


# The default tool, built from the generated checklist set. Consumers that want
# the built-in checklist read this; a caller extracting against a different field
# set calls build_checklist_tool with that set's value schemas.
CHECKLIST_TOOL = build_checklist_tool()

_PROMPT_TEMPLATE = """\
You are extracting structured data from a pathology report for the CAP/ICCR \
endometrial carcinoma synoptic checklist. Read the report text below and call \
{tool_name} with your findings.

For each field: only set a value if the report actually supports it; leave \
value and evidence_quote null rather than guessing when the report doesn't \
say. evidence_quote must be an exact, verbatim substring copied from the \
report -- do not paraphrase or summarize it.

REPORT TEXT:
{report_text}
"""


def get_client():
    """The LLM client every call site reaches the model through, from the provider seam
    (`endopath.llm`). The default is the hosted Anthropic API; `LLM_PROVIDER=openai_compatible` with
    `LLM_BASE_URL` swaps in a local, on-premise model. This stays the import path the call sites already
    use (induction, crosswalk, the two chat co-pilots, validation, resolution), so routing every call
    through the seam changed nothing at those sites.

    Config is read lazily inside the seam, so importing this module still leaks no key and leaves the
    skipif(not ANTHROPIC_API_KEY) real-API tests skipped in an ordinary run."""
    from endopath import llm

    return llm.get_client()


# --- The model seam --------------------------------------------------------
# Extraction reaches the model through an ExtractionModel, a provider-agnostic
# seam. The forced, schema-constrained tool call lives inside the implementation,
# so a local on-premise model (or a test fake) plugs in by implementing one
# method, and swapping it needs no change where extract_raw is called (#99). The
# call site names no provider; it hands the seam the tool and the prompt.


@runtime_checkable
class ExtractionModel(Protocol):
    """Runs one report's extraction: given the checklist tool and the prompt,
    returns the assembled tool input (the `{field: {value, confidence,
    evidence_quote}}` mapping). The seam a local, on-premise model or a test fake
    plugs into; it never surfaces its provider to the call site."""

    def extract(self, *, tool: dict, prompt: str) -> dict: ...


class AnthropicExtractionModel:
    """The hosted extraction model: a forced tool call is the schema-constrained
    decoding mechanism, so the model cannot return anything except a structurally
    valid checklist. Names no provider at the call site (it is reached through
    the ExtractionModel seam)."""

    def __init__(
        self,
        client: Optional[anthropic.Anthropic] = None,
        model: str = MODEL,
        # A full nine-field extraction, each with a verbatim evidence_quote
        # (regional-lymph-node and gross-description quotes can be long),
        # regularly needs well over 2048 tokens. Too small a budget truncates
        # the tool call mid-JSON, and the assembled tool input comes back empty
        # or partial -- which silently produced all-null extractions (and, when
        # a field was cut mid-object, a bare-string field value downstream).
        max_tokens: int = 4096,
    ) -> None:
        self._client = client
        self._model = model
        self._max_tokens = max_tokens

    def extract(self, *, tool: dict, prompt: str) -> dict:
        client = self._client or get_client()
        response = client.messages.create(
            model=self._model,
            max_tokens=self._max_tokens,
            tools=[tool],
            tool_choice={"type": "tool", "name": tool["name"]},
            messages=[{"role": "user", "content": prompt}],
        )
        # Fail loudly on truncation rather than returning a partial/empty tool
        # call that would be silently persisted as an all-not-found case.
        if response.stop_reason == "max_tokens":
            raise RuntimeError(
                "LLM extraction hit max_tokens before completing the tool call; "
                "the result would be truncated. Raise max_tokens."
            )
        tool_call = next(block for block in response.content if block.type == "tool_use")
        return tool_call.input


_EXTRACTION_MODEL: Optional[ExtractionModel] = None


def configure_extraction_model(model_impl: Optional[ExtractionModel]) -> None:
    """Register the extraction model every call uses by default, or reset to the
    hosted default with None. The one place a provider is named: a local
    on-premise model (or a test fake) registers here and every call site keeps
    working unchanged (#99)."""
    global _EXTRACTION_MODEL
    _EXTRACTION_MODEL = model_impl


def _resolve_extraction_model(
    client: Optional[anthropic.Anthropic], model: str, model_impl: Optional[ExtractionModel]
) -> ExtractionModel:
    """The model for one extraction call. An explicit `model_impl` wins; then an
    explicit Anthropic `client` (the back-compatible injection the pipeline and
    tests use); then the globally configured seam; then the hosted default."""
    if model_impl is not None:
        return model_impl
    if client is not None:
        return AnthropicExtractionModel(client=client, model=model)
    if _EXTRACTION_MODEL is not None:
        return _EXTRACTION_MODEL
    return AnthropicExtractionModel(model=model)


def extract_raw(
    report_text: str,
    client: Optional[anthropic.Anthropic] = None,
    model: str = MODEL,
    *,
    model_impl: Optional[ExtractionModel] = None,
    field_value_schemas: Optional[dict[str, dict]] = None,
) -> dict:
    """Extract one report against a field set, returning the model's raw tool
    input. The tool is built from `field_value_schemas` (the generated checklist
    set by default), and the model is reached through the ExtractionModel seam."""
    tool = build_checklist_tool(field_value_schemas)
    prompt = _PROMPT_TEMPLATE.format(tool_name=tool["name"], report_text=report_text)
    return _resolve_extraction_model(client, model, model_impl).extract(tool=tool, prompt=prompt)


def _normalize_for_match(text: str) -> tuple[str, list[int]]:
    """Fold OCR noise for evidence verification (#213): lower-case, and collapse every run of
    non-alphanumeric characters (whitespace, stray punctuation, fragmented-table gaps) to one space.

    Returns the normalized text and an index map, where `index_map[i]` is the original index the i-th
    normalized character came from, so a match found in the normalized text resolves back to a real offset
    in the original for the evidence highlight. This tolerates OCR artifacts, not wrong quotes: a quote
    whose words are genuinely absent from the page still will not be found.
    """
    norm: list[str] = []
    index_map: list[int] = []
    prev_sep = True  # skip leading separators so the normalized text has no leading space
    for i, ch in enumerate(text):
        if ch.isalnum():
            norm.append(ch.lower())
            index_map.append(i)
            prev_sep = False
        elif not prev_sep:
            norm.append(" ")
            index_map.append(i)
            prev_sep = True
    return "".join(norm), index_map


def build_checklist_fields(
    llm_output: dict,
    report_text: str,
    *,
    field_value_schemas: Optional[dict[str, dict]] = None,
    enum_fields: Optional[dict[str, Any]] = None,
) -> dict[str, ChecklistField]:
    """Turn one model tool call into ChecklistFields, over the same field set the
    tool was built from. `field_value_schemas` names the fields to read (the
    generated set by default); `enum_fields` names which of them coerce back to a
    Python Enum, so a dictionary-driven field set parses against its own enums."""
    schemas = field_value_schemas if field_value_schemas is not None else _FIELD_VALUE_SCHEMAS
    enums = enum_fields if enum_fields is not None else _ENUM_FIELDS
    fields: dict[str, ChecklistField] = {}

    for name in schemas:
        raw = llm_output.get(name)
        # Defensive: schema-constrained decoding should make every field a
        # {value, confidence, evidence_quote} object, but don't crash the whole
        # extraction if the model (or a truncated tool call) hands back a
        # non-object for one field -- treat it as not-extracted, which a human
        # reviews anyway, rather than guessing at a malformed value.
        if not isinstance(raw, dict):
            raw = {}
        value = raw.get("value")
        confidence = raw.get("confidence")
        quote = raw.get("evidence_quote")

        if name in enums and value is not None:
            try:
                value = enums[name](value)
            except ValueError:
                # The model returned something outside its own schema's enum
                # (shouldn't happen with tool_choice forcing, but don't trust
                # a value we can't actually represent).
                value = None
                confidence = None
                quote = None

        evidence = None
        if quote:
            char_start: Optional[int] = None
            char_end: Optional[int] = None
            if quote in report_text:
                verified = True
                char_start = report_text.find(quote)
                char_end = char_start + len(quote)
            else:
                # OCR-tolerant fallback (#213): the model reads the page image, so a quote that is correct
                # will not exact-match OCR text mangled by stray punctuation or fragmented tables. Match
                # modulo that noise, then map the hit back to a real offset for the highlight. A quote whose
                # words are genuinely absent still fails, so the invariant tolerates OCR artifacts, not
                # wrong quotes.
                norm_report, index_map = _normalize_for_match(report_text)
                norm_quote = _normalize_for_match(quote)[0].strip()
                pos = norm_report.find(norm_quote) if norm_quote else -1
                verified = pos != -1
                if verified:
                    char_start = index_map[pos]
                    char_end = index_map[pos + len(norm_quote) - 1] + 1
            evidence = EvidenceSpan(
                quote=quote,
                char_start=char_start,
                char_end=char_end,
                source="llm" if verified else "llm_unverified_quote",
            )
            if not verified and confidence is not None:
                # Don't silently trust a citation that isn't actually in the
                # source text -- a human still sees the value, but flagged
                # as much less trustworthy than the model claimed.
                confidence = min(confidence, 0.2)

        fields[name] = ChecklistField(
            value=value,
            confidence=confidence,
            evidence=evidence,
            status=FieldStatus.NEEDS_REVIEW,
        )

    return fields


def extract_case_with_llm(
    case: Case,
    report_text: str,
    client: Optional[anthropic.Anthropic] = None,
    *,
    model_impl: Optional[ExtractionModel] = None,
) -> dict[str, ChecklistField]:
    """Populate case.checklist from an LLM extraction pass. Fields are set
    onto the case in place; applicability rules (#3, #43) are re-applied so an
    LLM-guessed grade on a non-endometrioid/mucinous histotype, or a
    molecular classification on a pre-2023 report, gets marked not-applicable.
    The drafted value and its evidence are preserved under that state (#43),
    so a reviewer who corrects the gating histotype gets the draft back.

    This is the extraction step the upload/process path invokes (#25, pipeline.py).
    `report_text` is already resolved: the ingest path reads it through the
    pluggable text source (#90, `endopath.textsource`), so swapping the text
    source needs no change here. `model_impl` (or `configure_extraction_model`)
    swaps the model the same way.
    """
    # Forward model_impl only when set, so a test (or caller) that replaced
    # extract_raw with a narrower-signature stub still works.
    extra = {"model_impl": model_impl} if model_impl is not None else {}
    # Tag this case's model calls, so their provenance records link to it (#116).
    from endopath import llm

    with llm.call_context(f"case:{case.case_barcode}"):
        llm_output = extract_raw(report_text, client=client, **extra)
    fields = build_checklist_fields(llm_output, report_text)
    for name, field in fields.items():
        setattr(case.checklist, name, field)
    case.checklist.apply_applicability_rules(case.staging_edition)
    # Re-read from the checklist: apply_applicability_rules mutates the gated
    # fields' state in place, so the returned dict reflects any not-applicable
    # marking it made.
    return case.checklist.all_fields()