evaluate_questions.py

# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates

#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at

#      http://www.apache.org/licenses/LICENSE-2.0

#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""
AInsteinBench Unified Evaluator

This script evaluates questions from the AInsteinBench dataset 
"""

import os
import sys
import json
import logging
import argparse
import docker
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Any, Optional

# Add project root to path
sys.path.insert(0, os.path.dirname(__file__))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'curation', 'et'))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'curation', 'msb'))

from ainsteinbench.question import Question

# Import ET evaluator components
try:
    from ET_evaluator import (
        ConfigManager as ETConfigManager,
        DockerManager,
        EinsteinToolkitTester,
        EinsteinToolkitScorer
    )
    ET_AVAILABLE = True
except ImportError as e:
    print(f"Warning: ET evaluator not available: {e}")
    ET_AVAILABLE = False

# Import MSB log parser
try:
    from log_parser import parse_log as msb_parse_log
    MSB_PARSER_AVAILABLE = True
except ImportError as e:
    print(f"Warning: MSB log parser not available: {e}")
    MSB_PARSER_AVAILABLE = False


class UnifiedEvaluator:
    """Main unified evaluator for AInsteinBench questions"""
    
    def __init__(self, questions_file: str, answers_dir: str, config_file: Optional[str] = None):
        """
        Initialize the unified evaluator.
        
        Args:
            questions_file: Path to JSONL file with questions
            answers_dir: Directory containing reference answers
            config_file: Optional config file (defaults to ET config)
        """
        self.questions_file = questions_file
        self.answers_dir = Path(answers_dir)
        self.config_file = config_file or os.path.join(
            os.path.dirname(__file__), 'curation', 'et', 'config_server.json'
        )
        
        # Setup logging
        self.setup_logging()
        
        # Load questions
        self.questions = self.load_questions()
        logging.info(f"Loaded {len(self.questions)} questions from {questions_file}")
        
    def setup_logging(self, output_dir: str = "."):
        """Setup logging configuration"""
        os.makedirs(output_dir, exist_ok=True)
        
        log_filename = f"ainsteinbench_eval_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
        log_filepath = os.path.join(output_dir, log_filename)
        
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.StreamHandler(sys.stdout),
                logging.FileHandler(log_filepath)
            ]
        )
        logging.info(f"Logging to {log_filepath}")
    
    def load_questions(self) -> List[Question]:
        """Load questions from JSONL file"""
        questions = []
        
        with open(self.questions_file, 'r') as f:
            for line_num, line in enumerate(f, 1):
                try:
                    data = json.loads(line)
                    question = Question.from_dict(data)
                    questions.append(question)
                except Exception as e:
                    logging.error(f"Error loading question at line {line_num}: {e}")
        
        return questions
    
    def get_answer_path(self, question: Question) -> Optional[Path]:
        """
        Get the path to the answer file for a question in the answers directory.
        
        Matches files by question_id prefix, ignoring extensions.
        
        Args:
            question: Question object
            
        Returns:
            Path to answer file, or None if not found
        """
        question_id = question.question_id
        
        # List all files in answers directory
        if not self.answers_dir.exists():
            return None
        
        # Find files that start with the question_id
        for file_path in self.answers_dir.iterdir():
            if file_path.is_file() and file_path.stem == question_id:
                return file_path
        
        return None
    
    def evaluate(self, max_examples: Optional[int] = None) -> List[Dict[str, Any]]:
        """
        Run evaluation on all questions.
        
        Args:
            max_examples: Maximum number of examples to evaluate
            
        Returns:
            List of evaluation results
        """
        questions_to_eval = self.questions[:max_examples] if max_examples else self.questions
        
        logging.info(f"Starting evaluation of {len(questions_to_eval)} questions")
        
        results = []
        for i, question in enumerate(questions_to_eval):
            logging.info(f"\n{'='*70}")
            logging.info(f"Evaluating question {i+1}/{len(questions_to_eval)}: {question.question_id}")
            logging.info(f"Type: {question.question_type}")
            logging.info(f"{'='*70}")
            
            result = self.evaluate_question(question, i)
            results.append(result)
        
        return results
    
    def evaluate_question(self, question: Question, index: int) -> Dict[str, Any]:
        """
        Evaluate a single question by routing to appropriate evaluator.
        
        Args:
            question: Question object to evaluate
            index: Question index
            
        Returns:
            Evaluation result dictionary
        """
        # Get answer file
        answer_path = self.get_answer_path(question)
        if not answer_path:
            logging.error(f"Answer file not found for {question.question_id}")
            return self._create_error_result(question, index, "Answer file not found")
        
        logging.info(f"Answer file: {answer_path}")
        
        # Route to appropriate evaluator
        if question.question_type == 'einstein_toolkit':
            return self.evaluate_et_question(question, answer_path, index)
        elif question.question_type == 'multi_swe_bench':
            return self.evaluate_msb_question(question, answer_path, index)
        else:
            logging.error(f"Unknown question type: {question.question_type}")
            return self._create_error_result(question, index, f"Unknown question type: {question.question_type}")
    
    def evaluate_et_question(self, question: Question, answer_path: Path, index: int) -> Dict[str, Any]:
        """
        Evaluate an Einstein Toolkit question.
        
        Args:
            question: ET question object
            answer_path: Path to answer source file
            index: Question index
            
        Returns:
            Evaluation result dictionary
        """
        if not ET_AVAILABLE:
            return self._create_error_result(question, index, "ET evaluator not available")
        
        logging.info("Using Einstein Toolkit evaluator")
        
        try:
            # Load answer code
            with open(answer_path, 'r') as f:
                answer_code = f.read()
            
            # Setup ET evaluator components
            # Note: Save the original directory to return to it later
            original_dir = os.getcwd()
            et_dir = Path(__file__).parent / 'curation' / 'et'
            os.chdir(et_dir)
            
            config = ETConfigManager(self.config_file)
            docker_mgr = DockerManager(config, f"eval_{index}", worker_id=index)
            scorer = EinsteinToolkitScorer(config)
            
            # Start Docker container
            logging.info("Starting Docker container...")
            if not docker_mgr.start_container():
                return self._create_error_result(question, index, "Failed to start Docker container")
            
            try:
                # Create tester
                tester = EinsteinToolkitTester(config, docker_mgr, worker_id=index)
                
                # Extract ET-specific info from question using Question attributes
                thorn_name = question.content.get('thorn_name')
                src_filename = question.content.get('src_filename')
                
                if not thorn_name or not src_filename:
                    return self._create_error_result(
                        question, index, 
                        "Missing thorn_name or src_filename in question content"
                    )
                
                logging.info(f"Deploying code for {thorn_name}/{src_filename}")
                
                # Deploy answer code
                if not tester.deploy_code(thorn_name, src_filename, answer_code):
                    return self._create_error_result(question, index, "Failed to deploy code")
                
                # Build
                logging.info(f"Building thorn {thorn_name}")
                build_result = tester.build_thorn(thorn_name)
                
                # Test (only if build succeeded)
                if build_result['success']:
                    logging.info(f"Running tests for {thorn_name}")
                    test_result = tester.run_tests_for_thorn(thorn_name)
                else:
                    logging.warning("Build failed, skipping tests")
                    test_result = {
                        'total_tests': 0,
                        'passed_tests': 0,
                        'failed_tests': 0,
                        'test_results': []
                    }
                
                # Score
                score_result = scorer.calculate_score(build_result, test_result)
                
                # Cleanup
                for src_path in list(tester.backed_up_files.keys()):
                    tester.restore_source_file(src_path)
                
                result = {
                    'question_id': question.question_id,
                    'question_type': question.question_type,
                    'index': index,
                    'thorn_name': thorn_name,
                    'src_filename': src_filename,
                    'build_result': build_result,
                    'test_result': test_result,
                    'score_result': score_result,
                    'timestamp': datetime.now().isoformat(),
                    'success': True
                }
                
                logging.info(f"Score: {score_result['overall_score']:.1f}/100")
                return result
                
            finally:
                # Stop Docker container
                docker_mgr.stop_container()
                # Return to original directory
                os.chdir(original_dir)
                
        except Exception as e:
            logging.error(f"Error evaluating ET question: {e}", exc_info=True)
            return self._create_error_result(question, index, f"Evaluation error: {e}")
    
    def evaluate_msb_question(self, question: Question, answer_path: Path, index: int) -> Dict[str, Any]:
        """
        Evaluate a Multi-SWE-bench question.
        
        Args:
            question: MSB question object
            answer_path: Path to answer patch file
            index: Question index
            
        Returns:
            Evaluation result dictionary
        """
        if not MSB_PARSER_AVAILABLE:
            return self._create_error_result(question, index, "MSB log parser not available")
        
        logging.info("Using Multi-SWE-bench evaluator")
        
        try:
            # Load patch file
            with open(answer_path, 'r') as f:
                patch_content = f.read()
            
            # Extract MSB-specific info from question
            org = question.content.get('org', '')
            repo = question.content.get('repo', '')
            pr_number = question.content.get('pr_number', 0)
            docker_image = question.environment.get('docker_image', '')
            
            if not docker_image:
                return self._create_error_result(question, index, "Docker image not specified in question")
            
            logging.info(f"Docker image: {docker_image}")
            logging.info(f"Repository: {org}/{repo} PR#{pr_number}")
            
            # Initialize Docker client
            client = docker.from_env()
            
            # Start Docker container
            container_name = f"msb_eval_{index}"
            logging.info(f"Starting container: {container_name}")
            
            try:
                # Remove existing container if it exists
                try:
                    old_container = client.containers.get(container_name)
                    old_container.remove(force=True)
                except docker.errors.NotFound:
                    pass
                
                # Create container
                # Mount the patch file to /home/fix.patch in the container
                container = client.containers.run(
                    docker_image,
                    name=container_name,
                    detach=True,
                    tty=True,
                    stdin_open=True,
                    remove=False
                )
                
                # Copy patch file into container
                import tarfile
                import io
                
                # Create tar archive with patch file
                tar_stream = io.BytesIO()
                tar = tarfile.open(fileobj=tar_stream, mode='w')
                
                patch_data = patch_content.encode('utf-8')
                tarinfo = tarfile.TarInfo(name='fix.patch')
                tarinfo.size = len(patch_data)
                tar.addfile(tarinfo, io.BytesIO(patch_data))
                tar.close()
                tar_stream.seek(0)
                
                # Put tar archive into container
                container.put_archive('/home/', tar_stream)
                
                logging.info("Patch file copied to /home/fix.patch")
                
                # Execute the test run script
                # The container should have /home/fix-run.sh that applies patch and runs tests
                logging.info("Running tests in container...")
                exec_result = container.exec_run(
                    "bash /home/fix-run.sh",
                    stdout=True,
                    stderr=True,
                    stream=False
                )
                
                # Get the logs
                log_output = exec_result.output.decode('utf-8', errors='replace')
                exit_code = exec_result.exit_code
                
                logging.info(f"Test execution completed with exit code: {exit_code}")
                
                # Parse the log using the appropriate parser
                test_result = msb_parse_log(org, repo, log_output)
                
                # Calculate score
                if hasattr(test_result, 'passed_count'):
                    passed = test_result.passed_count
                    failed = test_result.failed_count
                    total = passed + failed
                else:
                    passed = test_result.get('passed_count', 0)
                    failed = test_result.get('failed_count', 0)
                    total = passed + failed
                
                # Score based on pass rate
                if total > 0:
                    score = (passed / total) * 100
                else:
                    score = 0.0
                
                success = failed == 0 and passed > 0
                
                result = {
                    'question_id': question.question_id,
                    'question_type': question.question_type,
                    'index': index,
                    'org': org,
                    'repo': repo,
                    'pr_number': pr_number,
                    'docker_image': docker_image,
                    'test_result': {
                        'exit_code': exit_code,
                        'passed_count': passed,
                        'failed_count': failed,
                        'total_count': total,
                        'passed_tests': list(test_result.passed_tests) if hasattr(test_result, 'passed_tests') else list(test_result.get('passed_tests', [])),
                        'failed_tests': list(test_result.failed_tests) if hasattr(test_result, 'failed_tests') else list(test_result.get('failed_tests', []))
                    },
                    'score_result': {
                        'overall_score': score,
                        'tests_passed': passed,
                        'tests_failed': failed,
                        'tests_total': total
                    },
                    'timestamp': datetime.now().isoformat(),
                    'success': success
                }
                
                logging.info(f"Score: {score:.1f}/100 ({passed}/{total} tests passed)")
                return result
                
            finally:
                # Clean up container
                try:
                    container.stop(timeout=5)
                    container.remove()
                    logging.info(f"Stopped and removed container: {container_name}")
                except Exception as e:
                    logging.warning(f"Error cleaning up container: {e}")
            
        except Exception as e:
            logging.error(f"Error evaluating MSB question: {e}", exc_info=True)
            return self._create_error_result(question, index, f"Evaluation error: {e}")
    
    def _create_error_result(self, question: Question, index: int, error_msg: str) -> Dict[str, Any]:
        """Create error result"""
        return {
            'question_id': question.question_id,
            'question_type': question.question_type,
            'index': index,
            'error': error_msg,
            'timestamp': datetime.now().isoformat(),
            'success': False
        }
    
    def save_results(self, results: List[Dict[str, Any]], output_file: str):
        """Save evaluation results to file"""
        with open(output_file, 'w') as f:
            json.dump(results, f, indent=2)
        logging.info(f"Results saved to {output_file}")


def main():
    """Main entry point"""
    parser = argparse.ArgumentParser(
        description='AInsteinBench Unified Evaluator'
    )
    
    parser.add_argument(
        '--questions',
        type=str,
        required=True,
        help='Path to questions JSONL file'
    )
    parser.add_argument(
        '--answers',
        type=str,
        required=True,
        help='Directory containing reference answers'
    )
    parser.add_argument(
        '--config',
        type=str,
        help='Path to config file (optional, defaults to ET config)'
    )
    parser.add_argument(
        '--output',
        type=str,
        default='evaluation_results.json',
        help='Output file for results (default: evaluation_results.json)'
    )
    parser.add_argument(
        '--max-examples',
        type=int,
        help='Maximum number of examples to evaluate'
    )
    
    args = parser.parse_args()
    
    # Create evaluator
    evaluator = UnifiedEvaluator(
        args.questions,
        args.answers,
        args.config
    )
    
    # Run evaluation
    results = evaluator.evaluate(max_examples=args.max_examples)
    
    # Save results
    evaluator.save_results(results, args.output)
    
    # Print summary
    print(f"\n{'='*70}")
    print("Evaluation Summary")
    print(f"{'='*70}")
    print(f"Total questions: {len(results)}")
    print(f"Successful: {sum(1 for r in results if r.get('success', False))}")
    print(f"Failed: {sum(1 for r in results if not r.get('success', False))}")
    
    # Print scores for successful evaluations
    successful = [r for r in results if r.get('success', False)]
    if successful:
        avg_score = sum(r.get('score_result', {}).get('overall_score', 0) for r in successful) / len(successful)
        print(f"Average score: {avg_score:.1f}/100")
    
    print(f"{'='*70}\n")


if __name__ == '__main__':
    main()