Update base_agent.py

base_agent.py

@@ -4,6 +4,7 @@ import logging
 from datetime import datetime
 from pathlib import Path
 from typing import Any, Dict, Optional
+import re
 from openai import OpenAI
 from config import settings
 
@@ -74,10 +75,29 @@ class BaseAgent:
             
             # Parse output (assuming JSON format)
             try:
+                # Attempt direct JSON parsing first
                 output_data = json.loads(output_text)
             except json.JSONDecodeError:
-                # If not JSON, wrap in a generic output
-                output_data = {"output": output_text}
+                # If direct parsing fails, try to extract JSON from markdown or embedded text
+                json_match = re.search(r'```json\s*(.*?)\s*```', output_text, re.DOTALL)
+                if json_match:
+                    try:
+                        output_data = json.loads(json_match.group(1))
+                    except json.JSONDecodeError:
+                        # Fallback if markdown-wrapped JSON is malformed
+                        output_data = {"output": output_text}
+                else:
+                    # Try to find any JSON object embedded in the text
+                    json_obj_match = re.search(r'(\{.*\})', output_text, re.DOTALL)
+                    if json_obj_match:
+                        try:
+                            output_data = json.loads(json_obj_match.group(1))
+                        except json.JSONDecodeError:
+                            # Fallback if embedded JSON is malformed
+                            output_data = {"output": output_text}
+                    else:
+                        # If no JSON found, wrap the entire output
+                        output_data = {"output": output_text}
             
             logger.info(f"{self.agent_name}: Processing completed")