Update app.py

app.py

@@ -4,7 +4,6 @@ import gradio as gr
 from transformers import ReactCodeAgent, HfEngine
 import pandas as pd
 
-from gradio import Chatbot
 from transformers.agents import stream_to_gradio
 from huggingface_hub import login
 
@@ -45,7 +44,7 @@ def get_images_in_directory(directory):
 
 def interact_with_agent(file_input, prompt):
     if file_input is None:
-        yield [gr.ChatMessage(role="assistant", content="Please upload a CSV file before running the analysis.")]
+        yield [["assistant", "Please upload a CSV file before running the analysis."]]
         return
 
     shutil.rmtree("./figures", ignore_errors=True)
@@ -54,7 +53,7 @@ def interact_with_agent(file_input, prompt):
     try:
         data_file = pd.read_csv(file_input.name)
     except Exception as e:
-        yield [gr.ChatMessage(role="assistant", content=f"Error reading CSV file: {str(e)}")]
+        yield [["assistant", f"Error reading CSV file: {str(e)}"]]
         return
 
     data_structure_notes = f"""- Description (output of .describe()):
@@ -66,25 +65,24 @@ def interact_with_agent(file_input, prompt):
     if prompt:
         full_prompt += f"\nAdditional notes: {prompt}"
 
-    messages = [gr.ChatMessage(role="user", content=full_prompt)]
-    yield messages + [
-        gr.ChatMessage(role="assistant", content="⏳ Starting task...")
-    ]
+    messages = [["user", full_prompt]]
+    yield messages + [["assistant", "⏳ Starting task..."]]
 
     plot_image_paths = {}
     for msg in stream_to_gradio(agent, full_prompt, data_file=data_file):
-        messages.append(msg)
-        for image_path in get_images_in_directory("./figures"):
+        if isinstance(msg.content, str):
+            messages.append(["assistant", msg.content])
+        elif isinstance(msg.content, dict) and 'path' in msg.content:
+            # Handle image messages
+            image_path = msg.content['path']
             if image_path not in plot_image_paths:
-                image_message = gr.ChatMessage(
-                    role="assistant",
-                    content=gr.Image.update(value=image_path),
-                )
+                messages.append(["assistant", (image_path,)])
                 plot_image_paths[image_path] = True
-                messages.append(image_message)
-        yield messages + [
-            gr.ChatMessage(role="assistant", content="⏳ Still processing...")
-        ]
+        
+        yield messages + [["assistant", "⏳ Still processing..."]]
+    
+    # Remove the last "Still processing..." message
+    messages = messages[:-1]
     yield messages
 
 with gr.Blocks(