Update app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 import time
- 
+
 # FIRST: Set page config before ANY other Streamlit command
 st.set_page_config(page_title="Spirituality Q&A")
 
@@ -27,9 +27,31 @@ if 'last_answer' not in st.session_state:
     st.session_state.last_answer = None
 
 # THEN: Import your modules
-from rag_engine import process_query, load_model
+from rag_engine import process_query, load_model, cached_load_data_files
 from utils import setup_all_auth
 
+# Preload resources during initialization
+init_message = st.empty()
+if not st.session_state.initialized:
+    init_message.info("Hang in there! We are setting the system up for you. 😊")
+    try:
+        # Setup authentication and preload heavy resources
+        setup_all_auth()
+        load_model()  # This uses cached_load_model via alias
+        cached_load_data_files()  # Preload FAISS index, text chunks, and metadata
+        st.session_state.initialized = True
+        st.session_state.init_time = time.time()
+        init_message.success("System initialized successfully!")
+        time.sleep(2)
+        init_message.empty()
+    except Exception as e:
+        init_message.error(f"Error initializing: {str(e)}")
+elif st.session_state.init_time is not None:
+    elapsed_time = time.time() - st.session_state.init_time
+    if elapsed_time >= 2.0:
+        init_message.empty()
+        st.session_state.init_time = None
+
 # Custom styling (pure CSS)
 st.markdown("""
 <style>
@@ -133,11 +155,9 @@ def group_buttons(questions, max_chars_per_row=100):
     rows = []
     current_row = []
     current_length = 0
-    
     for q in questions:
         # Add some buffer for button padding/margins
         q_length = len(q) + 5
-        
         if current_length + q_length > max_chars_per_row and current_row:
             rows.append(current_row)
             current_row = [q]
@@ -145,36 +165,10 @@ def group_buttons(questions, max_chars_per_row=100):
         else:
             current_row.append(q)
             current_length += q_length
-    
     if current_row:
         rows.append(current_row)
-    
     return rows
 
-# Handle initialization and success message timing
-init_message = st.empty()
-
-if not st.session_state.initialized:
-    init_message.info("Hang in there! We are setting the system up for you. 😊")
-
-    try:
-        setup_all_auth()
-        load_model()
-        st.session_state.initialized = True
-        st.session_state.init_time = time.time()
-        init_message.success("System initialized successfully!")
-        time.sleep(2)
-        init_message.empty()
-    except Exception as e:
-        init_message.error(f"Error initializing: {str(e)}")
-
-# Handle timing of success message disappearance
-elif st.session_state.init_time is not None:
-    elapsed_time = time.time() - st.session_state.init_time
-    if elapsed_time >= 2.0:
-        init_message.empty()
-        st.session_state.init_time = None
-
 # All common questions in a single list
 common_questions = [
     "What is the Atman or the soul?",
@@ -192,10 +186,8 @@ common_questions = [
 # Display heading for common questions
 st.markdown("### Common questions to try:")
 
-# Group questions into rows
+# Group questions into rows and create buttons (disabled if processing)
 question_rows = group_buttons(common_questions, max_chars_per_row=80)
-
-# Create buttons for each row. They are disabled if a query is processing.
 for row_idx, row in enumerate(question_rows):
     cols = st.columns(len(row))
     for i, (col, q) in enumerate(zip(cols, row)):
@@ -215,14 +207,13 @@ def handle_form_submit():
         # Increment the form key to force a reset
         st.session_state.form_key += 1
 
-# Create a form with dynamic key for resetting.
-# The form's submit button is also disabled when processing.
+# Create a form with a dynamic key (to allow resetting)
 with st.form(key=f"query_form_{st.session_state.form_key}"):
     query = st.text_input("Ask your question:", key="query_input", 
-                         placeholder="Press enter to submit your question", disabled=st.session_state.is_processing)
+                          placeholder="Press enter to submit your question", disabled=st.session_state.is_processing)
     submit_button = st.form_submit_button("Get Answer", on_click=handle_form_submit, disabled=st.session_state.is_processing)
 
-# Display the current question if there is one
+# Display the current question if available
 if st.session_state.last_query:
     st.markdown("### Current Question:")
     st.info(st.session_state.last_query)
@@ -234,19 +225,17 @@ with col1:
 with col2:
     word_limit = st.slider("Word limit:", 50, 500, 200)
 
-# Process the query only if explicitly submitted
+# 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
-
     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 the result in session state
+            st.session_state.last_answer = result  # Store result in session state
         except Exception as e:
             st.session_state.last_answer = {"answer_with_rag": f"Error processing query: {str(e)}", "citations": ""}
-    # Reset debouncing after processing
+    # Reset debouncing after processing and force a rerun to re-enable buttons
     st.session_state.is_processing = False
-    # Force a rerun so that buttons are re-rendered in enabled state and the answer is re-displayed
     st.experimental_rerun()
 
 # Display the answer if available