Update app.py

app.py

@@ -11,6 +11,8 @@ import pyodbc
 import openai
 from langchain_openai import AzureChatOpenAI
 
+from tabulate import tabulate
+
 from utils import SQLDatabaseChainPatched, table_search, extract_question_type, extract_table_name
 
 openai.api_key = os.environ['OPENAI_API_KEY']
@@ -83,7 +85,7 @@ template = """You are a database expert. Given an input question, first create a
             """
 
 if __name__ == '__main__':
-    connection_string = ("Driver=FreeTDS;Server=crawlersdb.c3pzpntwjvdf.us-east-1.rds.amazonaws.com;Database=SmartClever;PORT=1433;UID=CleverData;PWD={};TrustServerCertificate=yes;".format(os.environ['DB_PWD'])
+    connection_string = ("Driver=FreeTDS;Server=crawlersdb.c3pzpntwjvdf.us-east-1.rds.amazonaws.com;Database=SmartCleverST;PORT=1433;UID=CleverData;PWD={};TrustServerCertificate=yes;".format(os.environ['DB_PWD'])
     )
     connection_url = URL.create(
         "mssql+pyodbc", 
@@ -105,8 +107,22 @@ if __name__ == '__main__':
             '4k': llm
         })
 
-    question = st.text_input("Ask a question in natural language and press enter")
+    #question = st.text_input("Ask a question in natural language and press enter")
+
+    if question := st.chat_input("tion in natural language and press enter"):
+        # Display user message in chat message container
+        with st.chat_message("user"):
+            st.markdown(prompt)
+        # Add user message to chat history
+        st.session_state.messages.append({"role": "user", "content": question})
+    
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
 
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+        
     if 'questions' not in st.session_state:
         st.session_state['questions'] = []
 
@@ -130,19 +146,28 @@ if __name__ == '__main__':
                     pass
             
             if db_chain.intermediate_steps.get("sql_cmd_unchecked",'') == '':
-                print (db_chain)
-                st.markdown("Sorry I cannot answer that. Please try again later.")
+                #print (db_chain)
+                #st.markdown("Sorry I cannot answer that. Please try again later.")
+                response = "Sorry I cannot answer that. Please try again later."
             elif 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'])
+                #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'])
+                #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')
             else:
-                st.markdown("Sorry I cannot answer that. Please try again later.")
+                #st.markdown("Sorry I cannot answer that. Please try again later.")
+                response = "Sorry I cannot answer that. Please try again later."
+            
+            with st.chat_message("assistant"):
+                st.markdown(response)
+            # Add assistant response to chat history
+            st.session_state.messages.append({"role": "assistant", "content": response})
 
             st.session_state['questions'].append(question)