import streamlit as st
from rag import RAGinit, RAG_proximity_search

# Wrap RAGinit() inside a function that shows a spinner
def load_resources():
    with st.spinner("Loading resources..."):
        client, model, emb, chroma_collection, vector_index_properties, top_n = RAGinit()
    return client, model, emb, chroma_collection, vector_index_properties, top_n

# Initialize everything once at startup
client, model, emb, chroma_collection, vector_index_properties, top_n = load_resources()

def main():
    st.title("RAG-based QA App")

    question = st.text_input("Ask a question:")
    
    if st.button("Search"):
        if question.strip():
            answer = RAG_proximity_search(
                question,
                client,
                model,
                emb,
                chroma_collection,
                vector_index_properties,
                top_n
            )
            st.markdown("**Answer:**")
            st.write(answer)
        else:
            st.warning("Please enter a question before searching.")

if __name__ == "__main__":
    main()