import streamlit as st
from streamlit_chat import message
import requests
from nomic.gpt4all import GPT4All

st.set_page_config(
    page_title="Streamlit Chat - Demo",
    page_icon=":robot:"
)

@st.cache(allow_output_mutation=True)
def get_model():
    m = GPT4All()
    m.open()
    return m

m = get_model()

st.header("GPT4All Chatbot")

if 'generated' not in st.session_state:
    st.session_state['generated'] = []

if 'past' not in st.session_state:
    st.session_state['past'] = []

def get_text():
    input_text = st.text_input("You: ","", key="input")
    return input_text 


user_input = get_text()

if user_input:
    output = m.prompt(user_input)
    st.session_state.past.append(user_input)
    st.session_state.generated.append(output)


if st.session_state['generated']:
    for i in range(len(st.session_state['generated'])-1, -1, -1):
        message(st.session_state["generated"][i], key=str(i))
        message(st.session_state['past'][i], is_user=True, key=str(i) + '_user')