Blame for src/llm/classify.py

"""LLM classification into Narrow/Exclude with an optional Possible category."""

from __future__ import annotations

from dataclasses import dataclass, field
from collections.abc import Callable, Sequence
import json
from typing import Protocol

from src.data.loaders import EventDefinitionForm
from src.llm.client import call_jev_decisions, call_llm_json, resolve_provider
from src.llm.prompts import build_classify_prompt, classify_system_prompt
from src.utils.logging import get_logger

logger = get_logger(__name__)

VALID_LABELS = {"Narrow", "Possible", "Exclude"}


class ClassificationCandidate(Protocol):
    """Minimal candidate interface consumed by the classifier."""

    code: str
    description: str
    vocabulary: str


@dataclass(frozen=True, slots=True)
class AdaptiveStoppingConfig:
    """Stop classification after Narrow results remain sparse for several batches."""

    sparse_narrow_threshold: int = 1
    consecutive_sparse_batches: int = 3
    minimum_batches: int = 3

    def __post_init__(self) -> None:
        if self.sparse_narrow_threshold < 0:
            raise ValueError("sparse_narrow_threshold must be at least 0")
        if self.consecutive_sparse_batches < 1:
            raise ValueError("consecutive_sparse_batches must be at least 1")
        if self.minimum_batches < 1:
            raise ValueError("minimum_batches must be at least 1")


@dataclass(frozen=True, slots=True)
class ClassificationRunInfo:
    """Summary of how much of the ranked candidate list was classified."""

    candidates_available: int
    candidates_classified: int
    batches_completed: int
    narrow_counts_by_batch: tuple[int, ...]
    stopped_early: bool
    stop_reason: str


@dataclass(slots=True)
class ClassifiedCandidate:
    code: str
    description: str
    vocabulary: str
    label: str
    confidence: float
    explanation: str
    relevance_score: float = 0.0
    retrieval_score: float = 0.0
    decision_origin: str = "llm"
    mapping_provenance: list[dict] = field(default_factory=list)
    prior_label: str | None = None
    prior_confidence: float | None = None


@dataclass(frozen=True, slots=True)
class GPTReviewRunInfo:
    """Summary of final-stage GPT adjudication of uncertain classifications."""

    confidence_threshold: float
    max_candidates: int | None
    eligible_candidates: int
    candidates_reviewed: int
    labels_changed: int
    narrow_to_exclude: int
    exclude_to_narrow: int
    reasoning_effort: str


def select_gpt_review_candidates(
    classified: Sequence[ClassifiedCandidate],
    *,
    confidence_threshold: float,
    max_candidates: int | None = None,
) -> list[ClassifiedCandidate]:
    """Select the lowest-confidence eligible decisions for GPT review."""
    if not 0.0 <= confidence_threshold <= 1.0:
        raise ValueError("confidence_threshold must be between 0 and 1")
    if max_candidates is not None and max_candidates < 1:
        raise ValueError("max_candidates must be at least 1 when provided")

    eligible = [item for item in classified if item.confidence < confidence_threshold]
    eligible.sort(key=lambda item: item.confidence)
    return eligible[:max_candidates] if max_candidates is not None else eligible


def gpt_review_low_confidence(
    classified: Sequence[ClassifiedCandidate],
    edf: EventDefinitionForm,
    *,
    confidence_threshold: float = 0.7,
    max_candidates: int | None = None,
    model: str | None = None,
    max_retries: int = 3,
    batch_size: int = 10,
    progress_callback: Callable[[int, int], None] | None = None,
    use_possible_category: bool = False,
    reasoning_effort: str = "medium",
    provider: str = "openai",
) -> tuple[list[ClassifiedCandidate], GPTReviewRunInfo]:
    """Use GPT to adjudicate only classifications below the confidence cutoff.

    Results at or above the threshold are preserved exactly. Reviewed results
    replace the corresponding stage-2 labels while retaining their original
    list positions and recording the prior label/confidence in the explanation.
    """
    eligible_count = sum(item.confidence < confidence_threshold for item in classified)
    uncertain = select_gpt_review_candidates(
        classified,
        confidence_threshold=confidence_threshold,
        max_candidates=max_candidates,
    )
    if not uncertain:
        return list(classified), GPTReviewRunInfo(
            confidence_threshold=confidence_threshold,
            max_candidates=max_candidates,
            eligible_candidates=eligible_count,
            candidates_reviewed=0,
            labels_changed=0,
            narrow_to_exclude=0,
            exclude_to_narrow=0,
            reasoning_effort=reasoning_effort,
        )

    reviewed = llm_classify(
        uncertain,
        edf,
        provider=provider,
        model=model,
        max_retries=max_retries,
        batch_size=batch_size,
        progress_callback=progress_callback,
        use_possible_category=use_possible_category,
        adaptive_stopping=None,
        reasoning_effort=reasoning_effort,
    )
    reviewed_by_key = {(item.code, item.vocabulary): item for item in reviewed}

    merged: list[ClassifiedCandidate] = []
    labels_changed = 0
    narrow_to_exclude = 0
    exclude_to_narrow = 0
    for original in classified:
        replacement = reviewed_by_key.get((original.code, original.vocabulary))
        if replacement is None:
            merged.append(original)
            continue

        if replacement.label != original.label:
            labels_changed += 1
            if original.label == "Narrow" and replacement.label == "Exclude":
                narrow_to_exclude += 1
            elif original.label == "Exclude" and replacement.label == "Narrow":
                exclude_to_narrow += 1
        merged.append(
            ClassifiedCandidate(
                code=original.code,
                description=original.description,
                vocabulary=original.vocabulary,
                label=replacement.label,
                confidence=replacement.confidence,
                explanation=(
                    "GPT review of low-confidence stage-2 result "
                    f"(previous label={original.label}, confidence={original.confidence:.2f}). "
                    + replacement.explanation
                ),
                relevance_score=original.relevance_score,
                retrieval_score=original.retrieval_score,
            )
        )

    return merged, GPTReviewRunInfo(
        confidence_threshold=confidence_threshold,
        max_candidates=max_candidates,
        eligible_candidates=eligible_count,
        candidates_reviewed=len(uncertain),
        labels_changed=labels_changed,
        narrow_to_exclude=narrow_to_exclude,
        exclude_to_narrow=exclude_to_narrow,
        reasoning_effort=reasoning_effort,
    )


def llm_classify(
    candidates: Sequence[ClassificationCandidate],
    edf: EventDefinitionForm,
    provider: str = "auto",
    model: str | None = None,
    max_retries: int = 3,
    batch_size: int = 10,
    progress_callback: Callable[[int, int], None] | None = None,
    use_possible_category: bool = False,
    adaptive_stopping: AdaptiveStoppingConfig | None = None,
    run_info_callback: Callable[[ClassificationRunInfo], None] | None = None,
    reasoning_effort: str | None = None,
    thinking_budget: int | None = None,
) -> list[ClassifiedCandidate]:
    """Classify candidates as Narrow/Exclude, optionally allowing Possible.

    Parameters
    ----------
    candidates : sequence
        Similarity-ranked hybrid or LLM-ranked candidates.
    edf : EventDefinitionForm
        Structured phenotype definition providing semantic context.
    provider : str
        "auto", "anthropic", or "openai".
    model : str, optional
        Model identifier. Defaults to a provider-specific default.
    max_retries : int
        Number of LLM call attempts before failing.
    batch_size : int
        Maximum number of candidates to include in a single LLM request.
    progress_callback : callable, optional
        Callback invoked after each completed batch as ``callback(completed, total)``.
    use_possible_category : bool
        Enable three-way Narrow/Possible/Exclude labeling. Defaults to False.
    adaptive_stopping : AdaptiveStoppingConfig, optional
        Stop after Narrow results remain at or below the configured threshold
        for the configured number of consecutive batches.
    run_info_callback : callable, optional
        Invoked once with the completed/stopped run summary.
    reasoning_effort : str, optional
        OpenAI reasoning effort, such as ``"none"`` or ``"medium"``.

    Returns
    -------
    list[ClassifiedCandidate]
    """
    if not candidates:
        if run_info_callback is not None:
            run_info_callback(
                ClassificationRunInfo(0, 0, 0, (), False, "No candidates to classify.")
            )
        return []

    if batch_size < 1:
        raise ValueError("batch_size must be at least 1")

    resolved_provider = resolve_provider(provider)
    if resolved_provider in {"jev", "openai_decisions", "laya", "jpt"}:
        from src.llm.decisions import DEFAULT_DECISION_MODELS
        return _jev_classify(
            candidates,
            edf,
            model=model or DEFAULT_DECISION_MODELS[resolved_provider],
            max_retries=max_retries,
            batch_size=batch_size,
            progress_callback=progress_callback,
            use_possible_category=use_possible_category,
            adaptive_stopping=adaptive_stopping,
            run_info_callback=run_info_callback,
            provider=resolved_provider,
        )

    results: list[ClassifiedCandidate] = []
    narrow_counts: list[int] = []
    sparse_streak = 0

    total = len(candidates)
    for start in range(0, total, batch_size):
        batch = candidates[start : start + batch_size]
        batch_result_start = len(results)
        lookup = {(c.code, c.vocabulary): c for c in batch}
        payload = [
            {"code": c.code, "vocabulary": c.vocabulary, "description": c.description}
            for c in batch
        ]
        prompt = build_classify_prompt(edf, payload, use_possible_category=use_possible_category)
        response = call_llm_json(
            system_prompt=classify_system_prompt(use_possible_category),
            user_prompt=prompt,
            provider=resolved_provider,
            model=model,
            max_retries=max_retries,
            reasoning_effort=reasoning_effort,
            **({'thinking_budget':thinking_budget} if thinking_budget is not None else {}),
        )

        classifications_raw = response.get("classifications", [])
        batch_seen: set[tuple[str, str]] = set()
        for item in classifications_raw:
            key = (item.get("code", ""), item.get("vocabulary", ""))
            source = lookup.get(key)
            if source is None or key in batch_seen:
                continue
            fallback_label = "Possible" if use_possible_category else "Exclude"
            valid_labels = VALID_LABELS if use_possible_category else {"Narrow", "Exclude"}
            label = item.get("label", fallback_label)
            if label not in valid_labels:
                logger.warning(
                    "LLM returned invalid label %r for code %s; defaulting to %r",
                    label,
                    key,
                    fallback_label,
                )
                label = fallback_label
            batch_seen.add(key)
            results.append(
                ClassifiedCandidate(
                    code=source.code,
                    description=source.description,
                    vocabulary=source.vocabulary,
                    label=label,
                    confidence=float(item.get("confidence", 0.0)),
                    explanation=str(item.get("explanation", "")),
                    relevance_score=float(getattr(source, "relevance_score", 0.0)),
                    retrieval_score=float(
                        getattr(source, "retrieval_score", getattr(source, "score", 0.0))
                    ),
                )
            )

        # Keep every candidate in the result. A missing response receives the
        # least-committal enabled label and zero confidence so review selection
        # can annotate it without filtering or altering any other result.
        for c in batch:
            key = (c.code, c.vocabulary)
            if key not in batch_seen:
                results.append(
                    ClassifiedCandidate(
                        code=c.code,
                        description=c.description,
                        vocabulary=c.vocabulary,
                        label="Possible" if use_possible_category else "Exclude",
                        confidence=0.0,
                        explanation="No classification returned by LLM; routed for manual review.",
                        relevance_score=float(getattr(c, "relevance_score", 0.0)),
                        retrieval_score=float(
                            getattr(c, "retrieval_score", getattr(c, "score", 0.0))
                        ),
                    )
                )

        if progress_callback is not None:
            progress_callback(min(start + len(batch), total), total)

        narrow_count = sum(
            item.label == "Narrow" for item in results[batch_result_start:]
        )
        narrow_counts.append(narrow_count)
        sparse_streak = (
            sparse_streak + 1
            if adaptive_stopping is not None
            and narrow_count <= adaptive_stopping.sparse_narrow_threshold
            else 0
        )
        if _adaptive_stop_reached(adaptive_stopping, len(narrow_counts), sparse_streak):
            break

    _report_run_info(
        run_info_callback,
        total=total,
        results=results,
        narrow_counts=narrow_counts,
        adaptive_stopping=adaptive_stopping,
    )
    return results


def _jev_classify(
    candidates: Sequence[ClassificationCandidate],
    edf: EventDefinitionForm,
    model: str,
    max_retries: int,
    batch_size: int,
    progress_callback: Callable[[int, int], None] | None,
    use_possible_category: bool,
    adaptive_stopping: AdaptiveStoppingConfig | None,
    run_info_callback: Callable[[ClassificationRunInfo], None] | None,
    provider: str = "jev",
) -> list[ClassifiedCandidate]:
    """Classify candidates with Jev typed choice questions."""
    effective_batch_size = min(batch_size, 20)
    labels = {
        "Narrow": "The code is sufficiently specific to the phenotype and meets the Narrow criteria.",
        "Exclude": "The code is irrelevant, insufficiently specific, historical, or explicitly excluded.",
    }
    if use_possible_category:
        labels = {
            "Narrow": labels["Narrow"],
            "Possible": "The code is related but not sufficiently specific in every context.",
            "Exclude": labels["Exclude"],
        }

    results: list[ClassifiedCandidate] = []
    narrow_counts: list[int] = []
    sparse_streak = 0
    total = len(candidates)
    state = "Phenotype definition:\n" + edf.to_prompt_context(
        include_possible=use_possible_category
    )
    for start in range(0, total, effective_batch_size):
        batch = candidates[start : start + effective_batch_size]
        batch_result_start = len(results)
        questions = {
            f"candidate_{offset}": {
                "type": "choice",
                "instructions": (
                    "Classify this clinical terminology candidate for the phenotype: "
                    f"code={candidate.code}; vocabulary={candidate.vocabulary}; "
                    f"description={candidate.description}"
                ),
                "criteria": labels,
            }
            for offset, candidate in enumerate(batch)
        }
        if provider == "jev":
            decision_call = call_jev_decisions
        else:
            from src.llm.decisions import call_decisions
            decision_call = lambda **kwargs: call_decisions(provider=provider, **kwargs)
        response = decision_call(
            state=state,
            questions=questions,
            model=model,
            max_retries=max_retries,
        )
        answers = response.get("answers", {})
        for offset, candidate in enumerate(batch):
            answer = answers.get(f"candidate_{offset}", {})
            probabilities = answer.get("probabilities", {})
            fallback_label = "Possible" if use_possible_category else "Exclude"
            label = answer.get("choice", fallback_label)
            if label not in labels:
                label = fallback_label
            confidence = answer.get("confidence")
            if confidence is None:
                confidence = max((float(value) for value in probabilities.values()), default=0.0)
            results.append(
                ClassifiedCandidate(
                    code=candidate.code,
                    description=candidate.description,
                    vocabulary=candidate.vocabulary,
                    label=label,
                    confidence=float(confidence),
                    explanation=f"{provider} choice probabilities: " + json.dumps(probabilities, sort_keys=True)
                    + "; confidence source: " + answer.get("confidence_source", "provider" if answer.get("confidence") is not None else "max_probability"),
                    relevance_score=float(getattr(candidate, "relevance_score", 0.0)),
                    retrieval_score=float(
                        getattr(
                            candidate,
                            "retrieval_score",
                            getattr(candidate, "score", 0.0),
                        )
                    ),
                )
            )
        if progress_callback is not None:
            progress_callback(min(start + len(batch), total), total)
        narrow_count = sum(
            item.label == "Narrow" for item in results[batch_result_start:]
        )
        narrow_counts.append(narrow_count)
        sparse_streak = (
            sparse_streak + 1
            if adaptive_stopping is not None
            and narrow_count <= adaptive_stopping.sparse_narrow_threshold
            else 0
        )
        if _adaptive_stop_reached(adaptive_stopping, len(narrow_counts), sparse_streak):
            break

    _report_run_info(
        run_info_callback,
        total=total,
        results=results,
        narrow_counts=narrow_counts,
        adaptive_stopping=adaptive_stopping,
    )
    return results


def _adaptive_stop_reached(
    config: AdaptiveStoppingConfig | None,
    batches_completed: int,
    sparse_streak: int,
) -> bool:
    return bool(
        config is not None
        and batches_completed >= config.minimum_batches
        and sparse_streak >= config.consecutive_sparse_batches
    )


def _report_run_info(
    callback: Callable[[ClassificationRunInfo], None] | None,
    *,
    total: int,
    results: list[ClassifiedCandidate],
    narrow_counts: list[int],
    adaptive_stopping: AdaptiveStoppingConfig | None,
) -> None:
    if callback is None:
        return
    stopped_early = len(results) < total
    if stopped_early and adaptive_stopping is not None:
        reason = (
            f"Stopped after {adaptive_stopping.consecutive_sparse_batches} consecutive "
            f"batches with at most {adaptive_stopping.sparse_narrow_threshold} Narrow "
            "result(s) per batch."
        )
    else:
        reason = "All available candidates were classified."
    callback(
        ClassificationRunInfo(
            candidates_available=total,
            candidates_classified=len(results),
            batches_completed=len(narrow_counts),
            narrow_counts_by_batch=tuple(narrow_counts),
            stopped_early=stopped_early,
            stop_reason=reason,
        )
    )