Update app.py

app.py

@@ -457,7 +457,7 @@ if __name__ == '__main__':
         if 'id_chat' not in st.session_state:
             st.session_state.id_chat = 'New Chat'
 
-        def options_list():
+        def options_list(conversations: Dict[str, list]):
             if st.session_state.id_chat == 'New Chat':
                 return [st.session_state.id_chat] + list(conversations.keys())
             else:
@@ -465,10 +465,10 @@ if __name__ == '__main__':
         
         with st.sidebar:
             print('SELECTBOX')
-            print(st.session_state.id_chat)
+            print(st.session_state.id_chat
             st.session_state.id_chat = st.selectbox(
                 label='Choose a conversation',
-                options=options_list(),
+                options=options_list(conversations),
                 index=0,
                 placeholder='_',
                 key='chat_id'
@@ -476,43 +476,21 @@ if __name__ == '__main__':
 
             st.session_state.messages = conversations[st.session_state.id_chat] if st.session_state.id_chat != 'New Chat' else []
             
-            def update_selectbox_remove():
-                tmp = st.session_state.chat_id
-                
-                if len(list(conversations.keys())) < 2:
-                    st.session_state.chat_id = 'New Chat'
-                    st.session_state.messages = []
-                else:
-                    for conversation in list(conversations.keys()):
-                        if conversation != conversations[st.session_state.chat_id]:
-                            st.session_state.chat_id = conversation
-                            break
-                    
-                conversations.pop(tmp)
+            def update_selectbox_remove(conversations_path, conversations):
+                conversations.pop(st.session_state.chat_id)
                 with open(conversations_path, 'wb') as fp:
                     packed_bytes = msgpack.packb(conversations, use_bin_type=True)
                     fp.write(packed_bytes)
+                st.session_state.chat_id = 'New Chat'
                     
             
-            if st.button(
+            st.button(
                 'Delete Chat', 
                 use_container_width=True, 
-                disabled=False if st.session_state.id_chat != 'New Chat' else True
-            ):
-                tmp = st.session_state.id_chat
-                    
-                if len(list(conversations.keys())) < 2:
-                    st.session_state.id_chat = 'New Chat'
-                else:
-                    for conversation in list(conversations.keys()):
-                        if conversation != conversations[st.session_state.id_chat]:
-                            st.session_state.id_chat = conversation
-                            break
-                                
-                conversations.pop(tmp)
-                with open(conversations_path, 'wb') as fp:
-                    packed_bytes = msgpack.packb(conversations, use_bin_type=True)
-                    fp.write(packed_bytes)
+                disabled=False if st.session_state.id_chat != 'New Chat' else True,
+                on_click=update_selectbox_remove,
+                args=(conversations_path, conversations)
+            )
 
         def generate_conv_title(llm: vllm.LLM, tokenizer: AutoTokenizer):
             if st.session_state.chat_id == 'New Chat':