Update app.py

app.py

@@ -42,7 +42,6 @@ db = SQLDatabase.from_uri(f"sqlite:///{db_loc}")
 # Retrieve the schema information of the database tables
 database_schema = db.get_table_info()
 
-st.write(database_schema)
 
 #=================================Logging=====================================#
 
@@ -395,21 +394,16 @@ def chat_with_agent():
     if 'conversation_history' not in st.session_state:
         st.session_state.conversation_history = [{"role": "assistant", "content": "Hey, how can I help you?"}]
 
-    # Display the initial greeting
+    # Display chat messages from history on app rerun
     for message in st.session_state.conversation_history:
-        if message['role'] == 'assistant':
-            st.write(f"Chatbot: {message['content']}")
-        else:
-            st.write(f"You: {message['content']}")
-
-    user_input = st.text_input("You: ", key="user_input")
-
-    if st.button("Send"):
-        if user_input.lower() == "exit":
-            st.write("Chatbot: Thank you for chatting. Goodbye!")
-            return
-
-        # Add user input to conversation history
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+
+    # React to user input
+    if user_input := st.chat_input("You: "):
+        # Display user message in chat message container
+        st.chat_message("user").markdown(user_input)
+        # Add user message to conversation history
         st.session_state.conversation_history.append({"role": "user", "content": user_input})
 
         conversation_input = "\n".join(
@@ -423,7 +417,11 @@ def chat_with_agent():
         st.session_state.conversation_history.append({"role": "assistant", "content": chatbot_response})
 
         # Display the chatbot's response
-        st.write(f"Chatbot: {chatbot_response}")
+        with st.chat_message("assistant"):
+            st.markdown(chatbot_response)
+
+if __name__ == "__main__":
+    chat_with_agent()
 
 if __name__ == "__main__":
     chat_with_agent()