Update app.py

app.py

@@ -57,13 +57,19 @@ st.title("💬 AI Chat Assistant")
 st.divider()
 
 # Display chat history
-for message in memory_storage.messages:
+for idx, message in enumerate(memory_storage.messages):
     role = "user" if message.type == "human" else "assistant"
     avatar = user_avatar if role == "user" else ai_avatar  # Assign appropriate avatar
 
     with st.chat_message(role, avatar=f"data:image/jpeg;base64,{avatar}"):
         st.markdown(message.content)
 
+        # Add copy button only for AI messages
+        if role == "assistant":
+            copy_key = f"copy_btn_{idx}"  # Unique key for each copy button
+            st.code(message.content, language="text")  # Display code block with text formatting
+            st.button("📋 Copy", key=copy_key, on_click=lambda text=message.content: st.session_state.update({"copied_text": text}))
+
 # User input at the bottom
 user_input = st.chat_input("Type your message here...")
 
@@ -83,6 +89,12 @@ if user_input:
 
     memory_storage.add_ai_message(answer)
 
-    # Display AI response
+    # Display AI response with copy button
     with st.chat_message("assistant", avatar=f"data:image/jpeg;base64,{ai_avatar}"):
         st.markdown(answer.content)
+
+        # Copy button for AI response
+        st.code(answer.content, language="text")
+        if st.button("📋 Copy", key="copy_latest"):
+            st.session_state["copied_text"] = answer.content
+            st.success("✅ Copied to clipboard!")