diff --git a/webui.py b/webui.py index 5cfd650..2406c19 100644 --- a/webui.py +++ b/webui.py @@ -21,17 +21,20 @@ if __name__ == "__main__": ) if "chat_list" not in st.session_state: - st.session_state["chat_list"] = {"对话1": {"need_rename": True}} + st.session_state["chat_list"] = {"对话1": {"need_rename": True, "chat_no": 1}} if "cur_chat_name" not in st.session_state: st.session_state["cur_chat_name"] = list(st.session_state["chat_list"].keys())[0] if "need_chat_name" not in st.session_state: st.session_state["need_chat_name"] = True + if "chat_count" not in st.session_state: + st.session_state["chat_count"] = 1 chat_list = [{"name": k, "chat_no": v.get("chat_no", 0)} for k, v in st.session_state.chat_list.items()] - pages = {i["name"]: { + chat_list = [x["name"] for x in sorted(chat_list, key=lambda x: x["chat_no"])] + pages1 = {i: { "icon": "chat", "func": dialogue_page, - } for i in sorted(chat_list, key=lambda x: x["chat_no"])} + } for i in chat_list} pages2 = { "新建对话": { @@ -47,41 +50,43 @@ if __name__ == "__main__": "func": knowledge_base_page, }, } - pages.update(pages2) - - - def on_change(key): - selection = st.session_state[key] - st.write(f"Selection changed to {selection}") + pages = {**pages1, **pages2} with st.sidebar: + options = chat_list + list(pages2) + icons = ["chat"] * len(chat_list) + [x["icon"] for x in pages2.values()] + 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()], + options=options, + icons=icons, menu_icon="chat-quote", default_index=default_index, ) if selected_page == "新建对话": - cur_chat_name = st.session_state["cur_chat_name"] + cur_chat_name = st.session_state.get("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")): - chat_no = len(st.session_state.chat_list) + 1 + st.session_state.chat_count += 1 + chat_no = st.session_state.chat_count new_chat_name = f"对话{chat_no}" st.session_state.chat_list[new_chat_name] = {"need_rename": True, "chat_no": chat_no} 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 + else: + 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 + st.experimental_rerun() + if st.session_state.get("delete_chat"): + st.session_state.delete_chat = False + st.experimental_rerun() elif selected_page in st.session_state.chat_list: - st.session_state["cur_chat_name"] = selected_page + st.session_state["selected_page"] = 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 41d844d..1708efd 100644 --- a/webui_pages/dialogue/dialogue.py +++ b/webui_pages/dialogue/dialogue.py @@ -32,7 +32,7 @@ def get_messages_history(history_len: int) -> List[Dict]: def dialogue_page(api: ApiRequest): chat_box.init_session() - chat_box.use_chat_name(st.session_state.cur_chat_name) + chat_box.use_chat_name(st.session_state.selected_page) with st.sidebar: # TODO: 对话模型与会话绑定 @@ -82,7 +82,7 @@ def dialogue_page(api: ApiRequest): chat_box.output_messages() - if st.session_state.chat_list.get(st.session_state.cur_chat_name, {}).get("need_rename"): + if st.session_state.chat_list.get(st.session_state.selected_page, {}).get("need_rename"): chat_input_placeholder = "请输入对话名称" else: chat_input_placeholder = "请输入对话内容，换行请使用Ctrl+Enter " @@ -91,15 +91,18 @@ def dialogue_page(api: ApiRequest): 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"): + if st.session_state.chat_list.get(st.session_state.selected_page, {}).get("need_rename"): if prompt in st.session_state.chat_list.keys(): st.toast("已有同名对话，请重新命名") else: - st.session_state.chat_list[prompt] = {"need_rename": False} - st.session_state.chat_list.pop(st.session_state.cur_chat_name) - chat_box.del_chat_name(st.session_state.cur_chat_name) + cur_chat_name = st.session_state.get("selected_page") + st.session_state.chat_list[prompt] = { + "need_rename": False, + "chat_no": st.session_state.chat_list[cur_chat_name]["chat_no"]} + st.session_state.chat_list.pop(cur_chat_name) + chat_box.del_chat_name(cur_chat_name) st.session_state.cur_chat_name = prompt - chat_box.use_chat_name(st.session_state.cur_chat_name) + chat_box.use_chat_name(prompt) st.experimental_rerun() else: history = get_messages_history(history_len) @@ -157,16 +160,17 @@ def dialogue_page(api: ApiRequest): 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) + cur_chat_name = st.session_state.get("selected_page") + chat_box.del_chat_name(cur_chat_name) + st.session_state.chat_list.pop(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)), - file_name=f"{now:%Y-%m-%d %H.%M}_{st.session_state.cur_chat_name}.md", + "".join(chat_box.export2md(st.session_state.selected_page)), + file_name=f"{now:%Y-%m-%d %H.%M}_{st.session_state.selected_page}.md", mime="text/markdown", use_container_width=True, )