diff --git a/src/lgtm_ai/ai/schemas.py b/src/lgtm_ai/ai/schemas.py index 71a39e4..536570f 100644 --- a/src/lgtm_ai/ai/schemas.py +++ b/src/lgtm_ai/ai/schemas.py @@ -39,9 +39,11 @@ "gemini-2.0-flash-lite-preview-02-05", "gemini-2.0-pro-exp-02-05", "gemini-2.5-flash-preview-04-17", + "gemini-2.5-pro", "gemini-2.5-pro-exp-03-25", "gemini-2.5-pro-preview-03-25", "gemini-2.5-pro-preview-05-06", + "gemini-2.5-flash", "gemini-2.5-flash-preview-05-20", "gemini-2.5-pro-preview-06-05", ] diff --git a/src/lgtm_ai/ai/utils.py b/src/lgtm_ai/ai/utils.py index ffbd50e..44779be 100644 --- a/src/lgtm_ai/ai/utils.py +++ b/src/lgtm_ai/ai/utils.py @@ -26,6 +26,16 @@ def select_latest_gemini_model(matches: list[SupportedGeminiModel]) -> Supported if len(matches) == 1: return matches[0] + # If one of them is not a `preview` or `exp` model, select it + # This is because out of several possible models given a wildcard, we assume preference for stable models. + # If the user wanted a preview or experimental model, they could have specified it in the wildcard itself: + # e.g.: `gemini-2.5-flash*` will select the latest stable, while `gemini-2.5-flash-p*` will select the latest preview. + non_preview_matches = [model for model in matches if all(word not in model for word in ("preview", "exp"))] + if non_preview_matches: + if len(non_preview_matches) > 1: + raise InvalidGeminiWildcard(matches) + return non_preview_matches[0] + model_to_date = {model: model.split("-")[-2:] for model in matches} def _looks_like_date(date: list[str]) -> bool: diff --git a/tests/ai/test_utils.py b/tests/ai/test_utils.py index 88d6392..018d8ff 100644 --- a/tests/ai/test_utils.py +++ b/tests/ai/test_utils.py @@ -51,6 +51,8 @@ def test_get_ai_model(model: str, model_url: str | None, ai_api_key: str, expect ("gemini-2.5-flash-*-*", None, pytest.raises(InvalidModelWildCard)), # Wildcards in the middle of the model name are not allowed ("gemini-*-pro-preview-06-05", None, pytest.raises(InvalidModelWildCard)), + # Wildcards that match preview, experimental and stable models must select the stable one + ("gemini-2.5-pro*", "gemini-2.5-pro", does_not_raise()), ], ) def test_get_ai_model_with_wildcard(model: str, expected_model_name: str, expectation: Any) -> None: