diff --git a/docs/docs/usage-guide/changing_a_model.md b/docs/docs/usage-guide/changing_a_model.md index 24c7b934..d05a73ab 100644 --- a/docs/docs/usage-guide/changing_a_model.md +++ b/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: diff --git a/pr_agent/algo/__init__.py b/pr_agent/algo/__init__.py index 62cd3f61..e3abb93c 100644 --- a/pr_agent/algo/__init__.py +++ b/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 = [ diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py index d717c087..61af7383 100644 --- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py +++ b/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