react agent with astream_events is very slow after the first response

[fatiiiiima]

I'm building a chatbot to chat with my SQL database. I followed this tutorial: https://python.langchain.com/docs/tutorials/sql_qa/, so I am using the pre built react agent from Langgraph. I am using astream_events to stream the responses. However, I have noticed that while the first response is always really fast (within 5 to 6 secs), any response after that is very slow (like over a minute). The reason why I went with streaming in the first place is because while using .invoke, my responses were slow too, so I wanted to build a better experience for the users. Has anyone faced a similar issue like this before?

Here is my code for reference:

load_dotenv()

connectionString = os.environ["py-connectionString"]
db_engine = create_engine(connectionString)
db = SQLDatabase(db_engine, schema="dbo", include_tables=["", "", ""])


llm = AzureChatOpenAI(
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    azure_deployment=os.environ["AZURE_OPENAI_DEPLOYMENT_NAME"],
    openai_api_version=os.environ["AZURE_OPENAI_API_VERSION"],
    api_key=os.environ["AZURE_OPENAI_KEY"], 
    max_tokens=1000,
)



toolkit = SQLDatabaseToolkit(db=db, llm=llm)
tools = toolkit.get_tools()


SQL_PREFIX = f"""

"""

system_message = SystemMessage(content=SQL_PREFIX)

class QueryRequest(BaseModel):
    question: str

memory = MemorySaver()
agent_executor = create_react_agent(llm, tools, state_modifier=system_message, checkpointer=memory)
config = {"configurable": {"thread_id": "test-thread"}}

async def ask_question():
    while True:
        query = input("Ask your question:")
        if query.lower() == "quit":
            break
        
        inputs = {"messages": [("human", query)]}
        async for event in agent_executor.astream_events(inputs, config, version="v2"):
            if event["event"] == "on_chat_model_stream":
                data = event["data"]
                if data["chunk"].content:
                    print(data["chunk"].content, end="|", flush=True)
        
            
if __name__ == "__main__":
    asyncio.run(ask_question())
  

I have not copied the prompt or the table names as I want to keep that private.