Update app.py

app.py

@@ -140,9 +140,13 @@ with col2:
 
 # Display chat messages from history on app rerun
 for message in st.session_state.messages:
-    avatar = "🤖" if message["role"] == "assistant" else "🕺"
-    with st.chat_message(message["role"], avatar=avatar):
-        st.markdown(message["content"])
+    if hasattr(message, "role"):
+        avatar = "🤖" if message.role == "assistant" else "🕺"
+        with st.chat_message(message.role, avatar=avatar):
+            st.markdown(message.content)
+    else:
+        st.warning("Message object does not have the 'role' attribute.")
+
 
 
 def generate_chat_responses(chat_completion) -> Generator[str, None, None]: