use crate::benchmarks; use crate::benchmarks::judge; use crate::config::AppConfig; use crate::metrics::{LatencySummary, Metrics, MetricsSummary}; use crate::report::{ BenchmarkReport, BenchmarkSummaryReport, DatasetReport, LatencyReport, RpmReport, RpmRunReport, RpmSummaryReport, RunReport, WrongCaseReport, write_benchmark_report, write_rpm_report, }; use crate::runner::{ModelRequest, run_model_request}; use anyhow::{Context, Result, bail}; use chrono::Utc; use clap::{Parser, Subcommand}; use futures::{StreamExt, stream}; use regex::Regex; use std::path::{Path, PathBuf}; use std::time::{Duration, Instant}; use tokio::time::{Instant as TokioInstant, sleep_until}; #[derive(Debug, Parser)] #[command( name = "lq_token_test", version, about = "Test LLM relay protocols, RPM, and benchmark accuracy" )] pub struct Cli { #[command(subcommand)] pub command: Command, } #[derive(Debug, Subcommand)] pub enum Command { Check { #[arg(long, default_value = "config.yaml")] config: PathBuf, #[arg(long)] provider: Option, #[arg(long)] model: Option, #[arg(long)] prompt: String, }, Dataset { #[command(subcommand)] command: DatasetCommand, }, Bench { #[command(subcommand)] command: BenchCommand, }, Rpm { #[arg(long, default_value = "config.yaml")] config: PathBuf, #[arg(long)] provider: Option, #[arg(long)] model: Option, #[arg(long)] rpm: u32, #[arg(long)] duration: String, #[arg(long)] prompt: String, }, } #[derive(Debug, Subcommand)] pub enum DatasetCommand { Fetch { dataset: String }, } #[derive(Debug, Subcommand)] pub enum BenchCommand { Aime2026 { #[arg(long, default_value = "config.yaml")] config: PathBuf, #[arg(long)] provider: Option, #[arg(long)] model: Option, #[arg(long, default_value_t = 4)] concurrency: usize, #[arg(long)] limit: Option, }, GpqaDiamond { #[arg(long, default_value = "config.yaml")] config: PathBuf, #[arg(long)] provider: Option, #[arg(long)] model: Option, #[arg(long, default_value_t = 4)] concurrency: usize, #[arg(long)] limit: Option, }, } pub async fn dispatch(cli: Cli) -> Result<()> { match cli.command { Command::Check { config, provider, model, prompt, } => { let config = AppConfig::load(&config)?; let provider = config.resolved_provider(provider.as_deref())?; let request = ModelRequest { base_url: provider.base_url.clone(), api_token: provider.api_token.clone(), model: model.unwrap_or_else(|| provider.default_model.clone()), prompt, temperature: 0.0, max_tokens: 1024, }; let response = run_model_request(provider.protocol, request).await?; println!("status: {}", response.status); println!("elapsed_ms: {}", response.elapsed_ms); println!("{}", response.text); Ok(()) } Command::Dataset { command: DatasetCommand::Fetch { dataset }, } => { let data_dir = dataset_data_dir(Path::new("config.yaml"))?; let path = benchmarks::fetch_dataset(&dataset, &data_dir).await?; println!("{}", path.display()); Ok(()) } Command::Bench { command } => dispatch_bench(command).await, Command::Rpm { config, provider, model, rpm, duration, prompt, } => run_rpm(config, provider, model, rpm, duration, prompt).await, } } async fn dispatch_bench(command: BenchCommand) -> Result<()> { match command { BenchCommand::Aime2026 { config, provider, model, concurrency, limit, } => run_aime_benchmark(config, provider, model, concurrency, limit).await, BenchCommand::GpqaDiamond { config, provider, model, concurrency, limit, } => run_gpqa_benchmark(config, provider, model, concurrency, limit).await, } } async fn run_aime_benchmark( config_path: PathBuf, provider: Option, model: Option, concurrency: usize, limit: Option, ) -> Result<()> { let config = AppConfig::load(&config_path)?; let provider_name = provider_name(&config, provider.as_deref())?; let provider_config = config.resolved_provider(Some(&provider_name))?; let model = model.unwrap_or_else(|| provider_config.default_model.clone()); let loaded = benchmarks::aime::load_cases(Path::new(&config.benchmarks.data_dir))?; let dataset = dataset_report( config .benchmarks .aime2026 .as_ref() .map(|dataset| (dataset.source.as_str(), dataset.split.as_str())), &loaded.local_path, ); let cases = apply_limit(loaded.cases, limit); let started_at = Utc::now(); let started = Instant::now(); let base_request = request_template(&provider_config, &model, 0.0, 1024); let protocol = provider_config.protocol; let results = stream::iter(cases) .map(|case| { let mut request = base_request.clone(); request.prompt = case.prompt(); async move { let result = run_model_request(protocol, request).await; (case, result) } }) .buffer_unordered(nonzero_concurrency(concurrency)) .collect::>() .await; let mut metrics = Metrics::new(); let mut wrong_cases = Vec::new(); for (case, result) in results { match result { Ok(response) => { metrics.record_success(response.status, response.elapsed_ms as u64); let actual = judge::extract_final_integer(&response.text) .unwrap_or_else(|| "no_answer".to_string()); let correct = judge::judge_integer(&response.text, &case.answer); metrics.record_judgement(correct); if !correct { wrong_cases.push(WrongCaseReport { id: case.id, question: case.problem, expected: case.answer, actual, raw_output: response.text, }); } } Err(error) => metrics.record_failure(error_code(&error)), } } let summary = metrics.summary(); let report = benchmark_report(BenchmarkReportInput { benchmark: "aime2026", provider: provider_name, model, dataset, started_at, duration_ms: started.elapsed().as_millis(), concurrency, limit, summary, wrong_cases, }); let report_path = write_benchmark_report(Path::new("."), &report)?; print_benchmark_report(&report, &report_path); Ok(()) } async fn run_gpqa_benchmark( config_path: PathBuf, provider: Option, model: Option, concurrency: usize, limit: Option, ) -> Result<()> { let config = AppConfig::load(&config_path)?; let provider_name = provider_name(&config, provider.as_deref())?; let provider_config = config.resolved_provider(Some(&provider_name))?; let model = model.unwrap_or_else(|| provider_config.default_model.clone()); let loaded = benchmarks::gpqa::load_cases(Path::new(&config.benchmarks.data_dir))?; let dataset = dataset_report( config .benchmarks .gpqa_diamond .as_ref() .map(|dataset| (dataset.source.as_str(), dataset.split.as_str())), &loaded.local_path, ); let cases = apply_limit(loaded.cases, limit); let started_at = Utc::now(); let started = Instant::now(); let base_request = request_template(&provider_config, &model, 0.0, 1024); let protocol = provider_config.protocol; let results = stream::iter(cases) .map(|case| { let mut request = base_request.clone(); request.prompt = case.prompt(); async move { let result = run_model_request(protocol, request).await; (case, result) } }) .buffer_unordered(nonzero_concurrency(concurrency)) .collect::>() .await; let mut metrics = Metrics::new(); let mut wrong_cases = Vec::new(); for (case, result) in results { match result { Ok(response) => { metrics.record_success(response.status, response.elapsed_ms as u64); let actual = judge::extract_choice(&response.text) .map(|choice| choice.to_string()) .unwrap_or_else(|| "no_answer".to_string()); let expected = case.answer.to_string(); let correct = judge::judge_choice(&response.text, case.answer); metrics.record_judgement(correct); if !correct { wrong_cases.push(WrongCaseReport { id: case.id, question: case.question, expected, actual, raw_output: response.text, }); } } Err(error) => metrics.record_failure(error_code(&error)), } } let summary = metrics.summary(); let report = benchmark_report(BenchmarkReportInput { benchmark: "gpqa-diamond", provider: provider_name, model, dataset, started_at, duration_ms: started.elapsed().as_millis(), concurrency, limit, summary, wrong_cases, }); let report_path = write_benchmark_report(Path::new("."), &report)?; print_benchmark_report(&report, &report_path); Ok(()) } async fn run_rpm( config_path: PathBuf, provider: Option, model: Option, rpm: u32, duration: String, prompt: String, ) -> Result<()> { if rpm == 0 { bail!("rpm must be greater than 0"); } let duration = parse_duration(&duration)?; let config = AppConfig::load(&config_path)?; let provider_name = provider_name(&config, provider.as_deref())?; let provider_config = config.resolved_provider(Some(&provider_name))?; let model = model.unwrap_or_else(|| provider_config.default_model.clone()); let request = ModelRequest { prompt, ..request_template(&provider_config, &model, 0.0, 1024) }; let schedule = rpm_start_schedule(duration, rpm); let started_at = Utc::now(); let started = Instant::now(); let tokio_started = TokioInstant::now(); let in_flight_limit = schedule.len().max(1); let mut metrics = Metrics::new(); let results = stream::iter(schedule.into_iter().map(|offset| { let request = request.clone(); async move { sleep_until(tokio_started + offset).await; run_model_request(provider_config.protocol, request).await } })) .buffer_unordered(in_flight_limit) .collect::>() .await; for result in results { match result { Ok(response) => metrics.record_success(response.status, response.elapsed_ms as u64), Err(error) => metrics.record_failure(error_code(&error)), } } let summary = metrics.summary(); let report = RpmReport { benchmark: "rpm".to_string(), provider: provider_name, model, run: RpmRunReport { started_at, duration_ms: started.elapsed().as_millis(), target_rpm: rpm, temperature: 0.0, max_tokens: 1024, }, summary: RpmSummaryReport { actual_requests: summary.total, success: summary.success, failure: summary.failed, latency_ms: latency_report(&summary.latency_ms), }, errors: summary.errors, }; let report_path = write_rpm_report(Path::new("."), &report)?; print_rpm_report(&report, &report_path); Ok(()) } fn provider_name(config: &AppConfig, provider: Option<&str>) -> Result { match provider { Some(provider) => Ok(provider.to_string()), None => config .default_provider .clone() .context("no provider specified and config has no default_provider"), } } fn request_template( provider: &crate::config::ProviderConfig, model: &str, temperature: f32, max_tokens: u32, ) -> ModelRequest { ModelRequest { base_url: provider.base_url.clone(), api_token: provider.api_token.clone(), model: model.to_string(), prompt: String::new(), temperature, max_tokens, } } fn apply_limit(cases: Vec, limit: Option) -> Vec { match limit { Some(limit) => cases.into_iter().take(limit).collect(), None => cases, } } fn dataset_report(config: Option<(&str, &str)>, local_path: &Path) -> DatasetReport { let (source, split) = config.unwrap_or(("local", "train")); DatasetReport { source: source.to_string(), split: split.to_string(), revision: None, local_path: local_path.display().to_string(), } } struct BenchmarkReportInput { benchmark: &'static str, provider: String, model: String, dataset: DatasetReport, started_at: chrono::DateTime, duration_ms: u128, concurrency: usize, limit: Option, summary: MetricsSummary, wrong_cases: Vec, } fn benchmark_report(input: BenchmarkReportInput) -> BenchmarkReport { BenchmarkReport { benchmark: input.benchmark.to_string(), provider: input.provider, model: input.model, dataset: input.dataset, run: RunReport { started_at: input.started_at, duration_ms: input.duration_ms, concurrency: input.concurrency, limit: input.limit, temperature: 0.0, max_tokens: 1024, }, summary: BenchmarkSummaryReport { accuracy: input.summary.accuracy, success: input.summary.success, total: input.summary.total, correct: input.summary.correct, wrong: input.summary.wrong, failed: input.summary.failed, latency_ms: latency_report(&input.summary.latency_ms), }, errors: input.summary.errors, wrong_cases: input.wrong_cases, } } fn latency_report(summary: &LatencySummary) -> LatencyReport { LatencyReport { p50: summary.p50, p95: summary.p95, p99: summary.p99, } } fn print_benchmark_report(report: &BenchmarkReport, report_path: &Path) { println!("benchmark: {}", report.benchmark); println!( "accuracy: {}", report .summary .accuracy .map(|accuracy| format!("{:.2}%", accuracy * 100.0)) .unwrap_or_else(|| "n/a".to_string()) ); println!( "success: {}/{} (failed: {})", report.summary.success, report.summary.total, report.summary.failed ); println!( "latency_ms: p50={} p95={} p99={}", format_optional_latency(report.summary.latency_ms.p50), format_optional_latency(report.summary.latency_ms.p95), format_optional_latency(report.summary.latency_ms.p99) ); println!("errors:"); if report.errors.is_empty() { println!(" none"); } else { for error in &report.errors { println!(" {}: {}", error.code, error.count); } } println!("wrong_cases:"); if report.wrong_cases.is_empty() { println!(" none"); } else { for case in &report.wrong_cases { println!( " {} expected={} actual={}", case.id, case.expected, case.actual ); } } println!("report: {}", report_path.display()); } fn print_rpm_report(report: &RpmReport, report_path: &Path) { println!("target_rpm: {}", report.run.target_rpm); println!("actual_requests: {}", report.summary.actual_requests); println!( "success: {} failed: {}", report.summary.success, report.summary.failure ); println!( "latency_ms: p50={} p95={} p99={}", format_optional_latency(report.summary.latency_ms.p50), format_optional_latency(report.summary.latency_ms.p95), format_optional_latency(report.summary.latency_ms.p99) ); println!("errors:"); if report.errors.is_empty() { println!(" none"); } else { for error in &report.errors { println!(" {}: {}", error.code, error.count); } } println!("report: {}", report_path.display()); } fn format_optional_latency(value: Option) -> String { value .map(|value| value.to_string()) .unwrap_or_else(|| "n/a".to_string()) } fn error_code(error: &anyhow::Error) -> String { let message = error.to_string(); let status_regex = Regex::new(r"status\s+(\d{3})").expect("valid status regex"); status_regex .captures(&message) .and_then(|captures| captures.get(1)) .map(|code| code.as_str().to_string()) .unwrap_or_else(|| "request_error".to_string()) } fn nonzero_concurrency(concurrency: usize) -> usize { concurrency.max(1) } fn parse_duration(value: &str) -> Result { let value = value.trim(); let Some(number) = value.strip_suffix('s') else { if let Some(number) = value.strip_suffix('m') { let minutes = number .parse::() .with_context(|| format!("invalid duration: {value}"))?; return Ok(Duration::from_secs(minutes * 60)); } bail!("invalid duration: expected values like 60s or 5m"); }; let seconds = number .parse::() .with_context(|| format!("invalid duration: {value}"))?; Ok(Duration::from_secs(seconds)) } fn rpm_start_schedule(duration: Duration, rpm: u32) -> Vec { let interval_nanos = 60_000_000_000u128 / u128::from(rpm); if interval_nanos == 0 { return Vec::new(); } let duration_nanos = duration.as_nanos(); let mut offsets = Vec::new(); let mut offset = 0u128; while offset < duration_nanos { offsets.push(Duration::from_nanos(offset.min(u128::from(u64::MAX)) as u64)); offset += interval_nanos; } offsets } fn dataset_data_dir(config_path: &Path) -> Result { if !config_path.exists() { return Ok(PathBuf::from("data/benchmarks")); } let config = AppConfig::load(config_path)?; Ok(PathBuf::from(config.benchmarks.data_dir)) } #[cfg(test)] mod tests { use super::*; #[test] fn dataset_data_dir_defaults_when_config_is_missing() { let temp_dir = tempfile::tempdir().expect("create temp dir"); let missing_config = temp_dir.path().join("missing-config.yaml"); let data_dir = dataset_data_dir(&missing_config).expect("default data dir"); assert_eq!(data_dir, PathBuf::from("data/benchmarks")); } #[test] fn dataset_data_dir_propagates_invalid_existing_config() { let temp_dir = tempfile::tempdir().expect("create temp dir"); let config_path = temp_dir.path().join("config.yaml"); std::fs::write(&config_path, "providers: [").expect("write invalid config"); let error = dataset_data_dir(&config_path).expect_err("invalid config should fail"); assert!(error.to_string().contains("failed to parse config")); } #[test] fn parses_duration_seconds_and_minutes() { assert_eq!(parse_duration("60s").expect("seconds").as_secs(), 60); assert_eq!(parse_duration("5m").expect("minutes").as_secs(), 300); } #[test] fn computes_rpm_start_schedule_from_run_start() { let schedule = rpm_start_schedule(Duration::from_secs(60), 120); assert_eq!(schedule.len(), 120); assert_eq!(schedule[0], Duration::ZERO); assert_eq!(schedule[1], Duration::from_millis(500)); assert_eq!(schedule[119], Duration::from_millis(59_500)); } #[test] fn rejects_invalid_duration() { let error = parse_duration("one hour").expect_err("invalid duration"); assert!(error.to_string().contains("duration")); } }