From 7944ce059415414a8c4867549bcb776561dff944 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=B0=8F=E7=BE=BD?= <xiaoyu@sulg.top>
Date: Thu, 8 Aug 2024 12:01:21 +0000
Subject: [PATCH] feat: wenxin add yi-34b-chat (#7117)

---
 .../model_providers/wenxin/llm/ernie_bot.py   |  4 ++-
 .../wenxin/llm/yi_34b_chat.yaml               | 30 +++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)
 create mode 100644 api/core/model_runtime/model_providers/wenxin/llm/yi_34b_chat.yaml

diff --git a/api/core/model_runtime/model_providers/wenxin/llm/ernie_bot.py b/api/core/model_runtime/model_providers/wenxin/llm/ernie_bot.py
index f585fa2c41..e345663d36 100644
--- a/api/core/model_runtime/model_providers/wenxin/llm/ernie_bot.py
+++ b/api/core/model_runtime/model_providers/wenxin/llm/ernie_bot.py
@@ -142,6 +142,7 @@ class ErnieBotModel:
         'ernie-character-8k-0321': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-char-8k',
         'ernie-4.0-turbo-8k': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-4.0-turbo-8k',
         'ernie-4.0-turbo-8k-preview': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-4.0-turbo-8k-preview',
+        'yi_34b_chat': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/yi_34b_chat',
     }
 
     function_calling_supports = [
@@ -154,7 +155,8 @@ class ErnieBotModel:
         'ernie-3.5-128k',
         'ernie-4.0-8k',
         'ernie-4.0-turbo-8k',
-        'ernie-4.0-turbo-8k-preview'
+        'ernie-4.0-turbo-8k-preview',
+        'yi_34b_chat'
     ]
 
     api_key: str = ''
diff --git a/api/core/model_runtime/model_providers/wenxin/llm/yi_34b_chat.yaml b/api/core/model_runtime/model_providers/wenxin/llm/yi_34b_chat.yaml
new file mode 100644
index 0000000000..0b247fbd22
--- /dev/null
+++ b/api/core/model_runtime/model_providers/wenxin/llm/yi_34b_chat.yaml
@@ -0,0 +1,30 @@
+model: yi_34b_chat
+label:
+  en_US: yi_34b_chat
+model_type: llm
+features:
+  - agent-thought
+model_properties:
+  mode: chat
+  context_size: 32000
+parameter_rules:
+  - name: temperature
+    use_template: temperature
+    min: 0.1
+    max: 1.0
+    default: 0.95
+  - name: top_p
+    use_template: top_p
+    min: 0
+    max: 1.0
+    default: 0.7
+  - name: max_tokens
+    use_template: max_tokens
+    default: 4096
+    min: 2
+    max: 4096
+  - name: presence_penalty
+    use_template: presence_penalty
+    default: 1.0
+    min: 1.0
+    max: 2.0