Update app.py

app.py

@@ -36,44 +36,6 @@ if 'page_loaded' not in st.session_state:
     st.session_state.last_query = ""
     st.session_state.last_answer = None
 
-# Add a flag to track if we need to check for navigation
-if 'check_navigation' not in st.session_state:
-    st.session_state.check_navigation = True
-
-# Check if we need to reset state on page load
-if st.session_state.check_navigation:
-    # Reset key query-related state variables
-    st.session_state.last_query = ""
-    st.session_state.last_answer = None
-    st.session_state.submit_clicked = False
-    # If you have any other state to reset, do it here
-    
-    # Mark that we've checked navigation for this page load
-    st.session_state.check_navigation = False
-    
-    # Add JavaScript to reset the flag when navigating away
-    components.html(
-        """
-        <script>
-        // Set up beforeunload event to prepare for next page load
-        window.addEventListener('beforeunload', function() {
-            // This will be executed when the user navigates away
-            // We'll use sessionStorage which persists across page loads
-            sessionStorage.setItem('resetOnReturn', 'true');
-        });
-        
-        // Check if we need to reload on return
-        if (sessionStorage.getItem('resetOnReturn') === 'true') {
-            // Clear the flag
-            sessionStorage.removeItem('resetOnReturn');
-            // Force page reload to reset state
-            window.location.reload();
-        }
-        </script>
-        """,
-        height=0
-    )
-
 # THEN: Import your modules
 from rag_engine import process_query, load_model, cached_load_data_files
 from utils import setup_all_auth
@@ -317,8 +279,6 @@ def set_query(query):
         return
     st.session_state.last_query = query
     st.session_state.submit_clicked = True
-    st.session_state.is_processing = True
-    st.rerun()
 
 # Function to group questions into rows based on length
 def group_buttons(questions, max_chars_per_row=100):
@@ -373,7 +333,6 @@ def handle_form_submit():
     if st.session_state.query_input and st.session_state.query_input.strip():
         st.session_state.last_query = st.session_state.query_input.strip()
         st.session_state.submit_clicked = True
-        st.session_state.is_processing = True
         # Increment the form key to force a reset
         st.session_state.form_key += 1
 
@@ -395,26 +354,50 @@ with col1:
 with col2:
     word_limit = st.slider("Word limit:", 50, 500, 200)
 
+# Create placeholders for answer sections
+answer_container = st.container()
+answer_header = answer_container.empty()
+answer_text = answer_container.empty()
+sources_header = answer_container.empty()
+sources_text = answer_container.empty()
+
 # Process the query only if it has been explicitly submitted
 if st.session_state.submit_clicked and st.session_state.last_query:
-    st.session_state.submit_clicked = False
+    # Set processing flag first
+    st.session_state.is_processing = True
+    
+    # Check if we already have an answer to display
+    if st.session_state.last_answer is not None:
+        answer_header.subheader("Answer:")
+        answer_text.write(st.session_state.last_answer["answer_with_rag"])
+        sources_header.subheader("Sources:")
+        sources_text.write(st.session_state.last_answer["citations"])
+    
+    # Process in the same page load without rerun
     with st.spinner("Processing your question..."):
         try:
             result = process_query(st.session_state.last_query, top_k=top_k, word_limit=word_limit)
-            st.session_state.last_answer = result  # Store result in session state
+            st.session_state.last_answer = result
+            
+            # Update the display with new results
+            answer_header.subheader("Answer:")
+            answer_text.write(result["answer_with_rag"])
+            sources_header.subheader("Sources:")
+            sources_text.write(result["citations"])
+            
         except Exception as e:
-            st.session_state.last_answer = {"answer_with_rag": f"Error processing query: {str(e)}", "citations": ""}
-    # Reset debouncing after processing and force a rerun to re-enable buttons
+            error_result = {"answer_with_rag": f"Error processing query: {str(e)}", "citations": ""}
+            st.session_state.last_answer = error_result
+            
+            # Display the error
+            answer_header.subheader("Answer:")
+            answer_text.write(error_result["answer_with_rag"])
+            sources_header.subheader("Sources:")
+            sources_text.write(error_result["citations"])
+    
+    # Reset flags after processing is complete
+    st.session_state.submit_clicked = False
     st.session_state.is_processing = False
-    st.rerun()
-
-# Display the answer if available
-if st.session_state.last_answer is not None:
-    st.subheader("Answer:")
-    st.write(st.session_state.last_answer["answer_with_rag"])
-    st.subheader("Sources:")
-    for citation in st.session_state.last_answer["citations"].split("\n"):
-        st.write(citation)
 
 # Add helpful information
 st.markdown("---")