Update app.py

app.py

@@ -57,18 +57,15 @@ st.title("💬 AI Chat Assistant")
 st.divider()
 
 # Display chat history
-for idx, message in enumerate(memory_storage.messages):
+for message in 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}))
+            st.code(message.content, language="markdown")  # Copiable markdown
+        else:
+            st.markdown(message.content)
 
 # User input at the bottom
 user_input = st.chat_input("Type your message here...")
@@ -89,12 +86,6 @@ if user_input:
 
     memory_storage.add_ai_message(answer)
 
-    # Display AI response with copy button
+    # Display AI response as a copiable markdown block
     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!")
+        st.code(answer.content, language="markdown")