Update app.py

app.py

@@ -39,8 +39,10 @@ for message in st.session_state.messages:
         
 if prompt := st.chat_input("Subject of the song"):
     st.session_state.messages.append({"role": "user", "content": prompt})
-    st.messages.chat_message("user").write(prompt)
+    with st.chat_message("user"):
+        st.markdown(prompt)
     prompt = f"Write a rap in the style of the artist SUM about {prompt}"
     output = query(prompt)
     st.session_state.messages.append({"role": "assistant", "content": output})
-    st.messages.chat_message("assistant").write(output)
+    with st.chat_message("assistant"):
+        st.markdown(output)