Merge pull request #1457 from KennyDizi/main

Support deepseek-chat model
qodo-ai · Jan 18, 2025 · 536d97b · 536d97b
2 parents 4923c8d + cea1b58
commit 536d97b
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 0 deletions.
diff --git a/docs/docs/usage-guide/changing_a_model.md b/docs/docs/usage-guide/changing_a_model.md
@@ -166,6 +166,25 @@ drop_params = true
 
 AWS session is automatically authenticated from your environment, but you can also explicitly set `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY` and `AWS_REGION_NAME` environment variables. Please refer to [this document](https://litellm.vercel.app/docs/providers/bedrock) for more details.
 
+### DeepSeek
+
+To use deepseek-chat model with DeepSeek, for example, set:
+
+```toml
+[config] # in configuration.toml
+model = "deepseek/deepseek-chat"
+fallback_models=["deepseek/deepseek-chat"]
+```
+
+and fill up your key
+
+```toml
+[deepseek] # in .secrets.toml
+key = ...
+```
+
+(you can obtain a deepseek-chat key from [here](https://platform.deepseek.com))
+
 ### 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
@@ -29,6 +29,7 @@
     'claude-instant-1': 100000,
     'claude-2': 100000,
     'command-nightly': 4096,
+    'deepseek/deepseek-chat': 128000,  # 128K, but may be limited by config.max_model_tokens
     'replicate/llama-2-70b-chat:2c1608e18606fad2812020dc541930f2d0495ce32eee50074220b87300bc16e1': 4096,
     'meta-llama/Llama-2-7b-chat-hf': 4096,
     'vertex_ai/codechat-bison': 6144,

diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py
@@ -90,6 +90,10 @@ def __init__(self):
         if get_settings().get("GOOGLE_AI_STUDIO.GEMINI_API_KEY", None):
           os.environ["GEMINI_API_KEY"] = get_settings().google_ai_studio.gemini_api_key
 
+        # Support deepseek models
+        if get_settings().get("DEEPSEEK.KEY", None):
+            os.environ['DEEPSEEK_API_KEY'] = get_settings().get("DEEPSEEK.KEY")
+
     def prepare_logs(self, response, system, user, resp, finish_reason):
         response_log = response.dict().copy()
         response_log['system'] = system

diff --git a/pr_agent/settings/.secrets_template.toml b/pr_agent/settings/.secrets_template.toml
@@ -91,3 +91,6 @@ pat = ""
 # Optional, uncomment if you want to use Azure devops webhooks. Value assinged when you create the webhook
 # webhook_username = "<basic auth user>"
 # webhook_password = "<basic auth password>"
+
+[deepseek]
+key = ""