Addressing async with issue affecting response streaming

app.py

@@ -157,15 +157,25 @@ async def start_chat():
 # (2) Invokes the RAG chain with the user's message
 # (3) Extracts the content from the response and sends it back to the user
 
-@cl.on_message 
+@cl.on_message
 async def handle_message(message: cl.Message):
     settings = cl.user_session.get("settings")
 
-    # 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)
+    # 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})
+
+    # Stream each response chunk to Chainlit
+    async for response_chunk in response_generator:
+        # Extract the content from the chunk
+        chunk_content = response_chunk.get("response", {}).get("content", "")
+        if chunk_content:
+            # Stream the chunk to Chainlit
+            await stream_msg.stream_token(chunk_content)
+
 
     ## Remove to stream the response
     # response = retrieval_augmented_qa_chain.invoke({"question": message.content})