Update app.py

app.py

@@ -100,18 +100,24 @@ import ast
 
 def extract_json_from_codeagent_output(raw_output):
     try:
-        # Case 1: Already a dictionary
+        # Case 1: If it's already a dict
         if isinstance(raw_output, dict):
+            # If there's a stringified JSON inside a dict key like 'output'
+            if "output" in raw_output and isinstance(raw_output["output"], str):
+                try:
+                    return json.loads(raw_output["output"])
+                except json.JSONDecodeError:
+                    pass  # Not JSON inside
             return raw_output
 
-        # Case 2: Attempt direct JSON parsing
+        # Case 2: Try parsing the whole string as JSON
         if isinstance(raw_output, str):
             try:
                 return json.loads(raw_output)
             except json.JSONDecodeError:
-                pass  # fallback to code block extraction
+                pass  # fallback to deeper extraction
 
-            # Case 3: Extract code blocks and search for JSON-like structures
+            # Case 3: Extract code blocks (supports json/py/python/empty labels)
             code_blocks = re.findall(r"```(?:json|py|python)?\n([\s\S]*?)```", raw_output, re.DOTALL)
 
             for block in code_blocks:
@@ -119,7 +125,8 @@ def extract_json_from_codeagent_output(raw_output):
                     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*\)"
+                    r"final_answer\s*\(\s*(\{[\s\S]*?\})\s*\)",
+                    r"^(\{[\s\S]*\})$"  # Direct raw JSON block
                 ]:
                     match = re.search(pattern, block, re.DOTALL)
                     if match:
@@ -128,7 +135,7 @@ def extract_json_from_codeagent_output(raw_output):
                         except json.JSONDecodeError:
                             return ast.literal_eval(match.group(1))
 
-            # Case 4: Fallback — match the first top-level dict-like structure in raw_output
+            # Case 4: Final fallback - any dict-like structure anywhere in output
             fallback = re.search(r"\{[\s\S]+?\}", raw_output)
             if fallback:
                 try:
@@ -139,7 +146,7 @@ def extract_json_from_codeagent_output(raw_output):
     except Exception as e:
         print(f"[extract_json] Error: {e}")
 
-    # Final fallback — structured error response
+    # Case 5: If everything fails
     return {"error": "Failed to extract structured JSON"}