feat(#54): add support for gemini model wildcards

Author: scastlara

README.md

@@ -147,8 +147,8 @@ These are the main supported models, though the CLI may support additional ones
 
 | Model name |  Description | 
 | --------   |  ----- | 
-| **gpt-4.1** * | Powerful and reliable model for detailed code analysis. Strong at coding tasks. | 
-| **gpt-4.1-mini** | Lightweight variant of GPT-4.1 offering faster responses and lower cost—ideal for iterative or high-volume reviews. | 
+| gpt4.1 | Powerful and reliable model for detailed code analysis. Strong at coding tasks. | 
+| gpt-4.1-mini | Lightweight variant of GPT-4.1 offering faster responses and lower cost—ideal for iterative or high-volume reviews. | 
 | gpt-4.1-nano | Ultra-light model focused on speed and affordability—best for basic code checks or initial feedback. | 
 | gpt-4o * | Cutting-edge model with strong reasoning and code capabilities—ideal for detailed, context-aware reviews. | 
 | gpt-4o-mini | Streamlined GPT-4o variant optimized for fast, cost-effective feedback on code. | 
@@ -174,16 +174,20 @@ To use Gemini LLMs, you need to provide lgtm an API Key, which can be generated
 These are the main supported models, though the CLI may support additional ones due to the use of [pydantic-ai](https://ai.pydantic.dev). Gemini timestamps models, so be sure to always use the latest model of each family, if possible.
 
 
+For Gemini models exclusively, you can provide a wildcard and lgtm will attempt to select the latest model (e.g., `gemini-2.5-pro*`)
+
+
 <details>
 
 <summary>Supported Google's Gemini models</summary>
 
 
 | Model name  | Description |
 | -----------  | --- |
-| gemini-2.5-pro-preview-05-06       | Most advanced publicly available Gemini model. Strong code reasoning and long-context support. Ideal for complex or large reviews. |
-| **gemini-2.0-pro-exp-02-05**            | High-performing general-purpose model. Balances accuracy and efficiency—ideal for robust reviews without 2.5's higher cost.        |
-| **gemini-2.0-flash**                    | Optimized for low-latency, lower-cost analysis. Ideal for iterative feedback and smaller reviews.                                  |
+| gemini-2.5-pro-preview-06-05       | Most advanced publicly available Gemini model. Strong code reasoning and long-context support. Ideal for complex or large reviews. |
+| gemini-2.5-pro-preview-06-05       | Deprecated. Most advanced publicly available Gemini model. Strong code reasoning and long-context support. Ideal for complex or large reviews. |
+| gemini-2.0-pro-exp-02-05            | High-performing general-purpose model. Balances accuracy and efficiency—ideal for robust reviews without 2.5's higher cost.        |
+| gemini-2.0-flash                    | Optimized for low-latency, lower-cost analysis. Ideal for iterative feedback and smaller reviews.                                  |
 | gemini-1.5-pro                      |  Proven performer with solid context and reasoning. Still excellent for general code understanding.                                 |
 | gemini-1.5-flash                    |  Lightweight and fast—suited for real-time or continuous code review loops.                                                         |
 
@@ -222,7 +226,7 @@ These are the main supported models, though the CLI may support additional ones
 
 | Model name         | Description                                                                                         |
 | ------------------ | --------------------------------------------------------------------------------------------------- |
-| **mistral-large-latest** |  Mistral's top-tier reasoning model for high-complexity tasks. |
+| mistral-large-latest |  Mistral's top-tier reasoning model for high-complexity tasks. |
 | mistral-small        |  Lightweight and fast. Best used for simple syntax or formatting checks where cost and speed are priorities. |
 | codestrallatest     |   Codestral specializes in low-latency, high-frequency tasks such as fill-in-the-middle (FIM), code correction and test generation.
 
@@ -242,8 +246,8 @@ To get an API key for DeepSeek, create one at [DeepSeek Platform](https://platfo
 
 | Model name         |  Description                                                                                         |
  ----------- | --------------------------------------------------------------------------------------------------- |
-| **deepseek-chat**     | General-purpose LLM optimized for chat and code assistance. Ideal for standard reviews and developer interactions. |
-| **deepseek-reasoner** | Advanced model specialized in reasoning and problem solving—ideal for complex code analysis and critical thinking tasks. |
+| deepseek-chat    | General-purpose LLM optimized for chat and code assistance. Ideal for standard reviews and developer interactions. |
+| deepseek-reasoner | Advanced model specialized in reasoning and problem solving—ideal for complex code analysis and critical thinking tasks. |
 
 </details>
 

src/lgtm_ai/ai/agent.py

@@ -1,7 +1,7 @@
 import logging
-from typing import Any, TypeGuard, get_args
+from typing import Any, TypeGuard, cast, get_args
 
-from lgtm_ai.ai.exceptions import MissingAIAPIKey, MissingModelUrl
+from lgtm_ai.ai.exceptions import InvalidModelName, MissingAIAPIKey, MissingModelUrl
 from lgtm_ai.ai.prompts import GUIDE_SYSTEM_PROMPT, REVIEWER_SYSTEM_PROMPT, SUMMARIZING_SYSTEM_PROMPT
 from lgtm_ai.ai.schemas import (
     AgentSettings,
@@ -14,6 +14,7 @@
     SupportedAIModelsList,
     SupportedGeminiModel,
 )
+from lgtm_ai.ai.utils import match_model_by_wildcard, select_latest_gemini_model
 from openai.types import ChatModel
 from pydantic_ai import Agent, RunContext
 from pydantic_ai.models import Model
@@ -32,7 +33,8 @@
 
 def get_ai_model(model_name: SupportedAIModels | str, api_key: str, model_url: str | None = None) -> Model:  # noqa: C901
     def _is_gemini_model(model_name: SupportedAIModels) -> TypeGuard[SupportedGeminiModel]:
-        return model_name in get_args(SupportedGeminiModel)
+        matched_model = match_model_by_wildcard(model_name, get_args(SupportedGeminiModel))
+        return bool(matched_model)
 
     def _is_openai_model(model_name: SupportedAIModels) -> TypeGuard[ChatModel]:
         return model_name in get_args(ChatModel)
@@ -54,7 +56,13 @@ def _is_deepseek_model(model_name: SupportedAIModels) -> TypeGuard[DeepSeekModel
         raise MissingAIAPIKey(model_name=model_name)
 
     if _is_gemini_model(model_name):
-        return GoogleModel(model_name, provider=GoogleProvider(api_key=api_key))
+        matches = match_model_by_wildcard(
+            model_name,
+            cast(tuple[SupportedGeminiModel, ...], get_args(SupportedGeminiModel)),
+        )
+        if not matches:
+            raise InvalidModelName(model_name=model_name)
+        return GoogleModel(select_latest_gemini_model(matches), provider=GoogleProvider(api_key=api_key))
     elif _is_openai_model(model_name):
         return OpenAIModel(model_name=model_name, provider=OpenAIProvider(api_key=api_key))
     elif _is_anthropic_model(model_name):

src/lgtm_ai/ai/exceptions.py

@@ -1,4 +1,5 @@
 import click
+from lgtm_ai.ai.schemas import SupportedGeminiModel
 
 
 class MissingModelUrl(click.BadParameter):  # not a LGTMException because we want click to handle it gracefully
@@ -15,3 +16,19 @@ class MissingAIAPIKey(click.BadParameter):
     def __init__(self, model_name: str) -> None:
         msg = f"Model '{model_name}' requires an AI API key to be provided"
         super().__init__(msg)
+
+
+class InvalidModelName(click.BadParameter):
+    """Exception raised when an invalid AI model name is provided."""
+
+    def __init__(self, model_name: str) -> None:
+        msg = f"Model '{model_name}' is not a valid AI model name"
+        super().__init__(msg)
+
+
+class InvalidGeminiWildcard(click.BadParameter):
+    """Exception raised when a Gemini model name with wildcard is invalid."""
+
+    def __init__(self, matches: list[SupportedGeminiModel]) -> None:
+        msg = f"The provided Gemini model name matches multiple models that cannot be narrowed down based on when they were released: {', '.join(matches)}. Please specify a more specific model name."
+        super().__init__(msg)

src/lgtm_ai/ai/schemas.py

@@ -46,16 +46,11 @@
     "gemini-2.5-pro-preview-06-05",
 ]
 
-LocalAIModel = str
+AnyModel = str
 """Users may use any model name in their local AI server, so we just allow any string."""
 
 SupportedAIModels = (
-    ChatModel
-    | SupportedGeminiModel
-    | LatestAnthropicModelNames
-    | LatestMistralModelNames
-    | DeepSeekModel
-    | LocalAIModel
+    ChatModel | SupportedGeminiModel | LatestAnthropicModelNames | LatestMistralModelNames | DeepSeekModel | AnyModel
 )
 """Type of all supported AI models in lgtm."""
 
@@ -65,7 +60,7 @@
     + get_args(LatestAnthropicModelNames)
     + get_args(LatestMistralModelNames)
     + get_args(DeepSeekModel)
-)  # Keep in sync with SupportedAIModels except for LocalAIModel
+)  # Keep in sync with SupportedAIModels except for AnyModel
 """Tuple of all known supported AI models in lgtm."""
 
 

src/lgtm_ai/ai/utils.py

@@ -0,0 +1,43 @@
+from typing import cast
+
+from lgtm_ai.ai.exceptions import InvalidGeminiWildcard
+from lgtm_ai.ai.schemas import SupportedGeminiModel
+
+
+def match_model_by_wildcard[T](model_name: str, model_list: tuple[T, ...]) -> list[T] | None:
+    """Match a model name against a list of models with wildcard support."""
+    if model_name in model_list:
+        return [cast(T, model_name)]
+
+    if "*" not in model_name:
+        return None
+
+    all_matches = [
+        model for model in model_list if str(model).startswith(model_name.replace("*", "")) or model_name == model
+    ]
+    if not all_matches:
+        return None
+
+    return all_matches
+
+
+def select_latest_gemini_model(matches: list[SupportedGeminiModel]) -> SupportedGeminiModel:
+    if len(matches) == 1:
+        return matches[0]
+
+    model_to_date = {model: model.split("-")[-2:] for model in matches}
+
+    # If the values that are supposed to be date-like (06-05, 04-17) are not in the format we expect,
+    # we just raise an error.
+    def _looks_like_date(date: list[str]) -> bool:
+        return len(date) == 2 and all(elem.isdigit() for elem in date)
+
+    if not all(_looks_like_date(date) for date in model_to_date.values()):
+        raise InvalidGeminiWildcard(matches)
+
+    # Get the latest model by date
+    latest_model = max(
+        matches,
+        key=lambda model: tuple(int(x) for x in model_to_date[model]),
+    )
+    return latest_model

tests/ai/test_utils.py

@@ -3,7 +3,7 @@
 
 import pytest
 from lgtm_ai.ai.agent import get_ai_model
-from lgtm_ai.ai.exceptions import MissingAIAPIKey, MissingModelUrl
+from lgtm_ai.ai.exceptions import InvalidGeminiWildcard, MissingAIAPIKey, MissingModelUrl
 from pydantic_ai.models.google import GoogleModel
 from pydantic_ai.models.openai import OpenAIModel
 
@@ -28,3 +28,23 @@ def test_get_ai_model(model: str, model_url: str | None, ai_api_key: str, expect
         ai_model = get_ai_model(model, ai_api_key, model_url=model_url)
         assert isinstance(ai_model, expected_type)
         assert ai_model.model_name == model
+
+
+@pytest.mark.parametrize(
+    ("model", "expected_model_name", "expectation"),
+    [
+        # Normal matches
+        ("gemini-2.5-flash-*", "gemini-2.5-flash-preview-05-20", does_not_raise()),
+        ("gemini-2.5-pro-preview-*", "gemini-2.5-pro-preview-06-05", does_not_raise()),
+        # Exact match
+        ("gemini-2.5-pro-preview-06-05", "gemini-2.5-pro-preview-06-05", does_not_raise()),
+        # No wildcard results in no attempt to actually match
+        ("gemini-2.5-pro-preview-", "gemini-2.5-pro-preview-06-05", pytest.raises(MissingModelUrl)),
+        # Multiple matches that we cannot narrow down by date
+        ("gemini-2*", None, pytest.raises(InvalidGeminiWildcard)),
+    ],
+)
+def test_get_ai_model_with_wildcard(model: str, expected_model_name: str, expectation: Any) -> None:
+    with expectation:
+        ai_model = get_ai_model(model, "fake", model_url=None)
+        assert ai_model.model_name == expected_model_name