Update main.py

main.py

@@ -112,10 +112,8 @@ async def on_chat_start():
         timeout="3600"
     ).send()
 
-    if res and res.get("value") == "continue":
-        await cl.Message(
-            content="On continue!",
-        ).send()
+    if res:
+        cl.user_session.set("selectRequest", res.get("value"))
         
     template = """Answer the question based only on the following context:
 
@@ -162,6 +160,7 @@ async def on_message(message: cl.Message):
                 self.sources.add(source_page_pair)  # Add unique pairs to the set
 
         def on_llm_end(self, response, *, run_id, parent_run_id, **kwargs):
+            cl.user_session.set("selectRequest","")
             if len(self.sources):
                 sources_text = "\n".join([f"{source}#page={page}" for source, page in self.sources])
                 self.msg.elements.append(
@@ -169,13 +168,23 @@ async def on_message(message: cl.Message):
                 )
 
     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)
+        if cl.user_session.get("selectRequest"):
+            async for chunk in runnable.astream(
+                cl.user_session.get("selectRequest"),
+                config=RunnableConfig(callbacks=[
+                    cl.LangchainCallbackHandler(),
+                    PostMessageHandler(msg)
+                ]),
+            ):
+                await msg.stream_token(chunk)
+        else:
+            async for chunk in runnable.astream(
+                message.content,
+                config=RunnableConfig(callbacks=[
+                    cl.LangchainCallbackHandler(),
+                    PostMessageHandler(msg)
+                ]),
+            ):
+                await msg.stream_token(chunk)
 
     await msg.send()