app.py

import streamlit as st
import time

# FIRST: Set page config before ANY other Streamlit command
st.set_page_config(page_title="Indian Spiritual RAG")

# Initialize ALL session state variables
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 'query_answered' not in st.session_state:
    st.session_state.query_answered = False  # Tracks if the query has been answered

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

# Create a placeholder for our initialization message
init_message = st.empty()

# ✅ Custom CSS to handle dynamic text box color change
st.markdown("""
<style>
/* Styling for the main title */
.main-title {
    font-size: 2.5rem;
    color: #FF5722;
    text-align: center;
    margin-bottom: 1rem;
}

/* ✅ Green by default */
input[type="text"] {
    border: 2px solid #4CAF50 !important;
    border-radius: 8px !important;
    transition: border 0.3s ease-in-out;
}

/* 🔴 Change to red while typing */
input[type="text"]:focus {
    border: 2px solid #FF5722 !important;
}

/* ✅ Change back to green AFTER submission */
.answered input[type="text"] {
    border: 2px solid #4CAF50 !important;
}

/* ✅ FULLY GREEN BUTTON */
.stButton>button {
    background-color: #4CAF50 !important;
    color: #E8F5E9 !important;
    border: 2px solid #4CAF50 !important;
    border-radius: 8px !important;
}
</style>
<div class="main-title">Indian Spiritual Texts Q&A</div>
""", unsafe_allow_html=True)

# Handle initialization and success message timing
if not st.session_state.initialized:
    init_message.info("Hang in there! We are setting the system up for you. 😊")

    try:
        # Setup authentication
        setup_all_auth()

        # Load the model
        load_model()

        # Mark as initialized and record time
        st.session_state.initialized = True
        st.session_state.init_time = time.time()

        # Show success message
        init_message.success("System initialized successfully!")

        # Force rerun to start the timer for removing the message
        time.sleep(0.1)  # Small delay to ensure message appears
        st.rerun()

    except Exception as e:
        init_message.error(f"Error initializing: {str(e)}")

# Check if we need to hide the success message (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  # Reset timer after clearing

# Function to 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_answered = False  # Mark as not yet answered
        st.session_state.query_input = ""  # Clear input box after submission

# ✅ Conditionally apply CSS class based on query answered status
input_css_class = "answered" if st.session_state.query_answered else ""

# ✅ Wrap the text input inside a div to allow dynamic styling
st.markdown(f'<div class="{input_css_class}">', unsafe_allow_html=True)

# ✅ Create a form for handling user input
with st.form(key="query_form"):
    # Input field with dynamic border behavior
    query = st.text_input(
        "Ask your question:",
        key="query_input",
        value="",
        placeholder="Type here..."
    )

    # Submit button
    submit_button = st.form_submit_button("Get Answer", on_click=handle_form_submit)

st.markdown("</div>", unsafe_allow_html=True)  # ✅ Close the div for styling

# ✅ Display the last question after submission
if st.session_state.submit_clicked and 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)

# ✅ Process the query if submitted
if st.session_state.submit_clicked and st.session_state.last_query:
    st.session_state.submit_clicked = False  # Reset flag

    with st.spinner("Processing your question..."):
        try:
            result = process_query(st.session_state.last_query, top_k=top_k, word_limit=word_limit)

            st.subheader("Answer:")
            st.write(result["answer_with_rag"])

            st.subheader("Sources:")
            for citation in result["citations"].split("\n"):
                st.write(citation)

            st.session_state.query_answered = True  # ✅ Mark query as answered
            st.rerun()  # ✅ Refresh UI so that text box turns green again

        except Exception as e:
            st.error(f"Error processing query: {str(e)}")

# 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.
""")