diff --git a/.beads/interactions.jsonl b/.beads/interactions.jsonl index 878e09c..bff30de 100644 --- a/.beads/interactions.jsonl +++ b/.beads/interactions.jsonl @@ -84,3 +84,4 @@ {"id":"int-686e2908","kind":"field_change","created_at":"2026-06-30T22:03:16.088748369Z","actor":"dawn","issue_id":"klbr-zue","extra":{"field":"status","new_value":"closed","old_value":"in_progress","reason":"Replaced shell stdout/stderr byte slicing with char-boundary truncation helper and added UTF-8 regression tests. Verified cargo test -p klbr-core shell_truncation, full cargo test -p klbr-core (168 passed, 1 ignored), cargo fmt --check, git diff --check, and rg confirms the unsafe stdout/stderr slices are gone."}} {"id":"int-eabe816d","kind":"field_change","created_at":"2026-06-30T22:04:39.33696926Z","actor":"dawn","issue_id":"klbr-41z","extra":{"field":"status","new_value":"closed","old_value":"in_progress","reason":"Changed daemon default websocket bind from 0.0.0.0:8765 to 127.0.0.1:8765 and replaced DumpMemories std::fs::write with tokio::fs::write(...).await. Verified rg for old patterns, cargo fmt --check, and cargo test -p klbr-daemon."}} {"id":"int-a150f81b","kind":"field_change","created_at":"2026-06-30T22:08:23.925241993Z","actor":"dawn","issue_id":"klbr-7e0","extra":{"field":"status","new_value":"closed","old_value":"in_progress","reason":"Removed unregistered old memory tool modules and updated memory docs"}} +{"id":"int-c575a8f2","kind":"field_change","created_at":"2026-06-30T22:20:51.577686369Z","actor":"dawn","issue_id":"klbr-x1i","extra":{"field":"status","new_value":"closed","old_value":"in_progress","reason":"Replaced klbr-bench raw argv dispatch with clap-derived commands and structured runner options"}} diff --git a/.beads/issues.jsonl b/.beads/issues.jsonl index 045e8aa..0f59678 100644 --- a/.beads/issues.jsonl +++ b/.beads/issues.jsonl @@ -43,7 +43,7 @@ {"_type":"issue","id":"klbr-wmz","title":"Finish memory architecture follow-through","description":"Tracks the remaining memory architecture work identified from docs/memory-arch.md, docs/memory-benches.md, docs/memory-implementation-status.md, and current klbr-core/klbr-bench source. Current status says pipeline, typed memory packets, markdown notes, edge mirroring, lifecycle projection, and benchmark runner integration exist; this epic is for gaps still present in source/docs.","acceptance_criteria":"Close when the child issues are complete, docs/memory-implementation-status.md is updated from current verification, and the architecture docs no longer point at missing or stale follow-up work.","status":"closed","priority":1,"issue_type":"epic","owner":"90008@klbr.net","created_at":"2026-06-26T17:52:54Z","created_by":"dawn","updated_at":"2026-06-26T23:45:35Z","closed_at":"2026-06-26T23:45:35Z","close_reason":"All 19 memory architecture child issues are closed; docs/status were updated from current verification; remaining official evaluator run is tracked separately as external blocked klbr-1yn.","labels":["architecture","memory"],"dependency_count":0,"dependent_count":0,"comment_count":0} {"_type":"issue","id":"klbr-u03","title":"bridge: refactor hardcoded /home/mayer paths and make binary path configurable","description":"antigravity-bridge/src/auth.rs contains hardcoded /home/mayer home directories and hardcoded binary path BINARY_PATH. Use dirs::home_dir() and make the binary path configurable via environment variables or configuration.","status":"closed","priority":2,"issue_type":"task","assignee":"dawn","owner":"90008@klbr.net","created_at":"2026-06-30T21:42:26Z","created_by":"dawn","updated_at":"2026-06-30T21:59:42Z","started_at":"2026-06-30T21:58:02Z","closed_at":"2026-06-30T21:59:42Z","close_reason":"Replaced hardcoded /home/mayer paths in antigravity bridge auth with dirs::home_dir()-based helpers, added ANTIGRAVITY_BRIDGE_CONFIG_DIR for token/cache files, and made the agy binary path configurable via ANTIGRAVITY_BRIDGE_BINARY_PATH or AGY_BINARY_PATH with ~/.local/bin/agy fallback. Verified cargo fmt --check and cargo test -p antigravity-bridge (9 passed).","dependency_count":0,"dependent_count":0,"comment_count":0} {"_type":"issue","id":"klbr-noi","title":"nix: fix broken flake.nix referencing non-existent default.nix","description":"flake.nix references ./default.nix to build the default package, but default.nix does not exist in the repository. Refactor flake.nix to output the derivation directly using nix-cargo-integration outputs.","status":"closed","priority":2,"issue_type":"task","assignee":"dawn","owner":"90008@klbr.net","created_at":"2026-06-30T21:42:19Z","created_by":"dawn","updated_at":"2026-06-30T22:01:41Z","started_at":"2026-06-30T22:00:06Z","closed_at":"2026-06-30T22:01:41Z","close_reason":"Replaced the missing ./default.nix package with config.nci.outputs.\"klbr-daemon\".packages.release. Verified nix eval .#packages.x86_64-linux.default.name returns klbr-daemon and nix flake show succeeds without the default.nix error.","dependency_count":0,"dependent_count":0,"comment_count":0} -{"_type":"issue","id":"klbr-x1i","title":"bench: replace manual CLI argument parsing with clap","description":"klbr-bench main.rs parses command-line arguments using raw std::env::args and index offsets. This is fragile and lacks auto-help or completions. Port CLI arguments parsing to clap.","status":"open","priority":2,"issue_type":"task","owner":"90008@klbr.net","created_at":"2026-06-30T21:42:13Z","created_by":"dawn","updated_at":"2026-06-30T21:42:13Z","dependency_count":0,"dependent_count":0,"comment_count":0} +{"_type":"issue","id":"klbr-x1i","title":"bench: replace manual CLI argument parsing with clap","description":"klbr-bench main.rs parses command-line arguments using raw std::env::args and index offsets. This is fragile and lacks auto-help or completions. Port CLI arguments parsing to clap.","status":"closed","priority":2,"issue_type":"task","assignee":"dawn","owner":"90008@klbr.net","created_at":"2026-06-30T21:42:13Z","created_by":"dawn","updated_at":"2026-06-30T22:20:52Z","started_at":"2026-06-30T22:09:02Z","closed_at":"2026-06-30T22:20:52Z","close_reason":"Replaced klbr-bench raw argv dispatch with clap-derived commands and structured runner options","dependency_count":0,"dependent_count":0,"comment_count":0} {"_type":"issue","id":"klbr-rds","title":"bench: split 6700-line main.rs monolith into logical submodules","description":"klbr-bench main.rs is currently over 6700 lines of code. It contains benchmark suites for passive recall, continuous loop, router, sweep, and command line parsing. Needs to be split into domain submodules (cli.rs, sweep.rs, passive_recall.rs, continuous_loop.rs, etc.).","status":"open","priority":2,"issue_type":"task","owner":"90008@klbr.net","created_at":"2026-06-30T21:42:03Z","created_by":"dawn","updated_at":"2026-06-30T21:42:03Z","dependency_count":0,"dependent_count":0,"comment_count":0} {"_type":"issue","id":"klbr-f8w","title":"discord: decompose monolithic lib.rs and deduplicate batch formatting","description":"klbr-discord lib.rs is a 2150-line monolith containing gateway loop, state database, formatting styles, and tools. Split into logical submodules, and deduplicate bracket/xml string formatting loops.","status":"open","priority":2,"issue_type":"task","owner":"90008@klbr.net","created_at":"2026-06-30T21:40:14Z","created_by":"dawn","updated_at":"2026-06-30T21:40:14Z","dependency_count":0,"dependent_count":0,"comment_count":0} {"_type":"issue","id":"klbr-41z","title":"daemon: fix blocking std::fs calls and secure default WS bind address","description":"DumpMemories handler in daemon.rs blocks the tokio runtime thread using sync std::fs::write. Fix to tokio::fs::write. Also, change default WS bind address to 127.0.0.1 for local-only safety.","status":"closed","priority":2,"issue_type":"task","assignee":"dawn","owner":"90008@klbr.net","created_at":"2026-06-30T21:40:08Z","created_by":"dawn","updated_at":"2026-06-30T22:04:39Z","started_at":"2026-06-30T22:03:35Z","closed_at":"2026-06-30T22:04:39Z","close_reason":"Changed daemon default websocket bind from 0.0.0.0:8765 to 127.0.0.1:8765 and replaced DumpMemories std::fs::write with tokio::fs::write(...).await. Verified rg for old patterns, cargo fmt --check, and cargo test -p klbr-daemon.","dependency_count":0,"dependent_count":0,"comment_count":0} diff --git a/Cargo.lock b/Cargo.lock index 6e0ef81..ce69ed1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -55,6 +55,56 @@ dependencies = [ "libc", ] +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.61.2", +] + [[package]] name = "antigravity-bridge" version = "0.1.0" @@ -326,6 +376,46 @@ dependencies = [ "windows-link", ] +[[package]] +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + [[package]] name = "cmake" version = "0.1.58" @@ -335,6 +425,12 @@ dependencies = [ "cc", ] +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + [[package]] name = "combine" version = "4.6.7" @@ -1156,6 +1252,12 @@ dependencies = [ "serde", ] +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + [[package]] name = "itoa" version = "1.0.18" @@ -1247,6 +1349,7 @@ dependencies = [ "blake3", "bytemuck", "chrono", + "clap", "klbr-core", "rand 0.8.5", "rusqlite", @@ -1553,6 +1656,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + [[package]] name = "openssl-probe" version = "0.2.1" @@ -2371,6 +2480,12 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b2231b7c3057d5e4ad0156fb3dc807d900806020c5ffa3ee6ff2c8c76fb8520" +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "subtle" version = "2.6.1" @@ -3033,6 +3148,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "uuid" version = "1.23.0" diff --git a/klbr-bench/Cargo.toml b/klbr-bench/Cargo.toml index 37f2120..3e22de9 100644 --- a/klbr-bench/Cargo.toml +++ b/klbr-bench/Cargo.toml @@ -16,4 +16,4 @@ rusqlite = { version = "0.39", features = ["bundled"] } rand = "0.8" blake3 = "1.8.5" bytemuck = "1.25.0" - +clap = { version = "4", features = ["derive"] } diff --git a/klbr-bench/src/longmemeval.rs b/klbr-bench/src/longmemeval.rs index 36e13fc..d823ece 100644 --- a/klbr-bench/src/longmemeval.rs +++ b/klbr-bench/src/longmemeval.rs @@ -157,6 +157,55 @@ struct SampleReport { question_ids: Vec, } +#[derive(Debug, Clone)] +pub struct PipelineRunOptions { + pub suite: String, + pub data_path: String, + pub profile: String, + pub out_dir: Option, + pub top_k: usize, + pub budget_read: usize, + pub graph_depth: usize, + pub limit: Option, + pub sample: Option, + pub sample_mode: String, + pub sample_seed: Option, + pub question_id_filter: Option, + pub retrieval_only: bool, + pub diagnostic: Option, + pub official_eval_cmd: Option, + pub packet_rerank: bool, + pub llm_url: Option, + pub embed_url: Option, + pub embed_model: Option, + pub embed_dim: Option, +} + +#[derive(Debug, Clone, Copy)] +pub enum LongMemEvalAction { + Ingest, + Retrieve, + Answer, + EvalRetrieval, + SynthReflink, + BenchExact, +} + +#[derive(Debug, Clone)] +pub struct LongMemEvalOptions { + pub action: LongMemEvalAction, + pub data: Option, + pub out: Option, + pub trace_out: Option, + pub db_dir: Option, + pub reader: Option, + pub retrieval_mode: String, + pub max_resolved_ref_tokens: usize, + pub top_k: usize, + pub batch_sizes: Vec, + pub graph_depth: usize, +} + fn load_questions_for_suite(suite: &str, data_path: &str) -> Result> { let bytes = fs::read(data_path).with_context(|| format!("failed to read {data_path}"))?; if suite.eq_ignore_ascii_case("locomo") { @@ -1247,184 +1296,100 @@ fn compute_retrieval_metrics( }) } -pub async fn run_command(args: &[String]) -> Result<()> { - let mut data = None; - let mut out = None; - let mut trace_out = None; - let mut db_dir = None; - let mut reader = None; - let mut retrieval_mode = Some("exact+semantic+graph+rerank".to_string()); - let mut max_resolved_ref_tokens = 2500; - let mut top_k = 10; - let mut batch_sizes = vec![1, 10, 100]; - let mut graph_depth = 1; - - let mut i = 3; - while i < args.len() { - match args[i].as_str() { - "--data" => { - data = Some(args[i + 1].clone()); - i += 2; - } - "--out" => { - out = Some(args[i + 1].clone()); - i += 2; - } - "--trace-out" => { - trace_out = Some(args[i + 1].clone()); - i += 2; - } - "--db-dir" => { - db_dir = Some(args[i + 1].clone()); - i += 2; - } - "--reader" => { - reader = Some(args[i + 1].clone()); - i += 2; - } - "--retrieval" => { - retrieval_mode = Some(args[i + 1].clone()); - i += 2; - } - "--max-resolved-ref-tokens" => { - max_resolved_ref_tokens = args[i + 1].parse().unwrap_or(2500); - i += 2; - } - "--top-k" => { - top_k = args[i + 1].parse().unwrap_or(10); - i += 2; - } - "--batch-sizes" => { - batch_sizes = args[i + 1] - .split(',') - .map(|s| s.parse().unwrap_or(1)) - .collect(); - i += 2; - } - "--graph-depth" => { - graph_depth = args[i + 1].parse().unwrap_or(1); - i += 2; - } - _ => { - i += 1; - } - } - } - - let subcommand = args.get(2).map(|s| s.as_str()).unwrap_or(""); - match subcommand { - "ingest" => { - let data_path = data.context("Missing --data")?; - let db_dir_path = db_dir.context("Missing --db-dir")?; +pub async fn run_command(options: LongMemEvalOptions) -> Result<()> { + match options.action { + LongMemEvalAction::Ingest => { + let data_path = options.data.context("Missing --data")?; + let db_dir_path = options.db_dir.context("Missing --db-dir")?; run_ingest(&data_path, &db_dir_path).await } - "retrieve" => { - let data_path = data.context("Missing --data")?; - let db_dir_path = db_dir.context("Missing --db-dir")?; - let trace_path = trace_out.context("Missing --trace-out")?; - let retrieval_str = - retrieval_mode.unwrap_or_else(|| "exact+semantic+graph+rerank".to_string()); + LongMemEvalAction::Retrieve => { + let data_path = options.data.context("Missing --data")?; + let db_dir_path = options.db_dir.context("Missing --db-dir")?; + let trace_path = options.trace_out.context("Missing --trace-out")?; run_retrieve( &data_path, &db_dir_path, &trace_path, - &retrieval_str, - max_resolved_ref_tokens, - top_k, + &options.retrieval_mode, + options.max_resolved_ref_tokens, + options.top_k, ) .await } - "answer" => { - let data_path = data.context("Missing --data")?; - let db_dir_path = db_dir.context("Missing --db-dir")?; - let out_path = out.context("Missing --out")?; - let trace_path = trace_out.context("Missing --trace-out")?; - let reader_str = reader.unwrap_or_else(|| "llama-local".to_string()); - let retrieval_str = - retrieval_mode.unwrap_or_else(|| "exact+semantic+graph+rerank".to_string()); + LongMemEvalAction::Answer => { + let data_path = options.data.context("Missing --data")?; + let db_dir_path = options.db_dir.context("Missing --db-dir")?; + let out_path = options.out.context("Missing --out")?; + let trace_path = options.trace_out.context("Missing --trace-out")?; + let reader_str = options.reader.unwrap_or_else(|| "llama-local".to_string()); run_answer( &data_path, &db_dir_path, &out_path, &trace_path, &reader_str, - &retrieval_str, - max_resolved_ref_tokens, - top_k, + &options.retrieval_mode, + options.max_resolved_ref_tokens, + options.top_k, ) .await } - "eval-retrieval" => { - let data_path = data.context("Missing --data")?; - let trace_path = trace_out.context("Missing --trace-out")?; + LongMemEvalAction::EvalRetrieval => { + let data_path = options.data.context("Missing --data")?; + let trace_path = options.trace_out.context("Missing --trace-out")?; run_eval_retrieval(&data_path, &trace_path).await } - "synth-reflink" => { - let data_path = data.context("Missing --data")?; - let db_dir_path = db_dir.context("Missing --db-dir")?; - let out_path = out.context("Missing --out")?; + LongMemEvalAction::SynthReflink => { + let data_path = options.data.context("Missing --data")?; + let db_dir_path = options.db_dir.context("Missing --db-dir")?; + let out_path = options.out.context("Missing --out")?; run_synth_reflink(&data_path, &db_dir_path, &out_path).await } - "bench-exact" => { - let data_path = data.context("Missing --data")?; - let db_dir_path = db_dir.context("Missing --db-dir")?; - let out_path = out.context("Missing --out")?; + LongMemEvalAction::BenchExact => { + let data_path = options.data.context("Missing --data")?; + let db_dir_path = options.db_dir.context("Missing --db-dir")?; + let out_path = options.out.context("Missing --out")?; run_bench_exact( &data_path, &db_dir_path, &out_path, - &batch_sizes, - graph_depth, + &options.batch_sizes, + options.graph_depth, ) .await } - other => bail!("Unknown longmemeval subcommand: {other}"), - } -} - -pub async fn run_pipeline_command(args: &[String]) -> Result<()> { - let suite = optional_arg(args, "--suite").unwrap_or_else(|| "longmemeval-s".to_string()); - let data_path = required_arg(args, "--data")?; - let profile = optional_arg(args, "--profile").unwrap_or_else(|| "klbr-full".to_string()); - let out_dir = optional_arg(args, "--out") - .map(PathBuf::from) - .unwrap_or_else(|| default_pipeline_run_dir(&suite, &profile)); - let top_k = optional_arg(args, "--top-k") - .and_then(|value| value.parse::().ok()) - .unwrap_or(8); - let budget_read = optional_arg(args, "--budget-read") - .and_then(|value| value.parse::().ok()) - .unwrap_or(5_000); - let graph_depth = optional_arg(args, "--graph-depth") - .and_then(|value| value.parse::().ok()) - .unwrap_or(1); - let limit = optional_arg(args, "--limit").and_then(|value| value.parse::().ok()); - let sample = optional_arg(args, "--sample") - .map(|value| { - value - .parse::() - .with_context(|| format!("invalid --sample value {value}")) - }) - .transpose()?; + } +} + +pub async fn run_pipeline_command(options: PipelineRunOptions) -> Result<()> { + let PipelineRunOptions { + suite, + data_path, + profile, + out_dir, + top_k, + budget_read, + graph_depth, + limit, + sample, + sample_mode, + sample_seed, + question_id_filter, + retrieval_only, + diagnostic, + official_eval_cmd, + packet_rerank, + llm_url, + embed_url, + embed_model, + embed_dim, + } = options; + let out_dir = out_dir.unwrap_or_else(|| default_pipeline_run_dir(&suite, &profile)); if sample == Some(0) { bail!("--sample must be greater than 0"); } - let sample_mode = optional_arg(args, "--sample-mode") - .map(|value| SampleMode::parse(&value)) - .transpose()? - .unwrap_or(SampleMode::Stratified); - let sample_seed = optional_arg(args, "--sample-seed") - .map(|value| { - value - .parse::() - .with_context(|| format!("invalid --sample-seed value {value}")) - }) - .transpose()?; - let question_id_filter = optional_arg(args, "--question-id"); - let retrieval_only = has_flag(args, "--retrieval-only"); - let diagnostic = optional_arg(args, "--diagnostic"); - let official_eval_cmd = optional_arg(args, "--official-eval-cmd"); - let packet_rerank = has_flag(args, "--rerank-packets"); + let sample_mode = SampleMode::parse(&sample_mode)?; let protocol_note = suite_protocol_note(&suite); fs::create_dir_all(&out_dir)?; @@ -1454,17 +1419,16 @@ pub async fn run_pipeline_command(args: &[String]) -> Result<()> { }; let mut config = Config::load_bench()?; - if let Some(llm_url) = optional_arg(args, "--llm-url") { + if let Some(llm_url) = llm_url { config.models.llm.url = normalize_base_url(&llm_url); } - if let Some(embed_url) = optional_arg(args, "--embed-url") { + if let Some(embed_url) = embed_url { config.models.embedder.url = normalize_base_url(&embed_url); } - if let Some(embed_model) = optional_arg(args, "--embed-model") { + if let Some(embed_model) = embed_model { config.models.embedder.model = embed_model; } - if let Some(embed_dim) = optional_arg(args, "--embed-dim").and_then(|value| value.parse().ok()) - { + if let Some(embed_dim) = embed_dim { config.models.embed_dim = embed_dim; } config.memory.top_k = top_k; @@ -1843,7 +1807,22 @@ pub async fn run_pipeline_command(args: &[String]) -> Result<()> { fs::write( out_dir.join("manifest.json"), serde_json::to_vec_pretty(&json!({ - "command": args, + "command": { + "suite": suite, + "data": data_path, + "profile": profile, + "top_k": top_k, + "budget_read": budget_read, + "graph_depth": graph_depth, + "limit": limit, + "sample_mode": sample_mode.as_str(), + "sample_seed": sample_seed, + "question_id": question_id_filter.clone(), + "retrieval_only": retrieval_only, + "diagnostic": diagnostic.clone(), + "official_eval_cmd": official_eval_cmd, + "packet_rerank": packet_rerank, + }, "suite": suite, "profile": profile, "models": { @@ -2133,20 +2112,6 @@ fn render_official_eval_result(result: &OfficialEvalResult) -> String { ) } -fn required_arg(args: &[String], name: &str) -> Result { - optional_arg(args, name).with_context(|| format!("missing required argument {name}")) -} - -fn optional_arg(args: &[String], name: &str) -> Option { - args.windows(2) - .find(|pair| pair[0] == name) - .map(|pair| pair[1].clone()) -} - -fn has_flag(args: &[String], name: &str) -> bool { - args.iter().any(|arg| arg == name) -} - fn normalize_base_url(url: &str) -> String { let trimmed = url.trim().trim_end_matches('/'); if trimmed.ends_with("/v1") { @@ -2594,13 +2559,12 @@ fn official_eval_labels_from_metrics(metrics: &serde_json::Value) -> HashMap Result<()> { - if args.len() < 2 { - bail!("usage: cargo run -p klbr-bench -- trace-diff [question_id ...]"); - } - let old_dir = PathBuf::from(&args[0]); - let new_dir = PathBuf::from(&args[1]); - let filters = args[2..].iter().cloned().collect::>(); +pub fn run_trace_diff_command( + old_dir: PathBuf, + new_dir: PathBuf, + filters: Vec, +) -> Result<()> { + let filters = filters.into_iter().collect::>(); let old_rows = load_trace_summary_rows(&old_dir)?; let new_rows = load_trace_summary_rows(&new_dir)?; let old_by_id = old_rows diff --git a/klbr-bench/src/main.rs b/klbr-bench/src/main.rs index 5ecbc9d..3aa5d91 100644 --- a/klbr-bench/src/main.rs +++ b/klbr-bench/src/main.rs @@ -9,6 +9,7 @@ use std::{ }; use anyhow::{bail, Context, Result}; +use clap::{Args, Parser, Subcommand}; use klbr_core::{ context::{format_recalled_memories, RecalledMemory}, memory::{MarkdownNoteInput, MemoryLane, MemoryStore, MemoryStoreStats}, @@ -209,176 +210,419 @@ struct SweepAttemptRecord { rerank_latency_ms: u128, } +#[derive(Debug, Parser)] +#[command(name = "klbr-bench")] +#[command(about = "klbr benchmark runner")] +struct Cli { + #[command(subcommand)] + command: BenchCommand, +} + +#[derive(Debug, Subcommand)] +enum BenchCommand { + Retrieval(PositionalBenchArgs), + #[command(name = "run")] + Run(PipelineRunArgs), + #[command(name = "trace-diff")] + TraceDiff(TraceDiffArgs), + #[command(name = "passive-recall")] + PassiveRecall(PositionalBenchArgs), + Router(PositionalBenchArgs), + #[command(name = "router-multi")] + RouterMulti(MultiRouterArgs), + #[command(name = "router-multi-linear")] + RouterMultiLinear(MultiRouterArgs), + #[command(name = "dump-tools")] + DumpTools(OutputPathArg), + #[command(name = "continuous-loop")] + ContinuousLoop(OutputDirArg), + Longmemeval { + #[command(subcommand)] + command: LongMemEvalCommand, + }, + Longmem(LongmemArgs), + #[command(name = "longmem-retrieval")] + LongmemRetrieval(LongmemRetrievalArgs), + Sweep(SweepArgs), +} + +#[derive(Debug, Args)] +struct PositionalBenchArgs { + dataset_json: String, + config_json: String, + output_dir: String, +} + +#[derive(Debug, Args)] +struct MultiRouterArgs { + config_json: String, + output_dir: String, + #[arg(required = true)] + dataset_json: Vec, +} + +#[derive(Debug, Args)] +struct OutputPathArg { + output_path: String, +} + +#[derive(Debug, Args)] +struct OutputDirArg { + output_dir: String, +} + +#[derive(Debug, Args)] +struct TraceDiffArgs { + old_run_dir: String, + new_run_dir: String, + question_id: Vec, +} + +#[derive(Debug, Args)] +struct PipelineRunArgs { + #[arg(long, default_value = "longmemeval-s")] + suite: String, + #[arg(long)] + data: String, + #[arg(long, default_value = "klbr-full")] + profile: String, + #[arg(long)] + out: Option, + #[arg(long = "top-k", default_value_t = 8)] + top_k: usize, + #[arg(long = "budget-read", default_value_t = 5_000)] + budget_read: usize, + #[arg(long = "graph-depth", default_value_t = 1)] + graph_depth: usize, + #[arg(long)] + limit: Option, + #[arg(long)] + sample: Option, + #[arg(long = "sample-mode", default_value = "stratified")] + sample_mode: String, + #[arg(long = "sample-seed")] + sample_seed: Option, + #[arg(long = "question-id")] + question_id: Option, + #[arg(long = "retrieval-only")] + retrieval_only: bool, + #[arg(long)] + diagnostic: Option, + #[arg(long = "official-eval-cmd")] + official_eval_cmd: Option, + #[arg(long = "rerank-packets")] + rerank_packets: bool, + #[arg(long = "llm-url")] + llm_url: Option, + #[arg(long = "embed-url")] + embed_url: Option, + #[arg(long = "embed-model")] + embed_model: Option, + #[arg(long = "embed-dim")] + embed_dim: Option, +} + +#[derive(Debug, Subcommand)] +enum LongMemEvalCommand { + Ingest(LongMemEvalArgs), + Retrieve(LongMemEvalArgs), + Answer(LongMemEvalArgs), + #[command(name = "eval-retrieval")] + EvalRetrieval(LongMemEvalArgs), + #[command(name = "synth-reflink")] + SynthReflink(LongMemEvalArgs), + #[command(name = "bench-exact")] + BenchExact(LongMemEvalArgs), +} + +#[derive(Debug, Args)] +struct LongMemEvalArgs { + #[arg(long)] + data: Option, + #[arg(long)] + out: Option, + #[arg(long = "trace-out")] + trace_out: Option, + #[arg(long = "db-dir")] + db_dir: Option, + #[arg(long)] + reader: Option, + #[arg(long = "retrieval")] + retrieval_mode: Option, + #[arg(long = "max-resolved-ref-tokens", default_value_t = 2500)] + max_resolved_ref_tokens: usize, + #[arg(long = "top-k", default_value_t = 10)] + top_k: usize, + #[arg(long = "batch-sizes")] + batch_sizes: Option, + #[arg(long = "graph-depth", default_value_t = 1)] + graph_depth: usize, +} + +#[derive(Debug, Args)] +struct LongmemArgs { + dataset_json: String, + config_json: String, + output_dir: String, + llm_url: Option, + #[arg(long = "official-eval-cmd")] + official_eval_cmd: Option, +} + +#[derive(Debug, Args)] +struct LongmemRetrievalArgs { + dataset_json: String, + config_json: String, + output_dir: String, + split_json: Option, + subset: Option, +} + +#[derive(Debug, Args)] +struct SweepArgs { + dataset_json: String, + config_json: String, + output_dir: String, + score_start: Option, + score_end: Option, + score_step: Option, + margin_start: Option, + margin_end: Option, + margin_step: Option, + support_start: Option, + support_end: Option, + support_step: Option, +} + #[tokio::main] async fn main() -> Result<()> { - let args: Vec = env::args().collect(); - if args.len() < 2 { - bail!( - "usage:\n cargo run -p klbr-bench -- run --suite longmemeval-s --data [--sample ] [--sample-seed ] [--sample-mode stratified|random] [--out ] [--retrieval-only]\n cargo run -p klbr-bench -- trace-diff [question_id ...]\n cargo run -p klbr-bench -- retrieval \n cargo run -p klbr-bench -- passive-recall \n cargo run -p klbr-bench -- router \n cargo run -p klbr-bench -- router-multi [dataset2.json ...]\n cargo run -p klbr-bench -- router-multi-linear [dataset2.json ...]\n cargo run -p klbr-bench -- longmem [llm_url]\n cargo run -p klbr-bench -- longmem-retrieval [split.json subset]\n cargo run -p klbr-bench -- continuous-loop \n cargo run -p klbr-bench -- dump-tools \n cargo run -p klbr-bench -- sweep [score_start score_end score_step margin_start margin_end margin_step [support_start support_end support_step]]" - ); - } + let cli = Cli::parse(); - match args[1].as_str() { - "retrieval" => { - if args.len() != 5 { - bail!( - "usage: cargo run -p klbr-bench -- retrieval " - ); - } - run_retrieval_command(&args[2], &args[3], &args[4]).await - } - "run" => longmemeval::run_pipeline_command(&args).await, - "trace-diff" => longmemeval::run_trace_diff_command(&args[2..]), - "passive-recall" => { - if args.len() != 5 { - bail!( - "usage: cargo run -p klbr-bench -- passive-recall " - ); - } - run_passive_recall_command(&args[2], &args[3], &args[4]).await + match cli.command { + BenchCommand::Retrieval(args) => { + run_retrieval_command(&args.dataset_json, &args.config_json, &args.output_dir).await } - "router" => { - if args.len() != 5 { - bail!( - "usage: cargo run -p klbr-bench -- router " - ); - } - run_router_command(&args[2], &args[3], &args[4]).await + BenchCommand::Run(args) => longmemeval::run_pipeline_command(args.into_options()?).await, + BenchCommand::TraceDiff(args) => longmemeval::run_trace_diff_command( + PathBuf::from(args.old_run_dir), + PathBuf::from(args.new_run_dir), + args.question_id, + ), + BenchCommand::PassiveRecall(args) => { + run_passive_recall_command(&args.dataset_json, &args.config_json, &args.output_dir) + .await } - "router-multi" => { - if args.len() < 6 { - bail!( - "usage: cargo run -p klbr-bench -- router-multi [dataset2.json ...]" - ); - } - run_router_multi_command(&args[2], &args[3], &args[4..]).await + BenchCommand::Router(args) => { + run_router_command(&args.dataset_json, &args.config_json, &args.output_dir).await } - "router-multi-linear" => { - if args.len() < 6 { - bail!( - "usage: cargo run -p klbr-bench -- router-multi-linear [dataset2.json ...]" - ); - } - run_router_multi_linear_command(&args[2], &args[3], &args[4..]).await + BenchCommand::RouterMulti(args) => { + run_router_multi_command(&args.config_json, &args.output_dir, &args.dataset_json).await } - "dump-tools" => { - if args.len() != 3 { - bail!("usage: cargo run -p klbr-bench -- dump-tools "); - } - run_dump_tools_command(&args[2]).await + BenchCommand::RouterMultiLinear(args) => { + run_router_multi_linear_command(&args.config_json, &args.output_dir, &args.dataset_json) + .await } - "continuous-loop" => { - if args.len() != 3 { - bail!("usage: cargo run -p klbr-bench -- continuous-loop "); - } - run_continuous_loop_command(&args[2]).await + BenchCommand::DumpTools(args) => run_dump_tools_command(&args.output_path).await, + BenchCommand::ContinuousLoop(args) => run_continuous_loop_command(&args.output_dir).await, + BenchCommand::Longmemeval { command } => { + longmemeval::run_command(command.into_options()?).await } - "longmemeval" => { - longmemeval::run_command(&args).await + BenchCommand::Longmem(args) => { + run_longmem_command( + &args.dataset_json, + &args.config_json, + &args.output_dir, + args.llm_url, + args.official_eval_cmd, + ) + .await } - "longmem" => { - // positional: dataset config outdir [llm_url] - // optional flags: --official-eval-cmd - let official_eval_cmd = args - .windows(2) - .find(|w| w[0] == "--official-eval-cmd") - .map(|w| w[1].clone()); - let mut positional: Vec<&str> = Vec::new(); - let mut skip_next = false; - for a in args.iter().skip(2) { - if skip_next { - skip_next = false; - continue; - } - if a.starts_with("--") { - skip_next = true; - continue; - } - positional.push(a.as_str()); - } - if positional.len() < 3 || positional.len() > 4 { + BenchCommand::LongmemRetrieval(args) => { + if args.split_json.is_some() != args.subset.is_some() { bail!( - "usage: cargo run -p klbr-bench -- longmem [llm_url] [--official-eval-cmd ]" + "longmem-retrieval split arguments must be provided as " ); } - let llm_url = positional.get(3).map(|s| s.to_string()); - run_longmem_command(positional[0], positional[1], positional[2], llm_url, official_eval_cmd).await + let split = args.split_json.as_deref().zip(args.subset.as_deref()); + run_longmem_retrieval_command( + &args.dataset_json, + &args.config_json, + &args.output_dir, + split, + ) + .await } - "longmem-retrieval" => { - if args.len() != 5 && args.len() != 7 { - bail!( - "usage: cargo run -p klbr-bench -- longmem-retrieval [split.json subset]" - ); - } - let split = if args.len() == 7 { - Some((args[5].as_str(), args[6].as_str())) - } else { - None - }; - run_longmem_retrieval_command(&args[2], &args[3], &args[4], split).await + BenchCommand::Sweep(args) => { + let grid = args.grid()?; + run_sweep_command( + &args.dataset_json, + &args.config_json, + &args.output_dir, + grid, + ) + .await } - "sweep" => { - if args.len() != 5 && args.len() != 11 && args.len() != 14 { - bail!( - "usage: cargo run -p klbr-bench -- sweep [score_start score_end score_step margin_start margin_end margin_step [support_start support_end support_step]]" - ); - } - let default_support = SweepRange { - start: 0.0, - end: 1.0, - step: 0.1, - }; - let grid = match args.len() { - 14 => SweepGrid { - score: SweepRange { - start: parse_f32_arg(&args[5], "score_start")?, - end: parse_f32_arg(&args[6], "score_end")?, - step: parse_f32_arg(&args[7], "score_step")?, - }, - margin: SweepRange { - start: parse_f32_arg(&args[8], "margin_start")?, - end: parse_f32_arg(&args[9], "margin_end")?, - step: parse_f32_arg(&args[10], "margin_step")?, - }, - support: SweepRange { - start: parse_f32_arg(&args[11], "support_start")?, - end: parse_f32_arg(&args[12], "support_end")?, - step: parse_f32_arg(&args[13], "support_step")?, - }, + } +} + +impl PipelineRunArgs { + fn into_options(self) -> Result { + if self.sample == Some(0) { + bail!("--sample must be greater than 0"); + } + Ok(longmemeval::PipelineRunOptions { + suite: self.suite, + data_path: self.data, + profile: self.profile, + out_dir: self.out, + top_k: self.top_k, + budget_read: self.budget_read, + graph_depth: self.graph_depth, + limit: self.limit, + sample: self.sample, + sample_mode: self.sample_mode, + sample_seed: self.sample_seed, + question_id_filter: self.question_id, + retrieval_only: self.retrieval_only, + diagnostic: self.diagnostic, + official_eval_cmd: self.official_eval_cmd, + packet_rerank: self.rerank_packets, + llm_url: self.llm_url, + embed_url: self.embed_url, + embed_model: self.embed_model, + embed_dim: self.embed_dim, + }) + } +} + +impl LongMemEvalCommand { + fn into_options(self) -> Result { + let (action, args) = match self { + Self::Ingest(args) => (longmemeval::LongMemEvalAction::Ingest, args), + Self::Retrieve(args) => (longmemeval::LongMemEvalAction::Retrieve, args), + Self::Answer(args) => (longmemeval::LongMemEvalAction::Answer, args), + Self::EvalRetrieval(args) => (longmemeval::LongMemEvalAction::EvalRetrieval, args), + Self::SynthReflink(args) => (longmemeval::LongMemEvalAction::SynthReflink, args), + Self::BenchExact(args) => (longmemeval::LongMemEvalAction::BenchExact, args), + }; + Ok(longmemeval::LongMemEvalOptions { + action, + data: args.data, + out: args.out, + trace_out: args.trace_out, + db_dir: args.db_dir, + reader: args.reader, + retrieval_mode: args + .retrieval_mode + .unwrap_or_else(|| "exact+semantic+graph+rerank".to_string()), + max_resolved_ref_tokens: args.max_resolved_ref_tokens, + top_k: args.top_k, + batch_sizes: parse_batch_sizes(args.batch_sizes)?, + graph_depth: args.graph_depth, + }) + } +} + +fn parse_batch_sizes(value: Option) -> Result> { + let Some(value) = value else { + return Ok(vec![1, 10, 100]); + }; + value + .split(',') + .map(str::trim) + .filter(|part| !part.is_empty()) + .map(|part| { + part.parse::() + .with_context(|| format!("failed to parse batch size '{part}' as usize")) + }) + .collect() +} + +impl SweepArgs { + fn grid(&self) -> Result { + let default_support = SweepRange { + start: 0.0, + end: 1.0, + step: 0.1, + }; + match ( + self.score_start, + self.score_end, + self.score_step, + self.margin_start, + self.margin_end, + self.margin_step, + self.support_start, + self.support_end, + self.support_step, + ) { + (None, None, None, None, None, None, None, None, None) => Ok(SweepGrid { + score: SweepRange { + start: -11.5, + end: 2.0, + step: 1.0, }, - 11 => SweepGrid { - score: SweepRange { - start: parse_f32_arg(&args[5], "score_start")?, - end: parse_f32_arg(&args[6], "score_end")?, - step: parse_f32_arg(&args[7], "score_step")?, - }, - margin: SweepRange { - start: parse_f32_arg(&args[8], "margin_start")?, - end: parse_f32_arg(&args[9], "margin_end")?, - step: parse_f32_arg(&args[10], "margin_step")?, - }, - support: default_support, + margin: SweepRange { + start: 0.0, + end: 12.0, + step: 1.0, }, - _ => SweepGrid { - score: SweepRange { - start: -11.5, - end: 2.0, - step: 1.0, - }, - margin: SweepRange { - start: 0.0, - end: 12.0, - step: 1.0, - }, - support: default_support, + support: default_support, + }), + ( + Some(score_start), + Some(score_end), + Some(score_step), + Some(margin_start), + Some(margin_end), + Some(margin_step), + None, + None, + None, + ) => Ok(SweepGrid { + score: SweepRange { + start: score_start, + end: score_end, + step: score_step, }, - }; - run_sweep_command(&args[2], &args[3], &args[4], grid).await + margin: SweepRange { + start: margin_start, + end: margin_end, + step: margin_step, + }, + support: default_support, + }), + ( + Some(score_start), + Some(score_end), + Some(score_step), + Some(margin_start), + Some(margin_end), + Some(margin_step), + Some(support_start), + Some(support_end), + Some(support_step), + ) => Ok(SweepGrid { + score: SweepRange { + start: score_start, + end: score_end, + step: score_step, + }, + margin: SweepRange { + start: margin_start, + end: margin_end, + step: margin_step, + }, + support: SweepRange { + start: support_start, + end: support_end, + step: support_step, + }, + }), + _ => bail!( + "sweep accepts either no grid values, 6 score/margin values, or 9 score/margin/support values" + ), } - other => bail!( - "unknown subcommand '{}'; expected 'retrieval', 'passive-recall', 'router', 'router-multi', 'router-multi-linear', 'continuous-loop', 'dump-tools', or 'sweep'", - other - ), } } @@ -2506,12 +2750,6 @@ async fn run_sweep_command( Ok(()) } -fn parse_f32_arg(value: &str, name: &str) -> Result { - value - .parse::() - .with_context(|| format!("failed to parse {}='{}' as f32", name, value)) -} - fn validate_benchmark_inputs( dataset: &InternalEvalDataset, experiment: &RetrievalExperimentConfig,