Update app.py

app.py

@@ -85,9 +85,9 @@ 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)
+        col2.button("", icon = ":material/edit:", key="button_edit_message_"+str(pos), args=[pos], on_click=edit_message)
         if message["role"] == "assistant":
-            col2.button(str(pos), icon = ":material/delete:", key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
+            col2.button("", icon = ":material/delete:", 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"):
@@ -96,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("edit-"+str(pos), key="button_edit_message_"+str(pos), args=[pos], on_click=edit_message)
+        col2.button("", icon = ":material/edit:", 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
@@ -106,12 +106,12 @@ 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(str(pos), icon = ":material/delete:", key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
+            col2.button("", icon = ":material/edit:", key="button_edit_message_"+str(pos), args=[pos], on_click=edit_message)
+            col2.button("", icon = ":material/delete:", key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
     else:
         with st.chat_message("assistant"):
             col1, col2 = st.columns([9,1])
             response = col1.write("Failure!")
-            col2.button(str(pos), icon = ":material/delete:", key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
+            col2.button("", icon = ":material/delete:", key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
 
     st.session_state.messages.append({"role": "assistant", "content": response})