diff --git a/docs/docs/usage-guide/changing_a_model.md b/docs/docs/usage-guide/changing_a_model.md index fd491012..e28bf5e4 100644 --- a/docs/docs/usage-guide/changing_a_model.md +++ b/docs/docs/usage-guide/changing_a_model.md @@ -193,6 +193,19 @@ key = ... (you can obtain a deepseek-chat key from [here](https://platform.deepseek.com)) + +### DeepInfra + +To use DeepSeek model with DeepInfra, for example, set: +``` +[config] # in configuration.toml +model = "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Llama-70B" +fallback_models = ["deepinfra/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B"] +[deepinfra] # in .secrets.toml +key = ... # your DeepInfra api key +``` +(you can obtain a DeepInfra key from [here](https://deepinfra.com/dash/api_keys)) + ### 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 34eac5e1..4487b04a 100644 --- a/pr_agent/algo/__init__.py +++ b/pr_agent/algo/__init__.py @@ -86,6 +86,9 @@ MAX_TOKENS = { "watsonx/ibm/granite-13b-chat-v2": 8191, "watsonx/ibm/granite-34b-code-instruct": 8191, "watsonx/mistralai/mistral-large": 32768, + "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": 128000, + "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": 128000, + "deepinfra/deepseek-ai/DeepSeek-R1": 128000, } 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 e2894f1b..fe367c9e 100644 --- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py +++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py @@ -96,6 +96,10 @@ class LiteLLMAIHandler(BaseAiHandler): if get_settings().get("DEEPSEEK.KEY", None): os.environ['DEEPSEEK_API_KEY'] = get_settings().get("DEEPSEEK.KEY") + # Support deepinfra models + if get_settings().get("DEEPINFRA.KEY", None): + os.environ['DEEPINFRA_API_KEY'] = get_settings().get("DEEPINFRA.KEY") + # Models that only use user meessage self.user_message_only_models = USER_MESSAGE_ONLY_MODELS diff --git a/pr_agent/settings/.secrets_template.toml b/pr_agent/settings/.secrets_template.toml index bee90616..ee383a37 100644 --- a/pr_agent/settings/.secrets_template.toml +++ b/pr_agent/settings/.secrets_template.toml @@ -94,3 +94,6 @@ pat = "" [deepseek] key = "" + +[deepinfra] +key = ""