Update app.py

app.py

@@ -405,14 +405,15 @@ def process_button_click(head, button_index, topic_desc, top_concept_index):
     topic_text_index = 'topic_' + str(button_index)
     topic_text = head.ca._df[topic_text_index][top_concept_index]
 
-    response = head.answering_agent.generate_response_topic(topic_desc, topic_text, head.get_conversation(), head.selected_mode)
-    image, caption = head.answering_agent.generate_image(topic_text)
-    topic = topic_desc
-
-    st.session_state.topic_matrix[top_concept_index][button_index] += 1
-
-    with st.chat_message("assistant"): st.write(response)
-    if image and caption != "": st.image(image, caption)
+    with st.chat_message("assistant"):
+        response = head.answering_agent.generate_response_topic(topic_desc, topic_text, head.get_conversation(), head.selected_mode)
+        image, caption = head.answering_agent.generate_image(topic_text)
+        topic = topic_desc
+    
+        st.session_state.topic_matrix[top_concept_index][button_index] += 1
+    
+        st.write(response)
+        if image and caption != "": st.image(image, caption)
 
     # ... (add response & image to message)
     st.session_state.messages.append({"role": "user", "content": topic_desc, "topic": "ml", "image": None})