|
- 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<String>,
- #[arg(long)]
- model: Option<String>,
- #[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<String>,
- #[arg(long)]
- model: Option<String>,
- #[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<String>,
- #[arg(long)]
- model: Option<String>,
- #[arg(long, default_value_t = 4)]
- concurrency: usize,
- #[arg(long)]
- limit: Option<usize>,
- },
- GpqaDiamond {
- #[arg(long, default_value = "config.yaml")]
- config: PathBuf,
- #[arg(long)]
- provider: Option<String>,
- #[arg(long)]
- model: Option<String>,
- #[arg(long, default_value_t = 4)]
- concurrency: usize,
- #[arg(long)]
- limit: Option<usize>,
- },
- }
-
- 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<String>,
- model: Option<String>,
- concurrency: usize,
- limit: Option<usize>,
- ) -> 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::<Vec<_>>()
- .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<String>,
- model: Option<String>,
- concurrency: usize,
- limit: Option<usize>,
- ) -> 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::<Vec<_>>()
- .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<String>,
- model: Option<String>,
- 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::<Vec<_>>()
- .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<String> {
- 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<T>(cases: Vec<T>, limit: Option<usize>) -> Vec<T> {
- 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<Utc>,
- duration_ms: u128,
- concurrency: usize,
- limit: Option<usize>,
- summary: MetricsSummary,
- wrong_cases: Vec<WrongCaseReport>,
- }
-
- 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<u64>) -> 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<Duration> {
- let value = value.trim();
- let Some(number) = value.strip_suffix('s') else {
- if let Some(number) = value.strip_suffix('m') {
- let minutes = number
- .parse::<u64>()
- .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::<u64>()
- .with_context(|| format!("invalid duration: {value}"))?;
- Ok(Duration::from_secs(seconds))
- }
-
- fn rpm_start_schedule(duration: Duration, rpm: u32) -> Vec<Duration> {
- 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<PathBuf> {
- 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"));
- }
- }
|