tests/test_term_routing.py

from __future__ import annotations

from pathlib import Path

import pytest

from gcmd_classifier.classification import (
    TermCandidate,
    build_term_candidates,
    build_topic_candidates,
    route_terms,
    route_topics,
    validate_term_candidate_relationship,
)
from gcmd_classifier.config import ModelSettings
from gcmd_classifier.errors import StructuredModelResponseError, UnknownCandidateIDError
from gcmd_classifier.llm import FakeModelClient
from gcmd_classifier.llm.prompts import PromptCandidate
from gcmd_classifier.models import ArticleRecord, SupportType
from gcmd_classifier.vocabulary import build_vocabulary_index, load_vocabulary

FIXTURE_PATH = Path("tests/fixtures/gcmd_hierarchy_small.json")
FULL_HIERARCHY_PATH = Path("data/gcmd_hierarchy.json")


def _index():
    return load_vocabulary(FIXTURE_PATH)


def _article(
    abstract: str = "Atmospheric chemistry and weather events are discussed.",
) -> ArticleRecord:
    return ArticleRecord(
        DOI="10.example/term-routing",
        Title="Atmospheric chemistry and weather observations",
        Year=2025,
        Abstract=abstract,
    )


def _decision(candidate_id: str, confidence: float | None = 0.82) -> dict:
    return {
        "candidate_id": candidate_id,
        "confidence": confidence,
        "evidence": "The article supports this Term.",
        "support_type": "explicit",
        "reason": "Primary subject is represented by this Term.",
    }


def _topic_branch(topic_index: int = 0):
    index = _index()
    topic_candidate = build_topic_candidates(index)[topic_index]
    client = FakeModelClient([{"selected": [_decision(topic_candidate.candidate_id)]}])
    return route_topics(
        article=_article(),
        vocabulary=index,
        model_client=client,
        settings=ModelSettings(),
    ).branches[0]


def _no_term_index():
    data = {
        "level": "Category",
        "name": "EARTH SCIENCE",
        "children": [
            {
                "level": "Topic",
                "name": "LAND SURFACE",
                "UUID": "topic-land-surface",
            }
        ],
    }
    return build_vocabulary_index(data, vocabulary_version="no-term-test")


def test_term_candidates_are_built_only_from_selected_topic_direct_children() -> None:
    index = _index()
    topic = index.get("topic-atmosphere")
    candidates = build_term_candidates(index, topic_uuid=topic.UUID)

    assert [candidate.term_uuid for candidate in candidates] == list(topic.child_uuids)
    assert [index.get(candidate.term_uuid).level for candidate in candidates] == ["Term", "Term"]


def test_terms_from_sibling_topics_are_not_included() -> None:
    index = _index()
    candidates = build_term_candidates(index, topic_uuid="topic-atmosphere")

    assert "term-ocean-chemistry" not in [candidate.term_uuid for candidate in candidates]


def test_term_candidate_ids_are_unique_within_prompt() -> None:
    candidates = build_term_candidates(_index(), topic_uuid="topic-atmosphere")
    candidate_ids = [candidate.candidate_id for candidate in candidates]

    assert candidate_ids == ["term_0001", "term_0002"]
    assert len(candidate_ids) == len(set(candidate_ids))


def test_term_candidate_ids_map_to_uuid_bearing_term_records() -> None:
    index = _index()

    for candidate in build_term_candidates(index, topic_uuid="topic-atmosphere"):
        record = index.get(candidate.term_uuid)
        assert record.UUID
        assert record.level == "Term"
        assert candidate.prompt_candidate.candidate_id == candidate.candidate_id


def test_term_candidate_construction_does_not_hard_code_term_names_or_uuids() -> None:
    index = _index()
    candidates = build_term_candidates(index, topic_uuid="topic-atmosphere")
    direct_terms = index.terms_for_topic("topic-atmosphere")

    assert [candidate.term_uuid for candidate in candidates] == [term.UUID for term in direct_terms]
    assert [candidate.prompt_candidate.name for candidate in candidates] == [
        term.name for term in direct_terms
    ]


def test_selected_topic_with_no_direct_term_children_is_handled_explicitly() -> None:
    index = _no_term_index()
    topic = index.topics()[0]

    candidates = build_term_candidates(index, topic_uuid=topic.UUID)

    assert candidates == ()


def test_one_term_selected_under_selected_topic() -> None:
    index = _index()
    topic_branch = _topic_branch()
    term_candidate = build_term_candidates(index, topic_uuid=topic_branch.topic_uuid)[0]
    client = FakeModelClient(
        [
            {
                "selected": [_decision(term_candidate.candidate_id)],
                "stop_at_parent": False,
                "ambiguous_alternatives": [],
            }
        ]
    )

    result = route_terms(
        article=_article(),
        topic_branch=topic_branch,
        vocabulary=index,
        model_client=client,
        settings=ModelSettings(),
    )

    assert result.selected_count == 1
    assert result.stopped_at_topic is False
    assert result.term_branches[0].term_uuid == term_candidate.term_uuid


def test_multiple_terms_selected_under_one_topic() -> None:
    index = _index()
    topic_branch = _topic_branch()
    term_candidates = build_term_candidates(index, topic_uuid=topic_branch.topic_uuid)
    client = FakeModelClient(
        [
            {
                "selected": [
                    _decision(term_candidates[0].candidate_id),
                    _decision(term_candidates[1].candidate_id, confidence=None),
                ],
                "stop_at_parent": False,
                "ambiguous_alternatives": [],
            }
        ]
    )

    result = route_terms(
        article=_article(),
        topic_branch=topic_branch,
        vocabulary=index,
        model_client=client,
        settings=ModelSettings(),
    )

    assert [branch.term_uuid for branch in result.term_branches] == [
        term_candidates[0].term_uuid,
        term_candidates[1].term_uuid,
    ]
    assert result.term_branches[1].confidence is None


def test_stop_at_topic_with_stop_at_parent_true() -> None:
    topic_branch = _topic_branch()
    client = FakeModelClient(
        [
            {
                "selected": [],
                "stop_at_parent": True,
                "stop_reason": "The Topic is supported, but no child Term is supported.",
                "ambiguous_alternatives": [],
            }
        ]
    )

    result = route_terms(
        article=_article(),
        topic_branch=topic_branch,
        vocabulary=_index(),
        model_client=client,
        settings=ModelSettings(),
    )

    assert result.term_branches == ()
    assert result.stopped_at_topic is True
    assert result.stop_at_topic is not None
    assert result.stop_at_topic.topic_uuid == topic_branch.topic_uuid
    assert result.stop_at_topic.stop_reason == (
        "The Topic is supported, but no child Term is supported."
    )


def test_topic_with_no_supported_term_returns_stop_at_topic_result() -> None:
    index = _no_term_index()
    topic = index.topics()[0]
    topic_branch = _manual_topic_branch(topic.UUID, topic.name, topic.canonical_path)
    client = FakeModelClient(
        [
            {
                "selected": [],
                "stop_at_parent": True,
                "stop_reason": "No direct Term candidates are available or supported.",
                "ambiguous_alternatives": [],
            }
        ]
    )

    result = route_terms(
        article=_article(),
        topic_branch=topic_branch,
        vocabulary=index,
        model_client=client,
        settings=ModelSettings(),
    )

    assert result.stopped_at_topic is True
    assert result.term_branches == ()


def test_invalid_term_candidate_id_is_rejected() -> None:
    client = FakeModelClient(
        [
            {
                "selected": [_decision("unknown-term")],
                "stop_at_parent": False,
                "ambiguous_alternatives": [],
            }
        ]
    )

    with pytest.raises(UnknownCandidateIDError):
        route_terms(
            article=_article(),
            topic_branch=_topic_branch(),
            vocabulary=_index(),
            model_client=client,
            settings=ModelSettings(),
        )


def test_duplicate_term_candidate_id_is_rejected() -> None:
    topic_branch = _topic_branch()
    candidate_id = build_term_candidates(
        _index(),
        topic_uuid=topic_branch.topic_uuid,
    )[0].candidate_id
    client = FakeModelClient(
        [
            {
                "selected": [_decision(candidate_id), _decision(candidate_id)],
                "stop_at_parent": False,
                "ambiguous_alternatives": [],
            }
        ]
    )

    with pytest.raises(UnknownCandidateIDError):
        route_terms(
            article=_article(),
            topic_branch=topic_branch,
            vocabulary=_index(),
            model_client=client,
            settings=ModelSettings(),
        )


def test_empty_or_malformed_term_candidate_id_is_rejected_by_structured_validation() -> None:
    client = FakeModelClient(
        [
            {
                "selected": [_decision("")],
                "stop_at_parent": False,
                "ambiguous_alternatives": [],
            }
        ]
    )

    with pytest.raises(StructuredModelResponseError):
        route_terms(
            article=_article(),
            topic_branch=_topic_branch(),
            vocabulary=_index(),
            model_client=client,
            settings=ModelSettings(),
        )


def test_selected_term_branch_seed_is_populated_from_vocabulary_index() -> None:
    index = _index()
    topic_branch = _topic_branch()
    term_candidate = build_term_candidates(index, topic_uuid=topic_branch.topic_uuid)[0]
    term = index.get(term_candidate.term_uuid)
    client = FakeModelClient(
        [
            {
                "selected": [_decision(term_candidate.candidate_id, confidence=0.74)],
                "stop_at_parent": False,
                "ambiguous_alternatives": [],
            }
        ]
    )

    result = route_terms(
        article=_article(),
        topic_branch=topic_branch,
        vocabulary=index,
        model_client=client,
        settings=ModelSettings(model_name="unit-test-model", prompt_version_term="term-test"),
    )
    branch = result.term_branches[0]

    assert branch.branch_id == f"{topic_branch.branch_id}/term:{term_candidate.candidate_id}"
    assert branch.parent_topic_uuid == topic_branch.topic_uuid
    assert branch.parent_topic_name == topic_branch.topic_name
    assert branch.term_uuid == term.UUID
    assert branch.term_name == term.name
    assert branch.term_level == "Term"
    assert branch.term_canonical_path == term.canonical_path
    assert branch.evidence == "The article supports this Term."
    assert branch.support_type is SupportType.EXPLICIT
    assert branch.confidence == 0.74
    assert branch.reason == "Primary subject is represented by this Term."
    assert branch.candidate_id == term_candidate.candidate_id
    assert branch.prompt_version == "term-test"
    assert branch.model_name == "unit-test-model"


def test_confidence_is_preserved_as_uncalibrated_metadata_only() -> None:
    index = _index()
    topic_branch = _topic_branch()
    term_candidate = build_term_candidates(index, topic_uuid=topic_branch.topic_uuid)[0]
    client = FakeModelClient(
        [
            {
                "selected": [_decision(term_candidate.candidate_id, confidence=0.01)],
                "stop_at_parent": False,
                "ambiguous_alternatives": [],
            }
        ]
    )

    result = route_terms(
        article=_article(),
        topic_branch=topic_branch,
        vocabulary=index,
        model_client=client,
        settings=ModelSettings(),
    )

    assert result.term_branches[0].confidence == 0.01
    assert result.term_branches[0].term_uuid == term_candidate.term_uuid


def test_no_fake_fallback_term_is_created() -> None:
    topic_branch = _topic_branch()
    client = FakeModelClient(
        [{"selected": [], "stop_at_parent": True, "stop_reason": "Stop here."}]
    )

    result = route_terms(
        article=_article(),
        topic_branch=topic_branch,
        vocabulary=_index(),
        model_client=client,
        settings=ModelSettings(),
    )

    assert result.term_branches == ()
    assert "FALLBACK" not in str(result.model_dump())


def test_no_variable_routing_is_performed() -> None:
    index = _index()
    topic_branch = _topic_branch()
    term_candidate = build_term_candidates(index, topic_uuid=topic_branch.topic_uuid)[0]
    client = FakeModelClient(
        [
            {
                "selected": [_decision(term_candidate.candidate_id)],
                "stop_at_parent": False,
            }
        ]
    )

    route_terms(
        article=_article(),
        topic_branch=topic_branch,
        vocabulary=index,
        model_client=client,
        settings=ModelSettings(),
    )

    assert len(client.requests) == 1
    assert client.requests[0].stage.value == "term"


def test_term_selected_from_correct_topic_succeeds() -> None:
    index = _index()
    topic_branch = _topic_branch()
    candidate = build_term_candidates(index, topic_uuid=topic_branch.topic_uuid)[0]

    validate_term_candidate_relationship(
        candidate,
        selected_topic_uuid=topic_branch.topic_uuid,
        index=index,
    )


def test_candidate_from_sibling_topic_cannot_be_accepted() -> None:
    client = FakeModelClient(
        [
            {
                "selected": [_decision("term-ocean-chemistry")],
                "stop_at_parent": False,
            }
        ]
    )

    with pytest.raises(UnknownCandidateIDError):
        route_terms(
            article=_article(),
            topic_branch=_topic_branch(),
            vocabulary=_index(),
            model_client=client,
            settings=ModelSettings(),
        )


def test_manually_invalid_candidate_to_record_mapping_is_rejected() -> None:
    index = _index()
    invalid_candidate = TermCandidate(
        candidate_id="term_0001",
        topic_uuid="topic-atmosphere",
        term_uuid="term-ocean-chemistry",
        prompt_candidate=PromptCandidate(
            candidate_id="term_0001",
            name="OCEAN CHEMISTRY",
            level="Term",
        ),
    )

    with pytest.raises(ValueError):
        validate_term_candidate_relationship(
            invalid_candidate,
            selected_topic_uuid="topic-atmosphere",
            index=index,
        )


def test_term_router_never_evaluates_all_terms_globally() -> None:
    index = _index()
    topic_branch = _topic_branch()
    term_candidate = build_term_candidates(index, topic_uuid=topic_branch.topic_uuid)[0]
    client = FakeModelClient(
        [
            {
                "selected": [_decision(term_candidate.candidate_id)],
                "stop_at_parent": False,
            }
        ]
    )

    route_terms(
        article=_article(),
        topic_branch=topic_branch,
        vocabulary=index,
        model_client=client,
        settings=ModelSettings(),
    )

    prompt = client.requests[0].prompt
    assert "ATMOSPHERIC CHEMISTRY" in prompt
    assert "WEATHER EVENTS" in prompt
    assert "OCEAN CHEMISTRY" not in prompt


def test_term_router_uses_provider_neutral_model_interface() -> None:
    index = _index()
    topic_branch = _topic_branch()
    term_candidate = build_term_candidates(index, topic_uuid=topic_branch.topic_uuid)[0]
    client = FakeModelClient(
        [
            {
                "selected": [_decision(term_candidate.candidate_id)],
                "stop_at_parent": False,
            }
        ]
    )

    route_terms(
        article=_article(),
        topic_branch=topic_branch,
        vocabulary=index,
        model_client=client,
        settings=ModelSettings(provider="fake-provider", model_name="fake-model"),
    )

    request = client.requests[0]
    assert request.provider == "fake-provider"
    assert request.model_name == "fake-model"
    assert request.response_schema.__name__ == "TermResponse"


def test_fake_model_receives_article_fields_parent_context_and_direct_term_candidates() -> None:
    index = _index()
    topic_branch = _topic_branch()
    term_candidate = build_term_candidates(index, topic_uuid=topic_branch.topic_uuid)[0]
    article = _article()
    client = FakeModelClient(
        [
            {
                "selected": [_decision(term_candidate.candidate_id)],
                "stop_at_parent": False,
            }
        ]
    )

    route_terms(
        article=article,
        topic_branch=topic_branch,
        vocabulary=index,
        model_client=client,
        settings=ModelSettings(),
    )

    prompt = client.requests[0].prompt
    assert article.DOI in prompt
    assert article.Title in prompt
    assert article.Abstract in prompt
    assert topic_branch.topic_name in prompt
    assert topic_branch.topic_canonical_path in prompt
    assert term_candidate.candidate_id in prompt
    assert index.get(term_candidate.term_uuid).name in prompt


def test_empty_abstract_remains_valid_and_is_passed_to_term_prompt() -> None:
    index = _index()
    topic_branch = _topic_branch()
    term_candidate = build_term_candidates(index, topic_uuid=topic_branch.topic_uuid)[0]
    client = FakeModelClient(
        [
            {
                "selected": [_decision(term_candidate.candidate_id)],
                "stop_at_parent": False,
            }
        ]
    )

    route_terms(
        article=_article(abstract=""),
        topic_branch=topic_branch,
        vocabulary=index,
        model_client=client,
        settings=ModelSettings(),
    )

    assert "<ABSTRACT>\n\n</ABSTRACT>" in client.requests[0].prompt


def test_full_data_current_vocabulary_has_144_direct_term_candidates() -> None:
    index = load_vocabulary(FULL_HIERARCHY_PATH)
    all_candidates = [
        candidate
        for topic in index.topics()
        for candidate in build_term_candidates(index, topic_uuid=topic.UUID)
    ]
    term_uuids = [candidate.term_uuid for candidate in all_candidates]

    assert len(all_candidates) == 144
    assert len(term_uuids) == len(set(term_uuids))
    assert all(index.get(candidate.term_uuid).UUID for candidate in all_candidates)
    assert all(index.get(candidate.term_uuid).level == "Term" for candidate in all_candidates)


def _manual_topic_branch(topic_uuid: str, topic_name: str, canonical_path: str):
    from gcmd_classifier.classification import TopicBranchSeed

    return TopicBranchSeed(
        branch_id="topic:manual",
        topic_uuid=topic_uuid,
        topic_name=topic_name,
        topic_level="Topic",
        topic_canonical_path=canonical_path,
        evidence="The Topic is supported.",
        support_type="explicit",
        confidence=0.8,
        reason="Manual test branch.",
        candidate_id="topic_manual",
        prompt_version="topic-test",
        model_provider="fake",
        model_name="fake-model",
        retry_count=0,
    )