import time

import streamlit as st
from agent.react_agent import ReactAgent

st.title("智扫通智能机器人客服")
st.divider()

if "agent" not in st.session_state:
    st.session_state["agent"] = ReactAgent()

if "message" not in st.session_state:
    st.session_state["message"] = []

history = st.session_state["message"]
agent = st.session_state["agent"]

for message in history:
    st.chat_message(message["role"]).write(message["content"])

# 要求用胡输入提示词
prompt = st.chat_input()

if prompt:
    st.chat_message("user").write(prompt)
    history.append({"role": "user", "content": prompt})

    response_messages = []
    with st.spinner("思考中"):
        stream = agent.excute_stream(prompt)

        def capture(generator, cache_list):
            for chunk in generator:
                cache_list.append(chunk)
                for char in chunk:
                    time.sleep(0.05)
                    yield char


        st.chat_message("assistant").write_stream(
            capture(stream, response_messages)
        )
        history.append({"role": "assistant", "content": response_messages[-1]})
        st.rerun()