Integrating Agent-Protocol-Client with Streamlit Chat UI: Seeking Suggestions for Continuous User Input Handling

Hi everyone, Iโ€™m currently trying to merge Agent-protocol-client with Streamlit chat UI. I hope to be able to accept user input again when step.is_last is False. Iโ€™ve tried some methods, but they all failed. Does anyone have any suggestions?

import asyncio
import streamlit as st

from assistants.protocols.models import StepRequestBody
from agent_protocol_client import (
    Configuration,
    ApiClient,
    StepRequestBody,
    TaskRequestBody,
    AgentApi,
)

st.title('AI Assistant')

# Defining the host is optional and defaults to http://localhost
# See configuration.py for a list of all supported configuration parameters.
configuration = Configuration(host="http://localhost:8000")


# Streamlit settings

if "openai_model" not in st.session_state:
    # st.session_state["openai_model"] = "gpt-3.5-turbo"
    st.session_state["openai_model"] = "gpt-4"

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

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message['content'])



########################################

async def main():
    # Enter a context with an instance of the API client
    async with ApiClient(configuration) as api_client:
        # Create an instance of the API class
        api_instance = AgentApi(api_client)
        prompts = []

        # Start a task
        if not prompts:
            if prompt := st.chat_input("Ask me"):
                prompts.append(prompt)
                st.session_state.messages.append({"role": "user", "content": prompt})

                with st.chat_message("user"):
                    st.markdown(prompt)

                # task_request_body = TaskRequestBody(input="Write 'Hello world!' to hi.txt.")
                task_request_body = TaskRequestBody(input=prompt)

                response = await api_instance.create_agent_task(
                    task_request_body=task_request_body
                )
                print("The response of AgentApi->create_agent_task:\n")
                print(response)
                print(prompts)
                print("\n\n")

                task_id = response.task_id
                i = 1

                step = await api_instance.execute_agent_task_step(
                    task_id=task_id, step_request_body=StepRequestBody(input=prompt))
                print(step)
                with st.chat_message("assistant"):
                    st.session_state.messages.append({"role": "assistant", "content": step.output})
                    st.markdown(step.output)

        if prompts:
            if step.is_last is False:
                while (
                    step := await api_instance.execute_agent_task_step(
                        task_id=task_id, step_request_body=StepRequestBody(input=prompt)
                    )
                ) and step.is_last is False:
                    with st.chat_message("assistant"):
                        st.session_state.messages.append({"role": "assistant", "content": step.output})
                        st.markdown(step.output, step.is_last)
                        print(step.is_last)

                print("Agent finished its work!")
            else:
                prompts = []

asyncio.run(main())