Update app.py

app.py

@@ -1,7 +1,5 @@
 import streamlit as st
 import time
-import uuid
-from datetime import datetime
 
 # FIRST: Set page config before ANY other Streamlit command
 st.set_page_config(page_title="Spirituality Q&A")
@@ -21,12 +19,6 @@ if 'init_time' not in st.session_state:
     st.session_state.init_time = None
 if 'form_key' not in st.session_state:
     st.session_state.form_key = 0  # This will help us reset the form
-if 'button_disabled' not in st.session_state:
-    st.session_state.button_disabled = {}
-if 'last_click_id' not in st.session_state:
-    st.session_state.last_click_id = None
-if 'answer_displayed' not in st.session_state:
-    st.session_state.answer_displayed = False
 
 # THEN: Import your modules
 from rag_engine import process_query, load_model
@@ -121,18 +113,10 @@ div.stInfo {
 """, unsafe_allow_html=True)
 
 # Function to handle query selection
-def set_query(query, button_id):
-    if button_id not in st.session_state.button_disabled or not st.session_state.button_disabled[button_id]:
-        # Disable this button and all other buttons
-        for k in st.session_state.button_disabled:
-            st.session_state.button_disabled[k] = True
-        st.session_state.button_disabled[button_id] = True
-        st.session_state.last_click_id = button_id
-        
-        st.session_state.last_query = query
-        st.session_state.submit_clicked = True
-        st.session_state.answer_displayed = False
-        st.rerun()
+def set_query(query):
+    st.session_state.last_query = query
+    st.session_state.submit_clicked = True
+    st.rerun()
 
 # Function to group questions into rows based on length
 def group_buttons(questions, max_chars_per_row=100):
@@ -207,31 +191,14 @@ for row_idx, row in enumerate(question_rows):
     cols = st.columns(len(row))
     for i, (col, q) in enumerate(zip(cols, row)):
         with col:
-            button_id = f"r{row_idx}_q{i}"
-            # Initialize this button's state if not already done
-            if button_id not in st.session_state.button_disabled:
-                st.session_state.button_disabled[button_id] = False
-                
-            if st.button(q, key=button_id, use_container_width=True, 
-                        disabled=st.session_state.button_disabled.get(button_id, False)):
-                set_query(q, button_id)
+            if st.button(q, key=f"r{row_idx}_q{i}", use_container_width=True):
+                set_query(q)
 
 # Function to handle form submission
 def handle_form_submit():
-    form_id = "form_submit"
-    if (st.session_state.query_input and st.session_state.query_input.strip() and 
-        (form_id not in st.session_state.button_disabled or 
-         not st.session_state.button_disabled[form_id])):
-        
-        # Disable all buttons
-        for k in st.session_state.button_disabled:
-            st.session_state.button_disabled[k] = True
-        st.session_state.button_disabled[form_id] = True
-        st.session_state.last_click_id = form_id
-        
+    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.answer_displayed = False
         # Increment the form key to force a reset
         st.session_state.form_key += 1
 
@@ -239,13 +206,7 @@ def handle_form_submit():
 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")
-    # Disable the button based on state
-    form_id = "form_submit"
-    if form_id not in st.session_state.button_disabled:
-        st.session_state.button_disabled[form_id] = False
-    
-    submit_button = st.form_submit_button("Get Answer", on_click=handle_form_submit, 
-                                         disabled=st.session_state.button_disabled.get(form_id, False))
+    submit_button = st.form_submit_button("Get Answer", on_click=handle_form_submit)
 
 # Display the current question if there is one
 if st.session_state.last_query:
@@ -266,21 +227,12 @@ if st.session_state.submit_clicked and st.session_state.last_query:
     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.answer_displayed = True
-            
-            # Re-enable all buttons after answer is displayed
-            for k in st.session_state.button_disabled:
-                st.session_state.button_disabled[k] = False
-                
             st.subheader("Answer:")
             st.write(result["answer_with_rag"])
             st.subheader("Sources:")
             for citation in result["citations"].split("\n"):
                 st.write(citation)
         except Exception as e:
-            # Re-enable all buttons on error too
-            for k in st.session_state.button_disabled:
-                st.session_state.button_disabled[k] = False
             st.error(f"Error processing query: {str(e)}")
 
 # Add helpful information