|
- use crate::benchmarks;
- use crate::benchmarks::judge;
- use crate::config::{AppConfig, ProtocolKind, ProviderThinkingConfig};
- use crate::metrics::{LatencySummary, Metrics, MetricsSummary};
- use crate::report::{
- BenchmarkParamsReport, BenchmarkReport, BenchmarkSummaryReport, CorrectCaseReport,
- DatasetReport, LatencyReport, LimiterInferenceKind, LimiterInferenceReport, PhaseSummaryReport,
- ProbeSecondReport, RequestDebugReport, RpmModeDetailReport, RpmParamsReport, RpmReport,
- RpmRunReport, RpmSummaryReport, RunReport, SentRequestParamsReport, ThinkingParamsReport,
- WindowBoundaryReport, WrongCaseReport, write_benchmark_report, write_rpm_report,
- };
- use crate::rpm_modes::{
- ProbePhase, RpmMode, ScheduledProbe, burst_schedule, sliding_window_schedule,
- sustained_schedule, token_bucket_schedule, window_boundary_plan,
- };
- use crate::runner::{
- ModelRequest, ModelResponse, RawDebugConfig, ThinkingConfig, run_model_request,
- };
- use anyhow::{Context, Result, bail};
- use chrono::Utc;
- use clap::{Parser, Subcommand};
- use futures::{StreamExt, stream};
- use indicatif::{ProgressBar, ProgressStyle};
- use regex::Regex;
- use serde_json::{Value, json};
- use std::collections::BTreeMap;
- 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,
- #[arg(long, num_args = 0..=1, default_missing_value = "true")]
- stream: Option<bool>,
- #[arg(long)]
- debug_raw: bool,
- #[arg(long, num_args = 0..=1, default_missing_value = "true")]
- thinking: Option<bool>,
- #[arg(long)]
- thinking_type: Option<String>,
- #[arg(long)]
- thinking_budget_tokens: Option<u32>,
- #[arg(long)]
- thinking_effort: Option<String>,
- #[arg(long)]
- thinking_display: Option<String>,
- #[arg(long)]
- reasoning_effort: Option<String>,
- #[arg(long)]
- reasoning_summary: Option<String>,
- },
- Dataset {
- #[command(subcommand)]
- command: DatasetCommand,
- },
- Bench {
- #[command(subcommand)]
- command: BenchCommand,
- },
- Rpm {
- #[arg(long, default_value = "config.yaml")]
- config: PathBuf,
- #[arg(long, value_enum, default_value_t = RpmMode::Sustained)]
- mode: RpmMode,
- #[arg(long)]
- provider: Option<String>,
- #[arg(long)]
- model: Option<String>,
- #[arg(long)]
- rpm: Option<u32>,
- #[arg(long, default_value = "60s")]
- duration: String,
- #[arg(long)]
- burst: Option<u32>,
- #[arg(long)]
- probe_seconds: Option<u64>,
- #[arg(long, default_value_t = 500)]
- window_offset_ms: u64,
- #[arg(long)]
- concurrency: Option<usize>,
- #[arg(long)]
- prompt: String,
- #[arg(long, num_args = 0..=1, default_missing_value = "true")]
- stream: Option<bool>,
- #[arg(long)]
- debug_raw: bool,
- #[arg(long, num_args = 0..=1, default_missing_value = "true")]
- thinking: Option<bool>,
- #[arg(long)]
- thinking_type: Option<String>,
- #[arg(long)]
- thinking_budget_tokens: Option<u32>,
- #[arg(long)]
- thinking_effort: Option<String>,
- #[arg(long)]
- thinking_display: Option<String>,
- #[arg(long)]
- reasoning_effort: Option<String>,
- #[arg(long)]
- reasoning_summary: Option<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>,
- #[arg(long, num_args = 0..=1, default_missing_value = "true")]
- stream: Option<bool>,
- #[arg(long, default_value_t = 64_000)]
- max_tokens: u32,
- #[arg(long)]
- debug_raw: bool,
- #[arg(long, num_args = 0..=1, default_missing_value = "true")]
- thinking: Option<bool>,
- #[arg(long)]
- thinking_type: Option<String>,
- #[arg(long)]
- thinking_budget_tokens: Option<u32>,
- #[arg(long)]
- thinking_effort: Option<String>,
- #[arg(long)]
- thinking_display: Option<String>,
- #[arg(long)]
- reasoning_effort: Option<String>,
- #[arg(long)]
- reasoning_summary: Option<String>,
- },
- 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>,
- #[arg(long, num_args = 0..=1, default_missing_value = "true")]
- stream: Option<bool>,
- #[arg(long, default_value_t = 64_000)]
- max_tokens: u32,
- #[arg(long)]
- debug_raw: bool,
- #[arg(long, num_args = 0..=1, default_missing_value = "true")]
- thinking: Option<bool>,
- #[arg(long)]
- thinking_type: Option<String>,
- #[arg(long)]
- thinking_budget_tokens: Option<u32>,
- #[arg(long)]
- thinking_effort: Option<String>,
- #[arg(long)]
- thinking_display: Option<String>,
- #[arg(long)]
- reasoning_effort: Option<String>,
- #[arg(long)]
- reasoning_summary: Option<String>,
- },
- }
-
- pub async fn dispatch(cli: Cli) -> Result<()> {
- match cli.command {
- Command::Check {
- config,
- provider,
- model,
- prompt,
- stream,
- debug_raw,
- thinking,
- thinking_type,
- thinking_budget_tokens,
- thinking_effort,
- thinking_display,
- reasoning_effort,
- reasoning_summary,
- } => {
- let config = AppConfig::load(&config)?;
- let provider_name = provider_name(&config, provider.as_deref())?;
- let provider = config.resolved_provider(Some(&provider_name))?;
- let model = model.unwrap_or_else(|| provider.default_model.clone());
- let request = ModelRequest {
- base_url: provider.base_url.clone(),
- api_token: provider.api_token.clone(),
- model: model.clone(),
- prompt,
- temperature: 0.0,
- max_tokens: 1024,
- stream: stream.unwrap_or(provider.stream),
- raw_debug: raw_debug_config(debug_raw, &provider_name, &model),
- thinking: merged_thinking_config(
- provider.thinking.as_ref(),
- ThinkingOverrides {
- thinking,
- thinking_type,
- thinking_budget_tokens,
- thinking_effort,
- thinking_display,
- reasoning_effort,
- reasoning_summary,
- },
- ),
- };
- let response = run_model_request(provider.protocol, request).await?;
-
- println!("status: {}", response.status);
- println!("elapsed_ms: {}", response.elapsed_ms);
- if let Some(ttft) = response.first_token_ms {
- println!("first_token_ms: {}", ttft);
- }
- 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,
- mode,
- provider,
- model,
- rpm,
- duration,
- burst,
- probe_seconds,
- window_offset_ms,
- concurrency,
- prompt,
- stream,
- debug_raw,
- thinking,
- thinking_type,
- thinking_budget_tokens,
- thinking_effort,
- thinking_display,
- reasoning_effort,
- reasoning_summary,
- } => {
- run_rpm(
- config,
- RpmCommandOptions {
- mode,
- provider,
- model,
- rpm,
- duration,
- burst,
- probe_seconds,
- window_offset_ms,
- concurrency,
- prompt,
- stream,
- debug_raw,
- thinking,
- thinking_type,
- thinking_budget_tokens,
- thinking_effort,
- thinking_display,
- reasoning_effort,
- reasoning_summary,
- },
- )
- .await
- }
- }
- }
-
- async fn dispatch_bench(command: BenchCommand) -> Result<()> {
- match command {
- BenchCommand::Aime2026 {
- config,
- provider,
- model,
- concurrency,
- limit,
- stream,
- max_tokens,
- debug_raw,
- thinking,
- thinking_type,
- thinking_budget_tokens,
- thinking_effort,
- thinking_display,
- reasoning_effort,
- reasoning_summary,
- } => {
- run_aime_benchmark(BenchmarkCommandOptions {
- config,
- provider,
- model,
- concurrency,
- limit,
- stream,
- max_tokens,
- debug_raw,
- thinking_overrides: ThinkingOverrides {
- thinking,
- thinking_type,
- thinking_budget_tokens,
- thinking_effort,
- thinking_display,
- reasoning_effort,
- reasoning_summary,
- },
- })
- .await
- }
- BenchCommand::GpqaDiamond {
- config,
- provider,
- model,
- concurrency,
- limit,
- stream,
- max_tokens,
- debug_raw,
- thinking,
- thinking_type,
- thinking_budget_tokens,
- thinking_effort,
- thinking_display,
- reasoning_effort,
- reasoning_summary,
- } => {
- run_gpqa_benchmark(BenchmarkCommandOptions {
- config,
- provider,
- model,
- concurrency,
- limit,
- stream,
- max_tokens,
- debug_raw,
- thinking_overrides: ThinkingOverrides {
- thinking,
- thinking_type,
- thinking_budget_tokens,
- thinking_effort,
- thinking_display,
- reasoning_effort,
- reasoning_summary,
- },
- })
- .await
- }
- }
- }
-
- struct BenchmarkCommandOptions {
- config: PathBuf,
- provider: Option<String>,
- model: Option<String>,
- concurrency: usize,
- limit: Option<usize>,
- stream: Option<bool>,
- max_tokens: u32,
- debug_raw: bool,
- thinking_overrides: ThinkingOverrides,
- }
-
- async fn run_aime_benchmark(options: BenchmarkCommandOptions) -> Result<()> {
- let config = AppConfig::load(&options.config)?;
- let provider_name = provider_name(&config, options.provider.as_deref())?;
- let provider_config = config.resolved_provider(Some(&provider_name))?;
- let model = options
- .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, options.limit);
- let total = cases.len() as u64;
- let started_at = Utc::now();
- let started = Instant::now();
- let mut base_request = request_template(&provider_config, &model, 0.0, options.max_tokens);
- base_request.stream = options.stream.unwrap_or(provider_config.stream);
- base_request.raw_debug = raw_debug_config(options.debug_raw, &provider_name, &model);
- base_request.thinking = merged_thinking_config(
- provider_config.thinking.as_ref(),
- options.thinking_overrides,
- );
- let protocol = provider_config.protocol;
-
- let pb = ProgressBar::new(total);
- pb.set_style(
- ProgressStyle::default_bar()
- .template("[{elapsed_precise}] {bar:40} {pos}/{len} ({eta})")
- .expect("valid template"),
- );
-
- let mut 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(options.concurrency));
-
- let mut metrics = Metrics::new();
- let mut wrong_cases = Vec::new();
- let mut correct_samples = Vec::new();
- let mut debug_requests = Vec::new();
- while let Some((case, result)) = results.next().await {
- pb.inc(1);
- match result {
- Ok(response) => {
- let prompt = case.prompt();
- if let Some(debug) =
- response_debug_report(Some(case.id.clone()), None, None, prompt, &response)
- {
- debug_requests.push(debug);
- }
- if response_has_no_content(&response) {
- metrics.record_failure("request_error");
- continue;
- }
- metrics.record_success(
- response.status,
- response.elapsed_ms as u64,
- response.first_token_ms.map(|ms| 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 {
- if correct_samples.len() < 5 {
- correct_samples.push(CorrectCaseReport {
- id: case.id,
- question: case.problem,
- expected: case.answer,
- raw_output: response.text,
- });
- }
- } else {
- 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)),
- }
- }
- pb.finish_and_clear();
-
- let summary = metrics.summary();
- let report = benchmark_report(BenchmarkReportInput {
- benchmark: "aime2026",
- provider: provider_name,
- model,
- stream: base_request.stream,
- thinking: thinking_report(base_request.thinking.as_ref()),
- request: sent_request_params(protocol, &base_request),
- dataset,
- started_at,
- duration_ms: started.elapsed().as_millis(),
- concurrency: options.concurrency,
- limit: options.limit,
- max_tokens: options.max_tokens,
- summary,
- correct_samples,
- wrong_cases,
- debug_requests,
- });
- let report_path = write_benchmark_report(Path::new("."), &report)?;
- print_benchmark_report(&report, &report_path);
- Ok(())
- }
-
- async fn run_gpqa_benchmark(options: BenchmarkCommandOptions) -> Result<()> {
- let config = AppConfig::load(&options.config)?;
- let provider_name = provider_name(&config, options.provider.as_deref())?;
- let provider_config = config.resolved_provider(Some(&provider_name))?;
- let model = options
- .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, options.limit);
- let total = cases.len() as u64;
- let started_at = Utc::now();
- let started = Instant::now();
- let mut base_request = request_template(&provider_config, &model, 0.0, options.max_tokens);
- base_request.stream = options.stream.unwrap_or(provider_config.stream);
- base_request.raw_debug = raw_debug_config(options.debug_raw, &provider_name, &model);
- base_request.thinking = merged_thinking_config(
- provider_config.thinking.as_ref(),
- options.thinking_overrides,
- );
- let protocol = provider_config.protocol;
-
- let pb = ProgressBar::new(total);
- pb.set_style(
- ProgressStyle::default_bar()
- .template("[{elapsed_precise}] {bar:40} {pos}/{len} ({eta})")
- .expect("valid template"),
- );
-
- let mut 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(options.concurrency));
-
- let mut metrics = Metrics::new();
- let mut wrong_cases = Vec::new();
- let mut correct_samples = Vec::new();
- let mut debug_requests = Vec::new();
- while let Some((case, result)) = results.next().await {
- pb.inc(1);
- match result {
- Ok(response) => {
- let prompt = case.prompt();
- if let Some(debug) =
- response_debug_report(Some(case.id.clone()), None, None, prompt, &response)
- {
- debug_requests.push(debug);
- }
- if response_has_no_content(&response) {
- metrics.record_failure("request_error");
- continue;
- }
- metrics.record_success(
- response.status,
- response.elapsed_ms as u64,
- response.first_token_ms.map(|ms| 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 {
- if correct_samples.len() < 5 {
- correct_samples.push(CorrectCaseReport {
- id: case.id,
- question: case.question,
- expected,
- raw_output: response.text,
- });
- }
- } else {
- wrong_cases.push(WrongCaseReport {
- id: case.id,
- question: case.question,
- expected,
- actual,
- raw_output: response.text,
- });
- }
- }
- Err(error) => metrics.record_failure(error_code(&error)),
- }
- }
- pb.finish_and_clear();
-
- let summary = metrics.summary();
- let report = benchmark_report(BenchmarkReportInput {
- benchmark: "gpqa-diamond",
- provider: provider_name,
- model,
- stream: base_request.stream,
- thinking: thinking_report(base_request.thinking.as_ref()),
- request: sent_request_params(protocol, &base_request),
- dataset,
- started_at,
- duration_ms: started.elapsed().as_millis(),
- concurrency: options.concurrency,
- limit: options.limit,
- max_tokens: options.max_tokens,
- summary,
- correct_samples,
- wrong_cases,
- debug_requests,
- });
- let report_path = write_benchmark_report(Path::new("."), &report)?;
- print_benchmark_report(&report, &report_path);
- Ok(())
- }
-
- struct RpmCommandOptions {
- mode: RpmMode,
- provider: Option<String>,
- model: Option<String>,
- rpm: Option<u32>,
- duration: String,
- burst: Option<u32>,
- probe_seconds: Option<u64>,
- window_offset_ms: u64,
- concurrency: Option<usize>,
- prompt: String,
- stream: Option<bool>,
- debug_raw: bool,
- thinking: Option<bool>,
- thinking_type: Option<String>,
- thinking_budget_tokens: Option<u32>,
- thinking_effort: Option<String>,
- thinking_display: Option<String>,
- reasoning_effort: Option<String>,
- reasoning_summary: Option<String>,
- }
-
- async fn run_rpm(config_path: PathBuf, options: RpmCommandOptions) -> Result<()> {
- let config = AppConfig::load(&config_path)?;
- let mode_plan = build_rpm_mode_plan(
- options.mode,
- options.rpm,
- &options.duration,
- options.burst,
- options.probe_seconds,
- options.window_offset_ms,
- )?;
- let provider_name = provider_name(&config, options.provider.as_deref())?;
- let provider_config = config.resolved_provider(Some(&provider_name))?;
- let model = options
- .model
- .unwrap_or_else(|| provider_config.default_model.clone());
- let stream_enabled = options.stream.unwrap_or(provider_config.stream);
- let concurrency = options.concurrency.unwrap_or(mode_plan.default_concurrency);
- let request = ModelRequest {
- prompt: options.prompt.clone(),
- stream: stream_enabled,
- raw_debug: raw_debug_config(options.debug_raw, &provider_name, &model),
- thinking: merged_thinking_config(
- provider_config.thinking.as_ref(),
- ThinkingOverrides {
- thinking: options.thinking,
- thinking_type: options.thinking_type,
- thinking_budget_tokens: options.thinking_budget_tokens,
- thinking_effort: options.thinking_effort,
- thinking_display: options.thinking_display,
- reasoning_effort: options.reasoning_effort,
- reasoning_summary: options.reasoning_summary,
- },
- ),
- ..request_template(&provider_config, &model, 0.0, 1024)
- };
- let thinking = thinking_report(request.thinking.as_ref());
- let started_at = Utc::now();
- let started = Instant::now();
- let mut metrics = Metrics::new();
- let mut mode_summary = RpmModeSummaryBuilder::default();
- let mut debug_requests = Vec::new();
- let prompt_for_debug = request.prompt.clone();
-
- let results = run_scheduled_requests(
- provider_config.protocol,
- request,
- mode_plan.probes,
- concurrency,
- )
- .await;
-
- for result in results {
- let success = result
- .result
- .as_ref()
- .is_ok_and(|response| !response_has_no_content(response));
- mode_summary.record(result.phase, result.second, success);
- match result.result {
- Ok(response) => {
- if let Some(debug) = response_debug_report(
- None,
- Some(phase_name(result.phase).to_string()),
- result.second,
- prompt_for_debug.clone(),
- &response,
- ) {
- debug_requests.push(debug);
- }
- if response_has_no_content(&response) {
- metrics.record_failure("request_error");
- continue;
- }
- metrics.record_success(
- response.status,
- response.elapsed_ms as u64,
- response.first_token_ms.map(|ms| 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,
- params: RpmParamsReport {
- prompt: options.prompt,
- stream: stream_enabled,
- thinking,
- duration: options.duration,
- burst: options.burst,
- concurrency,
- probe_seconds: options.probe_seconds,
- window_offset_ms: options.window_offset_ms,
- },
- run: RpmRunReport {
- started_at,
- duration_ms: started.elapsed().as_millis(),
- target_rpm: mode_plan.target_rpm,
- actual_rpm: actual_rpm(summary.total, started.elapsed()),
- 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),
- ttft_ms: latency_report(&summary.ttft_ms),
- },
- mode: mode_plan.mode_name.to_string(),
- mode_detail: mode_summary.into_report(options.mode),
- errors: summary.errors,
- debug_requests,
- };
- let report_path = write_rpm_report(Path::new("."), &report)?;
- print_rpm_report(&report, &report_path);
- Ok(())
- }
-
- #[derive(Debug)]
- struct RpmModePlan {
- mode_name: &'static str,
- target_rpm: u32,
- probes: Vec<ScheduledProbe>,
- default_concurrency: usize,
- }
-
- #[derive(Debug)]
- struct ScheduledResult {
- phase: ProbePhase,
- second: Option<u64>,
- result: Result<crate::runner::ModelResponse>,
- }
-
- fn build_rpm_mode_plan(
- mode: RpmMode,
- rpm: Option<u32>,
- duration: &str,
- burst: Option<u32>,
- probe_seconds: Option<u64>,
- window_offset_ms: u64,
- ) -> Result<RpmModePlan> {
- let target_rpm = rpm.unwrap_or(0);
- let burst = burst.unwrap_or(target_rpm);
- let probes = match mode {
- RpmMode::Sustained => {
- let rpm = require_positive("rpm", rpm)?;
- let duration = parse_duration(duration)?;
- sustained_schedule(duration, rpm)
- .into_iter()
- .map(|offset| ScheduledProbe {
- offset,
- phase: ProbePhase::RefillProbe,
- second: Some(offset.as_secs()),
- })
- .collect()
- }
- RpmMode::Burst => {
- let burst = require_positive_value("burst", burst)?;
- burst_schedule(burst)
- .into_iter()
- .map(|offset| ScheduledProbe {
- offset,
- phase: ProbePhase::Burst,
- second: Some(0),
- })
- .collect()
- }
- RpmMode::TokenBucket => {
- let rpm = require_positive("rpm", rpm)?;
- let burst = require_positive_value("burst", burst)?;
- token_bucket_schedule(rpm, burst, probe_seconds.unwrap_or(30))
- }
- RpmMode::SlidingWindow => {
- let burst = require_positive_value("burst", burst)?;
- sliding_window_schedule(burst, probe_seconds.unwrap_or(90))
- }
- RpmMode::WindowBoundary => {
- let burst = require_positive_value("burst", burst)?;
- window_boundary_plan(Utc::now(), burst, window_offset_ms).probes
- }
- RpmMode::Diagnose => {
- let rpm = require_positive("rpm", rpm)?;
- let burst = require_positive_value("burst", burst)?;
- let mut probes = token_bucket_schedule(rpm, burst, probe_seconds.unwrap_or(90));
- probes.extend(sliding_window_schedule(burst, probe_seconds.unwrap_or(90)));
- probes.extend(window_boundary_plan(Utc::now(), burst, window_offset_ms).probes);
- probes
- }
- };
-
- let default_concurrency = probes.len().max(1);
- Ok(RpmModePlan {
- mode_name: mode_name(mode),
- target_rpm,
- default_concurrency,
- probes,
- })
- }
-
- fn mode_name(mode: RpmMode) -> &'static str {
- match mode {
- RpmMode::Sustained => "sustained",
- RpmMode::Burst => "burst",
- RpmMode::TokenBucket => "token-bucket",
- RpmMode::SlidingWindow => "sliding-window",
- RpmMode::WindowBoundary => "window-boundary",
- RpmMode::Diagnose => "diagnose",
- }
- }
-
- #[derive(Default)]
- struct RpmModeSummaryBuilder {
- phases: BTreeMap<&'static str, PhaseAccumulator>,
- refill_seconds: BTreeMap<u64, PhaseAccumulator>,
- sliding_seconds: BTreeMap<u64, PhaseAccumulator>,
- }
-
- impl RpmModeSummaryBuilder {
- fn record(&mut self, phase: ProbePhase, second: Option<u64>, success: bool) {
- match phase {
- ProbePhase::Burst => self.phase("burst").record(success),
- ProbePhase::RefillProbe => {
- self.phase("refill_probe").record(success);
- self.refill_seconds
- .entry(second.unwrap_or(0))
- .or_default()
- .record(success);
- }
- ProbePhase::SlidingProbe => {
- self.phase("sliding_probe").record(success);
- self.sliding_seconds
- .entry(second.unwrap_or(0))
- .or_default()
- .record(success);
- }
- ProbePhase::BeforeBoundary => self.phase("before_boundary").record(success),
- ProbePhase::AfterBoundary => self.phase("after_boundary").record(success),
- }
- }
-
- fn into_report(self, mode: RpmMode) -> Option<RpmModeDetailReport> {
- if mode == RpmMode::Sustained {
- return None;
- }
-
- let burst = self.phases.get("burst").map(PhaseAccumulator::to_report);
- let refill_probe = probe_seconds_report(self.refill_seconds);
- let sliding_probe = probe_seconds_report(self.sliding_seconds);
- let window_boundary = match (
- self.phases.get("before_boundary"),
- self.phases.get("after_boundary"),
- ) {
- (Some(before), Some(after)) => Some(WindowBoundaryReport {
- before: before.to_report(),
- after: after.to_report(),
- }),
- _ => None,
- };
- let inference = if mode == RpmMode::Diagnose {
- Some(infer_limiter(
- burst.as_ref(),
- &refill_probe,
- &sliding_probe,
- window_boundary.as_ref(),
- ))
- } else {
- None
- };
-
- Some(RpmModeDetailReport {
- burst,
- refill_probe,
- sliding_probe,
- window_boundary,
- inference,
- })
- }
-
- fn phase(&mut self, name: &'static str) -> &mut PhaseAccumulator {
- self.phases.entry(name).or_default()
- }
- }
-
- #[derive(Default)]
- struct PhaseAccumulator {
- sent: u64,
- success: u64,
- failure: u64,
- }
-
- impl PhaseAccumulator {
- fn record(&mut self, success: bool) {
- self.sent += 1;
- if success {
- self.success += 1;
- } else {
- self.failure += 1;
- }
- }
-
- fn to_report(&self) -> PhaseSummaryReport {
- PhaseSummaryReport {
- sent: self.sent,
- success: self.success,
- failure: self.failure,
- }
- }
-
- fn success_rate(&self) -> f64 {
- if self.sent == 0 {
- 0.0
- } else {
- self.success as f64 / self.sent as f64
- }
- }
- }
-
- fn probe_seconds_report(seconds: BTreeMap<u64, PhaseAccumulator>) -> Vec<ProbeSecondReport> {
- seconds
- .into_iter()
- .map(|(second, accumulator)| ProbeSecondReport {
- second,
- sent: accumulator.sent,
- success: accumulator.success,
- failure: accumulator.failure,
- })
- .collect()
- }
-
- fn infer_limiter(
- _burst: Option<&PhaseSummaryReport>,
- refill_probe: &[ProbeSecondReport],
- sliding_probe: &[ProbeSecondReport],
- window_boundary: Option<&WindowBoundaryReport>,
- ) -> LimiterInferenceReport {
- let mut signals = Vec::new();
-
- if let Some(boundary) = window_boundary {
- let before_rate = phase_success_rate(&boundary.before);
- let after_rate = phase_success_rate(&boundary.after);
- if after_rate > before_rate + 0.3 {
- signals.push("after-boundary success rate was much higher than before-boundary".into());
- return LimiterInferenceReport {
- likely_limiter: LimiterInferenceKind::FixedWindow,
- confidence: "medium".to_string(),
- signals,
- };
- }
- }
-
- let refill_sent: u64 = refill_probe.iter().map(|probe| probe.sent).sum();
- let refill_success: u64 = refill_probe.iter().map(|probe| probe.success).sum();
- if refill_sent > 0 && refill_success as f64 / refill_sent as f64 >= 0.5 {
- signals.push("refill probes recovered at a steady rate".into());
- return LimiterInferenceReport {
- likely_limiter: LimiterInferenceKind::TokenBucket,
- confidence: "medium".to_string(),
- signals,
- };
- }
-
- let early = sliding_probe
- .iter()
- .filter(|probe| probe.second <= 30)
- .fold(PhaseAccumulator::default(), |mut acc, probe| {
- acc.sent += probe.sent;
- acc.success += probe.success;
- acc.failure += probe.failure;
- acc
- });
- let late = sliding_probe
- .iter()
- .filter(|probe| probe.second >= 60)
- .fold(PhaseAccumulator::default(), |mut acc, probe| {
- acc.sent += probe.sent;
- acc.success += probe.success;
- acc.failure += probe.failure;
- acc
- });
- if late.sent > 0 && late.success_rate() > early.success_rate() + 0.3 {
- signals.push("probe recovery improved near the 60 second rolling window".into());
- return LimiterInferenceReport {
- likely_limiter: LimiterInferenceKind::SlidingWindow,
- confidence: "medium".to_string(),
- signals,
- };
- }
-
- signals.push("signals did not clearly match a limiter model".into());
- LimiterInferenceReport {
- likely_limiter: LimiterInferenceKind::Unknown,
- confidence: "low".to_string(),
- signals,
- }
- }
-
- fn phase_success_rate(phase: &PhaseSummaryReport) -> f64 {
- if phase.sent == 0 {
- 0.0
- } else {
- phase.success as f64 / phase.sent as f64
- }
- }
-
- fn actual_rpm(total_requests: u64, elapsed: Duration) -> Option<f64> {
- let elapsed_seconds = elapsed.as_secs_f64();
- if elapsed_seconds == 0.0 {
- None
- } else {
- Some(total_requests as f64 / elapsed_seconds * 60.0)
- }
- }
-
- fn require_positive(name: &str, value: Option<u32>) -> Result<u32> {
- let value = value.with_context(|| format!("{name} is required for this rpm mode"))?;
- require_positive_value(name, value)
- }
-
- fn require_positive_value(name: &str, value: u32) -> Result<u32> {
- if value == 0 {
- bail!("{name} must be greater than 0");
- }
- Ok(value)
- }
-
- async fn run_scheduled_requests(
- protocol: crate::config::ProtocolKind,
- request: ModelRequest,
- starts: Vec<ScheduledProbe>,
- max_in_flight: usize,
- ) -> Vec<ScheduledResult> {
- let total = starts.len() as u64;
- let pb = ProgressBar::new(total);
- pb.set_style(
- ProgressStyle::default_bar()
- .template("[{elapsed_precise}] {bar:40} {pos}/{len} ({eta})")
- .expect("valid template"),
- );
-
- let tokio_started = TokioInstant::now();
- let mut results = Vec::with_capacity(starts.len());
- let mut s = stream::iter(starts.into_iter().map(|start| {
- let request = request.clone();
- async move {
- sleep_until(tokio_started + start.offset).await;
- ScheduledResult {
- phase: start.phase,
- second: start.second,
- result: run_model_request(protocol, request).await,
- }
- }
- }))
- .buffer_unordered(nonzero_concurrency(max_in_flight));
-
- while let Some(result) = s.next().await {
- pb.inc(1);
- results.push(result);
- }
- pb.finish_and_clear();
- results
- }
-
- 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 raw_debug_config(enabled: bool, provider: &str, model: &str) -> Option<RawDebugConfig> {
- Some(
- RawDebugConfig::new(
- PathBuf::from("outputs/debug"),
- format!("{provider}-{model}"),
- )
- .with_success_raw(enabled),
- )
- }
-
- #[derive(Debug, Default)]
- struct ThinkingOverrides {
- thinking: Option<bool>,
- thinking_type: Option<String>,
- thinking_budget_tokens: Option<u32>,
- thinking_effort: Option<String>,
- thinking_display: Option<String>,
- reasoning_effort: Option<String>,
- reasoning_summary: Option<String>,
- }
-
- fn merged_thinking_config(
- provider: Option<&ProviderThinkingConfig>,
- overrides: ThinkingOverrides,
- ) -> Option<ThinkingConfig> {
- let has_overrides = overrides.thinking.is_some()
- || overrides.thinking_type.is_some()
- || overrides.thinking_budget_tokens.is_some()
- || overrides.thinking_effort.is_some()
- || overrides.thinking_display.is_some()
- || overrides.reasoning_effort.is_some()
- || overrides.reasoning_summary.is_some();
- if provider.is_none() && !has_overrides {
- return None;
- }
-
- let enabled = overrides
- .thinking
- .or_else(|| provider.and_then(|thinking| thinking.enabled))
- .unwrap_or(true);
- if !enabled {
- return Some(ThinkingConfig {
- enabled: false,
- kind: overrides
- .thinking_type
- .or_else(|| provider.and_then(|thinking| thinking.kind.clone())),
- budget_tokens: overrides
- .thinking_budget_tokens
- .or_else(|| provider.and_then(|thinking| thinking.budget_tokens)),
- effort: overrides
- .thinking_effort
- .or_else(|| provider.and_then(|thinking| thinking.effort.clone())),
- display: overrides
- .thinking_display
- .or_else(|| provider.and_then(|thinking| thinking.display.clone())),
- reasoning_effort: overrides
- .reasoning_effort
- .or_else(|| provider.and_then(|thinking| thinking.reasoning_effort.clone())),
- reasoning_summary: overrides
- .reasoning_summary
- .or_else(|| provider.and_then(|thinking| thinking.reasoning_summary.clone())),
- });
- }
-
- Some(ThinkingConfig {
- enabled,
- kind: overrides
- .thinking_type
- .or_else(|| provider.and_then(|thinking| thinking.kind.clone()))
- .or_else(|| Some("enabled".to_string())),
- budget_tokens: overrides
- .thinking_budget_tokens
- .or_else(|| provider.and_then(|thinking| thinking.budget_tokens)),
- effort: overrides
- .thinking_effort
- .or_else(|| provider.and_then(|thinking| thinking.effort.clone())),
- display: overrides
- .thinking_display
- .or_else(|| provider.and_then(|thinking| thinking.display.clone())),
- reasoning_effort: overrides
- .reasoning_effort
- .or_else(|| provider.and_then(|thinking| thinking.reasoning_effort.clone())),
- reasoning_summary: overrides
- .reasoning_summary
- .or_else(|| provider.and_then(|thinking| thinking.reasoning_summary.clone())),
- })
- }
-
- 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,
- stream: provider.stream,
- raw_debug: None,
- thinking: merged_thinking_config(provider.thinking.as_ref(), ThinkingOverrides::default()),
- }
- }
-
- 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(),
- }
- }
-
- fn thinking_report(thinking: Option<&ThinkingConfig>) -> Option<ThinkingParamsReport> {
- thinking.map(|thinking| ThinkingParamsReport {
- enabled: thinking.enabled,
- kind: thinking.kind.clone(),
- budget_tokens: thinking.budget_tokens,
- effort: thinking.effort.clone(),
- display: thinking.display.clone(),
- reasoning_effort: thinking.reasoning_effort.clone(),
- reasoning_summary: thinking.reasoning_summary.clone(),
- })
- }
-
- struct BenchmarkReportInput {
- benchmark: &'static str,
- provider: String,
- model: String,
- stream: bool,
- thinking: Option<ThinkingParamsReport>,
- request: SentRequestParamsReport,
- dataset: DatasetReport,
- started_at: chrono::DateTime<Utc>,
- duration_ms: u128,
- concurrency: usize,
- limit: Option<usize>,
- max_tokens: u32,
- summary: MetricsSummary,
- correct_samples: Vec<CorrectCaseReport>,
- wrong_cases: Vec<WrongCaseReport>,
- debug_requests: Vec<RequestDebugReport>,
- }
-
- fn benchmark_report(input: BenchmarkReportInput) -> BenchmarkReport {
- BenchmarkReport {
- benchmark: input.benchmark.to_string(),
- provider: input.provider,
- model: input.model,
- params: BenchmarkParamsReport {
- stream: input.stream,
- thinking: input.thinking,
- request: input.request,
- },
- 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: input.max_tokens,
- },
- 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),
- ttft_ms: latency_report(&input.summary.ttft_ms),
- },
- errors: input.summary.errors,
- correct_samples: input.correct_samples,
- wrong_cases: input.wrong_cases,
- debug_requests: input.debug_requests,
- }
- }
-
- const NEAR_TTFT_LATENCY_THRESHOLD_MS: u128 = 1_000;
-
- fn response_debug_report(
- id: Option<String>,
- phase: Option<String>,
- second: Option<u64>,
- prompt: String,
- response: &ModelResponse,
- ) -> Option<RequestDebugReport> {
- let ttft_latency_delta_ms = response
- .first_token_ms
- .map(|ttft_ms| response.elapsed_ms.abs_diff(ttft_ms));
- let near_ttft_latency =
- ttft_latency_delta_ms.is_some_and(|delta_ms| delta_ms < NEAR_TTFT_LATENCY_THRESHOLD_MS);
- let has_think_tags = contains_think_tags(&response.text);
- let no_content = response_has_no_content(response);
-
- let mut reasons = Vec::new();
- if no_content {
- reasons.push("no_content");
- }
- if near_ttft_latency {
- reasons.push("near_ttft_latency");
- }
- if has_think_tags {
- reasons.push("contains_think_tags");
- }
- if reasons.is_empty() {
- return None;
- }
-
- Some(RequestDebugReport {
- reason: reasons.join(","),
- id,
- phase,
- second,
- prompt,
- output: response.text.clone(),
- latency_ms: response.elapsed_ms,
- ttft_ms: response.first_token_ms,
- ttft_latency_delta_ms,
- has_think_tags,
- })
- }
-
- fn response_has_no_content(response: &ModelResponse) -> bool {
- response.text.trim().is_empty()
- }
-
- fn contains_think_tags(text: &str) -> bool {
- let lower = text.to_ascii_lowercase();
- lower.contains("<think>") && lower.contains("</think>")
- }
-
- fn phase_name(phase: ProbePhase) -> &'static str {
- match phase {
- ProbePhase::Burst => "burst",
- ProbePhase::RefillProbe => "refill_probe",
- ProbePhase::SlidingProbe => "sliding_probe",
- ProbePhase::BeforeBoundary => "before_boundary",
- ProbePhase::AfterBoundary => "after_boundary",
- }
- }
-
- fn sent_request_params(protocol: ProtocolKind, request: &ModelRequest) -> SentRequestParamsReport {
- let protocol_name = match protocol {
- ProtocolKind::Openai => "openai",
- ProtocolKind::Anthropic => "anthropic",
- ProtocolKind::Google => "google",
- };
- SentRequestParamsReport {
- protocol: protocol_name.to_string(),
- body: sent_request_body_params(protocol, request),
- }
- }
-
- fn sent_request_body_params(protocol: ProtocolKind, request: &ModelRequest) -> Value {
- match protocol {
- ProtocolKind::Openai => {
- let mut body = serde_json::Map::new();
- body.insert("temperature".to_string(), json!(request.temperature));
- body.insert("max_tokens".to_string(), json!(request.max_tokens));
- if request.stream {
- body.insert("stream".to_string(), json!(true));
- }
- if let Some(thinking) = &request.thinking
- && thinking.enabled
- {
- if let Some(reasoning_effort) = thinking
- .reasoning_effort
- .as_ref()
- .or(thinking.effort.as_ref())
- {
- body.insert("reasoning_effort".to_string(), json!(reasoning_effort));
- }
- if let Some(reasoning_summary) = &thinking.reasoning_summary {
- body.insert("reasoning_summary".to_string(), json!(reasoning_summary));
- }
- }
- Value::Object(body)
- }
- ProtocolKind::Anthropic => {
- let mut body = serde_json::Map::new();
- body.insert("max_tokens".to_string(), json!(request.max_tokens));
- if request.stream {
- body.insert("stream".to_string(), json!(true));
- }
- if let Some(thinking) = &request.thinking
- && thinking.enabled
- {
- let mut thinking_body = serde_json::Map::new();
- let thinking_type = thinking.kind.as_deref().unwrap_or("enabled");
- thinking_body.insert("type".to_string(), json!(thinking_type));
- if thinking_type == "adaptive" {
- if let Some(effort) = &thinking.effort {
- thinking_body.insert("effort".to_string(), json!(effort));
- }
- } else if let Some(budget_tokens) = thinking.budget_tokens {
- thinking_body.insert("budget_tokens".to_string(), json!(budget_tokens));
- }
- if let Some(display) = &thinking.display {
- thinking_body.insert("display".to_string(), json!(display));
- }
- body.insert("thinking".to_string(), Value::Object(thinking_body));
- } else {
- body.insert("temperature".to_string(), json!(request.temperature));
- }
- Value::Object(body)
- }
- ProtocolKind::Google => {
- let mut generation_config = serde_json::Map::new();
- generation_config.insert("temperature".to_string(), json!(request.temperature));
- generation_config.insert("maxOutputTokens".to_string(), json!(request.max_tokens));
- if let Some(thinking) = &request.thinking
- && thinking.enabled
- {
- let mut thinking_config = serde_json::Map::new();
- if let Some(display) = &thinking.display {
- thinking_config.insert(
- "includeThoughts".to_string(),
- json!(display != "omitted" && display != "false"),
- );
- }
- if let Some(budget_tokens) = thinking.budget_tokens {
- thinking_config.insert("thinkingBudget".to_string(), json!(budget_tokens));
- }
- if let Some(effort) = &thinking.effort {
- thinking_config.insert("thinkingLevel".to_string(), json!(effort));
- }
- if !thinking_config.is_empty() {
- generation_config
- .insert("thinkingConfig".to_string(), Value::Object(thinking_config));
- }
- }
- let mut body = serde_json::Map::new();
- body.insert(
- "generationConfig".to_string(),
- Value::Object(generation_config),
- );
- Value::Object(body)
- }
- }
- }
-
- 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)
- );
- if report.summary.ttft_ms.p50.is_some() {
- println!(
- "ttft_ms: p50={} p95={} p99={}",
- format_optional_latency(report.summary.ttft_ms.p50),
- format_optional_latency(report.summary.ttft_ms.p95),
- format_optional_latency(report.summary.ttft_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!("mode: {}", report.mode);
- println!("target_rpm: {}", report.run.target_rpm);
- println!(
- "actual_rpm: {}",
- report
- .run
- .actual_rpm
- .map(|rpm| format!("{rpm:.2}"))
- .unwrap_or_else(|| "n/a".to_string())
- );
- 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)
- );
- if report.summary.ttft_ms.p50.is_some() {
- println!(
- "ttft_ms: p50={} p95={} p99={}",
- format_optional_latency(report.summary.ttft_ms.p50),
- format_optional_latency(report.summary.ttft_ms.p95),
- format_optional_latency(report.summary.ttft_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 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 = sustained_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"));
- }
-
- #[test]
- fn rpm_command_defaults_to_sustained_mode() {
- let cli = Cli::try_parse_from([
- "lq_token_test",
- "rpm",
- "--provider",
- "anthropic",
- "--rpm",
- "120",
- "--prompt",
- "hello",
- ])
- .expect("parse rpm");
-
- let Command::Rpm { mode, rpm, .. } = cli.command else {
- panic!("expected rpm command");
- };
-
- assert_eq!(mode, RpmMode::Sustained);
- assert_eq!(rpm, Some(120));
- }
-
- #[test]
- fn rpm_command_parses_token_bucket_mode() {
- let cli = Cli::try_parse_from([
- "lq_token_test",
- "rpm",
- "--mode",
- "token-bucket",
- "--provider",
- "anthropic",
- "--rpm",
- "120",
- "--burst",
- "120",
- "--probe-seconds",
- "30",
- "--prompt",
- "hello",
- ])
- .expect("parse token bucket rpm");
-
- let Command::Rpm {
- mode,
- burst,
- probe_seconds,
- ..
- } = cli.command
- else {
- panic!("expected rpm command");
- };
-
- assert_eq!(mode, RpmMode::TokenBucket);
- assert_eq!(burst, Some(120));
- assert_eq!(probe_seconds, Some(30));
- }
-
- #[test]
- fn check_command_parses_bare_stream_as_true() {
- let cli = Cli::try_parse_from([
- "lq_token_test",
- "check",
- "--provider",
- "anthropic",
- "--prompt",
- "hello",
- "--stream",
- ])
- .expect("parse check stream");
-
- let Command::Check { stream, .. } = cli.command else {
- panic!("expected check command");
- };
-
- assert_eq!(stream, Some(true));
- }
-
- #[test]
- fn bench_command_parses_explicit_stream_false() {
- let cli = Cli::try_parse_from([
- "lq_token_test",
- "bench",
- "aime2026",
- "--provider",
- "anthropic",
- "--stream",
- "false",
- ])
- .expect("parse bench stream");
-
- let Command::Bench {
- command: BenchCommand::Aime2026 { stream, .. },
- } = cli.command
- else {
- panic!("expected aime2026 bench command");
- };
-
- assert_eq!(stream, Some(false));
- }
-
- #[test]
- fn bench_command_defaults_max_tokens_to_64000() {
- let cli = Cli::try_parse_from(["lq_token_test", "bench", "gpqa-diamond"])
- .expect("parse gpqa bench");
-
- let Command::Bench {
- command: BenchCommand::GpqaDiamond { max_tokens, .. },
- } = cli.command
- else {
- panic!("expected gpqa-diamond bench command");
- };
-
- assert_eq!(max_tokens, 64_000);
- }
-
- #[test]
- fn check_command_parses_debug_raw() {
- let cli = Cli::try_parse_from([
- "lq_token_test",
- "check",
- "--provider",
- "anthropic",
- "--prompt",
- "hello",
- "--debug-raw",
- ])
- .expect("parse check debug raw");
-
- let Command::Check { debug_raw, .. } = cli.command else {
- panic!("expected check command");
- };
-
- assert!(debug_raw);
- }
-
- #[test]
- fn raw_debug_config_is_created_even_when_success_raw_is_disabled() {
- let raw_debug = raw_debug_config(false, "google", "gemini-test");
-
- assert!(raw_debug.is_some());
- }
-
- #[test]
- fn check_command_parses_thinking_overrides() {
- let cli = Cli::try_parse_from([
- "lq_token_test",
- "check",
- "--provider",
- "anthropic",
- "--prompt",
- "hello",
- "--thinking",
- "true",
- "--thinking-type",
- "adaptive",
- "--thinking-budget-tokens",
- "10000",
- "--thinking-effort",
- "high",
- "--thinking-display",
- "omitted",
- "--reasoning-effort",
- "high",
- "--reasoning-summary",
- "auto",
- ])
- .expect("parse thinking options");
-
- let Command::Check {
- thinking,
- thinking_type,
- thinking_budget_tokens,
- thinking_effort,
- thinking_display,
- reasoning_effort,
- reasoning_summary,
- ..
- } = cli.command
- else {
- panic!("expected check command");
- };
-
- assert_eq!(thinking, Some(true));
- assert_eq!(thinking_type.as_deref(), Some("adaptive"));
- assert_eq!(thinking_budget_tokens, Some(10000));
- assert_eq!(thinking_effort.as_deref(), Some("high"));
- assert_eq!(thinking_display.as_deref(), Some("omitted"));
- assert_eq!(reasoning_effort.as_deref(), Some("high"));
- assert_eq!(reasoning_summary.as_deref(), Some("auto"));
- }
-
- #[test]
- fn check_command_parses_bare_thinking_as_true() {
- let cli = Cli::try_parse_from([
- "lq_token_test",
- "check",
- "--provider",
- "anthropic",
- "--prompt",
- "hello",
- "--thinking",
- ])
- .expect("parse bare thinking option");
-
- let Command::Check { thinking, .. } = cli.command else {
- panic!("expected check command");
- };
-
- assert_eq!(thinking, Some(true));
- }
-
- #[test]
- fn disabled_thinking_report_keeps_cli_budget_override() {
- let provider = ProviderThinkingConfig {
- enabled: Some(true),
- kind: Some("enabled".to_string()),
- budget_tokens: Some(10000),
- effort: None,
- display: None,
- reasoning_effort: None,
- reasoning_summary: None,
- };
-
- let merged = merged_thinking_config(
- Some(&provider),
- ThinkingOverrides {
- thinking: Some(false),
- thinking_budget_tokens: Some(20000),
- ..ThinkingOverrides::default()
- },
- )
- .expect("disabled thinking config");
-
- assert!(!merged.enabled);
- assert_eq!(merged.budget_tokens, Some(20000));
- }
-
- #[test]
- fn sent_anthropic_adaptive_request_params_omit_temperature_and_nulls() {
- let request = ModelRequest {
- base_url: "https://example.test".to_string(),
- api_token: "secret".to_string(),
- model: "claude-test".to_string(),
- prompt: "hello".to_string(),
- temperature: 0.0,
- max_tokens: 32_768,
- stream: true,
- raw_debug: None,
- thinking: Some(ThinkingConfig {
- enabled: true,
- kind: Some("adaptive".to_string()),
- budget_tokens: None,
- effort: Some("high".to_string()),
- display: Some("summarized".to_string()),
- reasoning_effort: None,
- reasoning_summary: None,
- }),
- };
-
- let params = sent_request_params(ProtocolKind::Anthropic, &request);
-
- assert_eq!(params.protocol, "anthropic");
- assert_eq!(params.body["max_tokens"], 32_768);
- assert_eq!(params.body["stream"], true);
- assert_eq!(params.body["thinking"]["type"], "adaptive");
- assert_eq!(params.body["thinking"]["effort"], "high");
- assert_eq!(params.body["thinking"]["display"], "summarized");
- assert!(params.body.get("temperature").is_none());
- assert!(params.body["thinking"].get("budget_tokens").is_none());
- assert!(params.body["thinking"].get("reasoning_effort").is_none());
- }
-
- #[test]
- fn sent_google_request_params_use_generation_config_names() {
- let request = ModelRequest {
- base_url: "https://example.test".to_string(),
- api_token: "secret".to_string(),
- model: "gemini-test".to_string(),
- prompt: "hello".to_string(),
- temperature: 0.0,
- max_tokens: 32_768,
- stream: true,
- raw_debug: None,
- thinking: Some(ThinkingConfig {
- enabled: true,
- kind: None,
- budget_tokens: Some(5000),
- effort: Some("high".to_string()),
- display: Some("summarized".to_string()),
- reasoning_effort: None,
- reasoning_summary: None,
- }),
- };
-
- let params = sent_request_params(ProtocolKind::Google, &request);
-
- assert_eq!(params.protocol, "google");
- assert_eq!(params.body["generationConfig"]["temperature"], 0.0);
- assert_eq!(params.body["generationConfig"]["maxOutputTokens"], 32_768);
- assert_eq!(
- params.body["generationConfig"]["thinkingConfig"]["thinkingBudget"],
- 5000
- );
- assert_eq!(
- params.body["generationConfig"]["thinkingConfig"]["thinkingLevel"],
- "high"
- );
- assert_eq!(
- params.body["generationConfig"]["thinkingConfig"]["includeThoughts"],
- true
- );
- }
-
- #[test]
- fn response_debug_report_records_near_ttft_latency_with_full_prompt_and_output() {
- let response = ModelResponse {
- text: "final answer".to_string(),
- status: 200,
- elapsed_ms: 1_500,
- first_token_ms: Some(750),
- };
-
- let debug = response_debug_report(
- Some("case-1".to_string()),
- None,
- None,
- "full prompt".to_string(),
- &response,
- )
- .expect("near ttft/latency should be debugged");
-
- assert_eq!(debug.reason, "near_ttft_latency");
- assert_eq!(debug.id.as_deref(), Some("case-1"));
- assert_eq!(debug.prompt, "full prompt");
- assert_eq!(debug.output, "final answer");
- assert_eq!(debug.latency_ms, 1_500);
- assert_eq!(debug.ttft_ms, Some(750));
- assert_eq!(debug.ttft_latency_delta_ms, Some(750));
- assert!(!debug.has_think_tags);
- }
-
- #[test]
- fn response_debug_report_marks_think_tags() {
- let response = ModelResponse {
- text: "<think>hidden</think>\nanswer".to_string(),
- status: 200,
- elapsed_ms: 5_000,
- first_token_ms: Some(100),
- };
-
- let debug = response_debug_report(None, None, None, "prompt".to_string(), &response)
- .expect("think tags should be debugged");
-
- assert_eq!(debug.reason, "contains_think_tags");
- assert!(debug.has_think_tags);
- assert_eq!(debug.ttft_latency_delta_ms, Some(4_900));
- }
-
- #[test]
- fn response_debug_report_records_empty_content_as_request_debug() {
- let response = ModelResponse {
- text: " ".to_string(),
- status: 200,
- elapsed_ms: 2_000,
- first_token_ms: None,
- };
-
- let debug = response_debug_report(
- None,
- Some("burst".to_string()),
- Some(0),
- "prompt".to_string(),
- &response,
- )
- .expect("empty content should be debugged");
-
- assert_eq!(debug.reason, "no_content");
- assert_eq!(debug.phase.as_deref(), Some("burst"));
- assert_eq!(debug.second, Some(0));
- assert!(response_has_no_content(&response));
- }
-
- #[test]
- fn rpm_command_parses_window_boundary_offset() {
- let cli = Cli::try_parse_from([
- "lq_token_test",
- "rpm",
- "--mode",
- "window-boundary",
- "--provider",
- "anthropic",
- "--burst",
- "10",
- "--window-offset-ms",
- "250",
- "--prompt",
- "hello",
- ])
- .expect("parse window boundary rpm");
-
- let Command::Rpm {
- mode,
- window_offset_ms,
- ..
- } = cli.command
- else {
- panic!("expected rpm command");
- };
-
- assert_eq!(mode, RpmMode::WindowBoundary);
- assert_eq!(window_offset_ms, 250);
- }
-
- #[test]
- fn rpm_mode_validation_rejects_zero_values() {
- let zero_rpm = build_rpm_mode_plan(RpmMode::Sustained, Some(0), "60s", None, None, 500)
- .expect_err("zero rpm should fail");
- let zero_burst = build_rpm_mode_plan(RpmMode::Burst, None, "60s", Some(0), None, 500)
- .expect_err("zero burst should fail");
-
- assert!(zero_rpm.to_string().contains("rpm"));
- assert!(zero_burst.to_string().contains("burst"));
- }
- }
|