import streamlit as st

st.set_page_config(layout="wide")
st.title("Mini Project 2 Part 3: Streamlit Chatbot")
if "messages" not in st.session_state:
    # ... (initialize messages)
    st.session_state.messages = []

with st.spinner("Please wait while the chat bot is loading!"):
    from agents2 import *


# Display existing chat messages
for message in st.session_state.messages:
    if message["role"] != "developer":
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
# Wait for user input
a_h = Head_Agent(st.secrets["OPENAI_KEY"], st.secrets["PINECONE_KEY"])
if prompt := st.chat_input("What would you like to chat about?"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    with st.chat_message("user"):
        st.markdown(prompt)
    with st.chat_message("assistant"):
        ai_message = a_h.generate_response(prompt, st.session_state.messages)
        st.markdown(ai_message)

    st.session_state.messages.append({"role": "assistant", "content": ai_message})