import streamlit as st
import time
import LlamaIndex_helper as lch
from llama_index.core import Settings
from llama_index.llms.openai import OpenAI

st.set_page_config(
    page_title="Work history response generator !", page_icon=":flag-ca:")
st.title("Work history response generator ! 🍁")
st.header("Let's chat :star2:")

DATA_PATH = "data"

Settings.llm = OpenAI(temperature=0.1, model="gpt-4o")

@st.cache_resource(show_spinner = "Loading AI agent")
def prepare_top_agents():
    file_names = lch.get_pdf_file_names(DATA_PATH)
    relevant_docs = lch.get_relevant_docs(file_names)
    agents = lch.building_document_agents(file_names=file_names,
                                          relevant_docs=relevant_docs)
    obj_index = lch.define_all_tools(agents=agents, file_names=file_names)
    top_agent = lch.create_top_agent(obj_index=obj_index)
    return top_agent

top_agent = prepare_top_agents()

YES_MESSAGE = "Hello there, thank you for trying out the app. So, how can I help you today? "
if "messages" not in st.session_state.keys():
    st.session_state.messages = [
        {"role": "assistant", "content": YES_MESSAGE}
    ]

for message in st.session_state.messages: # Display the prior chat messages
    with st.chat_message(message["role"]):
        st.write(message["content"])

if prompt := st.chat_input("Your question"): # Prompt for user input and save to chat history
    st.session_state.messages.append({"role": "user", "content": prompt})
    st.chat_message("user").write(prompt)
    with st.chat_message("assistant"):
        with st.spinner("Thinking...please be patient"):
            response = top_agent.stream_chat(prompt)
            # st.write_stream(response.response_gen)
            def generate():
                for text in response.response_gen:
                    yield text
                    time.sleep(0.05)
            st.write_stream(generate)
            message = {"role": "assistant", "content": response.response}
            st.session_state.messages.append(message)