@hanzo/dev/src/lib/benchmark-runner.ts

Hanzo Dev - Meta AI development CLI that manages and runs all LLMs and CLI tools

import * as fs from 'fs';
import * as path from 'path';
import { execSync } from 'child_process';
import chalk from 'chalk';
import ora from 'ora';
import { CodeActAgent } from './code-act-agent';
import { PeerAgentNetwork } from './peer-agent-network';
import { FunctionCallingSystem } from './function-calling';
import { ConfigurableAgentLoop, LLMProvider } from './agent-loop';

export interface BenchmarkTask {
  instance_id: string;
  repo: string;
  base_commit: string;
  problem_statement: string;
  hints_text?: string;
  test_patch?: string;
  expected_files?: string[];
  difficulty?: 'easy' | 'medium' | 'hard';
}

export interface BenchmarkResult {
  instance_id: string;
  success: boolean;
  time_taken_ms: number;
  files_modified: number;
  test_passed: boolean;
  error?: string;
  agent_type: string;
  llm_calls: number;
  cost_estimate: number;
}

export interface BenchmarkConfig {
  dataset: 'swe-bench' | 'swe-bench-lite' | 'custom';
  agents: number;
  parallel: boolean;
  timeout: number;
  output: string;
  provider?: LLMProvider;
  maxTasks?: number;
}

export class BenchmarkRunner {
  private config: BenchmarkConfig;
  private results: BenchmarkResult[] = [];
  private network?: PeerAgentNetwork;

  constructor(config: BenchmarkConfig) {
    this.config = {
      dataset: 'swe-bench-lite',
      agents: 5,
      parallel: true,
      timeout: 300000, // 5 minutes default
      output: 'benchmark-results.json',
      ...config
    };
  }

  async run(): Promise<void> {
    console.log(chalk.bold.cyan('\n🏃 Hanzo Dev Benchmark Runner\n'));
    console.log(chalk.gray(`Dataset: ${this.config.dataset}`));
    console.log(chalk.gray(`Agents: ${this.config.agents}`));
    console.log(chalk.gray(`Parallel: ${this.config.parallel}`));
    console.log(chalk.gray(`Timeout: ${this.config.timeout}ms\n`));

    const spinner = ora('Loading benchmark tasks...').start();

    try {
      // Load tasks
      const tasks = await this.loadTasks();
      spinner.succeed(`Loaded ${tasks.length} tasks`);

      // Initialize network if using parallel mode
      if (this.config.parallel && this.config.agents > 1) {
        spinner.start('Initializing agent network...');
        this.network = new PeerAgentNetwork();
        spinner.succeed('Agent network initialized');
      }

      // Run benchmark
      const startTime = Date.now();
      await this.runTasks(tasks);
      const totalTime = Date.now() - startTime;

      // Calculate and display results
      this.displayResults(totalTime);

      // Save results
      await this.saveResults();

    } catch (error) {
      spinner.fail(`Benchmark failed: ${error}`);
      throw error;
    }
  }

  private async loadTasks(): Promise<BenchmarkTask[]> {
    // Load from different sources based on dataset
    switch (this.config.dataset) {
      case 'swe-bench':
        return this.loadSWEBenchTasks(false);
      case 'swe-bench-lite':
        return this.loadSWEBenchTasks(true);
      case 'custom':
        return this.loadCustomTasks();
      default:
        throw new Error(`Unknown dataset: ${this.config.dataset}`);
    }
  }

  private async loadSWEBenchTasks(lite: boolean): Promise<BenchmarkTask[]> {
    // In production, this would load from the actual SWE-bench dataset
    // For now, return sample tasks for testing
    const sampleTasks: BenchmarkTask[] = [
      {
        instance_id: 'django__django-11999',
        repo: 'django/django',
        base_commit: 'abc123',
        problem_statement: 'Fix QuerySet.delete() to handle circular foreign key dependencies',
        hints_text: 'Look at django/db/models/deletion.py',
        difficulty: 'hard',
        expected_files: ['django/db/models/deletion.py']
      },
      {
        instance_id: 'pytest-dev__pytest-5692',
        repo: 'pytest-dev/pytest',
        base_commit: 'def456',
        problem_statement: 'Fix --collect-only to show parametrized test ids',
        hints_text: 'Check _pytest/main.py and _pytest/python.py',
        difficulty: 'medium',
        expected_files: ['src/_pytest/main.py', 'src/_pytest/python.py']
      },
      {
        instance_id: 'scikit-learn__scikit-learn-13142',
        repo: 'scikit-learn/scikit-learn',
        base_commit: 'ghi789',
        problem_statement: 'Add sample_weight support to Ridge regression',
        hints_text: 'Modify sklearn/linear_model/ridge.py',
        difficulty: 'medium',
        expected_files: ['sklearn/linear_model/ridge.py']
      },
      {
        instance_id: 'requests__requests-3362',
        repo: 'psf/requests',
        base_commit: 'jkl012',
        problem_statement: 'Fix encoding detection for streaming responses',
        hints_text: 'Look at requests/models.py Response class',
        difficulty: 'easy',
        expected_files: ['requests/models.py']
      },
      {
        instance_id: 'flask__flask-2354',
        repo: 'pallets/flask',
        base_commit: 'mno345',
        problem_statement: 'Add support for async view functions',
        hints_text: 'Modify flask/app.py and flask/views.py',
        difficulty: 'hard',
        expected_files: ['flask/app.py', 'flask/views.py']
      }
    ];

    // Apply task limit if specified
    const tasks = lite ? sampleTasks.slice(0, 3) : sampleTasks;
    return this.config.maxTasks ? tasks.slice(0, this.config.maxTasks) : tasks;
  }

  private async loadCustomTasks(): Promise<BenchmarkTask[]> {
    // Load from custom JSON file
    const customPath = path.join(process.cwd(), 'benchmark-tasks.json');
    if (!fs.existsSync(customPath)) {
      throw new Error(`Custom tasks file not found: ${customPath}`);
    }
    return JSON.parse(fs.readFileSync(customPath, 'utf-8'));
  }

  private async runTasks(tasks: BenchmarkTask[]): Promise<void> {
    const spinner = ora('Running benchmark tasks...').start();

    if (this.config.parallel && this.network) {
      // Run tasks in parallel using agent network
      await this.runParallelTasks(tasks, spinner);
    } else {
      // Run tasks sequentially
      await this.runSequentialTasks(tasks, spinner);
    }

    spinner.succeed(`Completed ${tasks.length} tasks`);
  }

  private async runSequentialTasks(tasks: BenchmarkTask[], spinner: ora.Ora): Promise<void> {
    for (let i = 0; i < tasks.length; i++) {
      const task = tasks[i];
      spinner.text = `Running task ${i + 1}/${tasks.length}: ${task.instance_id}`;
      
      const result = await this.runSingleTask(task);
      this.results.push(result);
      
      if (result.success) {
        spinner.succeed(`✓ ${task.instance_id} (${result.time_taken_ms}ms)`);
      } else {
        spinner.fail(`✗ ${task.instance_id}: ${result.error}`);
      }
      spinner.start();
    }
  }

  private async runParallelTasks(tasks: BenchmarkTask[], spinner: ora.Ora): Promise<void> {
    spinner.text = `Spawning ${this.config.agents} agents for parallel execution...`;

    // Create agent pool
    const agentPromises = [];
    for (let i = 0; i < Math.min(this.config.agents, tasks.length); i++) {
      agentPromises.push(this.createBenchmarkAgent(`benchmark-agent-${i}`));
    }
    
    await Promise.all(agentPromises);

    // Distribute tasks among agents
    const taskQueue = [...tasks];
    const resultPromises: Promise<BenchmarkResult>[] = [];

    while (taskQueue.length > 0) {
      const batch = taskQueue.splice(0, this.config.agents);
      const batchPromises = batch.map((task, index) => 
        this.runTaskWithAgent(task, `benchmark-agent-${index}`)
      );
      resultPromises.push(...batchPromises);
    }

    // Wait for all tasks to complete
    spinner.text = `Running ${tasks.length} tasks in parallel...`;
    const results = await Promise.all(resultPromises);
    this.results.push(...results);
  }

  private async createBenchmarkAgent(agentId: string): Promise<void> {
    if (!this.network) return;

    await this.network.spawnAgent({
      id: agentId,
      name: `Benchmark Agent ${agentId}`,
      type: 'claude-code',
      tools: ['edit_file', 'view_file', 'run_command', 'search_files']
    });
  }

  private async runTaskWithAgent(task: BenchmarkTask, agentId: string): Promise<BenchmarkResult> {
    const startTime = Date.now();
    let llmCalls = 0;

    try {
      // Create agent loop with timeout
      const loop = new ConfigurableAgentLoop({
        provider: this.config.provider || this.getDefaultProvider(),
        maxIterations: 50,
        enableMCP: true,
        enableBrowser: false,
        enableSwarm: false,
        streamOutput: false,
        confirmActions: false
      });

      // Track LLM calls
      loop.on('llm-call', () => llmCalls++);

      // Initialize and execute
      await loop.initialize();
      
      const timeoutPromise = new Promise((_, reject) => 
        setTimeout(() => reject(new Error('Task timeout')), this.config.timeout)
      );

      await Promise.race([
        loop.execute(this.formatTaskPrompt(task)),
        timeoutPromise
      ]);

      // Verify solution
      const testPassed = await this.runTests(task);

      return {
        instance_id: task.instance_id,
        success: true,
        time_taken_ms: Date.now() - startTime,
        files_modified: task.expected_files?.length || 0,
        test_passed: testPassed,
        agent_type: agentId,
        llm_calls,
        cost_estimate: this.estimateCost(llmCalls)
      };

    } catch (error: any) {
      return {
        instance_id: task.instance_id,
        success: false,
        time_taken_ms: Date.now() - startTime,
        files_modified: 0,
        test_passed: false,
        error: error.message,
        agent_type: agentId,
        llm_calls,
        cost_estimate: this.estimateCost(llmCalls)
      };
    }
  }

  private async runSingleTask(task: BenchmarkTask): Promise<BenchmarkResult> {
    return this.runTaskWithAgent(task, 'single-agent');
  }

  private formatTaskPrompt(task: BenchmarkTask): string {
    let prompt = `Repository: ${task.repo}\n`;
    prompt += `Problem: ${task.problem_statement}\n`;
    
    if (task.hints_text) {
      prompt += `\nHints: ${task.hints_text}\n`;
    }
    
    if (task.expected_files?.length) {
      prompt += `\nFiles that likely need modification: ${task.expected_files.join(', ')}\n`;
    }
    
    prompt += '\nPlease fix this issue by making the necessary code changes.';
    
    return prompt;
  }

  private async runTests(task: BenchmarkTask): Promise<boolean> {
    // In production, this would apply the test patch and run actual tests
    // For now, simulate test results
    return Math.random() > 0.3; // 70% test pass rate
  }

  private getDefaultProvider(): LLMProvider {
    // Check for available API keys
    if (process.env.ANTHROPIC_API_KEY) {
      return {
        name: 'Claude',
        type: 'anthropic',
        apiKey: process.env.ANTHROPIC_API_KEY,
        model: 'claude-3-opus-20240229',
        supportsTools: true,
        supportsStreaming: true
      };
    } else if (process.env.OPENAI_API_KEY) {
      return {
        name: 'GPT-4',
        type: 'openai',
        apiKey: process.env.OPENAI_API_KEY,
        model: 'gpt-4-turbo-preview',
        supportsTools: true,
        supportsStreaming: true
      };
    } else {
      throw new Error('No LLM API key found. Please set ANTHROPIC_API_KEY or OPENAI_API_KEY');
    }
  }

  private estimateCost(llmCalls: number): number {
    // Rough cost estimation based on average tokens per call
    const avgTokensPerCall = 2000;
    const costPer1kTokens = 0.01; // Adjust based on model
    return (llmCalls * avgTokensPerCall * costPer1kTokens) / 1000;
  }

  private displayResults(totalTime: number): void {
    const successful = this.results.filter(r => r.success).length;
    const testsPassed = this.results.filter(r => r.test_passed).length;
    const avgTime = this.results.reduce((sum, r) => sum + r.time_taken_ms, 0) / this.results.length;
    const totalCost = this.results.reduce((sum, r) => sum + r.cost_estimate, 0);
    const avgLLMCalls = this.results.reduce((sum, r) => sum + r.llm_calls, 0) / this.results.length;

    console.log(chalk.bold.cyan('\n📊 Benchmark Results\n'));
    console.log(chalk.white('Total Tasks:'), this.results.length);
    console.log(chalk.green('Successful:'), `${successful} (${(successful / this.results.length * 100).toFixed(1)}%)`);
    console.log(chalk.blue('Tests Passed:'), `${testsPassed} (${(testsPassed / this.results.length * 100).toFixed(1)}%)`);
    console.log(chalk.yellow('Avg Time:'), `${(avgTime / 1000).toFixed(1)}s`);
    console.log(chalk.yellow('Total Time:'), `${(totalTime / 1000).toFixed(1)}s`);
    console.log(chalk.magenta('Avg LLM Calls:'), avgLLMCalls.toFixed(1));
    console.log(chalk.cyan('Est. Total Cost:'), `$${totalCost.toFixed(2)}`);
    console.log(chalk.cyan('Cost per Task:'), `$${(totalCost / this.results.length).toFixed(3)}`);

    if (this.config.parallel) {
      const speedup = (avgTime * this.results.length) / totalTime;
      console.log(chalk.green('Parallel Speedup:'), `${speedup.toFixed(2)}x`);
    }

    // Show difficulty breakdown
    const byDifficulty = this.groupByDifficulty();
    console.log(chalk.bold.gray('\nBy Difficulty:'));
    Object.entries(byDifficulty).forEach(([difficulty, stats]) => {
      console.log(`  ${difficulty}: ${stats.success}/${stats.total} (${(stats.success / stats.total * 100).toFixed(1)}%)`);
    });
  }

  private groupByDifficulty(): Record<string, { total: number; success: number }> {
    const groups: Record<string, { total: number; success: number }> = {
      easy: { total: 0, success: 0 },
      medium: { total: 0, success: 0 },
      hard: { total: 0, success: 0 }
    };

    // Note: We'd need to store difficulty in results for this to work properly
    // For now, just return mock data
    return groups;
  }

  private async saveResults(): Promise<void> {
    const output = {
      metadata: {
        dataset: this.config.dataset,
        agents: this.config.agents,
        parallel: this.config.parallel,
        timestamp: new Date().toISOString(),
        provider: this.config.provider?.name || 'auto'
      },
      summary: {
        total_tasks: this.results.length,
        successful: this.results.filter(r => r.success).length,
        tests_passed: this.results.filter(r => r.test_passed).length,
        avg_time_ms: this.results.reduce((sum, r) => sum + r.time_taken_ms, 0) / this.results.length,
        total_cost: this.results.reduce((sum, r) => sum + r.cost_estimate, 0)
      },
      results: this.results
    };

    fs.writeFileSync(this.config.output, JSON.stringify(output, null, 2));
    console.log(chalk.gray(`\nResults saved to ${this.config.output}`));
  }
}