Update app.py

app.py

@@ -94,36 +94,55 @@ def clean_data(df):
 
 # Add a extraction of JSON if CodeAgent Output is not in format
 
+import json
+import re
+import ast
+
 def extract_json_from_codeagent_output(raw_output):
     try:
-        # Correct regex pattern to extract code blocks that contain Python code
-        code_blocks = re.findall(r"```(?:py|python)?\n(.*?)```", raw_output, re.DOTALL)
-        
-        # Loop through each code block and look for patterns related to JSON or result definitions
-        for block in code_blocks:
-            for pattern in [
-                r"print\(\s*json\.dumps\(\s*(\{[\s\S]*?\})\s*\)\s*\)",
-                r"json\.dumps\(\s*(\{[\s\S]*?\})\s*\)",
-                r"result\s*=\s*(\{[\s\S]*?\})"
-            ]:
-                match = re.search(pattern, block, re.DOTALL)
-                if match:
-                    try:
-                        return json.loads(match.group(1))  # Try to parse the JSON directly
-                    except json.JSONDecodeError:
-                        return ast.literal_eval(match.group(1))  # Fallback to evaluating the dictionary
-
-        # Final fallback: Look for any dictionary-like structure in the entire output
-        fallback = re.search(r"\{[\s\S]+?\}", raw_output)
-        if fallback:
-            return json.loads(fallback.group(0))
-    
+        # Case 1: Already a dictionary
+        if isinstance(raw_output, dict):
+            return raw_output
+
+        # Case 2: Attempt direct JSON parsing
+        if isinstance(raw_output, str):
+            try:
+                return json.loads(raw_output)
+            except json.JSONDecodeError:
+                pass  # fallback to code block extraction
+
+            # Case 3: Extract code blocks and search for JSON-like structures
+            code_blocks = re.findall(r"```(?:json|py|python)?\n(.*?)```", raw_output, re.DOTALL)
+
+            for block in code_blocks:
+                for pattern in [
+                    r"print\(\s*json\.dumps\(\s*(\{[\s\S]*?\})\s*\)\s*\)",
+                    r"json\.dumps\(\s*(\{[\s\S]*?\})\s*\)",
+                    r"result\s*=\s*(\{[\s\S]*?\})",
+                    r"final_answer\s*\(\s*(\{[\s\S]*?\})\s*\)"
+                ]:
+                    match = re.search(pattern, block, re.DOTALL)
+                    if match:
+                        try:
+                            return json.loads(match.group(1))
+                        except json.JSONDecodeError:
+                            return ast.literal_eval(match.group(1))
+
+            # Case 4: Fallback — match the first top-level dict-like structure in raw_output
+            fallback = re.search(r"\{[\s\S]+?\}", raw_output)
+            if fallback:
+                try:
+                    return json.loads(fallback.group(0))
+                except json.JSONDecodeError:
+                    return ast.literal_eval(fallback.group(0))
+
     except Exception as e:
         print(f"[extract_json] Error: {e}")
-    
-    # Return an error if JSON extraction fails
+
+    # Final fallback — structured error response
     return {"error": "Failed to extract structured JSON"}
 
+
 import pandas as pd
 import tempfile