From 9974bf33a0657e8c72fe74c3f97226d5f44e1028 Mon Sep 17 00:00:00 2001
From: liunux4odoo <liunu@qq.com>
Date: Sun, 13 Aug 2023 12:59:06 +0800
Subject: [PATCH] fix webui: switch to created chat session automatically

---
 webui.py                         | 39 +++++++++++++++++---------------
 webui_pages/dialogue/dialogue.py |  8 ++++++-
 2 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/webui.py b/webui.py
index d15dc59..c5a00e8 100644
--- a/webui.py
+++ b/webui.py
@@ -52,31 +52,34 @@ if __name__ == "__main__":
     def on_change(key):
         selection = st.session_state[key]
         st.write(f"Selection changed to {selection}")
-
-    def on_page_change(key):
-        cur_chat_name = st.session_state["cur_chat_name"]
-        if (st.session_state[key] == "新建对话"
-            and cur_chat_name != "新建对话"
-            and not st.session_state.get("prompt")):
-            new_chat_name = f"对话{len(st.session_state.chat_list) + 1}"
-            st.session_state.chat_list[new_chat_name] = {"need_rename": True}
-            st.session_state["cur_chat_name"] = new_chat_name
-            st.session_state[key]  = new_chat_name
-        elif st.session_state[key] not in ["新建对话", "知识库管理"]:
-            if st.session_state.get("prompt"):
-                st.session_state["cur_chat_name"] = st.session_state.get("prompt")
-            else:
-                st.session_state["cur_chat_name"] = st.session_state[key]
         
     with st.sidebar:
+        default_index = list(pages).index(st.session_state["cur_chat_name"])
         selected_page = option_menu(
             "langchain-chatglm",
             options=list(pages.keys()),
             icons=[i["icon"] for i in pages.values()],
             menu_icon="chat-quote",
-            key="selected_page",
-            on_change=on_page_change,
+            default_index=default_index,
         )
 
-    if selected_page == "知识库管理" or selected_page in pages:
+    if selected_page == "新建对话":
+        cur_chat_name = st.session_state["cur_chat_name"]
+        if (not st.session_state.get("create_chat")
+            and not st.session_state.get("renamde_chat")
+            and not st.session_state.get("delete_chat")):
+            new_chat_name = f"对话{len(st.session_state.chat_list) + 1}"
+            st.session_state.chat_list[new_chat_name] = {"need_rename": True}
+            st.session_state["cur_chat_name"] = new_chat_name
+            st.experimental_rerun()
+        if st.session_state.get("create_chat"):
+            st.session_state.create_chat = False
+        if st.session_state.get("renamde_chat"):
+            st.session_state.renamde_chat = False
+        if st.session_state.get("delete_chat"):
+            st.session_state.delete_chat = False
+    elif selected_page in st.session_state.chat_list:
+        st.session_state["cur_chat_name"] = selected_page
+
+    if selected_page in pages:
         pages[selected_page]["func"](api)
diff --git a/webui_pages/dialogue/dialogue.py b/webui_pages/dialogue/dialogue.py
index 8b4e7c4..41d844d 100644
--- a/webui_pages/dialogue/dialogue.py
+++ b/webui_pages/dialogue/dialogue.py
@@ -88,7 +88,7 @@ def dialogue_page(api: ApiRequest):
         chat_input_placeholder = "请输入对话内容，换行请使用Ctrl+Enter "
     
     def on_prompt():
-        st.session_state["selected_page"] = prompt        
+        st.session_state.rename_chat = True
 
     if prompt := st.chat_input(chat_input_placeholder, key="prompt", on_submit=on_prompt):
         if st.session_state.chat_list.get(st.session_state.cur_chat_name, {}).get("need_rename"):
@@ -146,17 +146,23 @@ def dialogue_page(api: ApiRequest):
                 use_container_width=True,
         ):
             chat_box.reset_history()
+            st.experimental_rerun()
+
+        def on_delete_chat():
+            st.session_state.delete_chat = True
 
         if cols[2].button(
                 "删除对话",
                 disabled=len(st.session_state.chat_list) <= 1,
                 use_container_width=True,
+                on_click=on_delete_chat
         ):
             chat_box.del_chat_name(st.session_state.cur_chat_name)
             st.session_state.chat_list.pop(st.session_state.cur_chat_name)
             st.session_state.cur_chat_name = list(st.session_state.chat_list.keys())[0]
             chat_box.use_chat_name(st.session_state.cur_chat_name)
             st.experimental_rerun()
+
     export_btn.download_button(
         "导出记录",
         "".join(chat_box.export2md(st.session_state.cur_chat_name)),