tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 8 Packages Projects Releases Wiki Activity

feat: efficient benching (#1921)

Browse Source 
Co-authored-by: Tyler Rockwood <rockwotj@gmail.com>
Co-authored-by: Kalvin C <kalvinnchau@users.noreply.github.com>
Co-authored-by: Alice Hau <110418948+ahau-square@users.noreply.github.com>
This commit is contained in:
marcelle
2025-04-08 14:43:43 -04:00
committed by GitHub
parent 319f2301f3
commit 8fbd9eb327
37 changed files with 1162 additions and 444 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/cli.rs
+64 -110
View File
@@ -4,7 +4,7 @@ use clap::{Args, Parser, Subcommand};
use goose::config::Config;
use crate::commands::agent_version::AgentCommand;
use crate::commands::bench::{list_selectors, run_benchmark};
use crate::commands::bench::agent_generator;
use crate::commands::configure::handle_configure;
use crate::commands::info::handle_info;
use crate::commands::mcp::run_server;
@@ -12,6 +12,10 @@ use crate::commands::session::handle_session_list;
use crate::logging::setup_logging;
use crate::session;
use crate::session::build_session;
use goose_bench::bench_config::BenchRunConfig;
use goose_bench::runners::bench_runner::BenchRunner;
use goose_bench::runners::eval_runner::EvalRunner;
use goose_bench::runners::model_runner::ModelRunner;
use std::io::Read;
use std::path::PathBuf;
@@ -71,6 +75,47 @@ enum SessionCommand {
},
}
#[derive(Subcommand)]
pub enum BenchCommand {
#[command(name = "init-config", about = "Create a new starter-config")]
InitConfig {
#[arg(short, long, help = "filename with extension for generated config")]
name: String,
},
#[command(about = "Run all benchmarks from a config")]
Run {
#[arg(
short,
long,
help = "A config file generated by the config-init command"
)]
config: PathBuf,
},
#[command(about = "List all available selectors")]
Selectors {
#[arg(
short,
long,
help = "A config file generated by the config-init command"
)]
config: Option<PathBuf>,
},
#[command(name = "eval-model", about = "Run an eval of model")]
EvalModel {
#[arg(short, long, help = "A serialized config file for the model only.")]
config: String,
},
#[command(name = "exec-eval", about = "run a single eval")]
ExecEval {
#[arg(short, long, help = "A serialized config file for the eval only.")]
config: String,
},
}
#[derive(Subcommand)]
enum Command {
/// Configure Goose settings
@@ -255,63 +300,8 @@ enum Command {
},
Bench {
#[arg(
short = 's',
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 = ','
)]
selectors: Vec<String>,
#[arg(
short = 'i',
long = "include-dir",
value_name = "DIR_NAME",
action = clap::ArgAction::Append,
long_help = "Make one or more dirs available to all bench suites. Specify either a single dir-name, a comma-separated list of dir-names, or use this multiple instances of this flag to specify multiple dirs.",
value_delimiter = ','
)]
include_dirs: Vec<PathBuf>,
#[arg(
long = "repeat",
value_name = "QUANTITY",
long_help = "Number of times to repeat the benchmark run.",
default_value = "1"
)]
repeat: usize,
#[arg(
long = "list",
value_name = "LIST",
help = "List all selectors and the number of evaluations they select."
)]
list: bool,
#[arg(
long = "output",
short = 'o',
value_name = "FILE",
help = "Save benchmark results to a file"
)]
output: Option<PathBuf>,
#[arg(
long = "format",
value_name = "FORMAT",
help = "Output format (text, json)",
default_value = "text"
)]
format: String,
#[arg(
long = "summary",
help = "Show only summary results",
action = clap::ArgAction::SetTrue
)]
summary: bool,
#[command(subcommand)]
cmd: BenchCommand,
},
}
@@ -346,10 +336,10 @@ pub async fn cli() -> Result<()> {
remote_extension,
builtin,
}) => {
match command {
return match command {
Some(SessionCommand::List { verbose, format }) => {
handle_session_list(verbose, format)?;
return Ok(());
Ok(())
}
None => {
// Run session command by default
@@ -367,9 +357,9 @@ pub async fn cli() -> Result<()> {
None,
)?;
let _ = session.interactive(None).await;
return Ok(());
Ok(())
}
}
};
}
Some(Command::Run {
instructions,
@@ -438,58 +428,22 @@ pub async fn cli() -> Result<()> {
crate::commands::update::update(canary, reconfigure)?;
return Ok(());
}
Some(Command::Bench {
selectors,
include_dirs,
repeat,
list,
output,
format,
summary,
}) => {
if list {
return list_selectors().await;
}
let selectors = if selectors.is_empty() {
vec!["core".to_string()]
} else {
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(selectors.clone(), include_dirs.clone()).await?;
// Handle output based on format
let output_str = match format.as_str() {
"json" => serde_json::to_string_pretty(&results)?,
_ => results.to_string(), // Uses Display impl
};
// Save to file if specified
if let Some(path) = &output {
std::fs::write(current_dir.join(path), &output_str)?;
println!("Results saved to: {}", path.display());
} else {
// Print to console
if summary {
println!("{}", results.summary());
} else {
println!("{}", output_str);
}
Some(Command::Bench { cmd }) => {
match cmd {
BenchCommand::Selectors { config } => BenchRunner::list_selectors(config)?,
BenchCommand::InitConfig { name } => BenchRunConfig::default().save(name),
BenchCommand::Run { config } => BenchRunner::new(config)?.run()?,
BenchCommand::EvalModel { config } => ModelRunner::from(config)?.run()?,
BenchCommand::ExecEval { config } => {
EvalRunner::from(config)?.run(agent_generator).await?
}
}
return Ok(());
}
None => {
if !Config::global().exists() {
return if !Config::global().exists() {
let _ = handle_configure().await;
return Ok(());
Ok(())
} else {
// Run session command by default
let mut session = build_session(None, false, vec![], vec![], vec![], false).await;
@@ -498,8 +452,8 @@ pub async fn cli() -> Result<()> {
None,
)?;
let _ = session.interactive(None).await;
return Ok(());
}
Ok(())
};
}
}
Ok(())