feat: refactor register eval (#1713)

crates/goose-cli/src/commands/bench.rs

@@ -5,9 +5,8 @@ use async_trait::async_trait;
 use goose::config::Config;
 use goose::message::Message;
 use goose_bench::bench_work_dir::BenchmarkWorkDir;
-use goose_bench::eval_suites::{BenchAgent, BenchAgentError, Evaluation, EvaluationSuiteFactory};
+use goose_bench::eval_suites::{BenchAgent, BenchAgentError, Evaluation, EvaluationSuite};
 use goose_bench::reporting::{BenchmarkResults, EvaluationResult, SuiteResult};
-use std::collections::HashMap;
 use std::path::PathBuf;
 use std::sync::Arc;
 use tokio::sync::Mutex;
@@ -116,31 +115,10 @@ async fn run_eval(
     Ok(result)
 }
 
-async fn run_suite(suite: &str, work_dir: &mut BenchmarkWorkDir) -> anyhow::Result<SuiteResult> {
-    let mut suite_result = SuiteResult::new(suite.to_string());
-    let eval_work_dir_guard = Mutex::new(work_dir);
-
-    if let Some(evals) = EvaluationSuiteFactory::create(suite) {
-        for eval in evals {
-            let mut eval_work_dir = eval_work_dir_guard.lock().await;
-            eval_work_dir.set_eval(eval.name());
-            let eval_result = run_eval(eval, &mut eval_work_dir).await?;
-            suite_result.add_evaluation(eval_result);
-        }
-    }
-
-    Ok(suite_result)
-}
-
 pub async fn run_benchmark(
-    suites: Vec<String>,
+    selectors: Vec<String>,
     include_dirs: Vec<PathBuf>,
 ) -> anyhow::Result<BenchmarkResults> {
-    let suites = EvaluationSuiteFactory::available_evaluations()
-        .into_iter()
-        .filter(|&s| suites.contains(&s.to_string()))
-        .collect::<Vec<_>>();
-
     let config = Config::global();
     let goose_model: String = config
         .get_param("GOOSE_MODEL")
@@ -151,30 +129,45 @@ pub async fn run_benchmark(
 
     let mut results = BenchmarkResults::new(provider_name.clone());
 
-    let suite_work_dir = Mutex::new(BenchmarkWorkDir::new(
+    let work_dir = Mutex::new(BenchmarkWorkDir::new(
         format!("{}-{}", provider_name, goose_model),
         include_dirs.clone(),
     ));
 
-    for suite in suites {
-        let mut work_dir = suite_work_dir.lock().await;
-        work_dir.set_suite(suite);
-        let suite_result = run_suite(suite, &mut work_dir).await?;
+    for (suite, evals) in EvaluationSuite::select(selectors).iter() {
+        let mut suite_result = SuiteResult::new(suite.clone());
+        for eval_selector in evals {
+            if let Some(eval) = EvaluationSuite::from(eval_selector) {
+                let mut work_dir = work_dir.lock().await;
+                work_dir.set_eval(eval_selector);
+                let eval_result = run_eval(eval, &mut work_dir).await?;
+                suite_result.add_evaluation(eval_result);
+            }
+        }
+
         results.add_suite(suite_result);
     }
 
     Ok(results)
 }
 
-pub async fn list_suites() -> anyhow::Result<HashMap<String, usize>> {
-    let suites = EvaluationSuiteFactory::available_evaluations();
-    let mut suite_counts = HashMap::new();
-
-    for suite in suites {
-        if let Some(evals) = EvaluationSuiteFactory::create(suite) {
-            suite_counts.insert(suite.to_string(), evals.len());
-        }
+pub async fn list_selectors() -> anyhow::Result<()> {
+    let selector_eval_counts = EvaluationSuite::available_selectors();
+    let mut keys: Vec<_> = selector_eval_counts.keys().collect();
+    keys.sort();
+    let max_key_len = keys.iter().map(|k| k.len()).max().unwrap_or(0);
+    println!(
+        "selector {} => Eval Count",
+        " ".repeat(max_key_len - "selector".len())
+    );
+    println!("{}", "-".repeat(max_key_len + 6));
+    for selector in keys {
+        println!(
+            "{} {} => {}",
+            selector,
+            " ".repeat(max_key_len - selector.len()),
+            selector_eval_counts.get(selector).unwrap()
+        );
     }
-
-    Ok(suite_counts)
+    Ok(())
 }

crates/goose-cli/src/main.rs

@@ -4,7 +4,7 @@ use clap::{Args, Parser, Subcommand};
 use goose::config::Config;
 
 use goose_cli::commands::agent_version::AgentCommand;
-use goose_cli::commands::bench::{list_suites, run_benchmark};
+use goose_cli::commands::bench::{list_selectors, run_benchmark};
 use goose_cli::commands::configure::handle_configure;
 use goose_cli::commands::info::handle_info;
 use goose_cli::commands::mcp::run_server;
@@ -237,13 +237,13 @@ enum Command {
     Bench {
         #[arg(
             short = 's',
-            long = "suites",
-            value_name = "BENCH_SUITE_NAME",
+            long = "selectors",
+            value_name = "EVALUATIONS_SELECTOR",
             help = "Run this list of bench-suites.",
             long_help = "Specify a comma-separated list of evaluation-suite names to be run.",
             value_delimiter = ','
         )]
-        suites: Vec<String>,
+        selectors: Vec<String>,
 
         #[arg(
             short = 'i',
@@ -266,7 +266,7 @@ enum Command {
         #[arg(
             long = "list",
             value_name = "LIST",
-            help = "List all available bench suites."
+            help = "List all selectors and the number of evaluations they select."
         )]
         list: bool,
 
@@ -416,7 +416,7 @@ async fn main() -> Result<()> {
             return Ok(());
         }
         Some(Command::Bench {
-            suites,
+            selectors,
             include_dirs,
             repeat,
             list,
@@ -425,24 +425,22 @@ async fn main() -> Result<()> {
             summary,
         }) => {
             if list {
-                let suites = list_suites().await?;
-                for suite in suites.keys() {
-                    println!("{}: {}", suite, suites.get(suite).unwrap());
-                }
-                return Ok(());
+                return list_selectors().await;
             }
-            let suites = if suites.is_empty() {
+
+            let selectors = if selectors.is_empty() {
                 vec!["core".to_string()]
             } else {
-                suites
+                selectors
             };
+
             let current_dir = std::env::current_dir()?;
 
             for i in 0..repeat {
                 if repeat > 1 {
                     println!("\nRun {} of {}:", i + 1, repeat);
                 }
-                let results = run_benchmark(suites.clone(), include_dirs.clone()).await?;
+                let results = run_benchmark(selectors.clone(), include_dirs.clone()).await?;
 
                 // Handle output based on format
                 let output_str = match format.as_str() {