src/gcmd_classifier/classification/routers.py

"""Topic and Term routing through the provider-neutral model interface."""

from __future__ import annotations

from pydantic import BaseModel, ConfigDict, Field

from gcmd_classifier.classification.candidates import (
    TermCandidate,
    TopicCandidate,
    build_term_candidates,
    build_topic_candidates,
    validate_term_candidate_relationship,
)
from gcmd_classifier.config import ModelSettings
from gcmd_classifier.errors import UnknownCandidateIDError
from gcmd_classifier.llm.base import (
    ModelClient,
    ModelRequest,
    ModelStage,
    RetryPolicy,
    generate_with_retries,
)
from gcmd_classifier.llm.prompts import ParentContext, build_term_prompt, build_topic_prompt
from gcmd_classifier.llm.schemas import CandidateDecision, TermResponse, TopicResponse
from gcmd_classifier.models import ArticleRecord, OutputError, OutputWarning, SupportType
from gcmd_classifier.vocabulary.index import VocabularyIndex


class ModelCallMetadata(BaseModel):
    """Safe model-call metadata retained in routing results."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    provider: str
    model_name: str
    prompt_version: str
    retry_count: int = Field(ge=0)
    duration_seconds: float | None = Field(default=None, ge=0.0)
    input_tokens: int | None = Field(default=None, ge=0)
    output_tokens: int | None = Field(default=None, ge=0)
    total_tokens: int | None = Field(default=None, ge=0)
    estimated_cost: float | None = Field(default=None, ge=0.0)


class TopicBranchSeed(BaseModel):
    """Seed for a later Term-routing branch under one selected Topic."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    branch_id: str = Field(min_length=1)
    topic_uuid: str = Field(min_length=1)
    topic_name: str = Field(min_length=1)
    topic_level: str = Field(pattern="^Topic$")
    topic_canonical_path: str = Field(min_length=1)
    evidence: str = Field(min_length=1)
    support_type: SupportType
    confidence: float | None = Field(default=None, ge=0.0, le=1.0)
    reason: str | None = None
    candidate_id: str = Field(min_length=1)
    prompt_version: str
    model_provider: str
    model_name: str
    retry_count: int = Field(ge=0)


class TopicRoutingResult(BaseModel):
    """Result of Topic routing for one article."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    branches: tuple[TopicBranchSeed, ...] = Field(default_factory=tuple)
    no_selection_reason: str | None = None
    model_metadata: ModelCallMetadata
    warnings: tuple[OutputWarning, ...] = Field(default_factory=tuple)
    errors: tuple[OutputError, ...] = Field(default_factory=tuple)

    @property
    def selected_count(self) -> int:
        """Number of selected Topic branches."""
        return len(self.branches)

    @property
    def is_no_topic(self) -> bool:
        """Whether routing completed successfully with no selected Topics."""
        return not self.branches


class TermBranchSeed(BaseModel):
    """Seed for a later Variable-routing branch under one selected Term."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    branch_id: str = Field(min_length=1)
    parent_topic_uuid: str = Field(min_length=1)
    parent_topic_name: str = Field(min_length=1)
    term_uuid: str = Field(min_length=1)
    term_name: str = Field(min_length=1)
    term_level: str = Field(pattern="^Term$")
    term_canonical_path: str = Field(min_length=1)
    evidence: str = Field(min_length=1)
    support_type: SupportType
    confidence: float | None = Field(default=None, ge=0.0, le=1.0)
    reason: str | None = None
    candidate_id: str = Field(min_length=1)
    parent_branch_id: str = Field(min_length=1)
    prompt_version: str
    model_provider: str
    model_name: str
    retry_count: int = Field(ge=0)


class TopicStopResult(BaseModel):
    """Successful Term-routing stop at the selected Topic parent."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    branch_id: str = Field(min_length=1)
    topic_uuid: str = Field(min_length=1)
    topic_name: str = Field(min_length=1)
    topic_level: str = Field(pattern="^Topic$")
    topic_canonical_path: str = Field(min_length=1)
    stop_reason: str = Field(min_length=1)
    parent_evidence: str = Field(min_length=1)
    parent_support_type: SupportType
    parent_confidence: float | None = Field(default=None, ge=0.0, le=1.0)
    topic_candidate_id: str = Field(min_length=1)
    prompt_version: str
    model_provider: str
    model_name: str
    retry_count: int = Field(ge=0)


class TermRoutingResult(BaseModel):
    """Result of Term routing under one selected Topic branch."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    term_branches: tuple[TermBranchSeed, ...] = Field(default_factory=tuple)
    stop_at_topic: TopicStopResult | None = None
    model_metadata: ModelCallMetadata
    warnings: tuple[OutputWarning, ...] = Field(default_factory=tuple)
    errors: tuple[OutputError, ...] = Field(default_factory=tuple)

    @property
    def selected_count(self) -> int:
        """Number of selected Term branches."""
        return len(self.term_branches)

    @property
    def stopped_at_topic(self) -> bool:
        """Whether routing stopped successfully at the parent Topic."""
        return self.stop_at_topic is not None


def route_topics(
    *,
    article: ArticleRecord,
    vocabulary: VocabularyIndex,
    model_client: ModelClient,
    settings: ModelSettings,
    retry_policy: RetryPolicy | None = None,
) -> TopicRoutingResult:
    """Route an article to zero, one, or more valid GCMD Topics."""
    candidates = build_topic_candidates(vocabulary)
    candidates_by_id = {candidate.candidate_id: candidate for candidate in candidates}
    prompt = build_topic_prompt(
        article=article,
        candidates=tuple(candidate.prompt_candidate for candidate in candidates),
        prompt_version=settings.prompt_version_topic,
    )
    request = ModelRequest.from_settings(
        stage=ModelStage.TOPIC,
        prompt=prompt,
        response_schema=TopicResponse,
        settings=settings,
        metadata={
            "DOI": article.DOI,
            "candidate_ids": tuple(candidates_by_id),
        },
    )
    response = generate_with_retries(
        model_client,
        request,
        retry_policy or RetryPolicy.from_settings(settings),
    )
    _validate_selected_candidate_ids(response.parsed.selected, candidates_by_id, stage="Topic")
    metadata = _model_metadata(response)
    branches = tuple(
        _topic_branch_seed(
            decision=decision,
            candidate=candidates_by_id[decision.candidate_id],
            vocabulary=vocabulary,
            metadata=metadata,
        )
        for decision in response.parsed.selected
    )
    return TopicRoutingResult(
        branches=branches,
        no_selection_reason=response.parsed.no_selection_reason if not branches else None,
        model_metadata=metadata,
    )


def route_terms(
    *,
    article: ArticleRecord,
    topic_branch: TopicBranchSeed,
    vocabulary: VocabularyIndex,
    model_client: ModelClient,
    settings: ModelSettings,
    retry_policy: RetryPolicy | None = None,
) -> TermRoutingResult:
    """Route one selected Topic branch to zero, one, or more direct Term children."""
    candidates = build_term_candidates(vocabulary, topic_uuid=topic_branch.topic_uuid)
    candidates_by_id = {candidate.candidate_id: candidate for candidate in candidates}
    parent_context = ParentContext(
        candidate_id=topic_branch.candidate_id,
        name=topic_branch.topic_name,
        level="Topic",
        canonical_path=topic_branch.topic_canonical_path,
    )
    prompt = build_term_prompt(
        article=article,
        parent=parent_context,
        candidates=tuple(candidate.prompt_candidate for candidate in candidates),
        prompt_version=settings.prompt_version_term,
    )
    request = ModelRequest.from_settings(
        stage=ModelStage.TERM,
        prompt=prompt,
        response_schema=TermResponse,
        settings=settings,
        metadata={
            "DOI": article.DOI,
            "parent_topic_uuid": topic_branch.topic_uuid,
            "parent_branch_id": topic_branch.branch_id,
            "candidate_ids": tuple(candidates_by_id),
        },
    )
    response = generate_with_retries(
        model_client,
        request,
        retry_policy or RetryPolicy.from_settings(settings),
    )
    _validate_selected_candidate_ids(response.parsed.selected, candidates_by_id, stage="Term")
    metadata = _model_metadata(response)
    term_branches = tuple(
        _term_branch_seed(
            decision=decision,
            candidate=candidates_by_id[decision.candidate_id],
            topic_branch=topic_branch,
            vocabulary=vocabulary,
            metadata=metadata,
        )
        for decision in response.parsed.selected
    )
    stop_at_topic = None
    if response.parsed.stop_at_parent:
        stop_at_topic = _topic_stop_result(
            topic_branch=topic_branch,
            stop_reason=response.parsed.stop_reason,
            metadata=metadata,
        )
    return TermRoutingResult(
        term_branches=term_branches,
        stop_at_topic=stop_at_topic,
        model_metadata=metadata,
    )


def _validate_selected_candidate_ids(
    selected: list[CandidateDecision],
    candidates_by_id: dict[str, TopicCandidate] | dict[str, TermCandidate],
    *,
    stage: str,
) -> None:
    seen: set[str] = set()
    for decision in selected:
        if decision.candidate_id not in candidates_by_id:
            raise UnknownCandidateIDError(
                f"{stage} model selected unknown candidate_id {decision.candidate_id!r}."
            )
        if decision.candidate_id in seen:
            raise UnknownCandidateIDError(
                f"{stage} model selected duplicate candidate_id {decision.candidate_id!r}."
            )
        seen.add(decision.candidate_id)


def _topic_branch_seed(
    *,
    decision: CandidateDecision,
    candidate: TopicCandidate,
    vocabulary: VocabularyIndex,
    metadata: ModelCallMetadata,
) -> TopicBranchSeed:
    topic = vocabulary.get(candidate.topic_uuid)
    return TopicBranchSeed(
        branch_id=f"topic:{candidate.candidate_id}",
        topic_uuid=topic.UUID,
        topic_name=topic.name,
        topic_level=topic.level,
        topic_canonical_path=topic.canonical_path,
        evidence=decision.evidence,
        support_type=decision.support_type,
        confidence=decision.confidence,
        reason=decision.reason,
        candidate_id=decision.candidate_id,
        prompt_version=metadata.prompt_version,
        model_provider=metadata.provider,
        model_name=metadata.model_name,
        retry_count=metadata.retry_count,
    )


def _term_branch_seed(
    *,
    decision: CandidateDecision,
    candidate: TermCandidate,
    topic_branch: TopicBranchSeed,
    vocabulary: VocabularyIndex,
    metadata: ModelCallMetadata,
) -> TermBranchSeed:
    validate_term_candidate_relationship(
        candidate,
        selected_topic_uuid=topic_branch.topic_uuid,
        index=vocabulary,
    )
    term = vocabulary.get(candidate.term_uuid)
    return TermBranchSeed(
        branch_id=f"{topic_branch.branch_id}/term:{candidate.candidate_id}",
        parent_topic_uuid=topic_branch.topic_uuid,
        parent_topic_name=topic_branch.topic_name,
        term_uuid=term.UUID,
        term_name=term.name,
        term_level=term.level,
        term_canonical_path=term.canonical_path,
        evidence=decision.evidence,
        support_type=decision.support_type,
        confidence=decision.confidence,
        reason=decision.reason,
        candidate_id=decision.candidate_id,
        parent_branch_id=topic_branch.branch_id,
        prompt_version=metadata.prompt_version,
        model_provider=metadata.provider,
        model_name=metadata.model_name,
        retry_count=metadata.retry_count,
    )


def _topic_stop_result(
    *,
    topic_branch: TopicBranchSeed,
    stop_reason: str | None,
    metadata: ModelCallMetadata,
) -> TopicStopResult:
    if not stop_reason:
        raise ValueError("stop_at_parent responses require a stop_reason.")
    return TopicStopResult(
        branch_id=topic_branch.branch_id,
        topic_uuid=topic_branch.topic_uuid,
        topic_name=topic_branch.topic_name,
        topic_level=topic_branch.topic_level,
        topic_canonical_path=topic_branch.topic_canonical_path,
        stop_reason=stop_reason,
        parent_evidence=topic_branch.evidence,
        parent_support_type=topic_branch.support_type,
        parent_confidence=topic_branch.confidence,
        topic_candidate_id=topic_branch.candidate_id,
        prompt_version=metadata.prompt_version,
        model_provider=metadata.provider,
        model_name=metadata.model_name,
        retry_count=metadata.retry_count,
    )


def _model_metadata(response) -> ModelCallMetadata:
    token_usage = response.token_usage
    return ModelCallMetadata(
        provider=response.provider,
        model_name=response.model_name,
        prompt_version=response.prompt_version,
        retry_count=response.retry_count,
        duration_seconds=response.duration_seconds,
        input_tokens=None if token_usage is None else token_usage.input_tokens,
        output_tokens=None if token_usage is None else token_usage.output_tokens,
        total_tokens=None if token_usage is None else token_usage.total_tokens,
        estimated_cost=response.estimated_cost,
    )