app.py

import streamlit as st
import time

# FIRST: Set page config before ANY other Streamlit command
st.set_page_config(page_title="Spirituality Q&A")

# Initialize ALL session state variables right at the beginning
if 'initialized' not in st.session_state:
    st.session_state.initialized = False
if 'model' not in st.session_state:
    st.session_state.model = None
if 'tokenizer' not in st.session_state:
    st.session_state.tokenizer = None
if 'last_query' not in st.session_state:
    st.session_state.last_query = ""
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_key' not in st.session_state:
    st.session_state.form_key = 0  # This will help us reset the form
if 'processing' not in st.session_state:
    st.session_state.processing = False

# THEN: Import your modules
from rag_engine import process_query, load_model
from utils import setup_all_auth

# Custom styling (pure CSS)
st.markdown("""
<style>
.main-title {
    font-size: 2.5rem;
    color: #c0392b;
    text-align: center;
    margin-bottom: 1rem;
}
/* Button styling */
.stButton>button {
    background-color: #fff0f0 !important;
    color: #3f51b5 !important;
    border: 1px solid #e1e4f2 !important;
    border-radius: 20px !important;
    padding: 8px 16px !important;
    box-shadow: 0 1px 2px rgba(0,0,0,0.03) !important;
    white-space: nowrap !important;
    overflow: hidden !important;
    text-overflow: ellipsis !important;
}
/* Form submit button specific styling */
button[type="submit"], 
.stFormSubmit>button, 
[data-testid="stFormSubmitButton"]>button {
    background-color: #fff0f0 !important;
    color: #3f51b5 !important;
    border: 1px solid #e1e4f2 !important;
    border-radius: 8px !important;
}
.stButton>button:hover {
    background-color: #fafbff !important;
    border-color: #c5cae9 !important;
}
/* Input field styling */
div[data-baseweb="input"] {
    border: 1px solid #fff0f0 !important;
    border-radius: 8px !important;
    background-color: #ffffff !important;
}
div[data-baseweb="input"]:focus-within {
    border: 1px solid #3f51b5 !important;
}
div[data-baseweb="input"]:active {
    border: 1px solid #fff0f0 !important;
}
/* Style the st.info boxes */
div.stInfo {
    background-color: #f8faff !important;
    color: #3f51b5 !important;
    border: 1px solid #e1e4f2 !important;
    border-radius: 8px !important;
}
/* COMBINED SCROLL CONTAINER */
.questions-scroll-container {
    width: 100%;
    overflow-x: auto;
    scrollbar-width: none; /* Firefox */
    -ms-overflow-style: none; /* IE and Edge */
}
/* Hide scrollbar for Chrome, Safari and Opera */
.questions-scroll-container::-webkit-scrollbar {
    display: none;
}
/* Inner content that holds both rows */
.questions-content {
    display: inline-flex;
    flex-direction: column;
    min-width: max-content;
    gap: 10px;
    padding: 5px 0;
}
/* Individual rows */
.questions-row {
    display: flex;
    flex-direction: row;
    gap: 10px;
}
/* Placeholder for buttons */
.button-placeholder {
    min-height: 38px;
    min-width: 120px;
    margin: 0 5px;
}
/* Hide "Press Enter to submit form" text */
.stTextInput label + div + div small,
[data-testid="stForm"] small,
form small {
    display: none !important;
}
</style>
<div class="main-title">Spirituality Q&A</div>
""", unsafe_allow_html=True)

# Function to handle query selection
def set_query(query):
    if not st.session_state.processing:
        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):
    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]
            current_length = q_length
        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
# This will run on each rerender after initialization
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?",
    "Are there rebirths?",
    "What is Karma?",
    "What is the ultimate truth?",
    "What is Swami Vivekananda's opinion about the SELF?",
    "Explain moksha or salvation. Is that for real?",
    "Destiny or free will?",
    "What is the ultimate goal of human life?",
    "Do we really die?",
    "How can you attain self-realization?"
]

# Display heading for common questions
st.markdown("### Common questions to try:")

# Group questions into rows
question_rows = group_buttons(common_questions, max_chars_per_row=80)

# Create buttons for each row
for row_idx, row in enumerate(question_rows):
    cols = st.columns(len(row))
    for i, (col, q) in enumerate(zip(cols, row)):
        with col:
            if st.button(q, key=f"r{row_idx}_q{i}", use_container_width=True, disabled=st.session_state.processing):
                set_query(q)

# Function to handle form submission
def handle_form_submit():
    if st.session_state.query_input and st.session_state.query_input.strip() and not st.session_state.processing:
        st.session_state.last_query = st.session_state.query_input.strip()
        st.session_state.submit_clicked = True
        # Increment the form key to force a reset
        st.session_state.form_key += 1

# Create a form with dynamic key for 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")
    submit_button = st.form_submit_button("Get Answer", on_click=handle_form_submit, disabled=st.session_state.processing)

# Display the current question if there is one
if st.session_state.last_query:
    st.markdown("### Current Question:")
    st.info(st.session_state.last_query)

# Sliders for customization
col1, col2 = st.columns(2)
with col1:
    top_k = st.slider("Number of sources:", 3, 10, 5)
with col2:
    word_limit = st.slider("Word limit:", 50, 500, 200)

# Only process the query if explicitly submitted and not already processing
if st.session_state.submit_clicked and st.session_state.last_query and not st.session_state.processing:
    st.session_state.submit_clicked = False
    st.session_state.processing = True

    with st.spinner("Processing your question..."):
        try:
            progress_bar = st.progress(0)
            progress_bar.progress(25, "Generating embeddings...")
            
            # Call processing function with progress tracking
            progress_bar.progress(50, "Retrieving relevant passages...")
            
            result = process_query(st.session_state.last_query, top_k=top_k, word_limit=word_limit)
            
            progress_bar.progress(75, "Generating answer...")
            time.sleep(0.5)  # Small delay for visual feedback
            progress_bar.progress(100)
            
            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:
            st.error(f"Error processing query: {str(e)}")
        finally:
            st.session_state.processing = False

# Add helpful information
st.markdown("---")
st.markdown("""
### About this app
This application uses a Retrieval-Augmented Generation (RAG) system to answer questions about Indian spiritual texts.
It searches through a database of texts to find relevant passages and generates answers based on those passages.
""")