Merge branch 'main' into UI-new

.gitignore

@@ -140,7 +140,7 @@ dmypy.json
 api_key.txt
 config.json
 auth.json
-models/
+.models/
 lora/
 .idea
 templates/*

ChuanhuChatbot.py

@@ -31,7 +31,7 @@ def create_new_model():
 
 with gr.Blocks(theme=small_and_beautiful_theme) as demo:
     user_name = gr.State("")
-    promptTemplates = gr.State(load_template(get_template_names(plain=True)[0], mode=2))
+    promptTemplates = gr.State(load_template(get_template_names()[0], mode=2))
     user_question = gr.State("")
     assert type(my_api_key)==str
     user_api_key = gr.State(my_api_key)
@@ -64,7 +64,7 @@ with gr.Blocks(theme=small_and_beautiful_theme) as demo:
                             with gr.Column(scale=6, elem_id="history-select-wrap"):
                                 historyFileSelectDropdown = gr.Radio(
                                     label=i18n("从列表中加载对话"),
-                                    choices=get_history_names(plain=True),
+                                    choices=get_history_names(),
                                     # multiselect=False,
                                     container=False,
                                     elem_id="history-select-dropdown"
@@ -90,7 +90,7 @@ with gr.Blocks(theme=small_and_beautiful_theme) as demo:
                                 gr.Markdown(i18n("默认保存于history文件夹"))
                         with gr.Row():
                             with gr.Column():
-                                downloadFile = gr.File(interactive=True)
+                                downloadFile = gr.File(interactive=True, label="下载/上传历史记录")
 
             with gr.Column(elem_id="chuanhu-menu-footer"):
                 with gr.Row(elem_id="chuanhu-func-nav"):
@@ -179,9 +179,9 @@ with gr.Blocks(theme=small_and_beautiful_theme) as demo:
                                     with gr.Column(scale=6):
                                         templateFileSelectDropdown = gr.Dropdown(
                                             label=i18n("选择Prompt模板集合文件"),
-                                            choices=get_template_names(plain=True),
+                                            choices=get_template_names(),
                                             multiselect=False,
-                                            value=get_template_names(plain=True)[0],
+                                            value=get_template_names()[0],
                                             container=False,
                                         )
                                     with gr.Column(scale=1):
@@ -191,7 +191,7 @@ with gr.Blocks(theme=small_and_beautiful_theme) as demo:
                                         templateSelectDropdown = gr.Dropdown(
                                             label=i18n("从Prompt模板中加载"),
                                             choices=load_template(
-                                                get_template_names(plain=True)[0], mode=1
+                                                get_template_names()[0], mode=1
                                             ),
                                             multiselect=False,
                                             container=False,
@@ -399,7 +399,7 @@ with gr.Blocks(theme=small_and_beautiful_theme) as demo:
         current_model = get_model(model_name = MODELS[DEFAULT_MODEL], access_key = my_api_key)[0]
         current_model.set_user_identifier(user_name)
         chatbot = gr.Chatbot.update(label=MODELS[DEFAULT_MODEL])
-        return user_info, user_name, current_model, toggle_like_btn_visibility(DEFAULT_MODEL), *current_model.auto_load(), get_history_names(False, user_name), chatbot
+        return user_info, user_name, current_model, toggle_like_btn_visibility(DEFAULT_MODEL), *current_model.auto_load(), get_history_dropdown(user_name), chatbot
     demo.load(create_greeting, inputs=None, outputs=[user_info, user_name, current_model, like_dislike_area, systemPromptTxt, chatbot, historyFileSelectDropdown, chatbot], api_name="load")
     chatgpt_predict_args = dict(
         fn=predict,
@@ -446,7 +446,7 @@ with gr.Blocks(theme=small_and_beautiful_theme) as demo:
     )
 
     refresh_history_args = dict(
-        fn=get_history_names, inputs=[gr.State(False), user_name], outputs=[historyFileSelectDropdown]
+        fn=get_history_dropdown, inputs=[user_name], outputs=[historyFileSelectDropdown]
     )
 
 
@@ -524,8 +524,8 @@ with gr.Blocks(theme=small_and_beautiful_theme) as demo:
 
     # Template
     systemPromptTxt.change(set_system_prompt, [current_model, systemPromptTxt], None)
-    templateRefreshBtn.click(get_template_names, None, [templateFileSelectDropdown])
-    templateFileSelectDropdown.change(
+    templateRefreshBtn.click(get_template_dropdown, None, [templateFileSelectDropdown])
+    templateFileSelectDropdown.input(
         load_template,
         [templateFileSelectDropdown],
         [promptTemplates, templateSelectDropdown],
@@ -545,7 +545,7 @@ with gr.Blocks(theme=small_and_beautiful_theme) as demo:
         downloadFile,
         show_progress=True,
     )
-    saveHistoryBtn.click(get_history_names, [gr.State(False), user_name], [historyFileSelectDropdown])
+    saveHistoryBtn.click(get_history_dropdown, [user_name], [historyFileSelectDropdown])
     exportMarkdownBtn.click(
         export_markdown,
         [current_model, saveFileName, chatbot, user_name],

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/index_func.py

@@ -23,6 +23,7 @@ def get_documents(file_src):
         filename = os.path.basename(filepath)
         file_type = os.path.splitext(filename)[1]
         logging.info(f"loading file: {filename}")
+        texts = None
         try:
             if file_type == ".pdf":
                 logging.debug("Loading PDF...")
@@ -72,8 +73,9 @@ def get_documents(file_src):
             logging.error(f"Error loading file: {filename}")
             traceback.print_exc()
 
-        texts = text_splitter.split_documents(texts)
-        documents.extend(texts)
+        if texts is not None:
+            texts = text_splitter.split_documents(texts)
+            documents.extend(texts)
     logging.debug("Documents loaded.")
     return documents
 

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:
@@ -718,7 +724,7 @@ class BaseLLMModel:
             history_file_path = filename
         try:
             os.remove(history_file_path)
-            return i18n("删除对话历史成功"), get_history_names(False, user_name), []
+            return i18n("删除对话历史成功"), get_history_dropdown(user_name), []
         except:
             logging.info(f"删除对话历史失败 {history_file_path}")
             return i18n("对话历史")+filename+i18n("已经被删除啦"), gr.update(), gr.update()

modules/models/models.py

@@ -580,8 +580,8 @@ def get_model(
             logging.info(msg)
             lora_selector_visibility = True
             if os.path.isdir("lora"):
-                lora_choices = get_file_names(
-                    "lora", plain=True, filetypes=[""])
+                lora_choices = get_file_names_dropdown_by_pinyin(
+                    "lora", filetypes=[""])
             lora_choices = ["No LoRA"] + lora_choices
         elif model_type == ModelType.LLaMA and lora_model_path != "":
             logging.info(f"正在加载LLaMA模型: {model_name} + {lora_model_path}")
@@ -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/models/spark.py

@@ -0,0 +1,163 @@
+import _thread as thread
+import base64
+import datetime
+import hashlib
+import hmac
+import json
+from collections import deque
+from urllib.parse import urlparse
+import ssl
+from datetime import datetime
+from time import mktime
+from urllib.parse import urlencode
+from wsgiref.handlers import format_date_time
+from threading import Condition
+import websocket
+import logging
+
+from .base_model import BaseLLMModel, CallbackToIterator
+
+
+class Ws_Param(object):
+    # 来自官方 Demo
+    # 初始化
+    def __init__(self, APPID, APIKey, APISecret, Spark_url):
+        self.APPID = APPID
+        self.APIKey = APIKey
+        self.APISecret = APISecret
+        self.host = urlparse(Spark_url).netloc
+        self.path = urlparse(Spark_url).path
+        self.Spark_url = Spark_url
+
+    # 生成url
+    def create_url(self):
+        # 生成RFC1123格式的时间戳
+        now = datetime.now()
+        date = format_date_time(mktime(now.timetuple()))
+
+        # 拼接字符串
+        signature_origin = "host: " + self.host + "\n"
+        signature_origin += "date: " + date + "\n"
+        signature_origin += "GET " + self.path + " HTTP/1.1"
+
+        # 进行hmac-sha256进行加密
+        signature_sha = hmac.new(
+            self.APISecret.encode("utf-8"),
+            signature_origin.encode("utf-8"),
+            digestmod=hashlib.sha256,
+        ).digest()
+
+        signature_sha_base64 = base64.b64encode(
+            signature_sha).decode(encoding="utf-8")
+
+        authorization_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'
+
+        authorization = base64.b64encode(authorization_origin.encode("utf-8")).decode(
+            encoding="utf-8"
+        )
+
+        # 将请求的鉴权参数组合为字典
+        v = {"authorization": authorization, "date": date, "host": self.host}
+        # 拼接鉴权参数，生成url
+        url = self.Spark_url + "?" + urlencode(v)
+        # 此处打印出建立连接时候的url,参考本demo的时候可取消上方打印的注释，比对相同参数时生成的url与自己代码生成的url是否一致
+        return url
+
+
+class Spark_Client(BaseLLMModel):
+    def __init__(self, model_name, appid, api_key, api_secret, user_name="") -> None:
+        super().__init__(model_name=model_name, user=user_name)
+        self.api_key = api_key
+        self.appid = appid
+        self.api_secret = api_secret
+        if None in [self.api_key, self.appid, self.api_secret]:
+            raise Exception("请在配置文件或者环境变量中设置讯飞的API Key、APP ID和API Secret")
+        if "2.0" in self.model_name:
+            self.spark_url = "wss://spark-api.xf-yun.com/v2.1/chat"
+            self.domain = "generalv2"
+        else:
+            self.spark_url = "wss://spark-api.xf-yun.com/v1.1/chat"
+            self.domain = "general"
+
+    # 收到websocket错误的处理
+    def on_error(self, ws, error):
+        ws.iterator.callback("出现了错误:" + error)
+
+    # 收到websocket关闭的处理
+    def on_close(self, ws, one, two):
+        pass
+
+    # 收到websocket连接建立的处理
+    def on_open(self, ws):
+        thread.start_new_thread(self.run, (ws,))
+
+    def run(self, ws, *args):
+        data = json.dumps(
+            self.gen_params()
+        )
+        ws.send(data)
+
+    # 收到websocket消息的处理
+    def on_message(self, ws, message):
+        ws.iterator.callback(message)
+
+    def gen_params(self):
+        """
+        通过appid和用户的提问来生成请参数
+        """
+        data = {
+            "header": {"app_id": self.appid, "uid": "1234"},
+            "parameter": {
+                "chat": {
+                    "domain": self.domain,
+                    "random_threshold": self.temperature,
+                    "max_tokens": 4096,
+                    "auditing": "default",
+                }
+            },
+            "payload": {"message": {"text": self.history}},
+        }
+        return data
+
+    def get_answer_stream_iter(self):
+        wsParam = Ws_Param(self.appid, self.api_key, self.api_secret, self.spark_url)
+        websocket.enableTrace(False)
+        wsUrl = wsParam.create_url()
+        ws = websocket.WebSocketApp(
+            wsUrl,
+            on_message=self.on_message,
+            on_error=self.on_error,
+            on_close=self.on_close,
+            on_open=self.on_open,
+        )
+        ws.appid = self.appid
+        ws.domain = self.domain
+
+        # Initialize the CallbackToIterator
+        ws.iterator = CallbackToIterator()
+
+        # Start the WebSocket connection in a separate thread
+        thread.start_new_thread(
+            ws.run_forever, (), {"sslopt": {"cert_reqs": ssl.CERT_NONE}}
+        )
+
+        # Iterate over the CallbackToIterator instance
+        answer = ""
+        total_tokens = 0
+        for message in ws.iterator:
+            data = json.loads(message)
+            code = data["header"]["code"]
+            if code != 0:
+                ws.close()
+                raise Exception(f"请求错误: {code}, {data}")
+            else:
+                choices = data["payload"]["choices"]
+                status = choices["status"]
+                content = choices["text"][0]["content"]
+                if "usage" in data["payload"]:
+                    total_tokens = data["payload"]["usage"]["text"]["total_tokens"]
+                answer += content
+                if status == 2:
+                    ws.iterator.finish()  # Finish the iterator when the status is 2
+                    ws.close()
+                yield answer, total_tokens

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 = [

modules/utils.py

@@ -354,31 +354,50 @@ def save_file(filename, system, history, chatbot, user_name):
 def sorted_by_pinyin(list):
     return sorted(list, key=lambda char: lazy_pinyin(char)[0][0])
 
+def sorted_by_last_modified_time(list, dir):
+    return sorted(list, key=lambda char: os.path.getmtime(os.path.join(dir, char)), reverse=True)
 
-def get_file_names(dir, plain=False, filetypes=[".json"]):
-    logging.debug(f"获取文件名列表，目录为{dir}，文件类型为{filetypes}，是否为纯文本列表{plain}")
+def get_file_names_by_type(dir, filetypes=[".json"]):
+    logging.debug(f"获取文件名列表，目录为{dir}，文件类型为{filetypes}")
     files = []
     try:
         for type in filetypes:
             files += [f for f in os.listdir(dir) if f.endswith(type)]
     except FileNotFoundError:
-        files = []
-    files = sorted_by_pinyin(files)
-    if files == []:
         files = [""]
     logging.debug(f"files are:{files}")
-    if plain:
-        return files
-    else:
-        return gr.Dropdown.update(choices=files)
+    return files
+
+def get_file_names_by_pinyin(dir, filetypes=[".json"]):
+    files = get_file_names_by_type(dir, filetypes)
+    if files != [""]:
+        files = sorted_by_pinyin(files)
+    logging.debug(f"files are:{files}")
+    return files
+
+def get_file_names_dropdown_by_pinyin(dir, filetypes=[".json"]):
+    files = get_file_names_by_pinyin(dir, filetypes)
+    return gr.Dropdown.update(choices=files)
+
+def get_file_names_by_last_modified_time(dir, filetypes=[".json"]):
+    files = get_file_names_by_type(dir, filetypes)
+    if files != [""]:
+        files = sorted_by_last_modified_time(files, dir)
+    logging.debug(f"files are:{files}")
+    return files
 
 
-def get_history_names(plain=False, user_name=""):
+def get_history_names(user_name=""):
     logging.debug(f"从用户 {user_name} 中获取历史记录文件名列表")
     if user_name == "" and hide_history_when_not_logged_in:
         return ""
     else:
-        return get_file_names(os.path.join(HISTORY_DIR, user_name), plain)
+        history_files = get_file_names_by_last_modified_time(os.path.join(HISTORY_DIR, user_name))
+        return history_files
+
+def get_history_dropdown(user_name=""):
+    history_names = get_history_names(user_name)
+    return gr.Dropdown.update(choices=history_names)
 
 
 def load_template(filename, mode=0):
@@ -406,9 +425,14 @@ def load_template(filename, mode=0):
         )
 
 
-def get_template_names(plain=False):
+def get_template_names():
     logging.debug("获取模板文件名列表")
-    return get_file_names(TEMPLATES_DIR, plain, filetypes=[".csv", "json"])
+    return get_file_names_by_pinyin(TEMPLATES_DIR, filetypes=[".csv", "json"])
+
+def get_template_dropdown():
+    logging.debug("获取模板下拉菜单")
+    template_names = get_template_names()
+    return gr.Dropdown.update(choices=template_names)
 
 
 def get_template_content(templates, selection, original_system_prompt):

requirements.txt

@@ -26,3 +26,5 @@ unstructured
 google-api-python-client
 tabulate
 ujson
+python-docx
+websocket_client