# chat_bot.py

import streamlit as st
from streamlit_chat import message
import zhipuai

zhipuai.api_key = "169859478749e13a43a2c1b3d3f0bce0.Bnq1AThjEwwmK1k8"
def getAnswerByZhiPu(t):
    response = zhipuai.model_api.sse_invoke(
        model="chatglm_pro",
        prompt=[
            {"role": "user", "content": t}
        ],
        temperature=0.95,
        top_p=0.7,
        incremental=True
    )
    result = ''
    for event in response.events():
        result += event.data
    return result


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

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

def end_click():
    st.session_state['prompts'] = []
    st.session_state['past'] = []
    st.session_state['generated'] = []
    st.session_state['user'] = ""
def chat_click():
    if st.session_state['user'] != '':
        chat_input = st.session_state['user']
        output = getAnswerByZhiPu(chat_input)
        st.session_state['past'].append(chat_input)
        st.session_state['generated'].append(output)
        st.session_state['prompts'].append(output)
        st.session_state['user'] = ""
st.image("./logo.jpg", use_column_width='auto')
st.title("珠宝贝专属聊天机器人")

user_input = st.text_input("输入:", key="user")
chat_button = st.button("发送", on_click=chat_click)
end_button = st.button("新聊天", on_click=end_click)

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