Update app.py

app.py

@@ -126,6 +126,12 @@ for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
+if retry not in st.session_state:
+        st.session_state.retry= False
+
+def retry_click:
+    st.session_state.retry= True
+
 
 if st.session_state.retry:
     lastmessage = st.session_state.messages.pop()
@@ -153,4 +159,4 @@ 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="retry")
+st.button("retry", key="retry", on_click=retry_click)