tests/evaluation_harness.py

"""
Comprehensive evaluation harness for the Agentic AI System.
Implements reproducible testing, performance metrics, and cost analysis.
"""

import time
import json
import statistics
from typing import Dict, Any, List, Tuple, Optional
from dataclasses import dataclass, asdict
from pathlib import Path
import sys

# Add src to path
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

from src.agents.supervisor_agent import SupervisorAgent
from src.agents.worker_agents import QueryParserAgent, APIMatcherAgent, APIExecutorAgent, ResultFormatterAgent
from src.evaluators.reflection_evaluator import ReflectionEvaluator
from src.evaluators.quality_assessor import QualityAssessor
from src.api_clients.api_matcher import APIOperation
from src.utils.logging_config import logger


@dataclass
class TestCase:
    """Single test case for evaluation."""
    id: str
    query: str
    expected_intent: str
    expected_keywords: List[str]
    expected_api_matches: List[str]
    difficulty: str  # easy, medium, hard
    category: str  # search, create, update, delete, api


@dataclass
class EvaluationResult:
    """Result of evaluating a single test case."""
    test_case_id: str
    success: bool
    processing_time: float
    quality_score: float
    accuracy_score: float
    extracted_keywords: List[str]
    matched_apis: List[str]
    executed_apis: int
    successful_apis: int
    error_message: Optional[str]
    reflection_insights: Dict[str, Any]


@dataclass
class EvaluationReport:
    """Complete evaluation report."""
    timestamp: float
    total_tests: int
    passed_tests: int
    failed_tests: int
    average_processing_time: float
    average_quality_score: float
    average_accuracy_score: float
    performance_by_difficulty: Dict[str, Dict[str, float]]
    performance_by_category: Dict[str, Dict[str, float]]
    detailed_results: List[EvaluationResult]
    recommendations: List[str]


class EvaluationHarness:
    """Comprehensive evaluation harness for reproducible testing."""
    
    def __init__(self, config: Dict[str, Any] = None):
        """Initialize the evaluation harness."""
        self.config = config or {}
        self.test_cases = []
        self.quality_assessor = QualityAssessor()
        
        # Initialize system components
        self._setup_system()
        
        # Load test cases
        self._load_test_cases()
        
        logger.info(f"Evaluation harness initialized with {len(self.test_cases)} test cases")
    
    def _setup_system(self) -> None:
        """Set up the agentic AI system for testing."""
        # Initialize supervisor agent
        self.supervisor = SupervisorAgent(config={
            'max_retries': 2,
            'retry_delay': 0.5
        })
        
        # Initialize worker agents
        query_parser = QueryParserAgent()
        api_matcher = APIMatcherAgent()
        api_executor = APIExecutorAgent(config={'cdms': {'enabled': False}})
        result_formatter = ResultFormatterAgent()
        reflection_evaluator = ReflectionEvaluator()
        
        # Load sample API operations
        sample_operations = self._create_sample_api_operations()
        api_matcher.api_matcher.add_operations(sample_operations)
        
        # Register worker agents
        self.supervisor.register_worker('query_parser', query_parser)
        self.supervisor.register_worker('api_matcher', api_matcher)
        self.supervisor.register_worker('api_executor', api_executor)
        self.supervisor.register_worker('result_formatter', result_formatter)
        self.supervisor.register_worker('evaluator', reflection_evaluator)
        
        logger.info("System components initialized for evaluation")
    
    def _create_sample_api_operations(self) -> List[APIOperation]:
        """Create comprehensive sample API operations for testing."""
        return [
            # User management APIs
            APIOperation("getUserProfile", "GET", "/users/{id}", "Get user profile", [], ["users", "profile"]),
            APIOperation("createUser", "POST", "/users", "Create new user", [], ["users", "create"]),
            APIOperation("updateUser", "PUT", "/users/{id}", "Update user", [], ["users", "update"]),
            APIOperation("deleteUser", "DELETE", "/users/{id}", "Delete user", [], ["users", "delete"]),
            APIOperation("authenticateUser", "POST", "/auth/login", "Authenticate user", [], ["auth", "login"]),
            
            # Data and ML APIs
            APIOperation("searchDatasets", "GET", "/datasets/search", "Search ML datasets", [], ["datasets", "search", "ml"]),
            APIOperation("getDataset", "GET", "/datasets/{id}", "Get dataset details", [], ["datasets", "data"]),
            APIOperation("createModel", "POST", "/models", "Create ML model", [], ["models", "ml", "create"]),
            APIOperation("trainModel", "POST", "/models/{id}/train", "Train ML model", [], ["models", "train", "ml"]),
            APIOperation("predictModel", "POST", "/models/{id}/predict", "Make predictions", [], ["models", "predict", "ml"]),
            
            # CDMS APIs
            APIOperation("getCDMSLabels", "GET", "/cdms/labels", "Get CDMS labels", [], ["cdms", "labels"]),
            APIOperation("searchCDMSLabels", "GET", "/cdms/labels/search", "Search CDMS labels", [], ["cdms", "search", "labels"]),
            APIOperation("getCDMSDatasets", "GET", "/cdms/datasets", "Get CDMS datasets", [], ["cdms", "datasets"]),
            
            # File and storage APIs
            APIOperation("uploadFile", "POST", "/files/upload", "Upload file", [], ["files", "upload"]),
            APIOperation("downloadFile", "GET", "/files/{id}/download", "Download file", [], ["files", "download"]),
            APIOperation("listFiles", "GET", "/files", "List files", [], ["files", "list"]),
            
            # Analytics APIs
            APIOperation("getAnalytics", "GET", "/analytics", "Get analytics data", [], ["analytics", "data"]),
            APIOperation("generateReport", "POST", "/reports", "Generate report", [], ["reports", "analytics"]),
            APIOperation("exportData", "GET", "/data/export", "Export data", [], ["data", "export"]),
        ]
    
    def _load_test_cases(self) -> None:
        """Load test cases for evaluation."""
        self.test_cases = [
            # Easy test cases
            TestCase("easy_01", "Get user profile", "search", ["user", "profile"], ["getUserProfile"], "easy", "search"),
            TestCase("easy_02", "Create new user", "create", ["create", "user"], ["createUser"], "easy", "create"),
            TestCase("easy_03", "List all files", "search", ["list", "files"], ["listFiles"], "easy", "search"),
            TestCase("easy_04", "Download file", "search", ["download", "file"], ["downloadFile"], "easy", "search"),
            TestCase("easy_05", "Get analytics", "search", ["analytics"], ["getAnalytics"], "easy", "search"),
            
            # Medium test cases
            TestCase("med_01", "Find machine learning datasets for image classification", "search", 
                    ["machine learning", "datasets", "image", "classification"], ["searchDatasets"], "medium", "search"),
            TestCase("med_02", "Create and train a new classification model", "create", 
                    ["create", "train", "classification", "model"], ["createModel", "trainModel"], "medium", "create"),
            TestCase("med_03", "Search for CDMS labels related to natural language processing", "search",
                    ["CDMS", "labels", "natural language processing"], ["searchCDMSLabels"], "medium", "search"),
            TestCase("med_04", "Update user profile and authenticate", "update",
                    ["update", "user", "profile", "authenticate"], ["updateUser", "authenticateUser"], "medium", "update"),
            TestCase("med_05", "Upload training data and create ML model", "create",
                    ["upload", "training", "data", "create", "model"], ["uploadFile", "createModel"], "medium", "create"),
            
            # Hard test cases
            TestCase("hard_01", "Find datasets for computer vision, create a CNN model, and generate performance analytics", "create",
                    ["datasets", "computer vision", "CNN", "model", "analytics"], 
                    ["searchDatasets", "createModel", "getAnalytics"], "hard", "create"),
            TestCase("hard_02", "Retrieve CDMS metadata for biomedical datasets, train a classification model, and export results", "search",
                    ["CDMS", "metadata", "biomedical", "datasets", "train", "classification", "export"],
                    ["getCDMSLabels", "searchDatasets", "trainModel", "exportData"], "hard", "search"),
            TestCase("hard_03", "Authenticate user, search for their uploaded files, and generate a comprehensive report", "search",
                    ["authenticate", "user", "files", "report"],
                    ["authenticateUser", "listFiles", "generateReport"], "hard", "search"),
            TestCase("hard_04", "Create user account, upload dataset, create ML model, and predict outcomes", "create",
                    ["create", "user", "upload", "dataset", "model", "predict"],
                    ["createUser", "uploadFile", "createModel", "predictModel"], "hard", "create"),
            TestCase("hard_05", "Search multiple data sources, merge results, train ensemble model, and validate performance", "search",
                    ["search", "data", "merge", "ensemble", "model", "validate"],
                    ["searchDatasets", "getCDMSDatasets", "createModel", "trainModel"], "hard", "search"),
        ]
    
    def run_evaluation(self, test_categories: List[str] = None, difficulties: List[str] = None) -> EvaluationReport:
        """Run comprehensive evaluation."""
        logger.info("Starting comprehensive evaluation")
        
        # Filter test cases if specified
        test_cases_to_run = self.test_cases
        if test_categories:
            test_cases_to_run = [tc for tc in test_cases_to_run if tc.category in test_categories]
        if difficulties:
            test_cases_to_run = [tc for tc in test_cases_to_run if tc.difficulty in difficulties]
        
        results = []
        start_time = time.time()
        
        for i, test_case in enumerate(test_cases_to_run, 1):
            logger.info(f"Running test {i}/{len(test_cases_to_run)}: {test_case.id}")
            
            result = self._run_single_test(test_case)
            results.append(result)
            
            # Log progress
            if result.success:
                logger.info(f"✅ {test_case.id}: {result.quality_score:.3f} quality, {result.processing_time:.3f}s")
            else:
                logger.warning(f"❌ {test_case.id}: {result.error_message}")
        
        total_time = time.time() - start_time
        
        # Generate report
        report = self._generate_report(results, total_time)
        
        # Save report
        self._save_report(report)
        
        logger.info(f"Evaluation completed: {report.passed_tests}/{report.total_tests} passed")
        return report
    
    def _run_single_test(self, test_case: TestCase) -> EvaluationResult:
        """Run a single test case."""
        start_time = time.time()
        
        try:
            # Execute the query through the system
            result = self.supervisor.execute({
                'query': test_case.query,
                'metadata': {
                    'test_case_id': test_case.id,
                    'evaluation_mode': True
                }
            })
            
            processing_time = time.time() - start_time
            
            if result.success:
                # Extract results for evaluation
                formatted_output = result.data.get('formatted_output', result.data)
                
                # Assess quality
                quality_metrics = self.quality_assessor.assess_query_processing(test_case.query, formatted_output)
                
                # Calculate accuracy
                accuracy_score = self._calculate_accuracy(test_case, formatted_output)
                
                # Extract reflection insights
                reflection_insights = formatted_output.get('evaluation', {})
                
                return EvaluationResult(
                    test_case_id=test_case.id,
                    success=True,
                    processing_time=processing_time,
                    quality_score=quality_metrics.overall_score,
                    accuracy_score=accuracy_score,
                    extracted_keywords=formatted_output.get('query', {}).get('keywords', []),
                    matched_apis=self._extract_matched_api_names(formatted_output),
                    executed_apis=formatted_output.get('results', {}).get('executed_count', 0),
                    successful_apis=formatted_output.get('results', {}).get('successful_calls', 0),
                    error_message=None,
                    reflection_insights=reflection_insights
                )
            else:
                return EvaluationResult(
                    test_case_id=test_case.id,
                    success=False,
                    processing_time=processing_time,
                    quality_score=0.0,
                    accuracy_score=0.0,
                    extracted_keywords=[],
                    matched_apis=[],
                    executed_apis=0,
                    successful_apis=0,
                    error_message=result.error_message,
                    reflection_insights={}
                )
                
        except Exception as e:
            processing_time = time.time() - start_time
            return EvaluationResult(
                test_case_id=test_case.id,
                success=False,
                processing_time=processing_time,
                quality_score=0.0,
                accuracy_score=0.0,
                extracted_keywords=[],
                matched_apis=[],
                executed_apis=0,
                successful_apis=0,
                error_message=str(e),
                reflection_insights={}
            )
    
    def _calculate_accuracy(self, test_case: TestCase, results: Dict[str, Any]) -> float:
        """Calculate accuracy score for a test case."""
        accuracy_factors = []
        
        # Intent accuracy
        predicted_intent = results.get('query', {}).get('intent', '')
        intent_accuracy = 1.0 if predicted_intent == test_case.expected_intent else 0.0
        accuracy_factors.append(('intent', intent_accuracy, 0.3))
        
        # Keyword accuracy
        extracted_keywords = results.get('query', {}).get('keywords', [])
        keyword_overlap = len(set(extracted_keywords).intersection(set(test_case.expected_keywords)))
        keyword_accuracy = keyword_overlap / len(test_case.expected_keywords) if test_case.expected_keywords else 0.0
        accuracy_factors.append(('keywords', keyword_accuracy, 0.4))
        
        # API matching accuracy
        matched_apis = self._extract_matched_api_names(results)
        api_overlap = len(set(matched_apis).intersection(set(test_case.expected_api_matches)))
        api_accuracy = api_overlap / len(test_case.expected_api_matches) if test_case.expected_api_matches else 0.0
        accuracy_factors.append(('apis', api_accuracy, 0.3))
        
        # Calculate weighted average
        total_accuracy = sum(score * weight for _, score, weight in accuracy_factors)
        return total_accuracy
    
    def _extract_matched_api_names(self, results: Dict[str, Any]) -> List[str]:
        """Extract API names from results."""
        api_matches = results.get('api_matches', {}).get('matches', [])
        return [match.get('operation', {}).get('name', '') for match in api_matches]
    
    def _generate_report(self, results: List[EvaluationResult], total_time: float) -> EvaluationReport:
        """Generate comprehensive evaluation report."""
        
        # Basic statistics
        total_tests = len(results)
        passed_tests = len([r for r in results if r.success])
        failed_tests = total_tests - passed_tests
        
        successful_results = [r for r in results if r.success]
        
        if successful_results:
            avg_processing_time = statistics.mean([r.processing_time for r in successful_results])
            avg_quality_score = statistics.mean([r.quality_score for r in successful_results])
            avg_accuracy_score = statistics.mean([r.accuracy_score for r in successful_results])
        else:
            avg_processing_time = 0.0
            avg_quality_score = 0.0
            avg_accuracy_score = 0.0
        
        # Performance by difficulty
        performance_by_difficulty = {}
        for difficulty in ['easy', 'medium', 'hard']:
            difficulty_results = [r for r in results if any(tc.difficulty == difficulty and tc.id == r.test_case_id for tc in self.test_cases)]
            if difficulty_results:
                performance_by_difficulty[difficulty] = {
                    'success_rate': len([r for r in difficulty_results if r.success]) / len(difficulty_results),
                    'avg_quality': statistics.mean([r.quality_score for r in difficulty_results if r.success]) if any(r.success for r in difficulty_results) else 0.0,
                    'avg_accuracy': statistics.mean([r.accuracy_score for r in difficulty_results if r.success]) if any(r.success for r in difficulty_results) else 0.0
                }
        
        # Performance by category
        performance_by_category = {}
        for category in ['search', 'create', 'update', 'delete', 'api']:
            category_results = [r for r in results if any(tc.category == category and tc.id == r.test_case_id for tc in self.test_cases)]
            if category_results:
                performance_by_category[category] = {
                    'success_rate': len([r for r in category_results if r.success]) / len(category_results),
                    'avg_quality': statistics.mean([r.quality_score for r in category_results if r.success]) if any(r.success for r in category_results) else 0.0,
                    'avg_accuracy': statistics.mean([r.accuracy_score for r in category_results if r.success]) if any(r.success for r in category_results) else 0.0
                }
        
        # Generate recommendations
        recommendations = self._generate_recommendations(results)
        
        return EvaluationReport(
            timestamp=time.time(),
            total_tests=total_tests,
            passed_tests=passed_tests,
            failed_tests=failed_tests,
            average_processing_time=avg_processing_time,
            average_quality_score=avg_quality_score,
            average_accuracy_score=avg_accuracy_score,
            performance_by_difficulty=performance_by_difficulty,
            performance_by_category=performance_by_category,
            detailed_results=results,
            recommendations=recommendations
        )
    
    def _generate_recommendations(self, results: List[EvaluationResult]) -> List[str]:
        """Generate improvement recommendations based on results."""
        recommendations = []
        
        successful_results = [r for r in results if r.success]
        
        if not successful_results:
            return ["System failing on all test cases - requires immediate attention"]
        
        # Analyze performance patterns
        avg_quality = statistics.mean([r.quality_score for r in successful_results])
        avg_accuracy = statistics.mean([r.accuracy_score for r in successful_results])
        
        if avg_quality < 0.7:
            recommendations.append("Overall quality below threshold - improve core processing pipeline")
        
        if avg_accuracy < 0.7:
            recommendations.append("Accuracy below threshold - enhance NLP understanding and API matching")
        
        # Check processing time
        avg_time = statistics.mean([r.processing_time for r in successful_results])
        if avg_time > 1.0:
            recommendations.append("Processing time too high - optimize performance bottlenecks")
        
        # Check failure patterns
        failed_results = [r for r in results if not r.success]
        if len(failed_results) > len(successful_results) * 0.2:  # More than 20% failure rate
            recommendations.append("High failure rate - strengthen error handling and robustness")
        
        # Difficulty-specific recommendations
        hard_results = [r for r in results if any(tc.difficulty == 'hard' and tc.id == r.test_case_id for tc in self.test_cases)]
        if hard_results:
            hard_success_rate = len([r for r in hard_results if r.success]) / len(hard_results)
            if hard_success_rate < 0.5:
                recommendations.append("Poor performance on complex queries - enhance multi-step reasoning")
        
        if not recommendations:
            recommendations.append("System performing well - consider advanced optimizations and new features")
        
        return recommendations
    
    def _save_report(self, report: EvaluationReport) -> None:
        """Save evaluation report to file."""
        reports_dir = Path("evaluation_reports")
        reports_dir.mkdir(exist_ok=True)
        
        timestamp_str = time.strftime("%Y%m%d_%H%M%S", time.localtime(report.timestamp))
        report_file = reports_dir / f"evaluation_report_{timestamp_str}.json"
        
        # Convert to JSON-serializable format
        report_dict = asdict(report)
        
        with open(report_file, 'w') as f:
            json.dump(report_dict, f, indent=2)
        
        logger.info(f"Evaluation report saved to {report_file}")
    
    def run_performance_benchmark(self, iterations: int = 10) -> Dict[str, Any]:
        """Run performance benchmark with repeated executions."""
        logger.info(f"Running performance benchmark with {iterations} iterations")
        
        # Use a representative test case
        test_case = next((tc for tc in self.test_cases if tc.difficulty == 'medium'), self.test_cases[0])
        
        times = []
        quality_scores = []
        
        for i in range(iterations):
            result = self._run_single_test(test_case)
            if result.success:
                times.append(result.processing_time)
                quality_scores.append(result.quality_score)
        
        if times:
            benchmark_results = {
                'test_case': test_case.id,
                'iterations': iterations,
                'successful_runs': len(times),
                'avg_time': statistics.mean(times),
                'min_time': min(times),
                'max_time': max(times),
                'std_time': statistics.stdev(times) if len(times) > 1 else 0,
                'avg_quality': statistics.mean(quality_scores),
                'min_quality': min(quality_scores),
                'max_quality': max(quality_scores)
            }
        else:
            benchmark_results = {
                'test_case': test_case.id,
                'iterations': iterations,
                'successful_runs': 0,
                'error': 'All benchmark runs failed'
            }
        
        logger.info(f"Benchmark completed: {benchmark_results}")
        return benchmark_results


def main():
    """Run the evaluation harness."""
    print("🧪 Agentic AI System - Evaluation Harness")
    print("=" * 50)
    
    harness = EvaluationHarness()
    
    # Run full evaluation
    print("\\n📊 Running comprehensive evaluation...")
    report = harness.run_evaluation()
    
    # Print summary
    print(f"\\n📈 Evaluation Results:")
    print(f"   Total Tests: {report.total_tests}")
    print(f"   Passed: {report.passed_tests} ({report.passed_tests/report.total_tests*100:.1f}%)")
    print(f"   Failed: {report.failed_tests}")
    print(f"   Avg Quality: {report.average_quality_score:.3f}")
    print(f"   Avg Accuracy: {report.average_accuracy_score:.3f}")
    print(f"   Avg Time: {report.average_processing_time:.3f}s")
    
    print(f"\\n🎯 Performance by Difficulty:")
    for difficulty, metrics in report.performance_by_difficulty.items():
        print(f"   {difficulty.title()}: {metrics['success_rate']:.1%} success, {metrics['avg_quality']:.3f} quality")
    
    print(f"\\n💡 Recommendations:")
    for rec in report.recommendations[:3]:
        print(f"   • {rec}")
    
    # Run performance benchmark
    print(f"\\n⚡ Running performance benchmark...")
    benchmark = harness.run_performance_benchmark(iterations=5)
    if 'avg_time' in benchmark:
        print(f"   Average Time: {benchmark['avg_time']:.3f}s ± {benchmark['std_time']:.3f}s")
        print(f"   Quality Range: {benchmark['min_quality']:.3f} - {benchmark['max_quality']:.3f}")
    
    print(f"\\n✅ Evaluation completed successfully!")


if __name__ == "__main__":
    main()