feat: 增加多用户显示与多用户历史文件夹 (#489)

* 功能添加: 对每个用户添加独立history文件夹

* 修改: 无用户名时的兼容性 Fix Username

* UI美化

---------

Co-authored-by: Keldos <i@keldos.me>

ChuanhuChatbot.py

@@ -19,6 +19,7 @@ with open("assets/custom.css", "r", encoding="utf-8") as f:
     customCSS = f.read()
 
 with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
+    user_name = gr.State("")
     history = gr.State([])
     token_count = gr.State([])
     promptTemplates = gr.State(load_template(get_template_names(plain=True)[0], mode=2))
@@ -29,8 +30,18 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
 
     with gr.Row():
         gr.HTML(title)
+        user_ip = gr.Markdown(value="", elem_id="user_ip")
         status_display = gr.Markdown(get_geoip(), elem_id="status_display")
 
+        # https://github.com/gradio-app/gradio/pull/3296
+        def create_greeting(request: gr.Request):
+            if hasattr(request, "username") and request.username: # is not None or is not ""
+                logging.info(f"Get User Name: {request.username}")
+                return gr.Markdown.update(value=f"User: {request.username}"), request.username
+            else:
+                return gr.Markdown.update(value=f"User: 管理员"), ""
+        demo.load(create_greeting, inputs=None, outputs=[user_ip, user_name])
+
     with gr.Row().style(equal_height=True):
         with gr.Column(scale=5):
             with gr.Row():
@@ -323,27 +334,27 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
     # S&L
     saveHistoryBtn.click(
         save_chat_history,
-        [saveFileName, systemPromptTxt, history, chatbot],
+        [saveFileName, systemPromptTxt, history, chatbot, user_name],
         downloadFile,
         show_progress=True,
     )
-    saveHistoryBtn.click(get_history_names, None, [historyFileSelectDropdown])
+    saveHistoryBtn.click(get_history_names, [gr.State(False), user_name], [historyFileSelectDropdown])
     exportMarkdownBtn.click(
         export_markdown,
-        [saveFileName, systemPromptTxt, history, chatbot],
+        [saveFileName, systemPromptTxt, history, chatbot, user_name],
         downloadFile,
         show_progress=True,
     )
-    historyRefreshBtn.click(get_history_names, None, [historyFileSelectDropdown])
+    historyRefreshBtn.click(get_history_names, [gr.State(False), user_name], [historyFileSelectDropdown])
     historyFileSelectDropdown.change(
         load_chat_history,
-        [historyFileSelectDropdown, systemPromptTxt, history, chatbot],
+        [historyFileSelectDropdown, systemPromptTxt, history, chatbot, user_name],
         [saveFileName, systemPromptTxt, history, chatbot],
         show_progress=True,
     )
     downloadFile.change(
         load_chat_history,
-        [downloadFile, systemPromptTxt, history, chatbot],
+        [downloadFile, systemPromptTxt, history, chatbot, user_name],
         [saveFileName, systemPromptTxt, history, chatbot],
     )
 

assets/custom.css

@@ -18,6 +18,19 @@ footer {
     opacity: 0.85;
 }
 
+/* user_ip */
+#user_ip {
+    display: flex;
+    min-height: 2.5em;
+    align-items: flex-end;
+    justify-content: flex-end;
+}
+#user_ip p {
+    font-size: .85em;
+    font-family: monospace;
+    color: var(--body-text-color-subdued);
+}
+
 /* status_display */
 #status_display {
     display: flex;

modules/presets.py

@@ -1,5 +1,6 @@
 # -*- coding:utf-8 -*-
 import gradio as gr
+from pathlib import Path
 
 # ChatGPT 设置
 initial_prompt = "You are a helpful assistant."
@@ -7,7 +8,7 @@ API_HOST = "api.openai.com"
 COMPLETION_URL = "https://api.openai.com/v1/chat/completions"
 BALANCE_API_URL="https://api.openai.com/dashboard/billing/credit_grants"
 USAGE_API_URL="https://api.openai.com/dashboard/billing/usage"
-HISTORY_DIR = "history"
+HISTORY_DIR = Path("history")
 TEMPLATES_DIR = "templates"
 
 # 错误信息

modules/utils.py

@@ -190,46 +190,46 @@ def delete_last_conversation(chatbot, history, previous_token_count):
     )
 
 
-def save_file(filename, system, history, chatbot):
-    logging.info("保存对话历史中……")
-    os.makedirs(HISTORY_DIR, exist_ok=True)
+def save_file(filename, system, history, chatbot, user_name):
+    logging.info(f"{user_name} 保存对话历史中……")
+    os.makedirs(HISTORY_DIR / user_name, exist_ok=True)
     if filename.endswith(".json"):
         json_s = {"system": system, "history": history, "chatbot": chatbot}
         print(json_s)
-        with open(os.path.join(HISTORY_DIR, filename), "w") as f:
+        with open(os.path.join(HISTORY_DIR / user_name, filename), "w") as f:
             json.dump(json_s, f)
     elif filename.endswith(".md"):
         md_s = f"system: \n- {system} \n"
         for data in history:
             md_s += f"\n{data['role']}: \n- {data['content']} \n"
-        with open(os.path.join(HISTORY_DIR, filename), "w", encoding="utf8") as f:
+        with open(os.path.join(HISTORY_DIR / user_name, filename), "w", encoding="utf8") as f:
             f.write(md_s)
-    logging.info("保存对话历史完毕")
-    return os.path.join(HISTORY_DIR, filename)
+    logging.info(f"{user_name} 保存对话历史完毕")
+    return os.path.join(HISTORY_DIR / user_name, filename)
 
 
-def save_chat_history(filename, system, history, chatbot):
+def save_chat_history(filename, system, history, chatbot, user_name):
     if filename == "":
         return
     if not filename.endswith(".json"):
         filename += ".json"
-    return save_file(filename, system, history, chatbot)
+    return save_file(filename, system, history, chatbot, user_name)
 
 
-def export_markdown(filename, system, history, chatbot):
+def export_markdown(filename, system, history, chatbot, user_name):
     if filename == "":
         return
     if not filename.endswith(".md"):
         filename += ".md"
-    return save_file(filename, system, history, chatbot)
+    return save_file(filename, system, history, chatbot, user_name)
 
 
-def load_chat_history(filename, system, history, chatbot):
-    logging.info("加载对话历史中……")
+def load_chat_history(filename, system, history, chatbot, user_name):
+    logging.info(f"{user_name} 加载对话历史中……")
     if type(filename) != str:
         filename = filename.name
     try:
-        with open(os.path.join(HISTORY_DIR, filename), "r") as f:
+        with open(os.path.join(HISTORY_DIR / user_name, filename), "r") as f:
             json_s = json.load(f)
         try:
             if type(json_s["history"][0]) == str:
@@ -245,10 +245,10 @@ def load_chat_history(filename, system, history, chatbot):
         except:
             # 没有对话历史
             pass
-        logging.info("加载对话历史完毕")
+        logging.info(f"{user_name} 加载对话历史完毕")
         return filename, json_s["system"], json_s["history"], json_s["chatbot"]
     except FileNotFoundError:
-        logging.info("没有找到对话历史文件，不执行任何操作")
+        logging.info(f"{user_name} 没有找到对话历史文件，不执行任何操作")
         return filename, system, history, chatbot
 
 
@@ -267,15 +267,16 @@ def get_file_names(dir, plain=False, filetypes=[".json"]):
     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)
 
 
-def get_history_names(plain=False):
-    logging.info("获取历史记录文件名列表")
-    return get_file_names(HISTORY_DIR, plain)
+def get_history_names(plain=False, user_name=""):
+    logging.info(f"从用户 {user_name} 中获取历史记录文件名列表")
+    return get_file_names(HISTORY_DIR / user_name, plain)
 
 
 def load_template(filename, mode=0):