diff --git a/Cargo.lock b/Cargo.lock index 85dfeab3..0e8456e0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2236,6 +2236,30 @@ dependencies = [ "zbus", ] +[[package]] +name = "ironclaw-bench" +version = "0.1.0" +dependencies = [ + "anyhow", + "async-trait", + "chrono", + "clap", + "futures", + "ironclaw", + "regex", + "rust_decimal", + "serde", + "serde_json", + "tempfile", + "thiserror 2.0.18", + "tokio", + "tokio-stream", + "toml", + "tracing", + "tracing-subscriber", + "uuid", +] + [[package]] name = "is-docker" version = "0.2.0" diff --git a/Cargo.toml b/Cargo.toml index 6fc39003..a23752d4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,6 @@ +[workspace] +members = [".", "benchmarks"] + [package] name = "ironclaw" version = "0.1.0" diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml new file mode 100644 index 00000000..0ad3ef35 --- /dev/null +++ b/benchmarks/Cargo.toml @@ -0,0 +1,49 @@ +[package] +name = "ironclaw-bench" +version = "0.1.0" +edition = "2024" +rust-version = "1.85" +description = "Benchmarking harness for IronClaw agent" +license = "MIT OR Apache-2.0" + +[[bin]] +name = "ironclaw-bench" +path = "src/main.rs" + +[dependencies] +ironclaw = { path = ".." } + +# Async runtime +tokio = { version = "1", features = ["full"] } +tokio-stream = { version = "0.1", features = ["sync"] } +futures = "0.3" + +# Serialization +serde = { version = "1", features = ["derive"] } +serde_json = "1" +toml = "0.8" + +# CLI +clap = { version = "4", features = ["derive"] } + +# Core types +uuid = { version = "1", features = ["v4", "serde"] } +chrono = { version = "0.4", features = ["serde"] } +rust_decimal = { version = "1", features = ["serde", "serde-with-str"] } + +# Error handling +thiserror = "2" +anyhow = "1" + +# Async traits +async-trait = "0.1" + +# Logging +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } + +# Scoring +regex = "1" + +[dev-dependencies] +tempfile = "3" diff --git a/benchmarks/src/adapters/custom.rs b/benchmarks/src/adapters/custom.rs new file mode 100644 index 00000000..2b9c5272 --- /dev/null +++ b/benchmarks/src/adapters/custom.rs @@ -0,0 +1,237 @@ +use std::io::BufRead; +use std::path::PathBuf; + +use async_trait::async_trait; +use serde::Deserialize; + +use crate::error::BenchError; +use crate::scoring; +use crate::suite::{BenchScore, BenchSuite, BenchTask, TaskSubmission}; + +/// A single entry in the custom JSONL format. +#[derive(Debug, Deserialize)] +struct CustomEntry { + id: String, + prompt: String, + #[serde(default)] + context: Option, + #[serde(default)] + tags: Vec, + #[serde(default)] + expected: Option, + #[serde(default)] + expected_contains: Option, + #[serde(default)] + expected_regex: Option, + /// "exact", "contains", "regex", or "llm" (default: "exact") + #[serde(default = "default_scorer")] + scorer: String, +} + +fn default_scorer() -> String { + "exact".to_string() +} + +/// Custom JSONL benchmark suite. +/// +/// Each line of the JSONL file is a task with `id`, `prompt`, and scoring +/// criteria (`expected`, `expected_contains`, `expected_regex`). +pub struct CustomSuite { + dataset_path: PathBuf, +} + +impl CustomSuite { + pub fn new(dataset_path: impl Into) -> Self { + Self { + dataset_path: dataset_path.into(), + } + } +} + +#[async_trait] +impl BenchSuite for CustomSuite { + fn name(&self) -> &str { + "Custom JSONL" + } + + fn id(&self) -> &str { + "custom" + } + + async fn load_tasks(&self) -> Result, BenchError> { + let file = std::fs::File::open(&self.dataset_path).map_err(BenchError::Io)?; + let reader = std::io::BufReader::new(file); + let mut tasks = Vec::new(); + + for (line_num, line) in reader.lines().enumerate() { + let line = line?; + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + let entry: CustomEntry = serde_json::from_str(trimmed) + .map_err(|e| BenchError::Config(format!("line {}: {}", line_num + 1, e)))?; + + let mut metadata = serde_json::json!({ + "scorer": entry.scorer, + }); + if let Some(ref expected) = entry.expected { + metadata["expected"] = serde_json::Value::String(expected.clone()); + } + if let Some(ref expected_contains) = entry.expected_contains { + metadata["expected_contains"] = + serde_json::Value::String(expected_contains.clone()); + } + if let Some(ref expected_regex) = entry.expected_regex { + metadata["expected_regex"] = serde_json::Value::String(expected_regex.clone()); + } + + tasks.push(BenchTask { + id: entry.id, + prompt: entry.prompt, + context: entry.context, + resources: vec![], + tags: entry.tags, + expected_turns: None, + timeout: None, + metadata, + }); + } + + Ok(tasks) + } + + async fn score( + &self, + task: &BenchTask, + submission: &TaskSubmission, + ) -> Result { + let scorer = task + .metadata + .get("scorer") + .and_then(|v| v.as_str()) + .unwrap_or("exact"); + + match scorer { + "exact" => { + if let Some(expected) = task.metadata.get("expected").and_then(|v| v.as_str()) { + Ok(scoring::exact_match(expected, &submission.response)) + } else { + Err(BenchError::Scoring { + task_id: task.id.clone(), + reason: "no 'expected' field for exact scoring".to_string(), + }) + } + } + "contains" => { + if let Some(expected) = task + .metadata + .get("expected_contains") + .and_then(|v| v.as_str()) + { + Ok(scoring::contains_match(expected, &submission.response)) + } else { + Err(BenchError::Scoring { + task_id: task.id.clone(), + reason: "no 'expected_contains' field for contains scoring".to_string(), + }) + } + } + "regex" => { + if let Some(pattern) = task.metadata.get("expected_regex").and_then(|v| v.as_str()) + { + Ok(scoring::regex_match(pattern, &submission.response)) + } else { + Err(BenchError::Scoring { + task_id: task.id.clone(), + reason: "no 'expected_regex' field for regex scoring".to_string(), + }) + } + } + "llm" => { + // TODO: LLM-as-judge scoring + Ok(BenchScore::partial(0.5, "LLM scoring not yet implemented")) + } + other => Err(BenchError::Scoring { + task_id: task.id.clone(), + reason: format!("unknown scorer: {other}"), + }), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + + #[tokio::test] + async fn test_custom_load_tasks() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("tasks.jsonl"); + let mut file = std::fs::File::create(&path).unwrap(); + writeln!( + file, + r#"{{"id": "t1", "prompt": "What is 2+2?", "expected": "4"}}"# + ) + .unwrap(); + writeln!( + file, + r#"{{"id": "t2", "prompt": "Say hello", "expected_contains": "hello", "scorer": "contains"}}"# + ) + .unwrap(); + + let suite = CustomSuite::new(&path); + let tasks = suite.load_tasks().await.unwrap(); + assert_eq!(tasks.len(), 2); + assert_eq!(tasks[0].id, "t1"); + assert_eq!(tasks[1].id, "t2"); + } + + #[tokio::test] + async fn test_custom_exact_scoring() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("tasks.jsonl"); + let mut file = std::fs::File::create(&path).unwrap(); + writeln!( + file, + r#"{{"id": "t1", "prompt": "What is 2+2?", "expected": "4"}}"# + ) + .unwrap(); + + let suite = CustomSuite::new(&path); + let tasks = suite.load_tasks().await.unwrap(); + + let submission = TaskSubmission { + response: "4".to_string(), + conversation: vec![], + tool_calls: vec![], + }; + let score = suite.score(&tasks[0], &submission).await.unwrap(); + assert_eq!(score.value, 1.0); + assert_eq!(score.label, "pass"); + } + + #[tokio::test] + async fn test_custom_contains_scoring() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("tasks.jsonl"); + let mut file = std::fs::File::create(&path).unwrap(); + writeln!( + file, + r#"{{"id": "t1", "prompt": "Greet me", "expected_contains": "hello", "scorer": "contains"}}"# + ) + .unwrap(); + + let suite = CustomSuite::new(&path); + let tasks = suite.load_tasks().await.unwrap(); + + let submission = TaskSubmission { + response: "Hello there!".to_string(), + conversation: vec![], + tool_calls: vec![], + }; + let score = suite.score(&tasks[0], &submission).await.unwrap(); + assert_eq!(score.value, 1.0); + } +} diff --git a/benchmarks/src/adapters/gaia.rs b/benchmarks/src/adapters/gaia.rs new file mode 100644 index 00000000..82bf773d --- /dev/null +++ b/benchmarks/src/adapters/gaia.rs @@ -0,0 +1,183 @@ +use std::io::BufRead; +use std::path::PathBuf; + +use async_trait::async_trait; +use serde::Deserialize; + +use crate::error::BenchError; +use crate::scoring; +use crate::suite::{BenchScore, BenchSuite, BenchTask, TaskResource, TaskSubmission}; + +/// GAIA dataset entry (Hugging Face JSONL format). +#[derive(Debug, Deserialize)] +struct GaiaEntry { + task_id: String, + #[serde(alias = "Question")] + question: String, + #[serde(alias = "Final answer", alias = "final_answer")] + final_answer: String, + #[serde(alias = "Level", default)] + level: Option, + #[serde(alias = "file_name", default)] + file_name: Option, + #[serde(alias = "Annotator Metadata", default)] + annotator_metadata: Option, +} + +/// GAIA benchmark suite. +/// +/// Tasks are loaded from HuggingFace JSONL exports. Scoring uses normalized +/// exact match against the `final_answer` field. +pub struct GaiaSuite { + dataset_path: PathBuf, + attachments_dir: Option, +} + +impl GaiaSuite { + pub fn new( + dataset_path: impl Into, + attachments_dir: Option>, + ) -> Self { + Self { + dataset_path: dataset_path.into(), + attachments_dir: attachments_dir.map(|d| d.into()), + } + } +} + +#[async_trait] +impl BenchSuite for GaiaSuite { + fn name(&self) -> &str { + "GAIA" + } + + fn id(&self) -> &str { + "gaia" + } + + async fn load_tasks(&self) -> Result, BenchError> { + let file = std::fs::File::open(&self.dataset_path)?; + let reader = std::io::BufReader::new(file); + let mut tasks = Vec::new(); + + for (line_num, line) in reader.lines().enumerate() { + let line = line?; + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + let entry: GaiaEntry = serde_json::from_str(trimmed) + .map_err(|e| BenchError::Config(format!("GAIA line {}: {}", line_num + 1, e)))?; + + let mut resources = Vec::new(); + if let Some(ref file_name) = entry.file_name { + if !file_name.is_empty() { + if let Some(ref dir) = self.attachments_dir { + resources.push(TaskResource { + name: file_name.clone(), + path: dir.join(file_name).to_string_lossy().to_string(), + resource_type: crate::suite::ResourceType::File, + }); + } + } + } + + let mut tags = Vec::new(); + if let Some(level) = entry.level { + tags.push(format!("level-{level}")); + } + + let metadata = serde_json::json!({ + "expected": entry.final_answer, + "level": entry.level, + }); + + tasks.push(BenchTask { + id: entry.task_id, + prompt: entry.question, + context: None, + resources, + tags, + expected_turns: None, + timeout: None, + metadata, + }); + } + + Ok(tasks) + } + + async fn score( + &self, + task: &BenchTask, + submission: &TaskSubmission, + ) -> Result { + let expected = task + .metadata + .get("expected") + .and_then(|v| v.as_str()) + .ok_or_else(|| BenchError::Scoring { + task_id: task.id.clone(), + reason: "missing expected answer in metadata".to_string(), + })?; + + Ok(scoring::exact_match(expected, &submission.response)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + + #[tokio::test] + async fn test_gaia_load_tasks() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("gaia.jsonl"); + let mut file = std::fs::File::create(&path).unwrap(); + writeln!( + file, + r#"{{"task_id": "g1", "question": "What is the capital of France?", "final_answer": "Paris", "Level": 1}}"# + ) + .unwrap(); + + let suite = GaiaSuite::new(&path, None::); + let tasks = suite.load_tasks().await.unwrap(); + assert_eq!(tasks.len(), 1); + assert_eq!(tasks[0].id, "g1"); + assert!(tasks[0].tags.contains(&"level-1".to_string())); + } + + #[tokio::test] + async fn test_gaia_scoring() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("gaia.jsonl"); + let mut file = std::fs::File::create(&path).unwrap(); + writeln!( + file, + r#"{{"task_id": "g1", "question": "Capital of France?", "final_answer": "Paris"}}"# + ) + .unwrap(); + + let suite = GaiaSuite::new(&path, None::); + let tasks = suite.load_tasks().await.unwrap(); + + // Exact match (case insensitive) + let submission = TaskSubmission { + response: "paris".to_string(), + conversation: vec![], + tool_calls: vec![], + }; + let score = suite.score(&tasks[0], &submission).await.unwrap(); + assert_eq!(score.value, 1.0); + + // Wrong answer + let submission = TaskSubmission { + response: "London".to_string(), + conversation: vec![], + tool_calls: vec![], + }; + let score = suite.score(&tasks[0], &submission).await.unwrap(); + assert_eq!(score.value, 0.0); + } +} diff --git a/benchmarks/src/adapters/mod.rs b/benchmarks/src/adapters/mod.rs new file mode 100644 index 00000000..7665f96d --- /dev/null +++ b/benchmarks/src/adapters/mod.rs @@ -0,0 +1,110 @@ +pub mod custom; +pub mod gaia; +pub mod swe_bench; +pub mod tau_bench; + +use crate::config::BenchConfig; +use crate::error::BenchError; +use crate::suite::BenchSuite; + +/// List of all known suite IDs. +pub const KNOWN_SUITES: &[(&str, &str)] = &[ + ("custom", "Custom JSONL tasks"), + ("gaia", "GAIA benchmark (knowledge & reasoning)"), + ("tau_bench", "Tau-bench (multi-turn tool use)"), + ("swe_bench", "SWE-bench Pro (software engineering)"), +]; + +/// Create a suite adapter by name. +pub fn create_suite(name: &str, config: &BenchConfig) -> Result, BenchError> { + let suite_map = config.suite_config_map(); + match name { + "custom" => { + let dataset_path = suite_map + .get("dataset_path") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .ok_or_else(|| { + BenchError::Config( + "suite_config.dataset_path is required for 'custom' suite".to_string(), + ) + })?; + Ok(Box::new(custom::CustomSuite::new(dataset_path))) + } + "gaia" => { + let dataset_path = suite_map + .get("dataset_path") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .ok_or_else(|| { + BenchError::Config( + "suite_config.dataset_path is required for 'gaia' suite".to_string(), + ) + })?; + let attachments_dir = suite_map + .get("attachments_dir") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + Ok(Box::new(gaia::GaiaSuite::new( + dataset_path, + attachments_dir, + ))) + } + "tau_bench" => { + let dataset_path = suite_map + .get("dataset_path") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .ok_or_else(|| { + BenchError::Config( + "suite_config.dataset_path is required for 'tau_bench' suite".to_string(), + ) + })?; + let domain = suite_map + .get("domain") + .and_then(|v| v.as_str()) + .unwrap_or("retail") + .to_string(); + Ok(Box::new(tau_bench::TauBenchSuite::new( + dataset_path, + domain, + ))) + } + "swe_bench" => { + let dataset_path = suite_map + .get("dataset_path") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .ok_or_else(|| { + BenchError::Config( + "suite_config.dataset_path is required for 'swe_bench' suite".to_string(), + ) + })?; + let workspace_dir = suite_map + .get("workspace_dir") + .and_then(|v| v.as_str()) + .unwrap_or("/tmp/swe-bench") + .to_string(); + let use_docker = suite_map + .get("use_docker") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + Ok(Box::new(swe_bench::SweBenchSuite::new( + dataset_path, + workspace_dir, + use_docker, + ))) + } + _ => { + let available = KNOWN_SUITES + .iter() + .map(|(id, _)| *id) + .collect::>() + .join(", "); + Err(BenchError::SuiteNotFound { + name: name.to_string(), + available, + }) + } + } +} diff --git a/benchmarks/src/adapters/swe_bench.rs b/benchmarks/src/adapters/swe_bench.rs new file mode 100644 index 00000000..2d97129a --- /dev/null +++ b/benchmarks/src/adapters/swe_bench.rs @@ -0,0 +1,291 @@ +use std::io::BufRead; +use std::path::PathBuf; + +use async_trait::async_trait; +use serde::Deserialize; + +use crate::error::BenchError; +use crate::suite::{BenchScore, BenchSuite, BenchTask, TaskSubmission}; + +/// SWE-bench dataset entry. +#[derive(Debug, Deserialize)] +struct SweBenchEntry { + instance_id: String, + repo: String, + base_commit: String, + #[serde(default)] + problem_statement: String, + #[serde(default)] + hints_text: Option, + #[serde(default)] + test_patch: Option, + #[serde(default)] + patch: Option, +} + +/// SWE-bench Pro: real-world software engineering tasks. +/// +/// Each task clones a repo at a specific commit, presents the problem statement, +/// and expects the agent to produce a patch. Scoring runs the test suite. +pub struct SweBenchSuite { + dataset_path: PathBuf, + workspace_dir: PathBuf, + use_docker: bool, +} + +impl SweBenchSuite { + pub fn new( + dataset_path: impl Into, + workspace_dir: impl Into, + use_docker: bool, + ) -> Self { + Self { + dataset_path: dataset_path.into(), + workspace_dir: workspace_dir.into(), + use_docker, + } + } +} + +#[async_trait] +impl BenchSuite for SweBenchSuite { + fn name(&self) -> &str { + "SWE-bench Pro" + } + + fn id(&self) -> &str { + "swe_bench" + } + + async fn load_tasks(&self) -> Result, BenchError> { + let file = std::fs::File::open(&self.dataset_path)?; + let reader = std::io::BufReader::new(file); + let mut tasks = Vec::new(); + + for (line_num, line) in reader.lines().enumerate() { + let line = line?; + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + let entry: SweBenchEntry = serde_json::from_str(trimmed).map_err(|e| { + BenchError::Config(format!("swe_bench line {}: {}", line_num + 1, e)) + })?; + + let metadata = serde_json::json!({ + "repo": entry.repo, + "base_commit": entry.base_commit, + "test_patch": entry.test_patch, + "gold_patch": entry.patch, + "use_docker": self.use_docker, + "workspace_dir": self.workspace_dir.to_string_lossy(), + }); + + let prompt = if let Some(ref hints) = entry.hints_text { + format!("{}\n\nHints:\n{}", entry.problem_statement, hints) + } else { + entry.problem_statement + }; + + tasks.push(BenchTask { + id: entry.instance_id, + prompt, + context: Some(format!( + "Repository: {}, Commit: {}", + entry.repo, entry.base_commit + )), + resources: vec![], + tags: vec![format!("repo-{}", entry.repo.replace('/', "-"))], + expected_turns: None, + timeout: None, + metadata, + }); + } + + Ok(tasks) + } + + async fn setup_task(&self, task: &BenchTask) -> Result<(), BenchError> { + let repo = task + .metadata + .get("repo") + .and_then(|v| v.as_str()) + .ok_or_else(|| BenchError::TaskFailed { + task_id: task.id.clone(), + reason: "missing repo in metadata".to_string(), + })?; + let base_commit = task + .metadata + .get("base_commit") + .and_then(|v| v.as_str()) + .ok_or_else(|| BenchError::TaskFailed { + task_id: task.id.clone(), + reason: "missing base_commit in metadata".to_string(), + })?; + + let task_dir = self.workspace_dir.join(&task.id); + + // Clone repo if not already present + if !task_dir.exists() { + let repo_url = format!("https://github.com/{}.git", repo); + let output = tokio::process::Command::new("git") + .args([ + "clone", + "--depth", + "1", + &repo_url, + &task_dir.to_string_lossy(), + ]) + .output() + .await + .map_err(|e| BenchError::TaskFailed { + task_id: task.id.clone(), + reason: format!("git clone failed: {e}"), + })?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + return Err(BenchError::TaskFailed { + task_id: task.id.clone(), + reason: format!("git clone failed: {stderr}"), + }); + } + } + + // Checkout the base commit + let output = tokio::process::Command::new("git") + .args(["checkout", base_commit]) + .current_dir(&task_dir) + .output() + .await + .map_err(|e| BenchError::TaskFailed { + task_id: task.id.clone(), + reason: format!("git checkout failed: {e}"), + })?; + + if !output.status.success() { + // Shallow clone might not have the commit; fetch more history + let _ = tokio::process::Command::new("git") + .args(["fetch", "--unshallow"]) + .current_dir(&task_dir) + .output() + .await; + + let output = tokio::process::Command::new("git") + .args(["checkout", base_commit]) + .current_dir(&task_dir) + .output() + .await + .map_err(|e| BenchError::TaskFailed { + task_id: task.id.clone(), + reason: format!("git checkout retry failed: {e}"), + })?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + return Err(BenchError::TaskFailed { + task_id: task.id.clone(), + reason: format!("git checkout failed: {stderr}"), + }); + } + } + + Ok(()) + } + + async fn teardown_task(&self, task: &BenchTask) -> Result<(), BenchError> { + let task_dir = self.workspace_dir.join(&task.id); + if task_dir.exists() { + // Reset any changes + let _ = tokio::process::Command::new("git") + .args(["checkout", "."]) + .current_dir(&task_dir) + .output() + .await; + let _ = tokio::process::Command::new("git") + .args(["clean", "-fdx"]) + .current_dir(&task_dir) + .output() + .await; + } + Ok(()) + } + + async fn score( + &self, + task: &BenchTask, + submission: &TaskSubmission, + ) -> Result { + // For SWE-bench, scoring requires running the test patch against the agent's changes. + // This is a simplified version that checks if the agent produced any code changes. + + let test_patch = task.metadata.get("test_patch").and_then(|v| v.as_str()); + + if submission.response.is_empty() { + return Ok(BenchScore::fail("no response from agent")); + } + + // If we have a test patch, try to verify the submission + if let Some(_test_patch) = test_patch { + // TODO: Apply agent's patch, then apply test patch, then run tests. + // For now, give partial credit if the agent produced some output. + Ok(BenchScore::partial( + 0.25, + "test execution not yet implemented; partial credit for response", + )) + } else { + // No test patch available; can't automatically score + Ok(BenchScore::partial( + 0.25, + "no test_patch available for automated scoring", + )) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + + #[tokio::test] + async fn test_swe_bench_load() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("swe.jsonl"); + let mut file = std::fs::File::create(&path).unwrap(); + writeln!( + file, + r#"{{"instance_id": "django__django-12345", "repo": "django/django", "base_commit": "abc123", "problem_statement": "Fix the ORM bug"}}"# + ) + .unwrap(); + + let suite = SweBenchSuite::new(&path, "/tmp/swe-test", false); + let tasks = suite.load_tasks().await.unwrap(); + assert_eq!(tasks.len(), 1); + assert_eq!(tasks[0].id, "django__django-12345"); + assert!(tasks[0].tags.contains(&"repo-django-django".to_string())); + } + + #[tokio::test] + async fn test_swe_bench_scoring_no_response() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("swe.jsonl"); + let mut file = std::fs::File::create(&path).unwrap(); + writeln!( + file, + r#"{{"instance_id": "s1", "repo": "org/repo", "base_commit": "abc", "problem_statement": "Fix bug"}}"# + ) + .unwrap(); + + let suite = SweBenchSuite::new(&path, "/tmp/swe-test", false); + let tasks = suite.load_tasks().await.unwrap(); + + let submission = TaskSubmission { + response: String::new(), + conversation: vec![], + tool_calls: vec![], + }; + let score = suite.score(&tasks[0], &submission).await.unwrap(); + assert_eq!(score.value, 0.0); + } +} diff --git a/benchmarks/src/adapters/tau_bench.rs b/benchmarks/src/adapters/tau_bench.rs new file mode 100644 index 00000000..1fa2b391 --- /dev/null +++ b/benchmarks/src/adapters/tau_bench.rs @@ -0,0 +1,227 @@ +use std::io::BufRead; +use std::path::PathBuf; + +use async_trait::async_trait; +use serde::Deserialize; + +use crate::error::BenchError; +use crate::suite::{BenchScore, BenchSuite, BenchTask, ConversationTurn, TaskSubmission}; + +/// Tau-bench task entry. +#[derive(Debug, Deserialize)] +struct TauBenchEntry { + id: String, + #[serde(default)] + domain: String, + instruction: String, + #[serde(default)] + user_persona: Option, + #[serde(default)] + expected_state: Option, + #[serde(default)] + expected_actions: Vec, + #[serde(default)] + max_turns: Option, +} + +/// Tau-bench: multi-turn tool-calling dialog benchmark. +/// +/// Tests agent ability to handle customer service scenarios with simulated +/// domain APIs (retail, airline). Scoring compares final state against expected. +pub struct TauBenchSuite { + dataset_path: PathBuf, + domain: String, +} + +impl TauBenchSuite { + pub fn new(dataset_path: impl Into, domain: impl Into) -> Self { + Self { + dataset_path: dataset_path.into(), + domain: domain.into(), + } + } +} + +#[async_trait] +impl BenchSuite for TauBenchSuite { + fn name(&self) -> &str { + "Tau-bench" + } + + fn id(&self) -> &str { + "tau_bench" + } + + async fn load_tasks(&self) -> Result, BenchError> { + let file = std::fs::File::open(&self.dataset_path)?; + let reader = std::io::BufReader::new(file); + let mut tasks = Vec::new(); + + for (line_num, line) in reader.lines().enumerate() { + let line = line?; + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + let entry: TauBenchEntry = serde_json::from_str(trimmed).map_err(|e| { + BenchError::Config(format!("tau_bench line {}: {}", line_num + 1, e)) + })?; + + let domain = if entry.domain.is_empty() { + self.domain.clone() + } else { + entry.domain.clone() + }; + + let metadata = serde_json::json!({ + "domain": domain, + "user_persona": entry.user_persona, + "expected_state": entry.expected_state, + "expected_actions": entry.expected_actions, + }); + + tasks.push(BenchTask { + id: entry.id, + prompt: entry.instruction, + context: entry.user_persona.clone(), + resources: vec![], + tags: vec![format!("domain-{domain}")], + expected_turns: entry.max_turns, + timeout: None, + metadata, + }); + } + + Ok(tasks) + } + + async fn score( + &self, + task: &BenchTask, + submission: &TaskSubmission, + ) -> Result { + // Score based on expected actions completion + let expected_actions: Vec = task + .metadata + .get("expected_actions") + .and_then(|v| serde_json::from_value(v.clone()).ok()) + .unwrap_or_default(); + + if expected_actions.is_empty() { + // No expected actions defined; score based on whether agent responded + if submission.response.is_empty() { + return Ok(BenchScore::fail("no response")); + } + return Ok(BenchScore::partial( + 0.5, + "no expected_actions to evaluate against", + )); + } + + // Check which expected actions were actually called + let called: std::collections::HashSet<&str> = + submission.tool_calls.iter().map(|s| s.as_str()).collect(); + let matched = expected_actions + .iter() + .filter(|a| called.contains(a.as_str())) + .count(); + + let ratio = matched as f64 / expected_actions.len() as f64; + if ratio >= 1.0 { + Ok(BenchScore::pass()) + } else if ratio > 0.0 { + Ok(BenchScore::partial( + ratio, + format!( + "{}/{} expected actions completed", + matched, + expected_actions.len() + ), + )) + } else { + Ok(BenchScore::fail(format!( + "0/{} expected actions completed", + expected_actions.len() + ))) + } + } + + async fn next_user_message( + &self, + task: &BenchTask, + conversation: &[ConversationTurn], + ) -> Result, BenchError> { + // Check if we've exceeded max turns + if let Some(max) = task.expected_turns { + let user_turns = conversation + .iter() + .filter(|t| matches!(t.role, crate::suite::TurnRole::User)) + .count(); + if user_turns >= max { + return Ok(None); + } + } + + // For now, multi-turn simulation requires an LLM (not implemented yet). + // Return None to end after the first turn. + // TODO: Use LLM to simulate customer based on user_persona. + Ok(None) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + + #[tokio::test] + async fn test_tau_bench_load() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("tau.jsonl"); + let mut file = std::fs::File::create(&path).unwrap(); + writeln!( + file, + r#"{{"id": "t1", "instruction": "Return my order", "expected_actions": ["lookup_order", "process_return"], "max_turns": 3}}"# + ) + .unwrap(); + + let suite = TauBenchSuite::new(&path, "retail"); + let tasks = suite.load_tasks().await.unwrap(); + assert_eq!(tasks.len(), 1); + assert_eq!(tasks[0].expected_turns, Some(3)); + } + + #[tokio::test] + async fn test_tau_bench_scoring() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("tau.jsonl"); + let mut file = std::fs::File::create(&path).unwrap(); + writeln!( + file, + r#"{{"id": "t1", "instruction": "Return order", "expected_actions": ["lookup_order", "process_return"]}}"# + ) + .unwrap(); + + let suite = TauBenchSuite::new(&path, "retail"); + let tasks = suite.load_tasks().await.unwrap(); + + // Partial completion + let submission = TaskSubmission { + response: "I found your order.".to_string(), + conversation: vec![], + tool_calls: vec!["lookup_order".to_string()], + }; + let score = suite.score(&tasks[0], &submission).await.unwrap(); + assert_eq!(score.value, 0.5); + assert_eq!(score.label, "partial"); + + // Full completion + let submission = TaskSubmission { + response: "Return processed.".to_string(), + conversation: vec![], + tool_calls: vec!["lookup_order".to_string(), "process_return".to_string()], + }; + let score = suite.score(&tasks[0], &submission).await.unwrap(); + assert_eq!(score.value, 1.0); + } +} diff --git a/benchmarks/src/channel.rs b/benchmarks/src/channel.rs new file mode 100644 index 00000000..71ffe91e --- /dev/null +++ b/benchmarks/src/channel.rs @@ -0,0 +1,229 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use tokio::sync::{Mutex, mpsc}; +use tokio_stream::wrappers::ReceiverStream; + +use ironclaw::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse, StatusUpdate}; +use ironclaw::error::ChannelError; + +use crate::results::TraceToolCall; +use crate::suite::ConversationTurn; + +/// Captured state from a benchmark channel run. +#[derive(Debug, Default)] +pub struct ChannelCapture { + /// All responses the agent sent back. + pub responses: Vec, + /// Tool calls observed (name, success, duration_ms). + pub tool_calls: Vec, + /// Full conversation turns for multi-turn scoring. + pub conversation: Vec, + /// Status messages (for debugging). + pub status_log: Vec, +} + +/// A headless Channel implementation for benchmarking. +/// +/// Modeled after `ReplChannel`: uses mpsc to inject messages and captures +/// all responses and tool status events. Auto-approves tool execution +/// so benchmarks run without user interaction. +pub struct BenchChannel { + /// Sender to inject messages into the agent loop. + msg_tx: mpsc::Sender, + /// Receiver the agent loop reads from (taken once by `start()`). + msg_rx: Mutex>>, + /// Accumulated capture data. + capture: Arc>, +} + +impl BenchChannel { + pub fn new() -> (Self, mpsc::Sender) { + let (tx, rx) = mpsc::channel(64); + let channel = Self { + msg_tx: tx.clone(), + msg_rx: Mutex::new(Some(rx)), + capture: Arc::new(Mutex::new(ChannelCapture::default())), + }; + (channel, tx) + } + + /// Get a handle to the capture data. + pub fn capture(&self) -> Arc> { + Arc::clone(&self.capture) + } + + /// Get a clone of the message sender for injecting follow-up messages. + pub fn sender(&self) -> mpsc::Sender { + self.msg_tx.clone() + } +} + +#[async_trait] +impl Channel for BenchChannel { + fn name(&self) -> &str { + "bench" + } + + async fn start(&self) -> Result { + let rx = self + .msg_rx + .lock() + .await + .take() + .ok_or_else(|| ChannelError::StartupFailed { + name: "bench".to_string(), + reason: "start() already called".to_string(), + })?; + Ok(Box::pin(ReceiverStream::new(rx))) + } + + async fn respond( + &self, + _msg: &IncomingMessage, + response: OutgoingResponse, + ) -> Result<(), ChannelError> { + let mut cap = self.capture.lock().await; + cap.responses.push(response.content.clone()); + cap.conversation.push(ConversationTurn { + role: crate::suite::TurnRole::Assistant, + content: response.content, + }); + Ok(()) + } + + async fn send_status( + &self, + status: StatusUpdate, + _metadata: &serde_json::Value, + ) -> Result<(), ChannelError> { + let mut cap = self.capture.lock().await; + + match status { + StatusUpdate::ToolCompleted { ref name, success } => { + cap.tool_calls.push(TraceToolCall { + name: name.clone(), + duration_ms: 0, // We don't have precise per-tool timing here + success, + }); + cap.status_log + .push(format!("tool_completed: {name} success={success}")); + } + StatusUpdate::ApprovalNeeded { ref request_id, .. } => { + // Auto-approve all tools during benchmarks + cap.status_log.push(format!("auto_approved: {request_id}")); + drop(cap); // Release lock before sending + let approval = IncomingMessage::new("bench", "bench-user", "always"); + let _ = self.msg_tx.send(approval).await; + return Ok(()); + } + StatusUpdate::Thinking(ref msg) => { + cap.status_log.push(format!("thinking: {msg}")); + } + StatusUpdate::ToolStarted { ref name } => { + cap.status_log.push(format!("tool_started: {name}")); + } + StatusUpdate::ToolResult { + ref name, + ref preview, + } => { + cap.status_log.push(format!( + "tool_result: {name} -> {}", + &preview[..preview.len().min(100)] + )); + } + StatusUpdate::StreamChunk(_) => {} + StatusUpdate::Status(ref msg) => { + cap.status_log.push(format!("status: {msg}")); + } + } + Ok(()) + } + + async fn broadcast( + &self, + _user_id: &str, + response: OutgoingResponse, + ) -> Result<(), ChannelError> { + let mut cap = self.capture.lock().await; + cap.status_log.push(format!( + "broadcast: {}", + &response.content[..response.content.len().min(100)] + )); + Ok(()) + } + + async fn health_check(&self) -> Result<(), ChannelError> { + Ok(()) + } + + async fn shutdown(&self) -> Result<(), ChannelError> { + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_bench_channel_captures_responses() { + let (channel, _tx) = BenchChannel::new(); + let capture = channel.capture(); + + let msg = IncomingMessage::new("bench", "user", "hello"); + let response = OutgoingResponse::text("world"); + channel.respond(&msg, response).await.unwrap(); + + let cap = capture.lock().await; + assert_eq!(cap.responses.len(), 1); + assert_eq!(cap.responses[0], "world"); + assert_eq!(cap.conversation.len(), 1); + } + + #[tokio::test] + async fn test_bench_channel_auto_approves() { + let (channel, _tx) = BenchChannel::new(); + // start() to consume the receiver + let _stream = channel.start().await.unwrap(); + + let status = StatusUpdate::ApprovalNeeded { + request_id: "req-1".to_string(), + tool_name: "shell".to_string(), + description: "run ls".to_string(), + parameters: serde_json::json!({}), + }; + channel + .send_status(status, &serde_json::Value::Null) + .await + .unwrap(); + + // The approval message was sent through msg_tx, + // which means the stream would receive it. + // We can't easily read from the stream in this test without + // consuming it, but we can verify the status log. + let capture_arc = channel.capture(); + let cap = capture_arc.lock().await; + assert!(cap.status_log.iter().any(|s| s.contains("auto_approved"))); + } + + #[tokio::test] + async fn test_bench_channel_captures_tool_events() { + let (channel, _tx) = BenchChannel::new(); + + let status = StatusUpdate::ToolCompleted { + name: "echo".to_string(), + success: true, + }; + channel + .send_status(status, &serde_json::Value::Null) + .await + .unwrap(); + + let capture_arc = channel.capture(); + let cap = capture_arc.lock().await; + assert_eq!(cap.tool_calls.len(), 1); + assert_eq!(cap.tool_calls[0].name, "echo"); + assert!(cap.tool_calls[0].success); + } +} diff --git a/benchmarks/src/config.rs b/benchmarks/src/config.rs new file mode 100644 index 00000000..eb492e0b --- /dev/null +++ b/benchmarks/src/config.rs @@ -0,0 +1,203 @@ +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use serde::Deserialize; + +use crate::error::BenchError; + +/// Top-level bench configuration, loaded from TOML. +#[derive(Debug, Clone, Deserialize)] +pub struct BenchConfig { + /// Where to write results. Default: "./bench-results". + #[serde(default = "default_results_dir")] + pub results_dir: PathBuf, + + /// Per-task timeout. Default: "300s". + #[serde( + default = "default_task_timeout", + deserialize_with = "deserialize_duration" + )] + pub task_timeout: Duration, + + /// Max agent iterations per task. Default: 15. + #[serde(default = "default_max_iterations")] + pub max_iterations: usize, + + /// How many tasks to run in parallel. Default: 1. + #[serde(default = "default_parallelism")] + pub parallelism: usize, + + /// Model/config matrix entries. At least one required. + #[serde(default)] + pub matrix: Vec, + + /// Suite-specific configuration (passed through to adapter). + #[serde(default = "default_suite_config")] + pub suite_config: toml::Value, +} + +/// A single model/config combination to benchmark. +#[derive(Debug, Clone, Deserialize)] +pub struct MatrixEntry { + /// Label for this configuration (used in results). + pub label: String, + + /// Model identifier. + #[serde(default)] + pub model: Option, + + /// Optional tool allowlist. If set, only these tools are available. + #[serde(default)] + pub tools: Option>, +} + +impl BenchConfig { + /// Load from a TOML file. + pub fn from_file(path: &Path) -> Result { + if !path.exists() { + return Err(BenchError::ConfigNotFound { + path: path.to_path_buf(), + }); + } + let content = std::fs::read_to_string(path)?; + let config: BenchConfig = toml::from_str(&content)?; + Ok(config) + } + + /// Create a minimal config for when no config file is provided. + /// Uses defaults and optional CLI overrides. + pub fn minimal(model: Option) -> Self { + let label = model.as_deref().unwrap_or("default").to_string(); + Self { + results_dir: default_results_dir(), + task_timeout: default_task_timeout(), + max_iterations: default_max_iterations(), + parallelism: default_parallelism(), + matrix: vec![MatrixEntry { + label, + model, + tools: None, + }], + suite_config: toml::Value::Table(toml::map::Map::new()), + } + } + + /// Get the suite_config as a generic map for adapter use. + pub fn suite_config_map(&self) -> toml::map::Map { + match &self.suite_config { + toml::Value::Table(map) => map.clone(), + _ => toml::map::Map::new(), + } + } + + /// Get a string value from suite_config. + pub fn suite_config_str(&self, key: &str) -> Option { + self.suite_config_map() + .get(key) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + } +} + +fn default_suite_config() -> toml::Value { + toml::Value::Table(toml::map::Map::new()) +} + +fn default_results_dir() -> PathBuf { + PathBuf::from("./bench-results") +} + +fn default_task_timeout() -> Duration { + Duration::from_secs(300) +} + +fn default_max_iterations() -> usize { + 15 +} + +fn default_parallelism() -> usize { + 1 +} + +/// Deserialize a duration from a string like "300s", "5m", etc. +fn deserialize_duration<'de, D>(deserializer: D) -> Result +where + D: serde::Deserializer<'de>, +{ + let s = String::deserialize(deserializer)?; + parse_duration(&s).map_err(serde::de::Error::custom) +} + +fn parse_duration(s: &str) -> Result { + let s = s.trim(); + if let Some(secs) = s.strip_suffix('s') { + secs.trim() + .parse::() + .map(Duration::from_secs) + .map_err(|e| format!("invalid seconds: {e}")) + } else if let Some(mins) = s.strip_suffix('m') { + mins.trim() + .parse::() + .map(|m| Duration::from_secs(m * 60)) + .map_err(|e| format!("invalid minutes: {e}")) + } else { + // Assume seconds if no suffix + s.parse::() + .map(Duration::from_secs) + .map_err(|e| format!("invalid duration '{s}': {e}")) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_duration() { + assert_eq!(parse_duration("300s").unwrap(), Duration::from_secs(300)); + assert_eq!(parse_duration("5m").unwrap(), Duration::from_secs(300)); + assert_eq!(parse_duration("60").unwrap(), Duration::from_secs(60)); + } + + #[test] + fn test_minimal_config() { + let config = BenchConfig::minimal(Some("test-model".to_string())); + assert_eq!(config.matrix.len(), 1); + assert_eq!(config.matrix[0].label, "test-model"); + assert_eq!(config.max_iterations, 15); + assert_eq!(config.parallelism, 1); + } + + #[test] + fn test_config_from_toml() { + let toml_str = r#" +results_dir = "./my-results" +task_timeout = "60s" +max_iterations = 10 +parallelism = 2 + +[[matrix]] +label = "fast" +model = "gpt-4o-mini" + +[[matrix]] +label = "full" +model = "claude-3-5-sonnet" +tools = ["echo", "time"] + +[suite_config] +dataset_path = "./data/test.jsonl" +"#; + let config: BenchConfig = toml::from_str(toml_str).unwrap(); + assert_eq!(config.results_dir, PathBuf::from("./my-results")); + assert_eq!(config.task_timeout, Duration::from_secs(60)); + assert_eq!(config.max_iterations, 10); + assert_eq!(config.parallelism, 2); + assert_eq!(config.matrix.len(), 2); + assert_eq!(config.matrix[1].tools.as_ref().unwrap().len(), 2); + assert_eq!( + config.suite_config_str("dataset_path").unwrap(), + "./data/test.jsonl" + ); + } +} diff --git a/benchmarks/src/error.rs b/benchmarks/src/error.rs new file mode 100644 index 00000000..1eb27f22 --- /dev/null +++ b/benchmarks/src/error.rs @@ -0,0 +1,40 @@ +use std::path::PathBuf; + +#[derive(Debug, thiserror::Error)] +pub enum BenchError { + #[error("Config error: {0}")] + Config(String), + + #[error("Config file not found: {path}")] + ConfigNotFound { path: PathBuf }, + + #[error("Suite {name} not found. Available: {available}")] + SuiteNotFound { name: String, available: String }, + + #[error("Task {task_id} failed: {reason}")] + TaskFailed { task_id: String, reason: String }, + + #[error("Timeout after {seconds}s for task {task_id}")] + Timeout { task_id: String, seconds: u64 }, + + #[error("Scoring error for task {task_id}: {reason}")] + Scoring { task_id: String, reason: String }, + + #[error("IO error: {0}")] + Io(#[from] std::io::Error), + + #[error("JSON error: {0}")] + Json(#[from] serde_json::Error), + + #[error("TOML parse error: {0}")] + Toml(#[from] toml::de::Error), + + #[error("Agent error: {0}")] + Agent(#[from] ironclaw::Error), + + #[error("Results directory error: {0}")] + ResultsDir(String), + + #[error("Resume failed: no completed tasks found in {path}")] + ResumeEmpty { path: PathBuf }, +} diff --git a/benchmarks/src/instrumented_llm.rs b/benchmarks/src/instrumented_llm.rs new file mode 100644 index 00000000..2e4e920a --- /dev/null +++ b/benchmarks/src/instrumented_llm.rs @@ -0,0 +1,248 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::time::Instant; + +use async_trait::async_trait; +use rust_decimal::Decimal; +use tokio::sync::Mutex; + +use ironclaw::error::LlmError; +use ironclaw::llm::{ + CompletionRequest, CompletionResponse, LlmProvider, ToolCompletionRequest, + ToolCompletionResponse, +}; + +/// Recorded metrics from a single LLM call. +#[derive(Debug, Clone)] +pub struct LlmCallRecord { + pub input_tokens: u32, + pub output_tokens: u32, + pub duration_ms: u64, + pub had_tool_calls: bool, +} + +/// Wraps an `LlmProvider` to record per-call metrics. +/// +/// The wrapper is transparent to the agent: it delegates every call +/// to the inner provider and captures token counts and timings. +pub struct InstrumentedLlm { + inner: Arc, + records: Mutex>, + total_input_tokens: AtomicU32, + total_output_tokens: AtomicU32, + call_count: AtomicU32, +} + +impl InstrumentedLlm { + pub fn new(inner: Arc) -> Self { + Self { + inner, + records: Mutex::new(Vec::new()), + total_input_tokens: AtomicU32::new(0), + total_output_tokens: AtomicU32::new(0), + call_count: AtomicU32::new(0), + } + } + + /// Take all recorded call metrics, clearing the internal buffer. + pub async fn take_records(&self) -> Vec { + let mut records = self.records.lock().await; + std::mem::take(&mut *records) + } + + /// Snapshot of total tokens without clearing. + pub fn total_input_tokens(&self) -> u32 { + self.total_input_tokens.load(Ordering::Relaxed) + } + + pub fn total_output_tokens(&self) -> u32 { + self.total_output_tokens.load(Ordering::Relaxed) + } + + pub fn call_count(&self) -> u32 { + self.call_count.load(Ordering::Relaxed) + } + + /// Estimated cost using the inner provider's cost-per-token rates. + pub fn estimated_cost(&self) -> f64 { + let (input_rate, output_rate) = self.inner.cost_per_token(); + let input_cost = + input_rate * Decimal::from(self.total_input_tokens.load(Ordering::Relaxed)); + let output_cost = + output_rate * Decimal::from(self.total_output_tokens.load(Ordering::Relaxed)); + let total = input_cost + output_cost; + // Convert Decimal to f64 for the trace (benchmarks don't need exact precision) + total.to_string().parse::().unwrap_or(0.0) + } + + /// Reset all counters and records. + pub async fn reset(&self) { + self.records.lock().await.clear(); + self.total_input_tokens.store(0, Ordering::Relaxed); + self.total_output_tokens.store(0, Ordering::Relaxed); + self.call_count.store(0, Ordering::Relaxed); + } + + async fn record( + &self, + input_tokens: u32, + output_tokens: u32, + duration_ms: u64, + had_tool_calls: bool, + ) { + self.total_input_tokens + .fetch_add(input_tokens, Ordering::Relaxed); + self.total_output_tokens + .fetch_add(output_tokens, Ordering::Relaxed); + self.call_count.fetch_add(1, Ordering::Relaxed); + self.records.lock().await.push(LlmCallRecord { + input_tokens, + output_tokens, + duration_ms, + had_tool_calls, + }); + } +} + +#[async_trait] +impl LlmProvider for InstrumentedLlm { + fn model_name(&self) -> &str { + self.inner.model_name() + } + + fn cost_per_token(&self) -> (Decimal, Decimal) { + self.inner.cost_per_token() + } + + async fn complete(&self, request: CompletionRequest) -> Result { + let start = Instant::now(); + let response = self.inner.complete(request).await?; + let elapsed = start.elapsed().as_millis() as u64; + self.record( + response.input_tokens, + response.output_tokens, + elapsed, + false, + ) + .await; + Ok(response) + } + + async fn complete_with_tools( + &self, + request: ToolCompletionRequest, + ) -> Result { + let start = Instant::now(); + let response = self.inner.complete_with_tools(request).await?; + let elapsed = start.elapsed().as_millis() as u64; + let had_tool_calls = !response.tool_calls.is_empty(); + self.record( + response.input_tokens, + response.output_tokens, + elapsed, + had_tool_calls, + ) + .await; + Ok(response) + } + + async fn list_models(&self) -> Result, LlmError> { + self.inner.list_models().await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use ironclaw::llm::{ChatMessage, CompletionRequest, CompletionResponse, FinishReason}; + + /// Fake LLM that returns a canned response with known token counts. + struct FakeLlm; + + #[async_trait] + impl LlmProvider for FakeLlm { + fn model_name(&self) -> &str { + "fake-model" + } + + fn cost_per_token(&self) -> (Decimal, Decimal) { + ( + Decimal::new(3, 6), // $0.000003 per input token + Decimal::new(15, 6), // $0.000015 per output token + ) + } + + async fn complete( + &self, + _request: CompletionRequest, + ) -> Result { + Ok(CompletionResponse { + content: "test response".to_string(), + input_tokens: 100, + output_tokens: 50, + finish_reason: FinishReason::Stop, + }) + } + + async fn complete_with_tools( + &self, + _request: ToolCompletionRequest, + ) -> Result { + Ok(ToolCompletionResponse { + content: Some("tool response".to_string()), + tool_calls: vec![], + input_tokens: 200, + output_tokens: 100, + finish_reason: FinishReason::Stop, + }) + } + } + + #[tokio::test] + async fn test_instrumented_records_metrics() { + let inner = Arc::new(FakeLlm); + let instrumented = InstrumentedLlm::new(inner); + + let request = CompletionRequest::new(vec![ChatMessage::user("hello")]); + let _ = instrumented.complete(request).await.unwrap(); + + assert_eq!(instrumented.call_count(), 1); + assert_eq!(instrumented.total_input_tokens(), 100); + assert_eq!(instrumented.total_output_tokens(), 50); + + let records = instrumented.take_records().await; + assert_eq!(records.len(), 1); + assert_eq!(records[0].input_tokens, 100); + assert!(!records[0].had_tool_calls); + } + + #[tokio::test] + async fn test_instrumented_cost_calculation() { + let inner = Arc::new(FakeLlm); + let instrumented = InstrumentedLlm::new(inner); + + let request = CompletionRequest::new(vec![ChatMessage::user("hello")]); + let _ = instrumented.complete(request).await.unwrap(); + + // 100 * 0.000003 + 50 * 0.000015 = 0.0003 + 0.00075 = 0.00105 + let cost = instrumented.estimated_cost(); + assert!((cost - 0.00105).abs() < 0.0001); + } + + #[tokio::test] + async fn test_instrumented_reset() { + let inner = Arc::new(FakeLlm); + let instrumented = InstrumentedLlm::new(inner); + + let request = CompletionRequest::new(vec![ChatMessage::user("hello")]); + let _ = instrumented.complete(request).await.unwrap(); + assert_eq!(instrumented.call_count(), 1); + + instrumented.reset().await; + assert_eq!(instrumented.call_count(), 0); + assert_eq!(instrumented.total_input_tokens(), 0); + + let records = instrumented.take_records().await; + assert!(records.is_empty()); + } +} diff --git a/benchmarks/src/main.rs b/benchmarks/src/main.rs new file mode 100644 index 00000000..0e1afe36 --- /dev/null +++ b/benchmarks/src/main.rs @@ -0,0 +1,314 @@ +mod adapters; +mod channel; +mod config; +mod error; +mod instrumented_llm; +mod results; +mod runner; +mod scoring; +mod suite; + +use std::path::PathBuf; +use std::sync::Arc; + +use clap::{Parser, Subcommand}; +use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt}; +use uuid::Uuid; + +use crate::config::BenchConfig; + +#[derive(Parser)] +#[command(name = "ironclaw-bench", about = "IronClaw benchmarking harness")] +struct Cli { + #[command(subcommand)] + command: Commands, +} + +#[derive(Subcommand)] +enum Commands { + /// Run a benchmark suite. + Run { + /// Suite to run (custom, gaia, tau_bench, swe_bench). + #[arg(long)] + suite: String, + + /// Path to bench config TOML. + #[arg(long)] + config: Option, + + /// Override model for all matrix entries. + #[arg(long)] + model: Option, + + /// Max tasks to run in parallel. + #[arg(long)] + parallelism: Option, + + /// Sample N tasks from the suite (for quick testing). + #[arg(long)] + sample: Option, + + /// Only run these task IDs (comma-separated). + #[arg(long, value_delimiter = ',')] + task_ids: Option>, + + /// Only run tasks with these tags (comma-separated). + #[arg(long, value_delimiter = ',')] + tags: Option>, + + /// Per-task timeout in seconds. + #[arg(long)] + timeout_secs: Option, + + /// Override results directory. + #[arg(long)] + results_dir: Option, + + /// Resume a previous run by ID. + #[arg(long)] + resume: Option, + }, + + /// Show results for a run. + Results { + /// Run ID or "latest". + #[arg(default_value = "latest")] + run_id: String, + + /// Output format. + #[arg(long, default_value = "table")] + format: ResultsFormat, + + /// Override results directory. + #[arg(long)] + results_dir: Option, + }, + + /// Compare two runs. + Compare { + /// Baseline run ID. + baseline: Uuid, + + /// Comparison run ID. + comparison: Uuid, + + /// Override results directory. + #[arg(long)] + results_dir: Option, + }, + + /// List available benchmark suites. + List, +} + +#[derive(Clone, Debug, clap::ValueEnum)] +enum ResultsFormat { + Table, + Json, + Csv, +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let cli = Cli::parse(); + + tracing_subscriber::registry() + .with( + EnvFilter::try_from_default_env() + .unwrap_or_else(|_| EnvFilter::new("ironclaw_bench=info,ironclaw=warn")), + ) + .with(tracing_subscriber::fmt::layer().with_target(false)) + .init(); + + match cli.command { + Commands::List => { + println!("Available benchmark suites:\n"); + for (id, desc) in adapters::KNOWN_SUITES { + println!(" {:<15} {}", id, desc); + } + println!(); + } + Commands::Run { + suite, + config: config_path, + model, + parallelism, + sample, + task_ids, + tags, + timeout_secs, + results_dir, + resume, + } => { + // Load or create config + let mut bench_config = if let Some(ref path) = config_path { + BenchConfig::from_file(path)? + } else { + BenchConfig::minimal(model.clone()) + }; + + // Apply CLI overrides + if let Some(p) = parallelism { + bench_config.parallelism = p; + } + if let Some(t) = timeout_secs { + bench_config.task_timeout = std::time::Duration::from_secs(t); + } + if let Some(ref dir) = results_dir { + bench_config.results_dir = dir.clone(); + } + + // If model override specified and we have matrix entries, update them + if let Some(ref m) = model { + for entry in &mut bench_config.matrix { + entry.model = Some(m.clone()); + } + } + + // Create suite + let bench_suite = adapters::create_suite(&suite, &bench_config)?; + + // Initialize ironclaw LLM provider + let ironclaw_config = ironclaw::Config::from_env().map_err(|e| { + anyhow::anyhow!( + "Failed to load ironclaw config: {}. Make sure .env is configured.", + e + ) + })?; + + let session = ironclaw::llm::create_session_manager(ironclaw::llm::SessionConfig { + auth_base_url: ironclaw_config.llm.nearai.auth_base_url.clone(), + session_path: ironclaw_config.llm.nearai.session_path.clone(), + ..Default::default() + }) + .await; + session.ensure_authenticated().await?; + + let llm = ironclaw::llm::create_llm_provider(&ironclaw_config.llm, session)?; + let safety = Arc::new(ironclaw::safety::SafetyLayer::new(&ironclaw_config.safety)); + + let runner = runner::BenchRunner::new(bench_suite, bench_config.clone(), llm, safety); + + // Run for each matrix entry + for matrix_entry in &bench_config.matrix { + let run_id = runner + .run( + matrix_entry, + sample, + task_ids.as_deref(), + tags.as_deref(), + resume, + ) + .await?; + println!("Run complete: {}", run_id); + } + } + Commands::Results { + run_id, + format, + results_dir, + } => { + let base = results_dir.unwrap_or_else(|| PathBuf::from("./bench-results")); + let uuid = if run_id == "latest" { + results::find_latest_run(&base)? + .ok_or_else(|| anyhow::anyhow!("No runs found in {}", base.display()))? + } else { + Uuid::parse_str(&run_id)? + }; + + let json_path = results::run_json_path(&base, uuid); + let jsonl_path = results::tasks_jsonl_path(&base, uuid); + + let run = results::read_run_result(&json_path)?; + let tasks = results::read_task_results(&jsonl_path)?; + + match format { + ResultsFormat::Table => { + results::print_results_table(&tasks, &run); + } + ResultsFormat::Json => { + let output = serde_json::json!({ + "run": run, + "tasks": tasks, + }); + println!("{}", serde_json::to_string_pretty(&output)?); + } + ResultsFormat::Csv => { + println!("task_id,score,label,tokens,cost,turns,time_s"); + for task in &tasks { + println!( + "{},{:.3},{},{},{:.4},{},{:.1}", + task.task_id, + task.score.value, + task.score.label, + task.trace.input_tokens + task.trace.output_tokens, + task.trace.estimated_cost_usd, + task.trace.turns, + task.trace.wall_time_ms as f64 / 1000.0, + ); + } + } + } + } + Commands::Compare { + baseline, + comparison, + results_dir, + } => { + let base = results_dir.unwrap_or_else(|| PathBuf::from("./bench-results")); + + let baseline_run = results::read_run_result(&results::run_json_path(&base, baseline))?; + let comparison_run = + results::read_run_result(&results::run_json_path(&base, comparison))?; + + println!("\nComparison: {} vs {}\n", baseline, comparison); + println!( + "{:<20} {:>12} {:>12} {:>10}", + "Metric", "Baseline", "Comparison", "Delta" + ); + println!("{}", "-".repeat(58)); + + let pass_delta = comparison_run.pass_rate - baseline_run.pass_rate; + println!( + "{:<20} {:>11.1}% {:>11.1}% {:>+9.1}%", + "Pass rate", + baseline_run.pass_rate * 100.0, + comparison_run.pass_rate * 100.0, + pass_delta * 100.0, + ); + + let score_delta = comparison_run.avg_score - baseline_run.avg_score; + println!( + "{:<20} {:>12.3} {:>12.3} {:>+10.3}", + "Avg score", baseline_run.avg_score, comparison_run.avg_score, score_delta, + ); + + let cost_delta = comparison_run.total_cost_usd - baseline_run.total_cost_usd; + println!( + "{:<20} {:>11.4}$ {:>11.4}$ {:>+9.4}$", + "Total cost", + baseline_run.total_cost_usd, + comparison_run.total_cost_usd, + cost_delta, + ); + + let time_b = baseline_run.total_wall_time_ms as f64 / 1000.0; + let time_c = comparison_run.total_wall_time_ms as f64 / 1000.0; + println!( + "{:<20} {:>11.1}s {:>11.1}s {:>+9.1}s", + "Total time", + time_b, + time_c, + time_c - time_b, + ); + + println!( + "{:<20} {:>12} {:>12}", + "Model", baseline_run.model, comparison_run.model, + ); + println!(); + } + } + + Ok(()) +} diff --git a/benchmarks/src/results.rs b/benchmarks/src/results.rs new file mode 100644 index 00000000..7aa33363 --- /dev/null +++ b/benchmarks/src/results.rs @@ -0,0 +1,392 @@ +use std::collections::HashSet; +use std::io::{BufRead, Write}; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use uuid::Uuid; + +use crate::error::BenchError; +use crate::suite::BenchScore; + +/// Metrics from a single task run: LLM usage, timing, tool calls. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct Trace { + pub wall_time_ms: u64, + pub llm_calls: u32, + pub input_tokens: u32, + pub output_tokens: u32, + pub estimated_cost_usd: f64, + pub tool_calls: Vec, + pub turns: u32, + pub hit_iteration_limit: bool, + pub hit_timeout: bool, +} + +impl Trace { + pub fn wall_time(&self) -> Duration { + Duration::from_millis(self.wall_time_ms) + } +} + +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct TraceToolCall { + pub name: String, + pub duration_ms: u64, + pub success: bool, +} + +/// Result of running a single benchmark task. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct TaskResult { + pub task_id: String, + pub suite_id: String, + pub score: BenchScore, + pub trace: Trace, + pub response: String, + pub started_at: DateTime, + pub finished_at: DateTime, + pub config_label: String, + #[serde(default)] + pub error: Option, +} + +/// Aggregate results for a full benchmark run. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct RunResult { + pub run_id: Uuid, + pub suite_id: String, + pub config_label: String, + pub model: String, + pub pass_rate: f64, + pub avg_score: f64, + pub total_tasks: usize, + pub completed_tasks: usize, + pub total_cost_usd: f64, + pub total_wall_time_ms: u64, + pub started_at: DateTime, + pub finished_at: DateTime, +} + +impl RunResult { + /// Build aggregate from individual task results. + pub fn from_tasks( + run_id: Uuid, + suite_id: &str, + config_label: &str, + model: &str, + total_tasks: usize, + tasks: &[TaskResult], + started_at: DateTime, + ) -> Self { + let pass_count = tasks.iter().filter(|t| t.score.value >= 1.0).count(); + let pass_rate = if tasks.is_empty() { + 0.0 + } else { + pass_count as f64 / tasks.len() as f64 + }; + let avg_score = if tasks.is_empty() { + 0.0 + } else { + tasks.iter().map(|t| t.score.value).sum::() / tasks.len() as f64 + }; + let total_cost: f64 = tasks.iter().map(|t| t.trace.estimated_cost_usd).sum(); + let total_wall: u64 = tasks.iter().map(|t| t.trace.wall_time_ms).sum(); + + Self { + run_id, + suite_id: suite_id.to_string(), + config_label: config_label.to_string(), + model: model.to_string(), + pass_rate, + avg_score, + total_tasks, + completed_tasks: tasks.len(), + total_cost_usd: total_cost, + total_wall_time_ms: total_wall, + started_at, + finished_at: Utc::now(), + } + } + + pub fn total_wall_time(&self) -> Duration { + Duration::from_millis(self.total_wall_time_ms) + } +} + +/// Append a single task result as one JSON line to the JSONL file. +pub fn append_task_result(path: &Path, result: &TaskResult) -> Result<(), BenchError> { + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(path)?; + let line = serde_json::to_string(result)?; + writeln!(file, "{line}")?; + Ok(()) +} + +/// Read all task results from a JSONL file. +pub fn read_task_results(path: &Path) -> Result, BenchError> { + if !path.exists() { + return Ok(Vec::new()); + } + let file = std::fs::File::open(path)?; + let reader = std::io::BufReader::new(file); + let mut results = Vec::new(); + for line in reader.lines() { + let line = line?; + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + let result: TaskResult = serde_json::from_str(trimmed)?; + results.push(result); + } + Ok(results) +} + +/// Write the aggregate run result as JSON. +pub fn write_run_result(path: &Path, result: &RunResult) -> Result<(), BenchError> { + let json = serde_json::to_string_pretty(result)?; + std::fs::write(path, json)?; + Ok(()) +} + +/// Read the aggregate run result from JSON. +pub fn read_run_result(path: &Path) -> Result { + let json = std::fs::read_to_string(path)?; + let result: RunResult = serde_json::from_str(&json)?; + Ok(result) +} + +/// Get the set of already-completed task IDs from a JSONL file (for resume). +pub fn completed_task_ids(path: &Path) -> Result, BenchError> { + let results = read_task_results(path)?; + Ok(results.into_iter().map(|r| r.task_id).collect()) +} + +/// Get the results directory for a specific run. +pub fn run_dir(base: &Path, run_id: Uuid) -> PathBuf { + base.join(run_id.to_string()) +} + +/// Get the tasks JSONL path for a run. +pub fn tasks_jsonl_path(base: &Path, run_id: Uuid) -> PathBuf { + run_dir(base, run_id).join("tasks.jsonl") +} + +/// Get the run JSON path for a run. +pub fn run_json_path(base: &Path, run_id: Uuid) -> PathBuf { + run_dir(base, run_id).join("run.json") +} + +/// Find the latest run directory (by modification time). +pub fn find_latest_run(base: &Path) -> Result, BenchError> { + if !base.exists() { + return Ok(None); + } + let mut entries: Vec<_> = std::fs::read_dir(base)? + .filter_map(|e| e.ok()) + .filter(|e| e.file_type().map(|ft| ft.is_dir()).unwrap_or(false)) + .filter_map(|e| { + let name = e.file_name().to_string_lossy().to_string(); + let uuid = Uuid::parse_str(&name).ok()?; + let modified = e.metadata().ok()?.modified().ok()?; + Some((uuid, modified)) + }) + .collect(); + entries.sort_by(|a, b| b.1.cmp(&a.1)); + Ok(entries.first().map(|(uuid, _)| *uuid)) +} + +/// Print a summary table of task results. +pub fn print_results_table(tasks: &[TaskResult], run: &RunResult) { + println!(); + println!( + "Run: {} | Suite: {} | Config: {} | Model: {}", + run.run_id, run.suite_id, run.config_label, run.model + ); + println!( + "Pass rate: {:.1}% | Avg score: {:.3} | Tasks: {}/{} | Cost: ${:.4} | Time: {:.1}s", + run.pass_rate * 100.0, + run.avg_score, + run.completed_tasks, + run.total_tasks, + run.total_cost_usd, + run.total_wall_time_ms as f64 / 1000.0, + ); + println!(); + + // Header + println!( + "{:<30} {:>6} {:>7} {:>8} {:>10} {:>6} {:>8}", + "Task ID", "Score", "Label", "Tokens", "Cost", "Turns", "Time" + ); + println!("{}", "-".repeat(80)); + + for task in tasks { + let total_tokens = task.trace.input_tokens + task.trace.output_tokens; + let task_id_display = if task.task_id.len() > 28 { + format!("{}...", &task.task_id[..25]) + } else { + task.task_id.clone() + }; + println!( + "{:<30} {:>6.3} {:>7} {:>8} {:>10.4} {:>6} {:>7.1}s", + task_id_display, + task.score.value, + task.score.label, + total_tokens, + task.trace.estimated_cost_usd, + task.trace.turns, + task.trace.wall_time_ms as f64 / 1000.0, + ); + } + println!(); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_run_result_from_tasks() { + let tasks = vec![ + TaskResult { + task_id: "t1".to_string(), + suite_id: "custom".to_string(), + score: BenchScore { + value: 1.0, + label: "pass".to_string(), + details: None, + }, + trace: Trace { + wall_time_ms: 1000, + llm_calls: 2, + input_tokens: 100, + output_tokens: 50, + estimated_cost_usd: 0.01, + tool_calls: vec![], + turns: 1, + hit_iteration_limit: false, + hit_timeout: false, + }, + response: "answer".to_string(), + started_at: Utc::now(), + finished_at: Utc::now(), + config_label: "default".to_string(), + error: None, + }, + TaskResult { + task_id: "t2".to_string(), + suite_id: "custom".to_string(), + score: BenchScore { + value: 0.0, + label: "fail".to_string(), + details: Some("wrong".to_string()), + }, + trace: Trace { + wall_time_ms: 2000, + llm_calls: 3, + input_tokens: 200, + output_tokens: 100, + estimated_cost_usd: 0.02, + tool_calls: vec![], + turns: 2, + hit_iteration_limit: false, + hit_timeout: false, + }, + response: "wrong answer".to_string(), + started_at: Utc::now(), + finished_at: Utc::now(), + config_label: "default".to_string(), + error: None, + }, + ]; + + let run = RunResult::from_tasks( + Uuid::new_v4(), + "custom", + "default", + "test-model", + 2, + &tasks, + Utc::now(), + ); + + assert_eq!(run.pass_rate, 0.5); + assert_eq!(run.avg_score, 0.5); + assert_eq!(run.total_tasks, 2); + assert_eq!(run.completed_tasks, 2); + assert!((run.total_cost_usd - 0.03).abs() < f64::EPSILON); + assert_eq!(run.total_wall_time_ms, 3000); + } + + #[test] + fn test_jsonl_roundtrip() { + let dir = tempfile::tempdir().expect("temp dir"); + let path = dir.path().join("tasks.jsonl"); + + let result = TaskResult { + task_id: "round-trip-test".to_string(), + suite_id: "custom".to_string(), + score: BenchScore::pass(), + trace: Trace { + wall_time_ms: 500, + llm_calls: 1, + input_tokens: 10, + output_tokens: 5, + estimated_cost_usd: 0.001, + tool_calls: vec![], + turns: 1, + hit_iteration_limit: false, + hit_timeout: false, + }, + response: "hello".to_string(), + started_at: Utc::now(), + finished_at: Utc::now(), + config_label: "test".to_string(), + error: None, + }; + + append_task_result(&path, &result).expect("append"); + append_task_result(&path, &result).expect("append"); + + let loaded = read_task_results(&path).expect("read"); + assert_eq!(loaded.len(), 2); + assert_eq!(loaded[0].task_id, "round-trip-test"); + } + + #[test] + fn test_completed_task_ids() { + let dir = tempfile::tempdir().expect("temp dir"); + let path = dir.path().join("tasks.jsonl"); + + let result = TaskResult { + task_id: "unique-id-1".to_string(), + suite_id: "custom".to_string(), + score: BenchScore::pass(), + trace: Trace { + wall_time_ms: 100, + llm_calls: 1, + input_tokens: 10, + output_tokens: 5, + estimated_cost_usd: 0.0, + tool_calls: vec![], + turns: 1, + hit_iteration_limit: false, + hit_timeout: false, + }, + response: "x".to_string(), + started_at: Utc::now(), + finished_at: Utc::now(), + config_label: "test".to_string(), + error: None, + }; + append_task_result(&path, &result).expect("append"); + + let ids = completed_task_ids(&path).expect("ids"); + assert!(ids.contains("unique-id-1")); + assert!(!ids.contains("unique-id-2")); + } +} diff --git a/benchmarks/src/runner.rs b/benchmarks/src/runner.rs new file mode 100644 index 00000000..68e839e0 --- /dev/null +++ b/benchmarks/src/runner.rs @@ -0,0 +1,460 @@ +use std::collections::HashSet; +use std::sync::Arc; +use std::time::Instant; + +use chrono::Utc; +use tokio::sync::Mutex; +use uuid::Uuid; + +use ironclaw::agent::{Agent, AgentDeps}; +use ironclaw::channels::{ChannelManager, IncomingMessage}; +use ironclaw::config::AgentConfig; +use ironclaw::llm::LlmProvider; +use ironclaw::safety::SafetyLayer; +use ironclaw::tools::ToolRegistry; + +use crate::channel::BenchChannel; +use crate::config::{BenchConfig, MatrixEntry}; +use crate::error::BenchError; +use crate::instrumented_llm::InstrumentedLlm; +use crate::results::{ + RunResult, TaskResult, Trace, append_task_result, completed_task_ids, run_dir, run_json_path, + tasks_jsonl_path, write_run_result, +}; +use crate::suite::{BenchSuite, BenchTask, ConversationTurn, TaskSubmission, TurnRole}; + +/// Parameters for running a single task in isolation. +struct TaskRunParams<'a> { + task: &'a BenchTask, + suite_id: &'a str, + config_label: &'a str, + llm: Arc, + safety: Arc, + timeout: std::time::Duration, + additional_tools: &'a [Arc], +} + +/// Orchestrates benchmark execution: loads tasks, runs agent per task, +/// scores results, writes JSONL output. +pub struct BenchRunner { + suite: Box, + config: BenchConfig, + llm: Arc, + safety: Arc, +} + +impl BenchRunner { + pub fn new( + suite: Box, + config: BenchConfig, + llm: Arc, + safety: Arc, + ) -> Self { + Self { + suite, + config, + llm, + safety, + } + } + + /// Run the benchmark for one matrix entry. + /// + /// Returns the run_id for result retrieval. + pub async fn run( + &self, + matrix: &MatrixEntry, + sample: Option, + task_filter: Option<&[String]>, + tag_filter: Option<&[String]>, + resume_run_id: Option, + ) -> Result { + let run_id = resume_run_id.unwrap_or_else(Uuid::new_v4); + let results_base = &self.config.results_dir; + let dir = run_dir(results_base, run_id); + std::fs::create_dir_all(&dir)?; + + let jsonl_path = tasks_jsonl_path(results_base, run_id); + let json_path = run_json_path(results_base, run_id); + + // Load completed task IDs for resume support + let completed: HashSet = if resume_run_id.is_some() { + completed_task_ids(&jsonl_path)? + } else { + HashSet::new() + }; + + if !completed.is_empty() { + tracing::info!( + "Resuming run {}: {} tasks already completed", + run_id, + completed.len() + ); + } + + // Load and filter tasks + let mut tasks = self.suite.load_tasks().await?; + + if let Some(ids) = task_filter { + let id_set: HashSet<&str> = ids.iter().map(|s| s.as_str()).collect(); + tasks.retain(|t| id_set.contains(t.id.as_str())); + } + + if let Some(tags) = tag_filter { + let tag_set: HashSet<&str> = tags.iter().map(|s| s.as_str()).collect(); + tasks.retain(|t| t.tags.iter().any(|tag| tag_set.contains(tag.as_str()))); + } + + // Filter out already-completed tasks + tasks.retain(|t| !completed.contains(&t.id)); + + // Sample if requested + if let Some(n) = sample { + tasks.truncate(n); + } + + let total_tasks = tasks.len() + completed.len(); + tracing::info!( + "Running {} tasks for suite '{}' (config: {}, run: {})", + tasks.len(), + self.suite.id(), + matrix.label, + run_id + ); + + let started_at = Utc::now(); + let all_results: Arc>> = + Arc::new(Mutex::new(Vec::with_capacity(tasks.len()))); + + if self.config.parallelism <= 1 { + // Sequential execution + let additional_tools = self.suite.additional_tools(); + for (i, task) in tasks.iter().enumerate() { + tracing::info!( + "[{}/{}] Running task: {}", + i + 1 + completed.len(), + total_tasks, + task.id + ); + let params = TaskRunParams { + task, + suite_id: self.suite.id(), + config_label: &matrix.label, + llm: Arc::clone(&self.llm), + safety: Arc::clone(&self.safety), + timeout: task.timeout.unwrap_or(self.config.task_timeout), + additional_tools: &additional_tools, + }; + let result = run_task_isolated(params).await; + append_task_result(&jsonl_path, &result)?; + all_results.lock().await.push(result); + } + } else { + // Parallel execution with bounded concurrency + let semaphore = Arc::new(tokio::sync::Semaphore::new(self.config.parallelism)); + + let mut handles = Vec::new(); + for (i, task) in tasks.into_iter().enumerate() { + let sem = Arc::clone(&semaphore); + let suite_id = self.suite.id().to_string(); + let config_label = matrix.label.clone(); + let llm = Arc::clone(&self.llm); + let safety = Arc::clone(&self.safety); + let timeout = task.timeout.unwrap_or(self.config.task_timeout); + let results_ref = Arc::clone(&all_results); + let jsonl = jsonl_path.clone(); + let completed_count = completed.len(); + let total = total_tasks; + let additional_tools = self.suite.additional_tools(); + + handles.push(tokio::spawn(async move { + let _permit = sem.acquire().await.expect("semaphore closed"); + tracing::info!( + "[{}/{}] Running task: {}", + i + 1 + completed_count, + total, + task.id + ); + let params = TaskRunParams { + task: &task, + suite_id: &suite_id, + config_label: &config_label, + llm, + safety, + timeout, + additional_tools: &additional_tools, + }; + let result = run_task_isolated(params).await; + if let Err(e) = append_task_result(&jsonl, &result) { + tracing::error!("Failed to write result for {}: {}", task.id, e); + } + results_ref.lock().await.push(result); + })); + } + + for handle in handles { + if let Err(e) = handle.await { + tracing::error!("Task panicked: {}", e); + } + } + } + + // Score all results + let results = all_results.lock().await; + let mut scored: Vec = Vec::with_capacity(results.len()); + for result in results.iter() { + let task_opt = self + .suite + .load_tasks() + .await? + .into_iter() + .find(|t| t.id == result.task_id); + + if let Some(task) = task_opt { + let submission = TaskSubmission { + response: result.response.clone(), + conversation: vec![], + tool_calls: result + .trace + .tool_calls + .iter() + .map(|tc| tc.name.clone()) + .collect(), + }; + match self.suite.score(&task, &submission).await { + Ok(score) => { + let mut scored_result = result.clone(); + scored_result.score = score; + scored.push(scored_result); + } + Err(e) => { + tracing::warn!("Scoring failed for {}: {}", result.task_id, e); + scored.push(result.clone()); + } + } + } else { + scored.push(result.clone()); + } + } + + // Combine with any previously completed results for the aggregate + let mut all_for_aggregate = crate::results::read_task_results(&jsonl_path)?; + // De-duplicate (prefer the newer scored versions) + let scored_ids: HashSet = scored.iter().map(|r| r.task_id.clone()).collect(); + all_for_aggregate.retain(|r| !scored_ids.contains(&r.task_id)); + all_for_aggregate.extend(scored); + + let model_name = matrix.model.as_deref().unwrap_or(self.llm.model_name()); + + let run_result = RunResult::from_tasks( + run_id, + self.suite.id(), + &matrix.label, + model_name, + total_tasks, + &all_for_aggregate, + started_at, + ); + + write_run_result(&json_path, &run_result)?; + + tracing::info!( + "Run {} complete: {:.1}% pass rate, {:.3} avg score, ${:.4} cost", + run_id, + run_result.pass_rate * 100.0, + run_result.avg_score, + run_result.total_cost_usd, + ); + + Ok(run_id) + } +} + +/// Run a single benchmark task in complete isolation. +/// +/// Creates a fresh Agent + BenchChannel + InstrumentedLlm for the task, +/// injects the prompt, waits for the response, and returns the result. +async fn run_task_isolated(params: TaskRunParams<'_>) -> TaskResult { + let TaskRunParams { + task, + suite_id, + config_label, + llm, + safety, + timeout, + additional_tools, + } = params; + + let started_at = Utc::now(); + let start = Instant::now(); + + // Wrap LLM with instrumentation + let instrumented = Arc::new(InstrumentedLlm::new(llm)); + + // Create bench channel + let (bench_channel, msg_tx) = BenchChannel::new(); + let capture = bench_channel.capture(); + + // Build tool registry + let tools = Arc::new(ToolRegistry::new()); + tools.register_builtin_tools(); + + // Register additional suite-specific tools + for tool in additional_tools { + tools.register(Arc::clone(tool)).await; + } + + // Build agent config (minimal, headless) + let agent_config = AgentConfig { + name: format!("bench-{}", task.id), + max_parallel_jobs: 1, + job_timeout: timeout, + stuck_threshold: timeout, + repair_check_interval: timeout + std::time::Duration::from_secs(999), + max_repair_attempts: 0, + use_planning: false, + session_idle_timeout: timeout, + }; + + let deps = AgentDeps { + store: None, + llm: instrumented.clone() as Arc, + safety, + tools, + workspace: None, + extension_manager: None, + }; + + let mut channels = ChannelManager::new(); + channels.add(Box::new(bench_channel)); + + let agent = Agent::new(agent_config, deps, channels, None, None, None); + + // Build the full prompt with context + let full_prompt = if let Some(ref ctx) = task.context { + format!("{}\n\nContext:\n{}", task.prompt, ctx) + } else { + task.prompt.clone() + }; + + // Inject the task prompt + let incoming = IncomingMessage::new("bench", "bench-user", &full_prompt); + if msg_tx.send(incoming).await.is_err() { + return make_error_result( + task, + suite_id, + config_label, + started_at, + "failed to send prompt", + ); + } + + // Record prompt in conversation + { + let mut cap = capture.lock().await; + cap.conversation.push(ConversationTurn { + role: TurnRole::User, + content: full_prompt, + }); + } + + // Run agent with timeout. + // After the first response, send /quit to end the session. + let quit_tx = msg_tx.clone(); + let capture_for_quit = Arc::clone(&capture); + let quit_handle = tokio::spawn(async move { + // Poll for first response + loop { + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + let cap = capture_for_quit.lock().await; + if !cap.responses.is_empty() { + break; + } + } + // Give a small grace period for any final status events + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + let quit = IncomingMessage::new("bench", "bench-user", "/quit"); + let _ = quit_tx.send(quit).await; + }); + + let agent_result = tokio::time::timeout(timeout, agent.run()).await; + + quit_handle.abort(); + + let wall_time = start.elapsed(); + let hit_timeout = agent_result.is_err(); + + if let Ok(Err(e)) = &agent_result { + tracing::warn!("Agent error for task {}: {}", task.id, e); + } + + // Extract results from capture + let cap = capture.lock().await; + let response = cap.responses.last().cloned().unwrap_or_default(); + + let trace = Trace { + wall_time_ms: wall_time.as_millis() as u64, + llm_calls: instrumented.call_count(), + input_tokens: instrumented.total_input_tokens(), + output_tokens: instrumented.total_output_tokens(), + estimated_cost_usd: instrumented.estimated_cost(), + tool_calls: cap.tool_calls.clone(), + turns: cap.responses.len() as u32, + hit_iteration_limit: false, + hit_timeout, + }; + + let error = if hit_timeout { + Some(format!("timeout after {}s", timeout.as_secs())) + } else if let Ok(Err(e)) = &agent_result { + Some(e.to_string()) + } else { + None + }; + + TaskResult { + task_id: task.id.clone(), + suite_id: suite_id.to_string(), + score: crate::suite::BenchScore { + value: 0.0, + label: "pending".to_string(), + details: None, + }, + trace, + response, + started_at, + finished_at: Utc::now(), + config_label: config_label.to_string(), + error, + } +} + +fn make_error_result( + task: &BenchTask, + suite_id: &str, + config_label: &str, + started_at: chrono::DateTime, + reason: &str, +) -> TaskResult { + TaskResult { + task_id: task.id.clone(), + suite_id: suite_id.to_string(), + score: crate::suite::BenchScore::fail(reason), + trace: Trace { + wall_time_ms: 0, + llm_calls: 0, + input_tokens: 0, + output_tokens: 0, + estimated_cost_usd: 0.0, + tool_calls: vec![], + turns: 0, + hit_iteration_limit: false, + hit_timeout: false, + }, + response: String::new(), + started_at, + finished_at: Utc::now(), + config_label: config_label.to_string(), + error: Some(reason.to_string()), + } +} diff --git a/benchmarks/src/scoring.rs b/benchmarks/src/scoring.rs new file mode 100644 index 00000000..312899e3 --- /dev/null +++ b/benchmarks/src/scoring.rs @@ -0,0 +1,113 @@ +use regex::Regex; + +use crate::suite::BenchScore; + +/// Normalize an answer string for comparison: lowercase, trim whitespace, +/// strip trailing punctuation, collapse internal whitespace. +pub fn normalize_answer(s: &str) -> String { + let trimmed = s.trim().to_lowercase(); + let collapsed: String = trimmed.split_whitespace().collect::>().join(" "); + collapsed.trim_end_matches(['.', ',', ';', '!']).to_string() +} + +/// Exact match after normalization. +pub fn exact_match(expected: &str, actual: &str) -> BenchScore { + let norm_expected = normalize_answer(expected); + let norm_actual = normalize_answer(actual); + if norm_expected == norm_actual { + BenchScore::pass() + } else { + BenchScore::fail(format!( + "expected \"{norm_expected}\", got \"{norm_actual}\"" + )) + } +} + +/// Check if the actual answer contains the expected substring (normalized). +pub fn contains_match(expected_substring: &str, actual: &str) -> BenchScore { + let norm_expected = normalize_answer(expected_substring); + let norm_actual = normalize_answer(actual); + if norm_actual.contains(&norm_expected) { + BenchScore::pass() + } else { + BenchScore::fail(format!("response does not contain \"{norm_expected}\"")) + } +} + +/// Check if the actual answer matches a regex pattern. +pub fn regex_match(pattern: &str, actual: &str) -> BenchScore { + match Regex::new(pattern) { + Ok(re) => { + if re.is_match(actual) { + BenchScore::pass() + } else { + BenchScore::fail(format!("response does not match pattern /{pattern}/")) + } + } + Err(e) => BenchScore::fail(format!("invalid regex pattern: {e}")), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_normalize_answer() { + assert_eq!(normalize_answer(" Hello World. "), "hello world"); + assert_eq!(normalize_answer("Yes!"), "yes"); + assert_eq!(normalize_answer("42"), "42"); + assert_eq!(normalize_answer(" "), ""); + } + + #[test] + fn test_exact_match_pass() { + let score = exact_match("Hello World", " hello world. "); + assert_eq!(score.value, 1.0); + assert_eq!(score.label, "pass"); + } + + #[test] + fn test_exact_match_fail() { + let score = exact_match("hello", "world"); + assert_eq!(score.value, 0.0); + assert_eq!(score.label, "fail"); + } + + #[test] + fn test_contains_match_pass() { + let score = contains_match("world", "Hello World!"); + assert_eq!(score.value, 1.0); + } + + #[test] + fn test_contains_match_fail() { + let score = contains_match("xyz", "Hello World!"); + assert_eq!(score.value, 0.0); + } + + #[test] + fn test_regex_match_pass() { + let score = regex_match(r"\d{4}", "The year is 2024."); + assert_eq!(score.value, 1.0); + } + + #[test] + fn test_regex_match_fail() { + let score = regex_match(r"\d{4}", "No numbers here."); + assert_eq!(score.value, 0.0); + } + + #[test] + fn test_regex_match_invalid_pattern() { + let score = regex_match(r"[invalid", "anything"); + assert_eq!(score.value, 0.0); + assert!( + score + .details + .as_deref() + .unwrap_or("") + .contains("invalid regex") + ); + } +} diff --git a/benchmarks/src/suite.rs b/benchmarks/src/suite.rs new file mode 100644 index 00000000..2f433387 --- /dev/null +++ b/benchmarks/src/suite.rs @@ -0,0 +1,156 @@ +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; + +use crate::error::BenchError; + +/// A single task in a benchmark suite. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct BenchTask { + pub id: String, + pub prompt: String, + #[serde(default)] + pub context: Option, + #[serde(default)] + pub resources: Vec, + #[serde(default)] + pub tags: Vec, + #[serde(default)] + pub expected_turns: Option, + #[serde(default)] + pub timeout: Option, + #[serde(default)] + pub metadata: serde_json::Value, +} + +/// A resource attached to a benchmark task (file, URL, etc.). +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct TaskResource { + pub name: String, + pub path: String, + #[serde(default)] + pub resource_type: ResourceType, +} + +#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ResourceType { + #[default] + File, + Url, + Directory, +} + +/// What the agent produced for scoring. +#[derive(Debug, Clone)] +pub struct TaskSubmission { + pub response: String, + pub conversation: Vec, + pub tool_calls: Vec, +} + +/// A single turn in a multi-turn conversation. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct ConversationTurn { + pub role: TurnRole, + pub content: String, +} + +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum TurnRole { + User, + Assistant, + System, +} + +/// Score for a single task. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct BenchScore { + /// 0.0 to 1.0 (1.0 = perfect). + pub value: f64, + /// "pass" / "fail" / "partial". + pub label: String, + #[serde(default)] + pub details: Option, +} + +impl BenchScore { + pub fn pass() -> Self { + Self { + value: 1.0, + label: "pass".to_string(), + details: None, + } + } + + pub fn fail(details: impl Into) -> Self { + Self { + value: 0.0, + label: "fail".to_string(), + details: Some(details.into()), + } + } + + pub fn partial(value: f64, details: impl Into) -> Self { + Self { + value: value.clamp(0.0, 1.0), + label: "partial".to_string(), + details: Some(details.into()), + } + } +} + +/// Trait for benchmark suite adapters. +/// +/// Each suite (GAIA, Tau-bench, custom, etc.) implements this trait +/// to provide task loading, scoring, and optional lifecycle hooks. +#[async_trait] +pub trait BenchSuite: Send + Sync { + /// Human-readable name (e.g., "GAIA Validation"). + fn name(&self) -> &str; + + /// Machine ID (e.g., "gaia"). + fn id(&self) -> &str; + + /// Load all tasks from the suite's data source. + async fn load_tasks(&self) -> Result, BenchError>; + + /// Score the agent's submission against the expected answer. + async fn score( + &self, + task: &BenchTask, + submission: &TaskSubmission, + ) -> Result; + + /// Optional: set up environment before running a task (clone repo, init DB, etc.). + async fn setup_task(&self, _task: &BenchTask) -> Result<(), BenchError> { + Ok(()) + } + + /// Optional: tear down environment after a task completes. + async fn teardown_task(&self, _task: &BenchTask) -> Result<(), BenchError> { + Ok(()) + } + + /// Optional: additional tools to register for this suite's tasks. + fn additional_tools(&self) -> Vec> { + vec![] + } + + /// Optional: restrict which tools the agent can use (allowlist). + fn tool_whitelist(&self) -> Option> { + None + } + + /// Multi-turn: generate next simulated user message based on conversation so far. + /// Return `None` to end the conversation. + async fn next_user_message( + &self, + _task: &BenchTask, + _conversation: &[ConversationTurn], + ) -> Result, BenchError> { + Ok(None) + } +} diff --git a/src/channels/web/types.rs b/src/channels/web/types.rs index 85420211..e476a0f6 100644 --- a/src/channels/web/types.rs +++ b/src/channels/web/types.rs @@ -314,6 +314,7 @@ impl WsServerMessage { SseEvent::Thinking { .. } => "thinking", SseEvent::ToolStarted { .. } => "tool_started", SseEvent::ToolCompleted { .. } => "tool_completed", + SseEvent::ToolResult { .. } => "tool_result", SseEvent::StreamChunk { .. } => "stream_chunk", SseEvent::Status { .. } => "status", SseEvent::ApprovalNeeded { .. } => "approval_needed", diff --git a/src/llm/mod.rs b/src/llm/mod.rs index f2d5ffcb..4990da5f 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -13,7 +13,7 @@ pub mod session; pub use nearai::{ModelInfo, NearAiProvider}; pub use nearai_chat::NearAiChatProvider; pub use provider::{ - ChatMessage, CompletionRequest, CompletionResponse, LlmProvider, Role, ToolCall, + ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult, }; pub use reasoning::{ActionPlan, Reasoning, ReasoningContext, RespondResult, ToolSelection};