How to persist websocket connection in Streamlit

🎈 Using Streamlit
1

I’m creating a chatbot which display the responses from websocket API. I want the same connectionId for a user interactions until connection timeout (meaning a user and can send multiple inputs and receive corresponding response but connectionId should remain the same). Below is the code I have written but it’s not working just showing as Running. How to fix this?


import plotly.graph_objects as go
import asyncio
import websockets
import json
import streamlit as st
import uuid

async def main():
    #capturing_responses_status = []

    api_url = "websocket-url"  # Define api_url here
    async with websockets.connect(api_url) as websocket:
        while True:
            try:
                if prompt := st.chat_input(key=uuid.uuid4(), disabled=False):
                    #st.session_state.messages.append(({"role": "user", "content": prompt}, None))
                    with st.chat_message("user"):
                        st.write(prompt)

                    with st.chat_message("assistant"):
                        input_data = {"action": "sendmessage", "input_message": prompt}
                        await websocket.send(json.dumps(input_data))

                        with st.spinner('Generating SQL...'):
                            response = await websocket.recv()  # Wait for a response
                            res = json.loads(response)
                            st.write(res["sql"])

                            #capturing_responses_status.append("sql")
 
                        with st.spinner('Generating Plot...'):
                            response = await websocket.recv()  # Wait for a response
                            res = json.loads(response)
                            figure = go.Figure(json.loads(res["plot"]))
                            st.plotly_chart(figure)
 
                            #capturing_responses_status.append("plot")          

                        with st.spinner('Generating Summary ...'):
                            response = await websocket.recv()  # Wait for a response
                            res = json.loads(response)
                            st.write(res["summary"])

                            #capturing_responses_status.append("summary")
  
            except websockets.exceptions.ConnectionClosedError:
                print("Connection closed. Reconnecting...")
                break
            except asyncio.TimeoutError:
                print("Timeout occurred. Reconnecting...")
                break

if __name__ == "__main__":
    asyncio.run(main())