From 4f2551e0a670cea735d7f04ff93e33e9640827eb Mon Sep 17 00:00:00 2001
From: muhammad-asn <muhammad.a.s.nugroho@gdplabs.id>
Date: Thu, 6 Mar 2025 15:49:07 +0700
Subject: [PATCH] feat: add DeepInfra support

---
 docs/docs/usage-guide/changing_a_model.md       | 13 +++++++++++++
 pr_agent/algo/__init__.py                       |  3 +++
 pr_agent/algo/ai_handlers/litellm_ai_handler.py |  4 ++++
 pr_agent/settings/.secrets_template.toml        |  3 +++
 4 files changed, 23 insertions(+)

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 = ""