from __future__ import annotations import pytest from gcmd_classifier.config import ModelSettings from gcmd_classifier.errors import ( NonRetryableModelError, RetryableModelError, StructuredModelResponseError, ) from gcmd_classifier.llm import FakeModelClient, ModelRequest, ModelStage from gcmd_classifier.llm.base import RetryPolicy, generate_with_retries from gcmd_classifier.llm.schemas import TermResponse, TopicResponse, VariableResponse def _request(stage: ModelStage, schema: type) -> ModelRequest: settings = ModelSettings() return ModelRequest.from_settings( stage=stage, prompt="prompt", response_schema=schema, settings=settings, ) def _decision(candidate_id: str = "candidate-1") -> dict: return { "candidate_id": candidate_id, "confidence": 0.8, "evidence": "The article supports this candidate.", "support_type": "explicit", } def test_fake_model_returns_typed_topic_response() -> None: client = FakeModelClient([{"selected": [_decision()], "ambiguous_alternatives": []}]) response = client.generate_structured(_request(ModelStage.TOPIC, TopicResponse)) assert isinstance(response.parsed, TopicResponse) assert response.parsed.selected[0].candidate_id == "candidate-1" assert response.provider == "fake" def test_fake_model_returns_typed_term_response() -> None: client = FakeModelClient( [ { "selected": [_decision("term-1")], "stop_at_parent": False, "ambiguous_alternatives": [], } ] ) response = client.generate_structured(_request(ModelStage.TERM, TermResponse)) assert isinstance(response.parsed, TermResponse) assert response.parsed.selected[0].candidate_id == "term-1" def test_fake_model_returns_typed_variable_response() -> None: client = FakeModelClient( [ { "selected": [_decision("variable-1")], "stop_at_parent": False, "ambiguous_alternatives": [], } ] ) response = client.generate_structured(_request(ModelStage.VARIABLE, VariableResponse)) assert isinstance(response.parsed, VariableResponse) assert response.parsed.selected[0].candidate_id == "variable-1" def test_fake_model_can_return_no_selection() -> None: client = FakeModelClient( [ { "selected": [], "ambiguous_alternatives": [], "no_selection_reason": "No supplied Topic is supported.", } ] ) response = client.generate_structured(_request(ModelStage.TOPIC, TopicResponse)) assert response.parsed.selected == [] assert response.parsed.no_selection_reason == "No supplied Topic is supported." def test_fake_model_records_requests() -> None: client = FakeModelClient([{"selected": [], "ambiguous_alternatives": []}]) request = _request(ModelStage.TOPIC, TopicResponse) client.generate_structured(request) assert client.requests == [request] def test_fake_model_can_return_invalid_candidate_ids_for_later_validation() -> None: client = FakeModelClient([{"selected": [_decision("not-a-supplied-id")]}]) response = client.generate_structured(_request(ModelStage.TOPIC, TopicResponse)) assert response.parsed.selected[0].candidate_id == "not-a-supplied-id" def test_fake_model_rejects_contradictory_structured_response() -> None: client = FakeModelClient( [ { "selected": [_decision()], "stop_at_parent": True, "stop_reason": "Contradictory response.", "ambiguous_alternatives": [], } ] ) with pytest.raises(StructuredModelResponseError): client.generate_structured(_request(ModelStage.TERM, TermResponse)) def test_fake_model_can_raise_retryable_errors() -> None: client = FakeModelClient([RetryableModelError("temporary failure")]) with pytest.raises(RetryableModelError): client.generate_structured(_request(ModelStage.TOPIC, TopicResponse)) def test_fake_model_can_raise_non_retryable_errors() -> None: client = FakeModelClient([NonRetryableModelError("permanent failure")]) with pytest.raises(NonRetryableModelError): client.generate_structured(_request(ModelStage.TOPIC, TopicResponse)) def test_fake_model_can_fail_once_then_succeed_with_retry_wrapper() -> None: client = FakeModelClient( [ RetryableModelError("temporary failure"), {"selected": [_decision()], "ambiguous_alternatives": []}, ] ) response = generate_with_retries( client, _request(ModelStage.TOPIC, TopicResponse), RetryPolicy(max_retries=1), ) assert response.retry_count == 1 assert len(client.requests) == 2 assert response.parsed.selected[0].candidate_id == "candidate-1" def test_fake_model_can_exhaust_retries() -> None: client = FakeModelClient( [RetryableModelError("one"), RetryableModelError("two"), RetryableModelError("three")] ) with pytest.raises(Exception) as exc_info: generate_with_retries( client, _request(ModelStage.TOPIC, TopicResponse), RetryPolicy(max_retries=2), ) assert exc_info.value.__class__.__name__ == "ModelRetriesExhaustedError" assert len(client.requests) == 3