Update app.py

app.py

@@ -138,11 +138,10 @@ def run_with_timeout(func, args, timeout):
 # In your Streamlit app
 def generate_response(chain, query, context):
     timeout_seconds = 60
-    result = run_with_timeout(chain.run, ({"context": context, "question": query},), timeout_seconds)
+    result = run_with_timeout(chain.invoke, ({"question": query, "chat_history": st.session_state.messages},), 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:
     st.title("OHW Assistant")
@@ -151,8 +150,13 @@ with st.sidebar:
     st.button('Clear Chat History', on_click=clear_chat_history)
 
 # Main app
+# Initialize session state variables
 if "messages" not in st.session_state:
     st.session_state.messages = []
+if "context_sources" not in st.session_state:
+    st.session_state.context_sources = []
+if "context_content" not in st.session_state:
+    st.session_state.context_content = []
 
 
 for q, message in enumerate(st.session_state.messages):
@@ -180,27 +184,28 @@ if prompt := st.chat_input("How may I assist you today?"):
         st.markdown(prompt)
 
     with st.chat_message("assistant"):
-        query=st.session_state.messages[-1]['content']
+        query = st.session_state.messages[-1]['content']
         tab1, tab2 = st.tabs(["Answer", "Sources"])
         with tab1:
             with st.spinner("Generating answer..."):
                 chain = get_chain(temperature)
-                context = "\n".join(st.session_state.context_content)
                 start_time = time.time()
-                full_answer = generate_response(chain, query, context)
+                full_answer = generate_response(chain, query, "")  # Context is handled within the chain
                 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}'
-                with st.expander(name):
-                    st.markdown(f'{st.session_state.context_content[i]}')
-
-
 
+        with tab2:
+            if st.session_state.context_sources:
+                for i, source in enumerate(st.session_state.context_sources):
+                    name = f'{source}'
+                    with st.expander(name):
+                        st.markdown(f'{st.session_state.context_content[i]}')
+            else:
+                st.write("No sources available for this query.")
 
     st.session_state.messages.append({"role": "assistant", "content": full_answer})
     st.session_state.messages[-1]['sources'] = st.session_state.context_sources
     st.session_state.messages[-1]['context'] = st.session_state.context_content
+