Update app.py

app.py

@@ -137,10 +137,21 @@ if st.session_state.retry:
     lastmessage = st.session_state.messages.pop()
     st.toast("popped msg: " + lastmessage["content"] + " // model: " + model_links[selected_model])
     response = get_assistant_aswer(model_links[selected_model], st.session_state.messages, temp_values, max_token_value)
-    with st.chat_message("assistant"):
-        st.write(response)
     st.session_state.messages.append({"role": "assistant", "content": response})
 
+if "remove" not in st.session_state:
+        st.session_state.remove= False
+
+def remove_click():
+    st.session_state.remove= True
+
+
+if st.session_state.remove:
+    lastmessage = st.session_state.messages.pop()
+    prelastmessage = st.session_state.messages.pop()
+    st.toast("popped msg: " + lastmessage["content"] + " // model: " + model_links[selected_model])
+
+
 
 
 # Accept user input
@@ -158,4 +169,6 @@ if prompt := st.chat_input(f"Hi I'm {selected_model}, ask me a question"):
 
     st.session_state.messages.append({"role": "assistant", "content": response})
 
-st.button("retry", key="retryButton", on_click=retry_click)
+if len(session_state.messages)>0:
+    st.button("retry", key="retryButton", on_click=retry_click)
+    st.button("remove", key="removeButton", on_click=remove_click)