Update app.py

app.py

@@ -111,10 +111,17 @@ def extract_json_from_codeagent_output(raw_output):
             except json.JSONDecodeError:
                 pass  # fallback to code block extraction
 
-            # Case 3: Extract code blocks and search for JSON-like structures
+            # Case 3: Extract code blocks with any tag: json, py, python, or untagged
             code_blocks = re.findall(r"```(?:json|py|python)?\n(.*?)```", raw_output, re.DOTALL)
 
             for block in code_blocks:
+                # Try loading the block directly as JSON first
+                try:
+                    return json.loads(block)
+                except json.JSONDecodeError:
+                    pass  # continue with patterns
+
+                # Then look for common assignment patterns
                 for pattern in [
                     r"print\(\s*json\.dumps\(\s*(\{[\s\S]*?\})\s*\)\s*\)",
                     r"json\.dumps\(\s*(\{[\s\S]*?\})\s*\)",
@@ -143,6 +150,7 @@ def extract_json_from_codeagent_output(raw_output):
     return {"error": "Failed to extract structured JSON"}
 
 
+
 import pandas as pd
 import tempfile