Update app.py

app.py

@@ -187,7 +187,10 @@ if __name__ == '__main__':
             st.session_state.messages.append({"role": "assistant", "content": response})
 
         else:
-            if st.button("Looks like this question is not related to the database, but a generic. Do you want me to answer it from the table? Otherwise I will use my own knowledge."):
+            with st.chat_message("assistant"):
+                st.markdown("Looks like this question is not related to the database, but a generic. Do you want me to answer it from the table? Otherwise I will use my own knowledge.")
+            if st.button("Yes"):
+                question = st.session_state.messages[-1]['question']
                 with st.status("Retrieving results..."):
                     #top_table_names = table_search(question, topk=1)['table'].tolist()
                     questions = extract_question_list(llm, question)
@@ -261,7 +264,8 @@ if __name__ == '__main__':
                     st.markdown(response)
                 # Add assistant response to chat history
                 st.session_state.messages.append({"role": "assistant", "content": response})
-            else:
+            elif st.button("No"):
+                question = st.session_state.messages[-1]['question']
                 response = llm.invoke(question).content
                 with st.chat_message("assistant"):
                     st.markdown(response)