Update app.py

app.py

@@ -171,18 +171,17 @@ def page1():
         # Find the last assistant message
         last_assistant_message = next((message for message in reversed(st.session_state.messages.data) if message.role == "assistant"), None)
         if last_assistant_message:
-            with st.chat_message('assistant'):
-                for content_part in last_assistant_message.content:
-                    if content_part.type == 'text':
-                        st.markdown(content_part.text.value)
-                    elif hasattr(content_part, 'image_file') and content_part.image_file:
-                        image = handle_and_display_image(content_part.image_file.file_id)
-                        if image:
-                            st.image(image, caption="Generated Image")
-
             # Add assistant's response to chat history
             st.session_state.chat_history.append(("Assistant", last_assistant_message.content[0].text.value))
 
+    # Display chat history in Streamlit chat format
+    for role, message in st.session_state.chat_history:
+        if role == "User":
+            with st.chat_message("User"):
+                st.markdown(message)
+        elif role == "Assistant":
+            with st.chat_message("Assistant"):
+                st.markdown(message)
 
     # Check if 'run' object has 'status' attribute
     if hasattr(st.session_state.run, 'status'):