import streamlit as st
from knowledge_engine import KnowledgeManager

st.set_page_config(page_title="LangChain QA Assistant", layout="centered")

def initialize_knowledge_manager():
    if "km" not in st.session_state:
        try:
            st.session_state.km = KnowledgeManager()
            st.success("Knowledge base loaded successfully!")
        except Exception as e:
            st.error(f"Error loading knowledge base: {e}")
            st.session_state.km = None

def main():
    st.title("LangChain QA Assistant")

    initialize_knowledge_manager()

    if st.session_state.km:
        query = st.text_input("Ask something about your knowledge base:")
        if query:
            with st.spinner("Getting answer..."):
                answer = st.session_state.km.ask(query)
            st.markdown(f"**Answer:** {answer}")
    else:
        st.info("Please upload your `.txt` files in the root of the repo.")

if __name__ == "__main__":
    main()