feat: 加入讯飞星火大模型支持 #877

config_example.json

@@ -11,6 +11,9 @@
     "midjourney_proxy_api_secret": "", // 你的 MidJourney Proxy API Secret，用于鉴权访问 api，可选
     "midjourney_discord_proxy_url": "", // 你的 MidJourney Discord Proxy URL，用于对生成对图进行反代，可选
     "midjourney_temp_folder": "./tmp", // 你的 MidJourney 临时文件夹，用于存放生成的图片，填空则关闭自动下载切图（直接显示MJ的四宫格图）
+    "spark_appid": "", // 你的 讯飞星火大模型 API AppID，用于讯飞星火大模型对话模型
+    "spark_api_key": "", // 你的 讯飞星火大模型 API Key，用于讯飞星火大模型对话模型
+    "spark_api_secret": "", // 你的 讯飞星火大模型 API Secret，用于讯飞星火大模型对话模型
 
 
     //== Azure ==

modules/config.py

@@ -123,6 +123,13 @@ os.environ["MIDJOURNEY_DISCORD_PROXY_URL"] = midjourney_discord_proxy_url
 midjourney_temp_folder = config.get("midjourney_temp_folder", "")
 os.environ["MIDJOURNEY_TEMP_FOLDER"] = midjourney_temp_folder
 
+spark_api_key = config.get("spark_api_key", "")
+os.environ["SPARK_API_KEY"] = spark_api_key
+spark_appid = config.get("spark_appid", "")
+os.environ["SPARK_APPID"] = spark_appid
+spark_api_secret = config.get("spark_api_secret", "")
+os.environ["SPARK_API_SECRET"] = spark_api_secret
+
 load_config_to_environ(["openai_api_type", "azure_openai_api_key", "azure_openai_api_base_url",
                        "azure_openai_api_version", "azure_deployment_name", "azure_embedding_deployment_name", "azure_embedding_model_name"])
 

modules/models/base_model.py

@@ -142,6 +142,7 @@ class ModelType(Enum):
     GooglePaLM = 9
     LangchainChat = 10
     Midjourney = 11
+    Spark = 12
 
     @classmethod
     def get_type(cls, model_name: str):
@@ -171,6 +172,8 @@ class ModelType(Enum):
             model_type = ModelType.Midjourney
         elif "azure" in model_name_lower or "api" in model_name_lower:
             model_type = ModelType.LangchainChat
+        elif "星火大模型" in model_name_lower:
+            model_type = ModelType.Spark
         else:
             model_type = ModelType.Unknown
         return model_type
@@ -269,9 +272,12 @@ class BaseLLMModel:
         if display_append:
             display_append = '\n\n<hr class="append-display no-in-raw" />' + display_append
         partial_text = ""
+        token_increment = 1
         for partial_text in stream_iter:
+            if type(partial_text) == tuple:
+                partial_text, token_increment = partial_text
             chatbot[-1] = (chatbot[-1][0], partial_text + display_append)
-            self.all_token_counts[-1] += 1
+            self.all_token_counts[-1] += token_increment
             status_text = self.token_message()
             yield get_return_value()
             if self.interrupted:

modules/models/models.py

@@ -625,6 +625,9 @@ def get_model(
             from .midjourney import Midjourney_Client
             mj_proxy_api_secret = os.getenv("MIDJOURNEY_PROXY_API_SECRET")
             model = Midjourney_Client(model_name, mj_proxy_api_secret, user_name=user_name)
+        elif model_type == ModelType.Spark:
+            from .spark import Spark_Client
+            model = Spark_Client(model_name, os.getenv("SPARK_APPID"), os.getenv("SPARK_API_KEY"), os.getenv("SPARK_API_SECRET"), user_name=user_name)
         elif model_type == ModelType.Unknown:
             raise ValueError(f"未知模型: {model_name}")
         logging.info(msg)

modules/presets.py

@@ -69,7 +69,9 @@ ONLINE_MODELS = [
     "yuanai-1.0-rhythm_poems",
     "minimax-abab4-chat",
     "minimax-abab5-chat",
-    "midjourney"
+    "midjourney",
+    "讯飞星火大模型V2.0",
+    "讯飞星火大模型V1.5"
 ]
 
 LOCAL_MODELS = [

requirements.txt

@@ -27,3 +27,4 @@ google-api-python-client
 tabulate
 ujson
 python-docx
+websocket_client