Update app.py

app.py

@@ -47,9 +47,6 @@ async def start():
         # path = r'assets/ChatPDFAvatar.jpg'
     ).send()
 
-
-@cl.langchain_factory(use_async=True)
-async def init():
     files = None
 
     # Wait for the user to upload a file
@@ -120,14 +117,21 @@ async def init():
     cl.user_session.set("texts", pages)
 
     # Let the user know that the system is ready
-    await msg.update(content=f''' "{file.name}" processed. You can now ask questions!''')
-
+    msg.content = f''' "{file.name}" processed. You can now ask questions!'''
+    await msg.update()
 
-    return chain
+    cl.user_session.set('chain', chain)
 
 
-@cl.langchain_postprocess
+@cl.on_message
 async def process_response(res):
+    chain = cl.user_session.get("chain")
+    cb = cl.AsyncLangchainCallbackHandler(
+        stream_final_answer=True, answer_prefix_tokens=["FINAL", "ANSWER"]
+    )
+    cb.answer_reached = True
+    res = await chain.acall(message, callbacks=[cb])
+    
     answer = res["answer"]
     source_documents = res['source_documents']
     content = [source_documents[i].page_content for i in range(len(source_documents))]
@@ -141,5 +145,8 @@ async def process_response(res):
     else:
         answer += "\n\nNo sources found"
     
-    await cl.Message(content=answer, elements=source_elements).send()
-    # await cl.Message(content=answer).send()
+    if cb.has_streamed_final_answer:
+        cb.final_stream.elements = source_elements
+        await cb.final_stream.update()
+    else:
+        await cl.Message(content=answer, elements=source_elements).send()