Update app.py

app.py

@@ -46,9 +46,8 @@ if "system_message" not in st.session_state:
 if "starter_message" not in st.session_state:
     st.session_state.starter_message = "Hello, there! How can I help you today?"
 
-# Initialize session state for continue action
-if "need_continue" not in st.session_state:
-    st.session_state.need_continue = False
+# Initialize state for continue action
+need_continue = False
 
 # Initialize the last response
 if "last_response" not in st.session_state:
@@ -186,8 +185,11 @@ if st.session_state.user_text:
     # Clear the user input
     st.session_state.user_text = None
 
+if st.button("Continue"):
+    need_continue = True
+
 # If "Continue" button was pressed
-if st.session_state.get('need_continue', False):
+if need_continue:
     # Display a spinner while generating the continuation
     with st.chat_message("assistant", avatar=st.session_state.avatars['assistant']):
         with st.spinner("Continuing..."):
@@ -200,10 +202,4 @@ if st.session_state.get('need_continue', False):
                 continue_last=True
             )
             st.markdown(response)
-        # refresh_chat()
-
-    # Reset the flag
-    st.session_state.need_continue = False
-
-if st.button("Continue"):
-    st.session_state.need_continue = True
+        # refresh_chat()