Update app.py

app.py

@@ -52,7 +52,7 @@ additional_notes = "Please note: Perform a comprehensive analysis including visu
 agent = CodeAgent(
     tools=[],
     model=model,
-    additional_authorized_imports=["numpy", "pandas", "matplotlib.pyplot", "seaborn"]
+    additional_authorized_imports=["numpy", "pandas", "matplotlib.pyplot", "seaborn", "os"]
 )
 
 
@@ -60,69 +60,54 @@ agent = CodeAgent(
 # Gradio Gallery and visualization output
 visual_output = gr.Gallery(label="Generated Visualizations", columns=3, height=600, object_fit="contain")
 
-# Fix in the `run_agent` function to handle agent results correctly
 def run_agent(_):
+
     if df_global is None:
         return "Please upload a file first.", []
 
+    # Save the dataset temporarily
     from tempfile import NamedTemporaryFile
     temp_file = NamedTemporaryFile(delete=False, suffix=".csv")
     df_global.to_csv(temp_file.name, index=False)
     temp_file.close()
 
+    # Prompt for the agent
     prompt = """
     You are an expert data analyst.
-    1. Load the provided dataset and analyze the structure.
-    2. Automatically detect key numeric and categorical columns.
+    1. Load the provided dataset using: df = pd.read_csv(source_file)
+    2. Automatically detect numeric and categorical columns.
     3. Perform:
-       - Basic descriptive statistics.
-       - Null and duplicate checks.
-       - Insightful relationships between key columns.
-       - At least 3 visualizations showing important trends.
-    4. Derive at least 3 actionable real-world insights.
-    5. Save all visualizations to ./figures/ directory.
-    6. Ensure all visualizations are at least 8x6 and saved at 150+ dpi.
-    Return a JSON object with keys:
-        - 'insights': clean bullet-point insights.
-        - 'figures': list of file paths of generated visualizations.
+       - Basic statistics
+       - Null/duplicate checks
+       - Correlation analysis
+       - 3+ visualizations
+    4. Extract 3+ bullet-point insights.
+    5. Ensure all figures are saved to ./figures/ directory (create it if needed).
+       Use at least 8x6 inches at 150+ dpi.
+    6. Return a JSON with:
+       - 'insights': list of insights
+       - 'figures': list of figure file paths
     """
 
-    result = agent.run(
-        prompt,
-        additional_args={"source_file": temp_file.name}
-    )
-
-    os.makedirs("figures", exist_ok=True)
+    result = agent.run(prompt, additional_args={"source_file": temp_file.name})
 
-    insights = "No insights generated."
-    image_paths = []
+    # Parse and process output
+    insights = "No insights returned."
+    images = []
 
-    # Handle different return formats from agent
     if isinstance(result, str):
         try:
             result = json.loads(result)
-        except json.JSONDecodeError:
-            return "Failed to parse result from agent.", []
-
+        except Exception:
+            return "Agent returned invalid JSON.", []
 
     if isinstance(result, dict):
-        raw_insights = result.get("insights", "No insights generated.")
-        insights = "\n".join(raw_insights) if isinstance(raw_insights, list) else raw_insights
+        raw_insights = result.get("insights", [])
+        insights = "\n".join(raw_insights) if isinstance(raw_insights, list) else str(raw_insights)
+
         image_paths = result.get("figures", [])
-        
-        print("Received image paths:", image_paths)
-
-        images = []
-        for path in image_paths:
-            try:
-                img = Image.open(path)
-                images.append(img)
-            except Exception as e:
-                print(f"Error loading image {path}: {e}")
-        
-        return insights, images
-
-    return "Unexpected output format from agent.", []
+        print("🔍 Image paths
+