bugfix: 加入LLaMA.cpp

modules/models/OpenAI.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import json
 import logging
+import traceback
 
 import colorama
 import requests
@@ -85,11 +86,11 @@ class OpenAIClient(BaseLLMModel):
 
             # return i18n("**本月使用金额** ") + f"\u3000 ${rounded_usage}"
             return get_html("billing_info.html").format(
-                label=i18n("本月使用金额"),
-                usage_percent=usage_percent,
-                rounded_usage=rounded_usage,
-                usage_limit=usage_limit
-            )
+                    label = i18n("本月使用金额"),
+                    usage_percent = usage_percent,
+                    rounded_usage = rounded_usage,
+                    usage_limit = usage_limit
+                )
         except requests.exceptions.ConnectTimeout:
             status_text = (
                 STANDARD_ERROR_MSG + CONNECTION_TIMEOUT_MSG + ERROR_RETRIEVE_MSG
@@ -161,6 +162,7 @@ class OpenAIClient(BaseLLMModel):
                     timeout=timeout,
                 )
             except:
+                traceback.print_exc()
                 return None
         return response
 
@@ -170,6 +172,7 @@ class OpenAIClient(BaseLLMModel):
             "Authorization": f"Bearer {sensitive_id}",
         }
 
+
     def _get_billing_data(self, billing_url):
         with retrieve_proxy():
             response = requests.get(
@@ -240,6 +243,7 @@ class OpenAIClient(BaseLLMModel):
 
         return response
 
+
     def auto_name_chat_history(self, name_chat_method, user_question, chatbot, user_name, single_turn_checkbox):
         if len(self.history) == 2 and not single_turn_checkbox:
             user_question = self.history[0]["content"]
@@ -247,22 +251,19 @@ class OpenAIClient(BaseLLMModel):
                 ai_answer = self.history[1]["content"]
                 try:
                     history = [
-                        {"role": "system", "content": SUMMARY_CHAT_SYSTEM_PROMPT},
-                        {"role": "user", "content": f"Please write a title based on the following conversation:\n---\nUser: {user_question}\nAssistant: {ai_answer}"}
+                        { "role": "system", "content": SUMMARY_CHAT_SYSTEM_PROMPT},
+                        { "role": "user", "content": f"Please write a title based on the following conversation:\n---\nUser: {user_question}\nAssistant: {ai_answer}"}
                     ]
-                    response = self._single_query_at_once(
-                        history, temperature=0.0)
+                    response = self._single_query_at_once(history, temperature=0.0)
                     response = json.loads(response.text)
                     content = response["choices"][0]["message"]["content"]
                     filename = replace_special_symbols(content) + ".json"
                 except Exception as e:
                     logging.info(f"自动命名失败。{e}")
-                    filename = replace_special_symbols(user_question)[
-                        :16] + ".json"
+                    filename = replace_special_symbols(user_question)[:16] + ".json"
                 return self.rename_chat_history(filename, chatbot, user_name)
             elif name_chat_method == i18n("第一条提问"):
-                filename = replace_special_symbols(user_question)[
-                    :16] + ".json"
+                filename = replace_special_symbols(user_question)[:16] + ".json"
                 return self.rename_chat_history(filename, chatbot, user_name)
             else:
                 return gr.update()

modules/models/XMChat.py

@@ -16,7 +16,7 @@ from ..utils import *
 from .base_model import BaseLLMModel
 
 
-class XMChatClient(BaseLLMModel):
+class XMChat(BaseLLMModel):
     def __init__(self, api_key, user_name=""):
         super().__init__(model_name="xmchat", user=user_name)
         self.api_key = api_key
@@ -31,7 +31,7 @@ class XMChatClient(BaseLLMModel):
     def reset(self):
         self.session_id = str(uuid.uuid4())
         self.last_conv_id = None
-        return [], "已重置"
+        return super().reset()
 
     def image_to_base64(self, image_path):
         # 打开并加载图片

modules/models/midjourney.py

@@ -2,11 +2,10 @@ import base64
 import io
 import json
 import logging
+import os
 import pathlib
-import time
 import tempfile
-import os
-
+import time
 from datetime import datetime
 
 import requests
@@ -14,7 +13,7 @@ import tiktoken
 from PIL import Image
 
 from modules.config import retrieve_proxy
-from modules.models.models import XMChat
+from modules.models.XMChat import XMChat
 
 mj_proxy_api_base = os.getenv("MIDJOURNEY_PROXY_API_BASE")
 mj_discord_proxy_url = os.getenv("MIDJOURNEY_DISCORD_PROXY_URL")

modules/models/models.py

@@ -69,10 +69,10 @@ def get_model(
             model = LLaMA_Client(
                 model_name, lora_model_path, user_name=user_name)
         elif model_type == ModelType.XMChat:
-            from .XMChat import XMChatClient
+            from .XMChat import XMChat
             if os.environ.get("XMCHAT_API_KEY") != "":
                 access_key = os.environ.get("XMCHAT_API_KEY")
-            model = XMChatClient(api_key=access_key, user_name=user_name)
+            model = XMChat(api_key=access_key, user_name=user_name)
         elif model_type == ModelType.StableLM:
             from .StableLM import StableLM_Client
             model = StableLM_Client(model_name, user_name=user_name)