src/endopath/executors.py

"""Executors for the Intake follow-up worker (`worker.py`).

Each executor runs one enqueued job to completion and writes a `submission_artifacts` row that the
mapping surface and later projection read back. report_induction (#107) induces each ingested report's
reporting schema. dictionary_compilation (#108) drafts the submitted dictionary's crosswalk.

Both reach their model through an injectable seam (the `SchemaInducer` / `CrosswalkDrafter` Protocols),
so a test substitutes a deterministic one and the fast suite spends no key. Importing this module
registers the executors with the worker.
"""

from __future__ import annotations

import json
import time
from datetime import datetime, timezone
from typing import Any, Callable, Optional

from sqlalchemy import Connection

from endopath import crosswalk, dictionary, induction, jobs, mapping, storage, textsource, worker
from endopath.schema import CaseStatus

REPORT_INDUCTION_ARTIFACT = "report_induction"
CROSSWALK_DRAFT_ARTIFACT = "crosswalk_draft"

# The model seams, overridable for tests. None means the hosted default (the pass constructs an
# Anthropic drafter or inducer), so production spends the key and a test never does.
_inducer: Optional[induction.SchemaInducer] = None
_drafter: Optional[crosswalk.CrosswalkDrafter] = None


def configure_inducer(inducer: Optional[induction.SchemaInducer]) -> None:
    """Substitute the schema inducer, for tests. None restores the hosted default."""
    global _inducer
    _inducer = inducer


def configure_drafter(drafter: Optional[crosswalk.CrosswalkDrafter]) -> None:
    """Substitute the crosswalk drafter, for tests. None restores the hosted default."""
    global _drafter
    _drafter = drafter


def report_text_for(patient_filename: Optional[str]) -> Optional[str]:
    """The report's recognized text, looked up the way the extraction pipeline looks it up, so an
    induced report and its extracted checklist read the same text. Imported lazily to keep this module's
    import light (the pipeline pulls the retrieval and extraction stack)."""
    from endopath import pipeline

    return pipeline.report_text_for_case(patient_filename)


def report_induction(conn: Connection, job: dict) -> None:
    """Induce the reporting schema of each ingested report and persist the contracts, keyed by the
    submission's dictionary id (#92, #107). A report whose text cannot be read, or whose induced
    contract violates a rule, is recorded in the artifact's `errors` rather than failing the whole
    family, so one bad report does not sink the job."""
    dictionary_id = job["dictionary_id"]
    registry = induction.load_registry()
    standards = induction.known_standards()

    reports: list[dict] = []
    errors: list[dict] = []
    for row in storage.list_cases(conn):
        request = textsource.TextRequest(
            case_barcode=row["case_barcode"],
            patient_filename=row["patient_filename"],
            corpus_text=report_text_for(row["patient_filename"]),
        )
        try:
            report = induction.induce_report(
                request,
                source_id=textsource.CORPUS_OCR,
                inducer=_inducer,
                registry=registry,
                standards=standards,
            )
            reports.append(report.model_dump(mode="json"))
            _persist_induced_observations(conn, row["case_barcode"], report)
        except induction.InductionError as exc:
            errors.append({"case_barcode": row["case_barcode"], "error": str(exc)})

    payload = {"source_id": textsource.CORPUS_OCR, "reports": reports, "errors": errors}
    storage.upsert_submission_artifact(
        conn,
        dictionary_id=dictionary_id,
        kind=REPORT_INDUCTION_ARTIFACT,
        payload_json=json.dumps(payload),
    )


def _persist_induced_observations(
    conn: Connection, case_barcode: str, report: induction.InducedReport
) -> None:
    """Promote one induced contract to the durable observation atom (#111): one observations row per
    InducedField, beside the checklist observations in the same table, built through the shared
    `storage.Observation.from_induced_field` shape. The strict-JSON artifact stays the reviewable
    contract; the table is the queryable atom. Attached to the case's resolved report; an ambiguous
    multi-report case is skipped rather than guessed, so a bad grain never invents a provenance."""
    report_id = storage.resolved_report_id(conn, case_barcode)
    if report_id is None:
        return
    now = datetime.now(timezone.utc).isoformat()
    storage.delete_induced_observations(conn, case_barcode)
    for induced_field in report.fields:
        observation = storage.Observation.from_induced_field(
            induced_field, case_barcode=case_barcode, report_id=report_id, created_at=now
        )
        storage.insert_observation(conn, observation)


worker.register_executor(jobs.REPORT_INDUCTION, report_induction)


def dictionary_compilation(conn: Connection, job: dict) -> None:
    """Compile the submitted dictionary against the canonical concept list and draft an edge from every
    variable to the concept that answers it, then persist the drafted crosswalk keyed by the dictionary
    id for the mapping surface to review (#91, #93, #108).

    The built-in CAP checklist already carries a committed, reviewed reference crosswalk, so it is not
    re-drafted: its mapping surface serves crosswalk.reference_crosswalk() instead (#109)."""
    dictionary_id = job["dictionary_id"]
    if dictionary_id == dictionary.BUILTIN_DICTIONARY_ID:
        return
    data_dictionary = dictionary.get_registered(conn, dictionary_id)
    if data_dictionary is None:
        raise RuntimeError(
            f"dictionary {dictionary_id!r} was not stored, so its crosswalk cannot be drafted"
        )

    concepts = mapping.concept_list()
    concept_ids = frozenset(concept["concept_id"] for concept in concepts)
    source_fields = [variable.name for variable in data_dictionary.variables]
    # The drafter reads the field summaries; label and permitted values give it context beyond the name.
    field_summaries = [
        {"name": variable.name, "label": variable.label, "permitted_values": variable.permitted_values}
        for variable in data_dictionary.variables
    ]
    drafter = _drafter or crosswalk.AnthropicCrosswalkDrafter()
    raw = drafter.draft(source_std=dictionary_id, source_fields=field_summaries, concepts=concepts)
    drafted = crosswalk.build_crosswalk(
        raw, source_std=dictionary_id, source_fields=source_fields, concept_ids=concept_ids
    )
    storage.upsert_submission_artifact(
        conn,
        dictionary_id=dictionary_id,
        kind=CROSSWALK_DRAFT_ARTIFACT,
        payload_json=crosswalk.crosswalk_json(drafted),
    )


worker.register_executor(jobs.DICTIONARY_COMPILATION, dictionary_compilation)

COHORT_EXTRACTION_ARTIFACT = "cohort_extraction"


def _transient_llm_errors() -> tuple:
    """The transient failures a retry should ride out: a rate limit, a dropped connection, or a 5xx from
    the hosted API or a local server. Resolved lazily so importing this module stays light."""
    import anthropic
    import httpx

    return (
        anthropic.RateLimitError,
        anthropic.APIConnectionError,
        anthropic.InternalServerError,
        httpx.HTTPError,
    )


def with_retry(
    fn: Callable[[], Any],
    *,
    attempts: int = 3,
    base_delay: float = 0.5,
    transient: Optional[tuple] = None,
    sleep: Callable[[float], None] = time.sleep,
) -> Any:
    """Run fn, retrying a transient LLM failure with exponential backoff and re-raising the last one when
    the attempts run out. A non-transient error (a validation or logic failure) is not retried."""
    transient = transient if transient is not None else _transient_llm_errors()
    for attempt in range(attempts):
        try:
            return fn()
        except transient:
            if attempt == attempts - 1:
                raise
            sleep(base_delay * (2**attempt))


def cohort_extraction(conn: Connection, job: dict) -> None:
    """Extract every queued case through the shared pipeline (#115), retrying a transient LLM failure per
    case. The pipeline advances each case queued -> ready-for-review and persists it, so a reviewer picks
    it up on the worklist. A case that still fails after retries, or has no report text, is recorded in
    the artifact's errors rather than failing the whole cohort."""
    from endopath import pipeline

    extracted: list[str] = []
    errors: list[dict] = []
    for row in storage.list_cases(conn, status=CaseStatus.QUEUED.value):
        barcode = row["case_barcode"]
        try:
            with_retry(lambda: pipeline.process_case_in_store(conn, barcode))
            extracted.append(barcode)
        except Exception as exc:  # noqa: BLE001 - one case's failure is recorded, not fatal to the cohort
            errors.append({"case_barcode": barcode, "error": str(exc)})
    payload = {"extracted": extracted, "errors": errors}
    storage.upsert_submission_artifact(
        conn, dictionary_id="(cohort)", kind=COHORT_EXTRACTION_ARTIFACT, payload_json=json.dumps(payload)
    )


worker.register_executor(jobs.COHORT_EXTRACTION, cohort_extraction)