Update app.py

app.py

@@ -92,6 +92,9 @@ if __name__ == '__main__':
     if "messages" not in st.session_state:
         st.session_state.messages = []
 
+    if "last_message_failed" not in st.session_state:
+        st.session_state.last_message_failed = False
+        
     for message in st.session_state.messages:
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
@@ -112,6 +115,7 @@ if __name__ == '__main__':
         st.session_state.messages.append({"role": "user", "content": question})
     
         if 'yes' in q_relevant.lower():
+            st.session_state.last_message_failed = False
             with st.status("Retrieving results..."):
                 #top_table_names = table_search(question, topk=1)['table'].tolist()
                 questions = extract_question_list(llm, question)
@@ -189,86 +193,89 @@ if __name__ == '__main__':
         else:
             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)
-    
-                    if type(questions) == list:
-                        responses = []
-                        for q in questions:
-                            top_table_names = extract_table_name(llm, q) #[extract_table_name(llm, q)] 
-                            print (top_table_names)
-                            #history = st.session_state['questions']
-                            history = st.session_state['history']
-                            try:
-                                db_chain._call(inputs={'query': q, 'history': history, \
-                                       'table_names_to_use': top_table_names})
-                            except:
-                                pass
-    
-                            if db_chain.intermediate_steps.get("result",'') != '':
-                                response = db_chain.intermediate_steps.get("result",'')
-                            elif db_chain.intermediate_steps.get("sql_data",'') != '':
-                                out = pd.DataFrame.from_dict(db_chain.intermediate_steps['sql_data'])
-                                response = tabulate(out, headers='keys', tablefmt='psql')
-                            else:
-                                response = ""
-    
-                            if "SQLQuery" in response or "Answer:" in response:
-                                response = ""
-                                
-                            responses.append(response)
-                            st.session_state['history'].append(db_chain.intermediate_steps.get("result",''))
-    
-                        response = "\n\n".join(responses)
-                        if response == "":
-                            response = "Sorry I may not have answer to this question."
-                    else:
-                        top_table_names = extract_table_name(llm, question) #[extract_table_name(llm, question)] 
+                st.session_state.last_message_failed = True
+
+    if st.session_state.last_message_failed == True:
+        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)
+
+                if type(questions) == list:
+                    responses = []
+                    for q in questions:
+                        top_table_names = extract_table_name(llm, q) #[extract_table_name(llm, q)] 
                         print (top_table_names)
                         #history = st.session_state['questions']
                         history = st.session_state['history']
-                        #try:
-                        db_chain._call(inputs={'query': question, 'history': history, \
+                        try:
+                            db_chain._call(inputs={'query': q, 'history': history, \
                                    'table_names_to_use': top_table_names})
-                        #except:
-                        #    pass
-    
+                        except:
+                            pass
+
                         if db_chain.intermediate_steps.get("result",'') != '':
-                            #st.markdown("Answer to your question is - " + db_chain.intermediate_steps.get("result",''))
-                            #print (db_chain.intermediate_steps.get("result",''))
-                            #st.markdown("The SQL query is - " + db_chain.intermediate_steps['sql_cmd'])
-                            response = "Answer to your question is - " + db_chain.intermediate_steps.get("result",'')
+                            response = db_chain.intermediate_steps.get("result",'')
                         elif db_chain.intermediate_steps.get("sql_data",'') != '':
                             out = pd.DataFrame.from_dict(db_chain.intermediate_steps['sql_data'])
-                            #st.markdown("Here is your result in a table format")
-                            #st.table(out)
-                            #st.markdown("The SQL query is - " + db_chain.intermediate_steps['sql_cmd'])
                             response = tabulate(out, headers='keys', tablefmt='psql')
-                        elif db_chain.intermediate_steps.get("sql_cmd_unchecked",'') == '':
-                            #print (db_chain)
-                            #st.markdown("Sorry I cannot answer that. Please try again later.")
-                            response = "Sorry I may not have answer to this question."
                         else:
-                            #st.markdown("Sorry I cannot answer that. Please try again later.")
-                            response = "Sorry I may not have answer to this question."
-    
+                            response = ""
+
                         if "SQLQuery" in response or "Answer:" in response:
-                            response = "Sorry I may not have answer to this question."
+                            response = ""
                             
+                        responses.append(response)
                         st.session_state['history'].append(db_chain.intermediate_steps.get("result",''))
-                
-                with st.chat_message("assistant"):
-                    st.markdown(response)
-                # Add assistant response to chat history
-                st.session_state.messages.append({"role": "assistant", "content": response})
-            elif st.button("No"):
-                question = st.session_state.messages[-1]['question']
-                response = llm.invoke(question).content
-                with st.chat_message("assistant"):
-                    st.markdown(response)
+
+                    response = "\n\n".join(responses)
+                    if response == "":
+                        response = "Sorry I may not have answer to this question."
+                else:
+                    top_table_names = extract_table_name(llm, question) #[extract_table_name(llm, question)] 
+                    print (top_table_names)
+                    #history = st.session_state['questions']
+                    history = st.session_state['history']
+                    #try:
+                    db_chain._call(inputs={'query': question, 'history': history, \
+                               'table_names_to_use': top_table_names})
+                    #except:
+                    #    pass
+
+                    if db_chain.intermediate_steps.get("result",'') != '':
+                        #st.markdown("Answer to your question is - " + db_chain.intermediate_steps.get("result",''))
+                        #print (db_chain.intermediate_steps.get("result",''))
+                        #st.markdown("The SQL query is - " + db_chain.intermediate_steps['sql_cmd'])
+                        response = "Answer to your question is - " + db_chain.intermediate_steps.get("result",'')
+                    elif db_chain.intermediate_steps.get("sql_data",'') != '':
+                        out = pd.DataFrame.from_dict(db_chain.intermediate_steps['sql_data'])
+                        #st.markdown("Here is your result in a table format")
+                        #st.table(out)
+                        #st.markdown("The SQL query is - " + db_chain.intermediate_steps['sql_cmd'])
+                        response = tabulate(out, headers='keys', tablefmt='psql')
+                    elif db_chain.intermediate_steps.get("sql_cmd_unchecked",'') == '':
+                        #print (db_chain)
+                        #st.markdown("Sorry I cannot answer that. Please try again later.")
+                        response = "Sorry I may not have answer to this question."
+                    else:
+                        #st.markdown("Sorry I cannot answer that. Please try again later.")
+                        response = "Sorry I may not have answer to this question."
+
+                    if "SQLQuery" in response or "Answer:" in response:
+                        response = "Sorry I may not have answer to this question."
+                        
+                    st.session_state['history'].append(db_chain.intermediate_steps.get("result",''))
+            
+            with st.chat_message("assistant"):
+                st.markdown(response)
+            # Add assistant response to chat history
+            st.session_state.messages.append({"role": "assistant", "content": response})
+        elif st.button("No"):
+            question = st.session_state.messages[-1]['question']
+            response = llm.invoke(question).content
+            with st.chat_message("assistant"):
+                st.markdown(response)
 
     if st.button("Reset Chat History"):
         #st.session_state['questions'] = []