Update app.py

app.py

@@ -17,10 +17,8 @@ if 'submit_clicked' not in st.session_state:
     st.session_state.submit_clicked = False
 if 'init_time' not in st.session_state:
     st.session_state.init_time = None
-if 'form_needs_rerun' not in st.session_state:
-    st.session_state.form_needs_rerun = False
-if 'init_needs_rerun' not in st.session_state:
-    st.session_state.init_needs_rerun = False
+if 'form_submitted' not in st.session_state:
+    st.session_state.form_submitted = False
 
 # THEN: Import your modules
 from rag_engine import process_query, load_model
@@ -72,7 +70,7 @@ if not st.session_state.initialized:
         st.session_state.init_time = time.time()
         init_message.success("System initialized successfully!")
         time.sleep(0.1)
-        st.session_state.init_needs_rerun = True
+        st.experimental_rerun()  # For compatibility with older Streamlit versions
     except Exception as e:
         init_message.error(f"Error initializing: {str(e)}")
 
@@ -82,28 +80,27 @@ elif st.session_state.init_time is not None:
         init_message.empty()
         st.session_state.init_time = None
 
-# Handle rerun for initialization
-if st.session_state.init_needs_rerun:
-    st.session_state.init_needs_rerun = False
-    st.rerun()
-
-# Handle form submission
+# Function to handle form submission
 def handle_form_submit():
     if st.session_state.query_input:
         st.session_state.last_query = st.session_state.query_input
         st.session_state.submit_clicked = True
-        st.session_state.query_input = ""
-        st.session_state.form_needs_rerun = True
+        # Set form_submitted to True to trigger rerun
+        st.session_state.form_submitted = True
 
 # Create a form for handling the user input
 with st.form(key="query_form"):
     query = st.text_input("Ask your question:", key="query_input", placeholder="Press enter to submit question")
     submit_button = st.form_submit_button("Get Answer", on_click=handle_form_submit)
 
-# Handle rerun for form submission
-if st.session_state.form_needs_rerun:
-    st.session_state.form_needs_rerun = False
-    st.rerun()
+# Check if form was submitted and handle the rerun outside the callback
+if st.session_state.form_submitted:
+    # Reset the flag
+    st.session_state.form_submitted = False
+    # Clear the input
+    st.session_state.query_input = ""
+    # Force a rerun to update the UI
+    st.experimental_rerun()
 
 # Display the current question if there is one
 if st.session_state.last_query: