Update app.py

app.py

@@ -80,20 +80,9 @@ if "messages" not in st.session_state:
     st.session_state.messages = []
 
 
-# Display chat messages from history on app rerun
-pos = 0
-for message in st.session_state.messages:
-    pos=pos+1
-    with st.chat_message(message["role"]):
-        col1, col2 = st.columns([9,1])
-        col1.markdown(message["content"])
-        col2.button("remove-"+str(pos), key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
-
-
 def remove_message(position):
     st.toast("remove message no:" + position)
 
-
 if "remove" not in st.session_state:
     st.session_state.remove= False
 
@@ -108,6 +97,14 @@ if st.session_state.remove:
     st.rerun()
 
 
+# Display chat messages from history on app rerun
+pos = 0
+for message in st.session_state.messages:
+    pos=pos+1
+    with st.chat_message(message["role"]):
+        col1, col2 = st.columns([9,1])
+        col1.markdown(message["content"])
+        col2.button("remove-"+str(pos), key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
 
 # Accept user input
 if prompt := st.chat_input(f"Hi I'm {selected_model}, ask me a question"):