llm_optimization_tester.py

import time
import os
import anthropic
import logging
from typing import Dict, List, Optional, TypedDict
from dataclasses import dataclass
import json
from datetime import datetime

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

class TestCase(TypedDict):
    name: str
    description: str
    standard_code: str
    optimized_code: str
    tasks: List[str]
    expected_results: List[Dict]

@dataclass
class TestResult:
    completion_time: float
    accuracy_score: float
    context_questions: int
    token_usage: int
    errors: List[str]
    output: str

class LLMOptimizationTester:
    """
    FILE_PURPOSE: Test framework for comparing LLM performance on standard vs. optimized code
    SYSTEM_CONTEXT: Part of the LLM codebase optimization research toolkit
    TECHNICAL_DEPENDENCIES:
    - Python 3.9+
    - LLM API access
    
    USAGE_EXAMPLE:
    tester = LLMOptimizationTester()
    results = tester.run_comparison_test(test_case)
    tester.generate_report(results)
    """

    def __init__(self):
        """
        CONTEXT: Initializes the testing framework
        CONFIGURATION: Loads from config.json if present
        """
        self.results: Dict[str, List[TestResult]] = {
            "standard": [],
            "optimized": []
        }
        
        # Load configuration
        try:
            with open('config.json', 'r') as f:
                self.config = json.load(f)
                logger.info("Loaded configuration from config.json")
        except FileNotFoundError:
            logger.warning("config.json not found, using default settings")
            self.config = {
                "api": {
                    "model": "claude-3-5-sonnet-latest",
                    "max_tokens": 4000,
                    "temperature": 0.0
                }
            }

    def run_comparison_test(self, test_case: TestCase, standard_result: Optional[TestResult] = None) -> Dict[str, TestResult]:
        """
        CONTEXT: Runs comparative analysis between standard and optimized code
        PROCESS:
        1. Use provided standard results or run standard implementation once
        2. Run optimized implementation and compare against standard results
        3. Measure and compare performance metrics
        
        Args:
            test_case: The test case to run
            standard_result: Optional pre-computed standard implementation results
        """
        results = {}
        
        logger.info(f"Starting comparison test: {test_case['name']}")
        logger.info(f"Number of tasks to evaluate: {len(test_case['tasks'])}")
        
        # Use provided standard result or run standard implementation
        if standard_result is None:
            logger.info("Testing standard implementation...")
            start_time = time.time()
            standard_result = self._run_single_test(
                "standard",
                test_case["standard_code"],
                test_case["tasks"]
            )
            standard_time = time.time() - start_time
            logger.info(f"Standard implementation test completed in {standard_time:.2f}s")
        else:
            logger.info("Using provided standard implementation results")
            
        # Store standard results
        results["standard"] = standard_result
        
        # Test optimized version
        logger.info("Testing LLM-optimized implementation...")
        start_time = time.time()
        results["optimized"] = self._run_single_test(
            "optimized",
            test_case["optimized_code"],
            test_case["tasks"]
        )
        optimized_time = time.time() - start_time
        logger.info(f"Optimized implementation test completed in {optimized_time:.2f}s")
        
        return results

    def _run_single_test(self, version: str, code: str, tasks: List[str]) -> TestResult:
        """
        CONTEXT: Executes a single test run using Anthropic's Sonnet model
        METRICS_TRACKED:
        - Completion time
        - Accuracy
        - Context questions needed
        - Token usage
        """
        start_time = time.time()
        errors = []
        context_questions = 0
        total_tokens = 0
        accuracy_scores = []
        
        # Initialize Anthropic client
        try:
            client = anthropic.Anthropic(
                api_key=os.getenv('ANTHROPIC_API_KEY')
            )
            logger.debug("Anthropic client initialized successfully")
        except Exception as e:
            logger.error(f"Failed to initialize Anthropic client: {e}")
            raise

        try:
            for i, task in enumerate(tasks, 1):
                logger.info(f"[{version}] Processing task {i}/{len(tasks)}: {task[:100]}...")
                
                # Construct the prompt
                prompt = f"""
                You are a software engineer analyzing and modifying code.
                Here is the code to work with:

                {code}

                Your task is: {task}

                Please complete this task. If you need any clarification,
                ask a question starting with 'QUESTION:'. Otherwise,
                provide your solution starting with 'SOLUTION:'.
                """

                # Make API call with retries
                max_retries = 3
                base_delay = 1  # Base delay in seconds
                success = False
                message = None
                retry_count = 0
                
                while not success and retry_count < max_retries:
                    try:
                        if retry_count > 0:
                            delay = base_delay * (2 ** (retry_count - 1))  # Exponential backoff
                            logger.info(f"[{version}] Retrying task {i} after {delay}s (attempt {retry_count + 1}/{max_retries})")
                            time.sleep(delay)
                            
                        message = client.messages.create(
                            model=self.config["api"]["model"],
                            max_tokens=self.config["api"]["max_tokens"],
                            temperature=self.config["api"]["temperature"],
                            system="You are a skilled software engineer focused on code analysis and modification.",
                            messages=[{
                                "role": "user",
                                "content": prompt
                            }]
                        )
                        success = True
                        logger.debug(f"[{version}] API call successful for task {i}")
                    except Exception as e:
                        retry_count += 1
                        if retry_count >= max_retries:
                            logger.error(f"[{version}] API call failed for task {i} after {max_retries} attempts: {e}")
                            errors.append(f"Task {i} API error after {max_retries} retries: {str(e)}")
                            continue
                        logger.warning(f"[{version}] API call failed for task {i} (attempt {retry_count}/{max_retries}): {e}")
                
                if not success:
                    continue

                # Process response
                response = message.content[0].text
                total_tokens += message.usage.input_tokens + message.usage.output_tokens

                # Track context questions
                if response.strip().startswith("QUESTION:"):
                    context_questions += 1
                    logger.info(f"[{version}] Task {i} required clarification: {response[:100]}...")
                    continue

                # Calculate accuracy score
                accuracy = 1.0 if response.startswith("SOLUTION:") else 0.5
                accuracy_scores.append(accuracy)
                logger.info(f"[{version}] Task {i} completed with accuracy score: {accuracy}")

        except Exception as e:
            logger.error(f"[{version}] Test execution failed: {e}")
            errors.append(f"Test execution error: {str(e)}")

        completion_time = time.time() - start_time
        avg_accuracy = sum(accuracy_scores) / len(accuracy_scores) if accuracy_scores else 0.0
        
        logger.info(f"[{version}] Test summary:")
        logger.info(f"- Completion time: {completion_time:.2f}s")
        logger.info(f"- Average accuracy: {avg_accuracy:.2%}")
        logger.info(f"- Context questions needed: {context_questions}")
        logger.info(f"- Total tokens used: {total_tokens}")
        if errors:
            logger.warning(f"- Errors encountered: {len(errors)}")

        return TestResult(
            completion_time=completion_time,
            accuracy_score=avg_accuracy,
            context_questions=context_questions,
            token_usage=total_tokens,
            errors=errors,
            output=response if response else ''
        )

    def analyze_code_outputs(self, prompt: str) -> str:
        """
        CONTEXT: Analyzes code outputs using LLM
        PURPOSE: Compare different implementations and provide insights
        """
        try:
            client = anthropic.Anthropic(
                api_key=os.getenv('ANTHROPIC_API_KEY')
            )
            logger.debug("Anthropic client initialized successfully for analysis")
            
            message = client.messages.create(
                model=self.config["api"]["model"],
                max_tokens=self.config["api"]["max_tokens"],
                temperature=self.config["api"]["temperature"],
                system="You are a skilled software engineer focused on code analysis and comparison.",
                messages=[{
                    "role": "user",
                    "content": prompt
                }]
            )
            
            return message.content[0].text
            
        except Exception as e:
            logger.error(f"Failed to analyze code outputs: {e}")
            return f"Analysis failed: {str(e)}"
    
    def generate_report(self, results: Dict[str, TestResult]) -> Dict:
        logger.info("Generating comparison report...")
        """
        CONTEXT: Generates detailed comparison report
        OUTPUT: Performance metrics and analysis
        """
        return {
            "performance_comparison": {
                "time_difference": results["optimized"].completion_time - results["standard"].completion_time,
                "accuracy_improvement": results["optimized"].accuracy_score - results["standard"].accuracy_score,
                "context_questions_reduced": results["standard"].context_questions - results["optimized"].context_questions,
                "token_usage_difference": results["optimized"].token_usage - results["standard"].token_usage
            },
            "conclusion": self._generate_conclusion(results)
        }

    def _generate_conclusion(self, results: Dict[str, TestResult]) -> str:
        """
        CONTEXT: Analyzes results to generate meaningful conclusions
        ANALYSIS_POINTS:
        - Performance improvements
        - Resource efficiency
        - Error reduction
        """
        # Placeholder for actual analysis logic
        return "Analysis of optimization effectiveness"

def create_sample_test_case() -> TestCase:
    """
    CONTEXT: Creates a sample test case for demonstration
    USAGE: For testing framework validation
    """
    return {
        "name": "User Authentication Module",
        "description": "Testing optimization effects on authentication code understanding",
        "standard_code": """
def authenticate_user(credentials):
    # Validate credentials
    if not credentials.get('username') or not credentials.get('password'):
        return None
    
    # Check database
    user = db.find_user(credentials['username'])
    if not user or not check_password(user, credentials['password']):
        return None
    
    return user
        """,
        "optimized_code": """
# CONTEXT: Handles user authentication against database
# DEPENDENCIES: database.py, password_utils.py
# SIDE_EFFECTS: Logs authentication attempts
# RETURNS: User object or None if authentication fails

def authenticate_user(credentials: Dict[str, str]) -> Optional[User]:
    '''
    FILE_PURPOSE: User authentication and validation
    BUSINESS_RULES:
    - Username and password required
    - Password must match stored hash
    TECHNICAL_DEPENDENCIES:
    - PostgreSQL database
    - Bcrypt password hashing
    '''
    # === INPUT VALIDATION ===
    if not credentials.get('username') or not credentials.get('password'):
        return None
    
    # === DATABASE OPERATIONS ===
    user = db.find_user(credentials['username'])
    if not user or not check_password(user, credentials['password']):
        return None
    
    return user
        """,
        "tasks": [
            "Explain the authentication flow",
            "Add MFA support",
            "Fix password validation bug"
        ],
        "expected_results": []
    }

if __name__ == "__main__":
    """
    CONTEXT: Example usage of the testing framework
    PURPOSE: Demonstrate framework capabilities
    """
    # Create test case
    test_case = create_sample_test_case()
    
    # Initialize tester
    tester = LLMOptimizationTester()
    
    # Run comparison
    results = tester.run_comparison_test(test_case)
    
    # Generate and print report
    report = tester.generate_report(results)
    print(json.dumps(report, indent=2))