From ea7d9556d34f8a248aaeb0ca45f6a4e557ddcbcc Mon Sep 17 00:00:00 2001
From: -LAN- <laipz8200@outlook.com>
Date: Thu, 6 Feb 2025 12:40:34 +0800
Subject: [PATCH] feat(model_providers): Support deepseek-r1 for Nvidia Catalog

Signed-off-by: -LAN- <laipz8200@outlook.com>
---
 .../model_providers/nvidia/llm/_position.yaml |  1 +
 .../nvidia/llm/deepseek-r1.yaml               | 35 +++++++++++++++++++
 .../model_providers/nvidia/llm/llm.py         |  2 +-
 3 files changed, 37 insertions(+), 1 deletion(-)
 create mode 100644 api/core/model_runtime/model_providers/nvidia/llm/deepseek-r1.yaml

diff --git a/api/core/model_runtime/model_providers/nvidia/llm/_position.yaml b/api/core/model_runtime/model_providers/nvidia/llm/_position.yaml
index ad01d430d61c79..36ec25fa2d534d 100644
--- a/api/core/model_runtime/model_providers/nvidia/llm/_position.yaml
+++ b/api/core/model_runtime/model_providers/nvidia/llm/_position.yaml
@@ -1,3 +1,4 @@
+- deepseek-ai/deepseek-r1
 - google/gemma-7b
 - google/codegemma-7b
 - google/recurrentgemma-2b
diff --git a/api/core/model_runtime/model_providers/nvidia/llm/deepseek-r1.yaml b/api/core/model_runtime/model_providers/nvidia/llm/deepseek-r1.yaml
new file mode 100644
index 00000000000000..159941eca442db
--- /dev/null
+++ b/api/core/model_runtime/model_providers/nvidia/llm/deepseek-r1.yaml
@@ -0,0 +1,35 @@
+model: deepseek-ai/deepseek-r1
+label:
+  en_US: deepseek-ai/deepseek-r1
+model_type: llm
+features:
+  - agent-thought
+model_properties:
+  mode: chat
+  context_size: 128000
+parameter_rules:
+  - name: temperature
+    use_template: temperature
+    min: 0
+    max: 1
+    default: 0.5
+  - name: top_p
+    use_template: top_p
+    min: 0
+    max: 1
+    default: 1
+  - name: max_tokens
+    use_template: max_tokens
+    min: 1
+    max: 1024
+    default: 1024
+  - name: frequency_penalty
+    use_template: frequency_penalty
+    min: -2
+    max: 2
+    default: 0
+  - name: presence_penalty
+    use_template: presence_penalty
+    min: -2
+    max: 2
+    default: 0
diff --git a/api/core/model_runtime/model_providers/nvidia/llm/llm.py b/api/core/model_runtime/model_providers/nvidia/llm/llm.py
index 1c98c6be6ca72d..fb572027984926 100644
--- a/api/core/model_runtime/model_providers/nvidia/llm/llm.py
+++ b/api/core/model_runtime/model_providers/nvidia/llm/llm.py
@@ -83,7 +83,7 @@ def validate_credentials(self, model: str, credentials: dict) -> None:
     def _add_custom_parameters(self, credentials: dict, model: str) -> None:
         credentials["mode"] = "chat"
 
-        if self.MODEL_SUFFIX_MAP[model]:
+        if self.MODEL_SUFFIX_MAP.get(model):
             credentials["server_url"] = f"https://ai.api.nvidia.com/v1/{self.MODEL_SUFFIX_MAP[model]}"
             credentials.pop("endpoint_url")
         else: