Addressing asynf for error to enable streaming responses in Chainlit chatbot

app.py

@@ -161,21 +161,11 @@ async def start_chat():
 async def handle_message(message: cl.Message):
     settings = cl.user_session.get("settings")
 
-
-    # Initialize the stream message in Chainlit
-    stream_msg = cl.Message(content="")
-    await stream_msg.send()  # Send initial empty message to start the stream
-
-    # Create a generator from the RAG chain
-    response_generator = retrieval_augmented_qa_chain.stream({"question": message.content})
-
-    async for response_chunk in response_generator:
-        # Extract the content from the chunk
-        chunk_content = response_chunk.get("response", {}).get("content", "")
-        if chunk_content:
-            # Append the chunk to the streaming message content
-            stream_msg.content += chunk_content
-            await stream_msg.update()  # Update the message in Chainlit
+    # Stream the response as it is generated
+    async with cl.Message(content="") as msg:
+        async for token in retrieval_augmented_qa_chain.stream({"question": message.content}):
+            # Append each token to the message
+            await msg.stream_token(token)
 
     ## Remove to stream the response
     # response = retrieval_augmented_qa_chain.invoke({"question": message.content})