Update app.py

app.py

@@ -64,10 +64,14 @@ st.sidebar.button('Reset Chat', on_click=reset_conversation) #Reset button
 st.sidebar.write(f"You're now chatting with **{selected_model}**")
 st.sidebar.markdown("*Generated content may be inaccurate or false.*")
 
+def edit_message(position):
+    st.toast("try to edit  message no: " + str(position))
+
 def remove_message(position):
     st.toast("try to remove message no: " + str(position-1) + " and "+ str(position))
     del st.session_state.messages[position-2:position]
 
+
 st.subheader(f'{selected_model}')
 
 # Initialize chat history
@@ -81,6 +85,7 @@ for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         col1, col2 = st.columns([9,1])
         col1.markdown(message["content"])
+        col2.button("edit-"+str(pos), key="button_edit_message_"+str(pos), args=[pos], on_click=edit_message)
         if message["role"] == "assistant":
             col2.button("remove-"+str(pos), key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
 
@@ -91,7 +96,7 @@ if prompt := st.chat_input(f"Hi I'm {selected_model}, ask me a question"):
     with st.chat_message("user"):
         col1, col2 = st.columns([9,1])
         col1.markdown(prompt)
-        # col2.button("remove-"+str(pos), key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
+        col2.button("edit-"+str(pos), key="button_edit_message_"+str(pos), args=[pos], on_click=edit_message)
     st.session_state.messages.append({"role": "user", "content": prompt})
     
     # Display assistant response in chat message container
@@ -101,6 +106,7 @@ if prompt := st.chat_input(f"Hi I'm {selected_model}, ask me a question"):
         with st.chat_message("assistant"):
             col1, col2 = st.columns([9,1])
             response = col1.write_stream(assistant["stream"])
+            col2.button("edit-"+str(pos), key="button_edit_message_"+str(pos), args=[pos], on_click=edit_message)
             col2.button("remove-"+str(pos), key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
     else:
         with st.chat_message("assistant"):