updates

app.py

@@ -6,6 +6,11 @@ IMPORTS HERE
 import chainlit as cl
 from langchain.chat_models import ChatOpenAI
 from langchain.chains import ConversationChain
+from langchain.prompts import ChatPromptTemplate
+from langchain.schema import StrOutputParser
+from langchain.schema.runnable import Runnable
+from langchain.schema.runnable.config import RunnableConfig
+from typing import cast
 
 ### Global Section ###
 """
@@ -24,6 +29,7 @@ conversation_chain = ConversationChain(llm=llm)
 async def on_chat_start():
     """ SESSION SPECIFIC CODE HERE """
     await cl.Message(content="Welcome! How can I assist you today?").send()
+
 ### Rename Chains ###
 @cl.author_rename
 def rename(orig_author: str):
@@ -35,12 +41,15 @@ def rename(orig_author: str):
 
 ### On Message Section ###
 @cl.on_message
-async def main(message: cl.Message):
-    # Get the user input message
-    user_message = message.content
+async def on_message(message: cl.Message):
+    runnable = cast(Runnable, cl.user_session.get("runnable"))  
+
+    msg = cl.Message(content="")
 
-    # Interact with a pre-built LLM conversation chain or other models
-    response = conversation_chain.run(input=user_message)
+    async for chunk in runnable.astream(
+        {"question": message.content},
+        config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
+    ):
+        await msg.stream_token(chunk)
 
-    # Send the response back to the user
-    await cl.Message(content=response).send()
+    await msg.send()