Update main.py

main.py

@@ -156,13 +156,14 @@ async def on_message(message: cl.Message):
                     cl.Text(name="Sources", content=sources_text, display="inline")
                 )
 
-    async for chunk in runnable.astream(
-        message.content,
-        config=RunnableConfig(callbacks=[
-            cl.LangchainCallbackHandler(),
-            PostMessageHandler(msg)
-        ]),
-    ):
-        await msg.stream_token(chunk)
+    async with cl.Step(type="run", name="QA Assistant"):
+        async for chunk in runnable.astream(
+            message.content,
+            config=RunnableConfig(callbacks=[
+                cl.LangchainCallbackHandler(),
+                PostMessageHandler(msg)
+            ]),
+        ):
+            await msg.stream_token(chunk)
 
     await msg.send()