Update app.py

app.py

@@ -13,8 +13,6 @@ app.secret_key = 'YourSecretKey'  # Change this to a real secret key for product
 
 # OAuth setup with Authlib
 oauth = OAuth(app)
-
-# Assuming environment variables are set for OAuth
 oauth.register(
     name='oauth_provider',
     client_id=os.getenv("OAUTH_CLIENT_ID"),
@@ -43,6 +41,24 @@ def initialize_chainlit():
         )
     )
 
+    # Setup chainlit callbacks
+    @cl.on_chat_start
+    async def on_chat_start():
+        prompt = ChatPromptTemplate.from_messages([("human", "{question}")])
+        runnable = prompt | llm | StrOutputParser()
+        cl.user_session.set("runnable", runnable)
+
+    @cl.on_message
+    async def on_message(message: cl.Message):
+        runnable = cl.user_session.get("runnable")  # type: Runnable
+        msg = cl.Message(content="")
+        async for chunk in runnable.astream(
+            {"question": message.content},
+            config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
+        ):
+            await msg.stream_token(chunk)
+        await msg.send()
+
 @app.route('/')
 def home():
     return 'Home - <a href="/login">Login with OAuth Provider</a>'
@@ -55,10 +71,10 @@ def login():
 @app.route('/login/callback')
 def authorize():
     token = oauth.oauth_provider.authorize_access_token()
-    # You can use token to fetch user info or proceed directly if not needed
-    # Here, initialize ChainLit or perform actions based on the authenticated user
+    # Initialize ChainLit or perform actions based on the authenticated user
     initialize_chainlit()
     return 'Logged in and language model initialized. Proceed with operations.'
 
 if __name__ == "__main__":
     app.run(debug=True)
+