Update app.py

app.py

@@ -81,7 +81,7 @@ if "messages" not in st.session_state:
 
 
 def remove_message(position):
-    st.toast("try to remove message no:" + str(position))
+    st.toast("try to remove message no: " + str(position) + " and "+ str(position+1))
 
 if "remove" not in st.session_state:
     st.session_state.remove= False
@@ -104,7 +104,8 @@ 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("remove-"+str(pos), key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
+        if message["role"] == "assistant":
+            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"):
@@ -113,7 +114,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("remove-"+str(pos), key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
     st.session_state.messages.append({"role": "user", "content": prompt})
     
     # Display assistant response in chat message container