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": ""})