Update app.py

app.py

@@ -14,7 +14,8 @@ 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 time
 # Environment variables
 os.environ['LANGCHAIN_TRACING_V2'] = 'true'
 os.environ['LANGCHAIN_ENDPOINT'] = 'https://api.smith.langchain.com'
@@ -125,7 +126,14 @@ def clear_chat_history():
     st.session_state.context_sources = []
     st.session_state.key = 0
 
-# st.set_page_config(page_title='OHW AI')
+# 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:
+        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?"
+
 
 # Sidebar
 with st.sidebar:
@@ -168,11 +176,14 @@ if prompt := st.chat_input("How may I assist you today?"):
         tab1, tab2 = st.tabs(["Answer", "Sources"])
         with tab1:
             with st.spinner("Generating answer..."):
-                Generate the full answer at once
-                full_answer = chain.invoke({"question": query, "chat_history": st.session_state.messages})
-            
-            # Display the full answer
-            st.markdown(full_answer, unsafe_allow_html=True)
+                chain = get_chain(temperature)
+                context = "\n".join(st.session_state.context_content)
+                start_time = time.time()
+                full_answer = generate_response(chain, query, context)
+                end_time = time.time()
+                
+            st.markdown(full_answer)
+            st.caption(f"Response time: {end_time - start_time:.2f} seconds")
         with tab2:
             for i, source in enumerate(st.session_state.context_sources):
                 name = f'{source}'