Merge pull request #1707 from dst03106/feature/add-support-for-mistral-and-codestral

Add support for Mistral and Codestral models

docs/docs/usage-guide/changing_a_model.md

@@ -250,6 +250,34 @@ key = ... # your DeepInfra api key
 
 (you can obtain a DeepInfra key from [here](https://deepinfra.com/dash/api_keys))
 
+### Mistral
+
+To use models like Mistral or Codestral with Mistral, for example, set:
+
+```toml
+[config] # in configuration.toml
+model = "mistral/mistral-small-latest"
+fallback_models = ["mistral/mistral-medium-latest"]
+[mistral] # in .secrets.toml
+key = "..." # your Mistral api key
+```
+
+(you can obtain a Mistral key from [here](https://console.mistral.ai/api-keys))
+
+### Codestral
+
+To use Codestral model with Codestral, for example, set:
+
+```toml
+[config] # in configuration.toml
+model = "codestral/codestral-latest"
+fallback_models = ["codestral/codestral-2405"]
+[codestral] # in .secrets.toml
+key = "..." # your Codestral api key
+```
+
+(you can obtain a Codestral key from [here](https://console.mistral.ai/codestral))
+
 ### Custom models
 
 If the relevant model doesn't appear [here](https://github.com/Codium-ai/pr-agent/blob/main/pr_agent/algo/__init__.py), you can still use it as a custom model:

pr_agent/algo/__init__.py

@@ -112,6 +112,20 @@ MAX_TOKENS = {
     "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": 128000,
     "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": 128000,
     "deepinfra/deepseek-ai/DeepSeek-R1": 128000,
+    "mistral/mistral-small-latest": 8191,
+    "mistral/mistral-medium-latest": 8191,
+    "mistral/mistral-large-2407": 128000,
+    "mistral/mistral-large-latest": 128000,
+    "mistral/open-mistral-7b": 8191,
+    "mistral/open-mixtral-8x7b": 8191,
+    "mistral/open-mixtral-8x22b": 8191,
+    "mistral/codestral-latest": 8191,
+    "mistral/open-mistral-nemo": 128000,
+    "mistral/open-mistral-nemo-2407": 128000,
+    "mistral/open-codestral-mamba": 256000,
+    "mistral/codestral-mamba-latest": 256000,
+    "codestral/codestral-latest": 8191,
+    "codestral/codestral-2405": 8191,
 }
 
 USER_MESSAGE_ONLY_MODELS = [

pr_agent/algo/ai_handlers/litellm_ai_handler.py

@@ -97,6 +97,14 @@ class LiteLLMAIHandler(BaseAiHandler):
         if get_settings().get("DEEPINFRA.KEY", None):
             os.environ['DEEPINFRA_API_KEY'] = get_settings().get("DEEPINFRA.KEY")
 
+        # Support mistral models
+        if get_settings().get("MISTRAL.KEY", None):
+            os.environ["MISTRAL_API_KEY"] = get_settings().get("MISTRAL.KEY")
+        
+        # Support codestral models
+        if get_settings().get("CODESTRAL.KEY", None):
+            os.environ["CODESTRAL_API_KEY"] = get_settings().get("CODESTRAL.KEY")
+
         # Check for Azure AD configuration
         if get_settings().get("AZURE_AD.CLIENT_ID", None):
             self.azure = True