feat: add wenxin 4.0 model

zhayujie · Oct 27, 2023 · 49446d4 · 49446d4
1 parent ced560e
commit 49446d4
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 6 deletions.
diff --git a/bot/baidu/baidu_wenxin.py b/bot/baidu/baidu_wenxin.py
@@ -16,7 +16,10 @@ class BaiduWenxinBot(Bot):
 
     def __init__(self):
         super().__init__()
-        self.sessions = SessionManager(BaiduWenxinSession, model=conf().get("baidu_wenxin_model") or "eb-instant")
+        wenxin_model = conf().get("baidu_wenxin_model") or "eb-instant"
+        if conf().get("model") and conf().get("model") == "wenxin-4":
+            wenxin_model = "completions_pro"
+        self.sessions = SessionManager(BaiduWenxinSession, model=wenxin_model)
 
     def reply(self, query, context=None):
         # acquire reply content

diff --git a/bridge/bridge.py b/bridge/bridge.py
@@ -23,7 +23,7 @@ def __init__(self):
             self.btype["chat"] = const.OPEN_AI
         if conf().get("use_azure_chatgpt", False):
             self.btype["chat"] = const.CHATGPTONAZURE
-        if model_type in ["wenxin"]:
+        if model_type in ["wenxin", "wenxin-4"]:
             self.btype["chat"] = const.BAIDU
         if model_type in ["xunfei"]:
             self.btype["chat"] = const.XUNFEI

diff --git a/common/const.py b/common/const.py
@@ -5,8 +5,6 @@
 XUNFEI = "xunfei"
 CHATGPTONAZURE = "chatGPTOnAzure"
 LINKAI = "linkai"
-
-VERSION = "1.3.0"
-
 CLAUDEAI = "claude"
-MODEL_LIST = ["gpt-3.5-turbo", "gpt-3.5-turbo-16k", "gpt-4", "wenxin", "xunfei","claude"]
+
+MODEL_LIST = ["gpt-3.5-turbo", "gpt-3.5-turbo-16k", "gpt-4", "wenxin", "wenxin-4", "xunfei", "claude"]