tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 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(())
crates/goose-cli/src/commands/bench.rs
+26 -149
View File
@@ -1,88 +1,37 @@
use crate::logging;
use crate::session::build_session;
use crate::Session;
use crate::{logging, session, Session};
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, EvaluationSuite};
use goose_bench::reporting::{BenchmarkResults, EvaluationResult, SuiteResult};
use goose_bench::bench_session::{BenchAgent, BenchBaseSession};
use goose_bench::eval_suites::ExtensionRequirements;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::Mutex;
pub struct BenchSession {
session: Session,
errors: Arc<Mutex<Vec<BenchAgentError>>>,
}
impl BenchSession {
pub fn new(session: Session) -> Self {
let errors = Arc::new(Mutex::new(Vec::new()));
// Initialize logging with error capture
logging::setup_logging(Some("bench"), Some(errors.clone()))
.expect("Failed to initialize logging");
Self { session, errors }
}
}
// allow session obj to be used in benchmarking
#[async_trait]
impl BenchAgent for BenchSession {
async fn prompt(&mut self, p: String) -> anyhow::Result<Vec<Message>> {
// Clear previous errors
{
let mut errors = self.errors.lock().await;
errors.clear();
}
self.session.headless(p).await?;
Ok(self.session.message_history())
impl BenchBaseSession for Session {
async fn headless(&mut self, message: String) -> anyhow::Result<()> {
self.headless(message).await
}
async fn get_errors(&self) -> Vec<BenchAgentError> {
let errors = self.errors.lock().await;
errors.clone()
fn session_file(&self) -> PathBuf {
self.session_file()
}
async fn get_token_usage(&self) -> Option<i32> {
self.session.get_total_token_usage().ok().flatten()
fn message_history(&self) -> Vec<Message> {
self.message_history()
}
fn get_total_token_usage(&self) -> anyhow::Result<Option<i32>> {
self.get_total_token_usage()
}
}
pub async fn agent_generator(
requirements: ExtensionRequirements,
session_id: String,
) -> BenchAgent {
let identifier = Some(session::Identifier::Name(session_id));
// Wrapper struct to implement BenchAgent for Arc<Mutex<BenchSession>>
struct BenchAgentWrapper(Arc<Mutex<BenchSession>>);
#[async_trait]
impl BenchAgent for BenchAgentWrapper {
async fn prompt(&mut self, p: String) -> anyhow::Result<Vec<Message>> {
let mut session = self.0.lock().await;
session.prompt(p).await
}
async fn get_errors(&self) -> Vec<BenchAgentError> {
let session = self.0.lock().await;
session.get_errors().await
}
async fn get_token_usage(&self) -> Option<i32> {
let session = self.0.lock().await;
session.get_token_usage().await
}
}
async fn run_eval(
evaluation: Box<dyn Evaluation>,
work_dir: &mut BenchmarkWorkDir,
) -> anyhow::Result<EvaluationResult> {
let mut result = EvaluationResult::new(evaluation.name().to_string());
let requirements = evaluation.required_extensions();
// Create session with error capture
let base_session = build_session(
None,
identifier,
false,
requirements.external,
requirements.remote,
@@ -91,84 +40,12 @@ async fn run_eval(
)
.await;
let bench_session = Arc::new(Mutex::new(BenchSession::new(base_session)));
let bench_session_clone = bench_session.clone();
// package session obj into benchmark-compatible struct
let bench_agent = BenchAgent::new(Box::new(base_session));
if let Ok(metrics) = evaluation
.run(Box::new(BenchAgentWrapper(bench_session)), work_dir)
.await
{
for (name, metric) in metrics {
result.add_metric(name, metric);
}
// Initialize logging with error capture
let errors = Some(Arc::new(Mutex::new(bench_agent.get_errors().await)));
logging::setup_logging(Some("bench"), errors).expect("Failed to initialize logging");
// Add any errors that occurred
let agent = BenchAgentWrapper(bench_session_clone);
for error in agent.get_errors().await {
result.add_error(error);
}
}
let current_dir = std::env::current_dir()?;
let output_str = serde_json::to_string_pretty(&result)?;
std::fs::write(current_dir.join("eval_result.json"), &output_str)?;
Ok(result)
}
pub async fn run_benchmark(
selectors: Vec<String>,
include_dirs: Vec<PathBuf>,
) -> anyhow::Result<BenchmarkResults> {
let config = Config::global();
let goose_model: String = config
.get_param("GOOSE_MODEL")
.expect("No model configured. Run 'goose configure' first");
let provider_name: String = config
.get_param("GOOSE_PROVIDER")
.expect("No provider configured. Run 'goose configure' first");
let mut results = BenchmarkResults::new(provider_name.clone());
let work_dir = Mutex::new(BenchmarkWorkDir::new(
format!("{}-{}", provider_name, goose_model),
include_dirs.clone(),
));
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_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(())
bench_agent
}
crates/goose-cli/src/logging.rs
+1 -1
View File
@@ -12,8 +12,8 @@ use tracing_subscriber::{
};
use goose::tracing::langfuse_layer;
use goose_bench::bench_session::BenchAgentError;
use goose_bench::error_capture::ErrorCaptureLayer;
use goose_bench::eval_suites::BenchAgentError;
// Used to ensure we only set up tracing once
static INIT: Once = Once::new();