增加多账号机制 (#513)

ChuanhuChatbot.py

@@ -76,7 +76,10 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
                         visible=not HIDE_MY_KEY,
                         label="API-Key",
                     )
-                    usageTxt = gr.Markdown("**发送消息** 或 **提交key** 以显示额度", elem_id="usage_display")
+                    if multi_api_key:
+                        usageTxt = gr.Markdown("多账号模式已开启，无需输入key，可直接开始对话", elem_id="usage_display")
+                    else:
+                        usageTxt = gr.Markdown("**发送消息** 或 **提交key** 以显示额度", elem_id="usage_display")
                     model_select_dropdown = gr.Dropdown(
                         label="选择模型", choices=MODELS, multiselect=False, value=MODELS[0]
                     )

config_example.json

@@ -8,5 +8,11 @@
             "two_column": false,
             "formula_ocr": true
         }
-    }
+    },
+    "multi_api_key": false,
+    "api_key_list": [
+        "sk-xxxxxxxxxxxxxxxxxxxxxxxx1",
+        "sk-xxxxxxxxxxxxxxxxxxxxxxxx2",
+        "sk-xxxxxxxxxxxxxxxxxxxxxxxx3"
+    ]
 }

modules/chat_func.py

@@ -35,6 +35,7 @@ initial_prompt = "You are a helpful assistant."
 HISTORY_DIR = "history"
 TEMPLATES_DIR = "templates"
 
+@shared.state.switching_api_key # 在不开启多账号模式的时候，这个装饰器不会起作用
 def get_response(
     openai_api_key, system_prompt, history, temperature, top_p, stream, selected_model
 ):
@@ -330,7 +331,7 @@ def predict(
     else:
         display_reference = ""
 
-    if len(openai_api_key) != 51:
+    if len(openai_api_key) != 51 and not shared.state.multi_api_key:
         status_text = standard_error_msg + no_apikey_msg
         logging.info(status_text)
         chatbot.append((inputs, ""))

modules/config.py

@@ -16,7 +16,8 @@ __all__ = [
     "retrieve_proxy",
     "log_level",
     "advance_docs",
-    "update_doc_config"
+    "update_doc_config",
+    "multi_api_key",
 ]
 
 # 添加一个统一的config文件，避免文件过多造成的疑惑（优先级最低）
@@ -36,6 +37,15 @@ if os.environ.get("dockerrun") == "yes":
 my_api_key = config.get("openai_api_key", "") # 在这里输入你的 API 密钥
 my_api_key = os.environ.get("my_api_key", my_api_key)
 
+## 多账户机制
+multi_api_key = config.get("multi_api_key", False) # 是否开启多账户机制
+if multi_api_key:
+    api_key_list = config.get("api_key_list", [])
+    if len(api_key_list) == 0:
+        logging.error("多账号模式已开启，但api_key_list为空，请检查config.json")
+        sys.exit(1)
+    shared.state.set_api_key_queue(api_key_list)
+
 auth_list = config.get("users", []) # 实际上是使用者的列表
 authflag = len(auth_list) > 0  # 是否开启认证的状态值，改为判断auth_list长度
 

modules/shared.py

@@ -1,7 +1,10 @@
 from modules.presets import COMPLETION_URL, BALANCE_API_URL, USAGE_API_URL, API_HOST
 import os
+import queue
+
 class State:
     interrupted = False
+    multi_api_key = False
     completion_url = COMPLETION_URL
     balance_api_url = BALANCE_API_URL
     usage_api_url = USAGE_API_URL
@@ -28,5 +31,25 @@ class State:
     def reset_all(self):
         self.interrupted = False
         self.completion_url = COMPLETION_URL
+    
+    def set_api_key_queue(self, api_key_list):
+        self.multi_api_key = True
+        self.api_key_queue = queue.Queue()
+        for api_key in api_key_list:
+            self.api_key_queue.put(api_key)
+
+    def switching_api_key(self, func):
+        if not hasattr(self, "api_key_queue"):
+            return func
+        
+        def wrapped(*args, **kwargs):
+            api_key = self.api_key_queue.get()
+            args = list(args)[1:]
+            ret = func(api_key, *args, **kwargs)
+            self.api_key_queue.put(api_key)
+            return ret
+
+        return wrapped
+        
 
 state = State()