example_test.py

"""
FILE_PURPOSE: Example test script demonstrating the LLM optimization testing framework
SYSTEM_CONTEXT: Part of the research validation toolkit
"""

import json
import os
import logging
import shutil
from datetime import datetime
from pathlib import Path
from llm_optimization_tester import LLMOptimizationTester, TestResult, TestCase
from typing import Dict, List

# Set up base directories
BASE_DIR = Path("test_outputs")
LOG_DIR = BASE_DIR / "logs"
RESULTS_DIR = BASE_DIR / "results"

# Create directories
LOG_DIR.mkdir(parents=True, exist_ok=True)
RESULTS_DIR.mkdir(parents=True, exist_ok=True)

# Configure file logging
log_file = LOG_DIR / f"test_run_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
file_handler = logging.FileHandler(log_file)
file_handler.setFormatter(logging.Formatter('%(asctime)s [%(levelname)s] %(message)s'))
logging.getLogger().addHandler(file_handler)

logger = logging.getLogger(__name__)

def discover_implementations() -> List[str]:
    """Dynamically discover all implementation directories."""
    impl_dir = Path("implementations")
    return [d.name for d in impl_dir.iterdir() 
            if d.is_dir() and not d.name.startswith('__')]

def read_implementation(impl_dir: str) -> str:
    """Read the implementation from a specific directory."""
    impl_path = Path("implementations") / impl_dir / "user_manager.py"
    with open(impl_path, 'r') as f:
        return f.read()

def create_complex_test_case() -> TestCase:
    """
    CONTEXT: Creates a complex test case using the user management system
    DEMONSTRATES: Impact of documentation and structure on LLM comprehension
    """
    impl_dirs = discover_implementations()
    implementations = {impl_dir: read_implementation(impl_dir) for impl_dir in impl_dirs}

    return {
        "name": "User Management System",
        "description": "Complex authentication and user management system with caching and email notifications",
        "implementations": implementations,
        "tasks": [
            # Understanding Tasks
            "Explain the complete user authentication flow including all side effects and security measures",
            
            # Feature Addition Tasks
            "Add support for OAuth2 authentication while maintaining the existing email/password flow",

            # Bug Fixing Tasks
            "Fix a race condition in the login attempt counter that could allow bypassing the account lock",
            "Resolve a security vulnerability where password reset tokens don't expire",
            
            # Refactoring Tasks
            "Refactor the password hashing system to use a more secure algorithm with salt and pepper",
            
            # Integration Tasks
            "Implement an audit logging system that tracks all security-relevant events"
        ],
        "expected_results": []
    }

def save_task_results(all_results: Dict, task_name: str) -> Path:
    """Save consolidated results for a task across all implementations."""
    task_dir = RESULTS_DIR / task_name.replace(" ", "_").lower()
    task_dir.mkdir(parents=True, exist_ok=True)
    
    output_file = task_dir / "consolidated_results.json"
    with open(output_file, 'w') as f:
        json.dump(all_results, f, indent=2)
    return task_dir

def analyze_task_results(task_dir: Path, task_name: str, results: Dict, tester: LLMOptimizationTester) -> str:
    """Use LLM to analyze and compare code outputs from different implementations for a single task."""
    # Results are already loaded from consolidated_results.json
    
    # Prepare comparison prompt
    prompt = f"""Task: {task_name}

Please analyze the code outputs from different implementations and explain:
1. The key differences in approach between implementations
2. Any improvements or regressions in code quality
3. Specific optimizations that made certain implementations more effective
4. Security considerations and potential vulnerabilities
5. Performance Analysis:
   - Completion times
   - Token usage efficiency
   - Error rates and robustness
   - Overall performance comparison
6. Overall recommendation for which implementation is best, considering both code quality and performance metrics

Implementations and their outputs:
"""

    # Sort implementations to ensure consistent order with poor implementation first
    sorted_impls = sorted(results.items(), key=lambda x: (
        # Poor implementation first, then alphabetically
        0 if x[0] == 'poor' else 1,
        x[0]
    ))
    
    for impl_name, result in sorted_impls:
        prompt += f"\n{impl_name.upper()} IMPLEMENTATION:\n"
        prompt += f"Code output:\n{result['output']}\n"
        prompt += f"Performance Metrics:\n"
        prompt += f"- Accuracy: {result['accuracy']}\n"
        prompt += f"- Completion Time: {result['completion_time']}s\n"
        prompt += f"- Tokens Used: {result['tokens_used']}\n"
        prompt += f"- Context Questions: {result['context_questions']}\n"
        prompt += f"- Errors: {len(result['errors'])}\n"

    # Get LLM analysis
    analysis = tester.analyze_code_outputs(prompt)
    
    # Save analysis
    analysis_file = task_dir / "implementation_analysis.txt"
    with open(analysis_file, 'w') as f:
        f.write(analysis)
    
    return analysis

def generate_final_analysis(tester: LLMOptimizationTester):
    """Generate a final analysis comparing implementations across all tasks."""
    # Collect all task analyses
    analyses = []
    for task_dir in RESULTS_DIR.glob("*"):
        if not task_dir.is_dir():
            continue
            
        analysis_file = task_dir / "implementation_analysis.txt"
        if analysis_file.exists():
            with open(analysis_file) as f:
                analyses.append({
                    "task": task_dir.name.replace("_", " "),
                    "analysis": f.read()
                })
    
    # Prepare final analysis prompt
    prompt = """Please provide a comprehensive analysis of all implementations across all tasks. Consider:

1. Overall Performance Patterns:
   - Which implementation consistently performed best?
   - Were there specific types of tasks where certain implementations excelled?
   - What were the general trends in completion times and token usage?

2. Code Quality Assessment:
   - How did code quality compare across implementations?
   - What were the common strengths and weaknesses?
   - Which implementation had the best balance of readability and efficiency?
   - What specific examples illustrate these differences?

3. Reliability and Error Handling:
   - Compare error rates and robustness
   - Identify any patterns in where errors occurred
   - Assess overall reliability of each implementation
   - What specific examples illustrate these differences?

4. Final Recommendations:
   - Best overall implementation
   - Best for specific use cases
   - Areas for improvement in each implementation
   - What specific examples illustrate these differences?
Task Analyses:
"""
    
    for analysis in analyses:
        prompt += f"\n=== {analysis['task']} ===\n"
        prompt += analysis['analysis'] + "\n"
    
    # Get final analysis
    final_analysis = tester.analyze_code_outputs(prompt)
    
    # Save final analysis
    final_analysis_file = RESULTS_DIR / "final_analysis.txt"
    with open(final_analysis_file, 'w') as f:
        f.write(final_analysis)
    
    return final_analysis

def main():
    """
    CONTEXT: Main test execution
    PURPOSE: Demonstrate framework with complex real-world example
    """
    logger.info("Starting test execution")

    # Ensure API key is set
    if not os.getenv('ANTHROPIC_API_KEY'):
        logger.error("ANTHROPIC_API_KEY environment variable not set")
        print("Error: ANTHROPIC_API_KEY environment variable not set")
        return
    
    logger.info("API key verification successful")

    try:
        # Clean previous results
        if RESULTS_DIR.exists():
            shutil.rmtree(RESULTS_DIR)
        RESULTS_DIR.mkdir(exist_ok=True)

        # Create test case
        logger.info("Creating complex test case")
        test_case = create_complex_test_case()
        
        # Initialize tester
        logger.info("Initializing LLM Optimization Tester")
        tester = LLMOptimizationTester()
        
        # Discover implementations
        impl_dirs = discover_implementations()
        logger.info(f"Discovered implementations: {impl_dirs}")
        
        # First run standard implementation once
        logger.info("Running standard implementation baseline test")
        standard_test_case = test_case.copy()
        standard_test_case['standard_code'] = test_case['implementations']['standard']
        standard_test_case['optimized_code'] = test_case['implementations']['standard']
        standard_results = tester.run_comparison_test(standard_test_case)
        standard_baseline = standard_results['standard']
        
        # Run tests for each implementation
        task_results = {}
        
        # First add standard implementation results to task_results
        for task in test_case['tasks']:
            if task not in task_results:
                task_results[task] = {}
            task_results[task]['standard'] = {
                'accuracy': standard_baseline.accuracy_score,
                'completion_time': standard_baseline.completion_time,
                'tokens_used': standard_baseline.token_usage,
                'context_questions': standard_baseline.context_questions,
                'errors': standard_baseline.errors,
                'output': standard_baseline.output
            }
        
        # Run tests for other implementations
        for impl_name in impl_dirs:
            if impl_name == 'standard':
                continue  # Skip standard as we already ran it
            
            print(f"\nRunning tests for {impl_name} implementation...")
            logger.info(f"Starting tests for {impl_name} implementation")
            
            # Create a test case specific to this implementation
            impl_test_case = test_case.copy()
            impl_test_case['standard_code'] = test_case['implementations']['standard']
            impl_test_case['optimized_code'] = test_case['implementations'][impl_name]
            
            # Run comparison test with pre-computed standard results
            results = tester.run_comparison_test(impl_test_case, standard_result=standard_baseline)
            
            # Save results organized by task
            for task_idx, task in enumerate(test_case['tasks']):
                if task not in task_results:
                    task_results[task] = {}
                
                standard_result = results['standard']
                optimized_result = results['optimized']
                
                task_results[task][impl_name] = {
                    'accuracy': optimized_result.accuracy_score,
                    'completion_time': optimized_result.completion_time,
                    'tokens_used': optimized_result.token_usage,
                    'context_questions': optimized_result.context_questions,
                    'errors': optimized_result.errors,
                    'output': optimized_result.output
                }
        
        # Save and analyze results for each task
        for task, results in task_results.items():
            task_dir = save_task_results(results, task)
            analysis = analyze_task_results(task_dir, task, results, tester)
            print(f"\n=== Implementation Analysis for Task: {task} ===\n")
            print(analysis)
            
        # Generate final analysis across all tasks
        logger.info("Generating final analysis across all tasks")
        final_analysis = generate_final_analysis(tester)
        print("\n=== Final Analysis Across All Tasks ===\n")
        print(final_analysis)
        
        # Generate overall report
        logger.info("Generating test report")
        
        # Get the first task's results for comparison
        first_task = next(iter(task_results.keys()))
        standard_impl = task_results[first_task]['standard']
        llm_optimized_impl = task_results[first_task]['llm_optimized']
        
        # Create TestResult objects for the report
        standard_result = TestResult(
            completion_time=standard_impl['completion_time'],
            accuracy_score=standard_impl['accuracy'],
            context_questions=standard_impl['context_questions'],
            token_usage=standard_impl['tokens_used'],
            errors=standard_impl['errors'],
            output=standard_impl['output']
        )
        
        optimized_result = TestResult(
            completion_time=llm_optimized_impl['completion_time'],
            accuracy_score=llm_optimized_impl['accuracy'],
            context_questions=llm_optimized_impl['context_questions'],
            token_usage=llm_optimized_impl['tokens_used'],
            errors=llm_optimized_impl['errors'],
            output=llm_optimized_impl['output']
        )
        
        report = tester.generate_report({
            'standard': standard_result,
            'optimized': optimized_result
        })
        
        # Output results
        print("\n=== LLM Optimization Test Results ===")
        print("\nTest Case:", test_case["name"])
        print("Description:", test_case["description"])
        logger.info(f"Test completed for: {test_case['name']}")
    
        print("\nTasks Attempted:", len(test_case["tasks"]))
        logger.info(f"Total tasks attempted: {len(test_case['tasks'])}")
        
        print("Standard Implementation Metrics:")
        logger.info("Standard Implementation Results:")
        print(f"- Completion Time: {standard_result.completion_time:.2f}s")
        print(f"- Accuracy Score: {standard_result.accuracy_score:.2%}")
        print(f"- Context Questions: {standard_result.context_questions}")
        print(f"- Token Usage: {standard_result.token_usage}")
    
        print("\nOptimized Implementation Metrics:")
        logger.info("Optimized Implementation Results:")
        print(f"- Completion Time: {optimized_result.completion_time:.2f}s")
        print(f"- Accuracy Score: {optimized_result.accuracy_score:.2%}")
        print(f"- Context Questions: {optimized_result.context_questions}")
        print(f"- Token Usage: {optimized_result.token_usage}")
    
        print("\nPerformance Comparison:")
        logger.info("Performance Comparison Results:")
        comparison_json = json.dumps(report["performance_comparison"], indent=2)
        print(comparison_json)
        logger.info(f"Performance comparison: {comparison_json}")
        
        if results['standard']['errors'] or results['llm_optimized']['errors']:
            print("\nErrors Encountered:")
            logger.warning("Errors encountered during testing:")
            if results['standard']['errors']:
                print("Standard Implementation:", results['standard']['errors'])
                logger.warning(f"Standard Implementation errors: {results['standard']['errors']}")
            if results['llm_optimized']['errors']:
                print("Optimized Implementation:", results['llm_optimized']['errors'])
                logger.warning(f"Optimized Implementation errors: {results['llm_optimized']['errors']}")
        
        print("\nConclusion:", report["conclusion"])
        logger.info(f"Test conclusion: {report['conclusion']}")
        
        logger.info(f"Full test results saved to: {log_file}")
        print(f"\nDetailed logs saved to: {log_file}")
        
    except Exception as e:
        logger.error(f"Test execution failed: {e}", exc_info=True)
        print(f"\nError: Test execution failed. See {log_file} for details.")

if __name__ == "__main__":
    main()