Update app.py

app.py

@@ -73,21 +73,12 @@ if "prev_option" not in st.session_state:
 #    st.session_state.prev_option = selected_model
 #    reset_conversation()
 
-#Pull in the model we want to use
-#repo_id = model_links[selected_model]
-
 st.subheader(f'{selected_model}')
 
-# Set a default model
-#if selected_model not in st.session_state:
-#    st.session_state[selected_model] = model_links[selected_model] 
-
 # Initialize chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-def remove_message(position):
-    st.toast("remove message no:" + position)
 
 # Display chat messages from history on app rerun
 pos = 0
@@ -96,7 +87,11 @@ 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", 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)
+
+
+def remove_message(position):
+    st.toast("remove message no:" + position)
 
 
 if "remove" not in st.session_state:
@@ -117,18 +112,26 @@ if st.session_state.remove:
 # Accept user input
 if prompt := st.chat_input(f"Hi I'm {selected_model}, ask me a question"):
     # Display user message in chat message container and Add user message to chat history
+    pos = st.session_state.messages.length 
     with st.chat_message("user"):
-        st.markdown(prompt)
+        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)
     st.session_state.messages.append({"role": "user", "content": prompt})
     
     # Display assistant response in chat message container
     assistant = ask_assistant_stream(model_links[selected_model], st.session_state.messages, temp_values, max_token_value)
+    pos = st.session_state.messages.length
     if "stream" in assistant: 
         with st.chat_message("assistant"):
-            response = st.write_stream(assistant["stream"])
+            col1, col2 = st.columns([9,1])
+            response = col1.write_stream(assistant["stream"])
+            col2.button("remove-"+str(pos), key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
     else:
         with st.chat_message("assistant"):
-            response = st.write("Failure")
+            col1, col2 = st.columns([9,1])
+            response = col1.write("Failure!")
+            col2.button("remove-"+str(pos), key="button_remove_message_"+str(pos), args=[pos], on_click=remove_message)
 
     st.session_state.messages.append({"role": "assistant", "content": response})