run_eval.py

"""
Evaluation system for the multi-agent AI system.
Tests agent chaining, goal satisfaction, and iterative refinement.
"""

import os
import json
import logging
from datetime import datetime, timedelta
from typing import Dict, Any, List
from main import AgentOrchestrator

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class EvaluationMetrics:
    """Metrics for evaluating agent performance"""
    
    def __init__(self):
        self.goals_achieved = 0
        self.total_goals = 0
        self.agent_usage = {}
        self.iteration_counts = []
        self.execution_times = []
        self.error_counts = {}
        
    def record_goal_result(self, achieved: bool):
        """Record whether a goal was achieved"""
        self.total_goals += 1
        if achieved:
            self.goals_achieved += 1
            
    def record_agent_usage(self, agent_name: str):
        """Record agent usage"""
        self.agent_usage[agent_name] = self.agent_usage.get(agent_name, 0) + 1
        
    def record_iteration(self, count: int):
        """Record number of iterations"""
        self.iteration_counts.append(count)
        
    def record_execution_time(self, time_ms: float):
        """Record execution time"""
        self.execution_times.append(time_ms)
        
    def record_error(self, error_type: str):
        """Record error occurrence"""
        self.error_counts[error_type] = self.error_counts.get(error_type, 0) + 1
        
    def get_summary(self) -> Dict[str, Any]:
        """Get evaluation summary"""
        return {
            'goal_success_rate': self.goals_achieved / self.total_goals if self.total_goals > 0 else 0,
            'average_iterations': sum(self.iteration_counts) / len(self.iteration_counts) if self.iteration_counts else 0,
            'average_execution_time': sum(self.execution_times) / len(self.execution_times) if self.execution_times else 0,
            'agent_usage': self.agent_usage,
            'error_counts': self.error_counts
        }

class TestCase:
    """Test case for evaluating the system"""
    
    def __init__(self, goal: str, expected_agents: List[str], expected_data: Dict[str, List[str]]):
        self.goal = goal
        self.expected_agents = expected_agents
        self.expected_data = expected_data
        
    def validate_results(self, results: Dict[str, Any]) -> bool:
        """Validate if results meet expectations"""
        try:
            # Check if all expected agents were used
            for agent in self.expected_agents:
                if agent not in results:
                    return False
                    
            # Check if all expected data is present
            for agent, fields in self.expected_data.items():
                if agent in results:
                    for field in fields:
                        if field not in results[agent]:
                            return False
                            
            return True
            
        except Exception as e:
            logger.error(f"Error validating results: {str(e)}")
            return False

def run_evaluation():
    """Run the evaluation suite"""
    try:
        # Initialize metrics
        metrics = EvaluationMetrics()
        
        # Define test cases
        test_cases = [
            TestCase(
                goal="Check if next SpaceX launch will be delayed",
                expected_agents=['spacex', 'weather', 'news'],
                expected_data={
                    'spacex': ['launch_date', 'launch_site'],
                    'weather': ['temperature', 'conditions', 'wind_speed'],
                    'news': ['articles', 'delay_indicators']
                }
            ),
            TestCase(
                goal="Get weather conditions for next launch",
                expected_agents=['spacex', 'weather'],
                expected_data={
                    'spacex': ['launch_date', 'launch_site'],
                    'weather': ['temperature', 'conditions', 'wind_speed']
                }
            ),
            TestCase(
                goal="Find news about launch delays",
                expected_agents=['spacex', 'news'],
                expected_data={
                    'spacex': ['launch_date', 'launch_site'],
                    'news': ['articles', 'delay_indicators']
                }
            )
        ]
        
        # Initialize orchestrator
        orchestrator = AgentOrchestrator()
        orchestrator.initialize()
        
        # Run test cases
        for test_case in test_cases:
            try:
                logger.info(f"Running test case: {test_case.goal}")
                start_time = datetime.now()
                
                # Execute the test case
                plan, results = orchestrator.execute(test_case.goal)
                
                # Record metrics
                execution_time = (datetime.now() - start_time).total_seconds() * 1000
                metrics.record_execution_time(execution_time)
                metrics.record_iteration(orchestrator.current_iteration)
                
                # Record agent usage
                for agent in results.keys():
                    metrics.record_agent_usage(agent)
                    
                # Validate results
                goal_achieved = test_case.validate_results(results)
                metrics.record_goal_result(goal_achieved)
                
                logger.info(f"Test case completed: {'Success' if goal_achieved else 'Failed'}")
                
            except Exception as e:
                logger.error(f"Error in test case: {str(e)}")
                metrics.record_error(type(e).__name__)
                
        # Get and save evaluation summary
        summary = metrics.get_summary()
        with open('evaluation_results.json', 'w') as f:
            json.dump(summary, f, indent=2)
            
        # Print summary
        print("\nEvaluation Summary:")
        print(f"Goal Success Rate: {summary['goal_success_rate']:.2%}")
        print(f"Average Iterations: {summary['average_iterations']:.2f}")
        print(f"Average Execution Time: {summary['average_execution_time']:.2f}ms")
        print("\nAgent Usage:")
        for agent, count in summary['agent_usage'].items():
            print(f"- {agent}: {count}")
        print("\nErrors:")
        for error, count in summary['error_counts'].items():
            print(f"- {error}: {count}")
            
    except Exception as e:
        logger.error(f"Evaluation error: {str(e)}")
        raise
    finally:
        orchestrator.cleanup()

if __name__ == "__main__":
    run_evaluation()