Implementing Streaming Output with Autogen and Chainlit Integration

[lianghsun]

I have a question regarding an example on https://github.com/Chainlit/cookbook/blob/main/pyautogen/app.py (or async version: https://github.com/Chainlit/cookbook/blob/main/pyautogen/async_app.py) that combines autogen and chainlit. I've been trying to modify the code to enable streaming output, but so far, no results have appeared on the screen. My code is as follows:

class ChainlitAgent(AssistantAgent):
    def __init__(self, *args, **kwargs):
        super().__init__(
            system_message=SYSTEM_PROMPT,
            name="assistant",
            llm_config={
                "stream": True
            }
        )
        
    async def a_send(
        self,
        message: Union[Dict, str],
        recipient: Agent,
        request_reply: Optional[bool] = None,
        silent: Optional[bool] = False,
    ) -> bool:
            
        msg = cl.run_sync(
            cl.Message(content="", author="assistant")
        )
     
        async for part in message:
            if token := part.choices[0].delta.content or "":
                print(token)
                await msg.stream_token(token)
                
        await msg.send()

        await super(ChainlitAgent, self).a_send(
            message=message,
            recipient=recipient,
            request_reply=request_reply,
            silent=silent,
        )

Is there a way to implement streaming with autogen in a chainlit environment?