Update app.py

app.py

@@ -17,6 +17,10 @@ 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
 
 # THEN: Import your modules
 from rag_engine import process_query, load_model
@@ -68,7 +72,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.rerun()  # Ensure the UI updates correctly
+        st.session_state.init_needs_rerun = True
     except Exception as e:
         init_message.error(f"Error initializing: {str(e)}")
 
@@ -78,19 +82,29 @@ 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
 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.rerun()  # Forces UI to refresh, ensuring input box turns green
+        st.session_state.form_needs_rerun = 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()
+
 # Display the current question if there is one
 if st.session_state.last_query:
     st.markdown("### Current Question:")