Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import re
 import gradio as gr
 import pandas as pd
 import numpy as np
@@ -139,6 +140,23 @@ def format_insights(insights, visuals):
     ])
 
 
+def extract_json_from_codeagent_output(raw_output):
+    try:
+        # Try to extract just the json.dumps(...) output if the model followed instructions
+        match = re.search(r'json\.dumps\((.*?)\)\s*$', raw_output, re.DOTALL)
+        if match:
+            json_string = match.group(1)
+            # Attempt to load the inner Python dict
+            parsed_dict = eval(json_string)  # Safe-ish since it’s model-generated
+            return parsed_dict
+        else:
+            # As fallback, try loading the whole string as JSON
+            return json.loads(raw_output)
+    except Exception as e:
+        print(f"[Error] Failed to parse agent output as JSON: {e}")
+        return None
+
+
 def analyze_data(csv_file, additional_notes=""):
     start_time = time.time()
     process = psutil.Process(os.getpid())
@@ -213,7 +231,14 @@ Be concise and avoid any narrative outside this final dictionary.
         wandb.log({os.path.basename(viz): wandb.Image(viz)})
 
     run.finish()
-    return format_analysis_report(analysis_result, visuals)
+    parsed_result = extract_json_from_codeagent_output(analysis_result)
+
+    if parsed_result:
+        return format_analysis_report(parsed_result, visuals)
+    else:
+    # Fallback to showing the raw string if nothing worked
+        return f"<pre>{str(analysis_result)}</pre>", visuals
+
 
 def compare_models():
     import seaborn as sns