import streamlit as st
from langchain.schema import SystemMessage, HumanMessage, AIMessage
from langchain.chat_models import ChatOpenAI
from datasets import load_dataset

st.set_page_config(page_title="Langchain Chat-GPT", page_icon=":robot:")
st.header("Hey! I'm your Bot")

if 'sessionMessages' not in st.session_state:
    st.session_state.sessionMessages = [SystemMessage(content="You are a helpful assistant")]

def get_text():
    user_input = st.text_input("Type your query ...", key="input")
    return user_input

def get_answer(query):

    st.session_state.sessionMessages.append(HumanMessage(content=query))
    response = chat(st.session_state.sessionMessages)
    st.session_state.sessionMessages.append(AIMessage(content=response.content))

    return response.content

# If the dataset is gated/private, make sure you have run huggingface-cli login
#dataset = load_dataset("satkr7-org/chromadb_test")
#st.write(dataset)

chat = ChatOpenAI(model='gpt-3.5-turbo', temperature=0.0)

user_input = get_text()
submit = st.button("Generate")

if submit:
    response = get_answer(user_input)
    st.header("Your response is:")
    st.write(response)