Update app.py

app.py

@@ -33,6 +33,9 @@ from langchain_core.messages import BaseMessage
 from pydantic import BaseModel, Field
 from typing import List
 
+from concurrent.futures import ThreadPoolExecutor
+
+
 class InMemoryHistory(BaseChatMessageHistory, BaseModel):
     messages: List[BaseMessage] = Field(default_factory=list)
 
@@ -74,11 +77,9 @@ async def async_post_request(url, headers, data):
     return await asyncio.to_thread(requests.post, url, headers=headers, json=data)
 
 
-async def keepalive():
-    while True:
-        await asyncio.sleep(10)
-        await cl.Message(content="").send()
 
+
+        
 @cl.set_starters
 async def set_starters():
     return [
@@ -103,19 +104,13 @@ async def set_starters():
             icon="/public/idea.svg",
             )
         ]
-
-
-
-@cl.on_chat_start # <————————————————— Chainlit context |  cl.config.run.log_level = "INFO"
-def setup_multiple_chains():
-    cl.user_session.set("connection_attempts", 0)
-    cl.user_session.set("last_connection_time", datetime.now())
-
-    print("Session id:")
-    #cl.user_session.set("id", session_id)
     
-    #ws_session = WebsocketSession.get_by_id(session_id=session_id)  # find other methods for WebsocketSession() to prevent oytage
-    #init_ws_context(ws_session)
+@cl.on_chat_start
+async def setup():
+  
+    cl.user_session.set("socket_auth", True)
+    cl.user_session.set("max_retries", 3)
+    cl.user_session.set("connection_attempts", 0)
     
     llm = OpenAI(
         model="gpt-3.5-turbo-instruct",
@@ -140,7 +135,21 @@ def setup_multiple_chains():
     )
 
     cl.user_session.set("llm_chain", llm_chain)
+    
+    asyncio.create_task(keep_alive())
+    
 
+async def long_running_task(message_content: str):
+    loop = asyncio.get_running_loop()
+    with ThreadPoolExecutor() as pool:
+        llm_chain = cl.user_session.get("llm_chain")
+        return await loop.run_in_executor(
+            pool,
+            lambda: llm_chain.invoke({
+                "question": message_content,
+                "chat_history": ""
+            })
+        )
     
 @cl.on_message
 async def handle_message(message: cl.Message):
@@ -193,13 +202,7 @@ async def handle_message(message: cl.Message):
 
     else:
         try:
-            # --invoke LLM w/ user_message
-            response = await llm_chain.ainvoke({
-                "question": user_message,
-                "chat_history": ""
-            }, callbacks=[cl.AsyncLangchainCallbackHandler()])
-
-            await cl.make_async.Message(content=response["text"]).send()
-
+            response = await long_running_task(message.content)
+            await cl.Message(content=response["text"]).send()
         except Exception as e:
             await cl.Message(content=f"Error: {str(e)}").send()