Update app.py

app.py

@@ -14,7 +14,7 @@ from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnableLambda
 from datetime import date
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from langchain.callbacks import TimeoutCallback
+import threading
 import time
 # Environment variables
 os.environ['LANGCHAIN_TRACING_V2'] = 'true'
@@ -125,15 +125,23 @@ def clear_chat_history():
     st.session_state.messages = []
     st.session_state.context_sources = []
     st.session_state.key = 0
-
+def run_with_timeout(func, args, timeout):
+    result = [None]
+    def worker():
+        result[0] = func(*args)
+    thread = threading.Thread(target=worker)
+    thread.start()
+    thread.join(timeout)
+    if thread.is_alive():
+        return None
+    return result[0]
 # In your Streamlit app
 def generate_response(chain, query, context):
-    timeout = TimeoutCallback(timeout=60)  # 60 second timeout
-    try:
-        return chain.run(context=context, question=query, callbacks=[timeout])
-    except TimeoutError:
+    timeout_seconds = 60
+    result = run_with_timeout(chain.run, ({"context": context, "question": query},), timeout_seconds)
+    if result is None:
         return "I apologize, but I couldn't generate a response in time. The query might be too complex for me to process quickly. Could you try simplifying your question?"
-
+    return result
 
 # Sidebar
 with st.sidebar: