tests/test_structured_responses.py

from __future__ import annotations

import pytest
from pydantic import ValidationError

from gcmd_classifier.llm.schemas import (
    CandidateDecision,
    TermResponse,
    TopicResponse,
    VariableResponse,
)

CANDIDATE = {
    "candidate_id": "topic-atmosphere",
    "confidence": 0.91,
    "evidence": "The article discusses atmospheric composition.",
    "support_type": "explicit",
    "reason": "Primary subject is atmospheric science.",
}
SECOND_CANDIDATE = {
    "candidate_id": "topic-oceans",
    "confidence": None,
    "evidence": "Ocean observations are also central to the article.",
    "support_type": "mixed",
}


def test_valid_topic_response_with_one_selected_candidate() -> None:
    response = TopicResponse.model_validate(
        {
            "selected": [CANDIDATE],
            "ambiguous_alternatives": [],
            "no_selection_reason": None,
        }
    )

    assert response.selected[0].candidate_id == "topic-atmosphere"
    assert response.selected[0].confidence == 0.91


def test_valid_topic_response_with_multiple_selected_candidates() -> None:
    response = TopicResponse.model_validate(
        {
            "selected": [CANDIDATE, SECOND_CANDIDATE],
            "ambiguous_alternatives": ["topic-cryosphere"],
        }
    )

    assert [candidate.candidate_id for candidate in response.selected] == [
        "topic-atmosphere",
        "topic-oceans",
    ]
    assert response.selected[1].confidence is None


def test_valid_topic_response_with_no_selected_candidates_and_reason() -> None:
    response = TopicResponse.model_validate(
        {
            "selected": [],
            "ambiguous_alternatives": [],
            "no_selection_reason": "No supplied Topic is supported by the title or abstract.",
        }
    )

    assert response.selected == []
    assert response.no_selection_reason is not None


def test_valid_term_response_selecting_candidates() -> None:
    response = TermResponse.model_validate(
        {
            "selected": [CANDIDATE, SECOND_CANDIDATE],
            "stop_at_parent": False,
            "stop_reason": None,
            "ambiguous_alternatives": [],
        }
    )

    assert len(response.selected) == 2
    assert response.stop_at_parent is False


def test_valid_term_response_stopping_at_parent() -> None:
    response = TermResponse.model_validate(
        {
            "selected": [],
            "stop_at_parent": True,
            "stop_reason": "No child Term is adequately supported.",
            "ambiguous_alternatives": ["term-weather-events"],
        }
    )

    assert response.stop_at_parent is True
    assert response.selected == []


@pytest.mark.parametrize(
    "payload",
    [
        {
            "selected": [CANDIDATE],
            "stop_at_parent": True,
            "stop_reason": "Contradictory response.",
            "ambiguous_alternatives": [],
        },
        {
            "selected": [CANDIDATE],
            "stop_at_parent": True,
            "stop_reason": None,
            "ambiguous_alternatives": [],
        },
    ],
)
def test_invalid_term_response_with_stop_and_selected(payload: dict) -> None:
    with pytest.raises(ValidationError):
        TermResponse.model_validate(payload)


@pytest.mark.parametrize("stop_reason", [None, ""])
def test_invalid_term_response_stopping_without_stop_reason(stop_reason: str | None) -> None:
    with pytest.raises(ValidationError):
        TermResponse.model_validate(
            {
                "selected": [],
                "stop_at_parent": True,
                "stop_reason": stop_reason,
                "ambiguous_alternatives": [],
            }
        )


def test_valid_variable_response_selecting_candidates() -> None:
    response = VariableResponse.model_validate(
        {
            "selected": [CANDIDATE],
            "stop_at_parent": False,
            "stop_reason": None,
            "ambiguous_alternatives": [],
        }
    )

    assert response.selected[0].support_type.value == "explicit"


def test_valid_variable_response_stopping_at_parent() -> None:
    response = VariableResponse.model_validate(
        {
            "selected": [],
            "stop_at_parent": True,
            "stop_reason": "Variable children are too specific for the evidence.",
            "ambiguous_alternatives": [],
        }
    )

    assert response.stop_at_parent is True


@pytest.mark.parametrize(
    "payload",
    [
        {
            "selected": [CANDIDATE],
            "stop_at_parent": True,
            "stop_reason": "Contradictory response.",
            "ambiguous_alternatives": [],
        },
        {
            "selected": [CANDIDATE],
            "stop_at_parent": True,
            "stop_reason": None,
            "ambiguous_alternatives": [],
        },
    ],
)
def test_invalid_variable_response_with_stop_and_selected(payload: dict) -> None:
    with pytest.raises(ValidationError):
        VariableResponse.model_validate(payload)


@pytest.mark.parametrize("stop_reason", [None, ""])
def test_invalid_variable_response_stopping_without_stop_reason(stop_reason: str | None) -> None:
    with pytest.raises(ValidationError):
        VariableResponse.model_validate(
            {
                "selected": [],
                "stop_at_parent": True,
                "stop_reason": stop_reason,
                "ambiguous_alternatives": [],
            }
        )


@pytest.mark.parametrize("confidence", [-0.01, 1.01])
def test_invalid_confidence_range(confidence: float) -> None:
    payload = CANDIDATE | {"confidence": confidence}

    with pytest.raises(ValidationError):
        CandidateDecision.model_validate(payload)


def test_valid_missing_confidence() -> None:
    payload = CANDIDATE.copy()
    del payload["confidence"]

    decision = CandidateDecision.model_validate(payload)

    assert decision.confidence is None


def test_unknown_fields_rejected() -> None:
    with pytest.raises(ValidationError):
        TopicResponse.model_validate({"selected": [], "ambiguous_alternatives": [], "extra": "bad"})


def test_malformed_candidate_decision_rejected() -> None:
    payload = CANDIDATE.copy()
    del payload["evidence"]

    with pytest.raises(ValidationError):
        CandidateDecision.model_validate(payload)


def test_empty_candidate_id_rejected() -> None:
    with pytest.raises(ValidationError):
        CandidateDecision.model_validate(CANDIDATE | {"candidate_id": ""})