Update app.py

app.py

@@ -78,8 +78,12 @@ if "messages" not in st.session_state:
 
 # Display the existing chat messages via st.chat_message.
 for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
+    if message["role"] == "assistant":
+        with st.chat_message("assistant"):
+            st.markdown(message["content"])
+    elif message["role"] == "user":
+        with st.chat_message("user"):
+            st.markdown(message["content"])
 
 # Create a chat input field to allow the user to enter a message. This will display
 # automatically at the bottom of the page.