Update app.py

app.py

@@ -55,7 +55,7 @@ agent = CodeAgent(
 
 def run_agent(_):
     if df_global is None:
-        return "Please upload a file first."
+        return "Please upload a file first.", []
 
     from tempfile import NamedTemporaryFile
     temp_file = NamedTemporaryFile(delete=False, suffix=".csv")
@@ -78,24 +78,15 @@ def run_agent(_):
       - File paths of the generated visualizations.
     """
 
-    result = agent.run(
-        prompt,
-        additional_args={"source_file": temp_file.name}
-    )
-
-    # Extract image paths from output (assuming SmolAgent prints them)
-    image_paths = []
-    for line in result.splitlines():
-        if line.strip().endswith(".png"):
-            image_paths.append(line.strip())
+    result = agent.run(prompt, additional_args={"source_file": temp_file.name})
 
+    image_paths = [line.strip() for line in result.splitlines() if line.strip().endswith(".png")]
     insights = "\n".join([line for line in result.splitlines() if not line.strip().endswith(".png")])
 
     return insights, image_paths
 
 
 
-
 def train_model(_):
     wandb.login(key=os.environ.get("WANDB_API_KEY"))
     #wandb_run = wandb.init(project="huggingface-data-analysis", name="Optuna_Run", reinit=True)
@@ -223,7 +214,6 @@ def explainability(_):
 
 
 
-
 with gr.Blocks() as demo:
     gr.Markdown("## 📊 AI-Powered Data Analysis with Hyperparameter Optimization")
 
@@ -247,7 +237,13 @@ with gr.Blocks() as demo:
         shap_img = gr.Image(label="SHAP Summary Plot")
         lime_img = gr.Image(label="LIME Explanation")
 
-    agent_btn.click(fn=run_agent, inputs=df_output, outputs=insights_output)
+    with gr.Row():
+        agent_btn = gr.Button("Run AI Agent (5 Insights + 5 Visualizations)")
+        insights_output = gr.Textbox(label="Insights from SmolAgent", lines=15)
+        visual_output = gr.Gallery(label="Generated Visualizations").style(grid=3, height="auto")
+
+    #agent_btn.click(fn=run_agent, inputs=df_output, outputs=insights_output)
+    agent_btn.click(fn=run_agent, inputs=df_output, outputs=[insights_output, visual_output])
     train_btn.click(fn=train_model, inputs=df_output, outputs=[metrics_output, trials_output])
     explain_btn.click(fn=explainability, inputs=df_output, outputs=[shap_img, lime_img])