planner_agent.py

"""
Planner Agent for coordinating between other agents and managing the execution plan.
Responsible for parsing goals, creating plans, and coordinating agent interactions.
"""

import os
import logging
import json
from dataclasses import dataclass
from datetime import datetime
from typing import Dict, Any, List, Optional
from utils import (
    format_timestamp, APIError, DataValidationError
)

logger = logging.getLogger(__name__)

@dataclass
class PlanStep:
    """Data class for plan step information"""
    agent: str
    action: str
    parameters: Dict[str, Any]
    expected_output: Dict[str, Any]
    status: str = "pending"
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None

@dataclass
class ExecutionPlan:
    """Data class for execution plan information"""
    goal: str
    steps: List[PlanStep]
    status: str = "pending"
    start_time: Optional[datetime] = None
    end_time: Optional[datetime] = None
    error: Optional[str] = None

class PlannerAgent:
    """Agent for planning and coordinating execution"""
    
    def __init__(self):
        """Initialize PlannerAgent"""
        self.agents = {
            'spacex': None,
            'weather': None,
            'news': None
        }
        self.logger = logging.getLogger(__name__)
        self.logger.info("PlannerAgent initialized")

    def register_agent(self, agent_type: str, agent_instance: Any) -> None:
        """
        Register an agent instance
        Args:
            agent_type: Type of agent ('spacex', 'weather', or 'news')
            agent_instance: Instance of the agent
        Raises:
            ValueError: If agent type is invalid
        """
        try:
            if agent_type not in self.agents:
                raise ValueError(f"Invalid agent type: {agent_type}")
            
            self.agents[agent_type] = agent_instance
            self.logger.info(f"Registered {agent_type} agent")
        except Exception as e:
            self.logger.error(f"Error registering agent: {str(e)}")
            raise

    def create_plan(self, goal: str) -> ExecutionPlan:
        """
        Create an execution plan based on the goal
        Args:
            goal: Natural language goal
        Returns:
            ExecutionPlan object
        Raises:
            ValueError: If goal is invalid
            DataValidationError: If plan creation fails
        """
        try:
            # Validate goal
            if not goal or not isinstance(goal, str):
                raise ValueError("Invalid goal")
            
            # Parse goal into steps
            steps = self._parse_goal(goal)
            if not steps:
                raise DataValidationError("Failed to create plan steps")
            
            # Create execution plan
            plan = ExecutionPlan(
                goal=goal,
                steps=steps,
                start_time=datetime.now()
            )
            
            self.logger.info(f"Created execution plan with {len(steps)} steps")
            return plan
        except Exception as e:
            self.logger.error(f"Error creating plan: {str(e)}")
            raise

    def _parse_goal(self, goal: str) -> List[PlanStep]:
        """
        Parse goal into plan steps
        Args:
            goal: Natural language goal
        Returns:
            List of PlanStep objects
        """
        try:
            steps = []
            
            # Check for SpaceX launch information
            if any(keyword in goal.lower() for keyword in ['launch', 'spacex', 'mission']):
                steps.append(PlanStep(
                    agent='spacex',
                    action='get_next_launch',
                    parameters={},
                    expected_output={
                        'mission_name': str,
                        'launch_date': datetime,
                        'launch_site': str,
                        'coordinates': dict
                    }
                ))
            
            # Check for weather information
            if any(keyword in goal.lower() for keyword in ['weather', 'conditions', 'forecast']):
                steps.append(PlanStep(
                    agent='weather',
                    action='get_weather_data',
                    parameters={
                        'lat': float,
                        'lon': float,
                        'launch_date': datetime
                    },
                    expected_output={
                        'temperature': float,
                        'conditions': str,
                        'wind_speed': float,
                        'delay_probability': float
                    }
                ))
            
            # Check for news information
            if any(keyword in goal.lower() for keyword in ['news', 'articles', 'reports']):
                steps.append(PlanStep(
                    agent='news',
                    action='search_articles',
                    parameters={
                        'query': str,
                        'days_back': int
                    },
                    expected_output={
                        'articles': list,
                        'has_delay_indicators': bool,
                        'confidence': float
                    }
                ))
            
            return steps
        except Exception as e:
            self.logger.error(f"Error parsing goal: {str(e)}")
            return []

    def execute_plan(self, plan: ExecutionPlan) -> Dict[str, Any]:
        """
        Execute the plan and coordinate between agents
        Args:
            plan: ExecutionPlan object
        Returns:
            Dictionary with execution results (including failed steps as None)
        """
        try:
            # Validate plan
            if not plan or not plan.steps:
                raise ValueError("Invalid execution plan")
            
            results = {}
            for step in plan.steps:
                try:
                    # Get agent instance
                    agent = self.agents.get(step.agent)
                    if not agent:
                        raise ValueError(f"Agent not registered: {step.agent}")
                    
                    # Prepare parameters (if weather, try to get from previous step)
                    if step.agent == 'weather' and ('lat' in step.parameters or 'lon' in step.parameters):
                        # Try to get coordinates from previous spacex step
                        spacex_result = results.get('spacex')
                        if spacex_result and hasattr(spacex_result, 'coordinates'):
                            step.parameters['lat'] = spacex_result.coordinates.get('latitude')
                            step.parameters['lon'] = spacex_result.coordinates.get('longitude')
                            if hasattr(spacex_result, 'launch_date'):
                                step.parameters['launch_date'] = spacex_result.launch_date
                    
                    self.logger.info(f"Executing step: {step.agent}.{step.action}")
                    result = getattr(agent, step.action)(**step.parameters)
                    step.status = "completed"
                    step.result = result
                    results[step.agent] = result
                except Exception as e:
                    self.logger.error(f"Error executing step {step.agent}.{step.action}: {str(e)}")
                    step.status = "failed"
                    step.error = str(e)
                    results[step.agent] = None  # Mark as failed, but continue
                    continue  # Continue to next step
            # Update plan status
            plan.status = "completed"
            plan.end_time = datetime.now()
            return results
        except Exception as e:
            self.logger.error(f"Error executing plan: {str(e)}")
            plan.status = "failed"
            plan.error = str(e)
            raise

    def format_plan(self, plan: ExecutionPlan) -> str:
        """
        Format execution plan for display
        Args:
            plan: ExecutionPlan object
        Returns:
            Formatted string
        """
        try:
            formatted_steps = []
            for idx, step in enumerate(plan.steps, 1):
                # Convert step results to string representation
                result_str = str(step.result)
                if hasattr(step.result, 'to_dict'):
                    result_str = str(step.result.to_dict())
                
                formatted_steps.append(
                    f"Step {idx}: {step.agent}.{step.action}\n"
                    f"Status: {step.status}\n"
                    f"Result: {result_str}\n"
                )
            
            return "\n".join(formatted_steps)
        except Exception as e:
            self.logger.error(f"Error formatting plan: {str(e)}")
            return ""

    def add_feedback(self, results: Dict[str, Any], goal: str):
        """Add feedback from previous execution to improve next plan"""
        try:
            # Log feedback for debugging
            self.logger.info(f"Adding feedback for goal: {goal}")
            self.logger.info(f"Previous results: {list(results.keys())}")
            
            # Here you could implement more sophisticated feedback mechanisms
            # For now, we'll just log the feedback
            
        except Exception as e:
            self.logger.error(f"Error adding feedback: {str(e)}")
            # Don't raise the error as feedback is not critical

# Example usage
if __name__ == "__main__":
    try:
        # Initialize planner
        planner = PlannerAgent()
        
        # Create plan
        goal = "Find next SpaceX launch, check weather, and search for news"
        plan = planner.create_plan(goal)
        
        # Print plan
        print("Execution Plan:")
        print(planner.format_plan(plan))
    except Exception as e:
        logger.error(f"Error in example usage: {str(e)}")
        print(f"Error: {str(e)}")