diff --git a/src/cli.rs b/src/cli.rs index 5ada6dc..3a1cd53 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -4,10 +4,10 @@ use crate::config::AppConfig; use crate::metrics::{LatencySummary, Metrics, MetricsSummary}; use crate::report::{ BenchmarkParamsReport, BenchmarkReport, BenchmarkSummaryReport, CorrectCaseReport, - DatasetReport, LatencyReport, LimiterInferenceKind, LimiterInferenceReport, - PhaseSummaryReport, ProbeSecondReport, RpmModeDetailReport, RpmParamsReport, RpmReport, - RpmRunReport, RpmSummaryReport, RunReport, WindowBoundaryReport, WrongCaseReport, - write_benchmark_report, write_rpm_report, + DatasetReport, LatencyReport, LimiterInferenceKind, LimiterInferenceReport, PhaseSummaryReport, + ProbeSecondReport, RpmModeDetailReport, RpmParamsReport, RpmReport, RpmRunReport, + RpmSummaryReport, RunReport, WindowBoundaryReport, WrongCaseReport, write_benchmark_report, + write_rpm_report, }; use crate::rpm_modes::{ ProbePhase, RpmMode, ScheduledProbe, burst_schedule, sliding_window_schedule, @@ -47,7 +47,7 @@ pub enum Command { model: Option, #[arg(long)] prompt: String, - #[arg(long)] + #[arg(long, num_args = 0..=1, default_missing_value = "true")] stream: Option, }, Dataset { @@ -81,7 +81,7 @@ pub enum Command { concurrency: Option, #[arg(long)] prompt: String, - #[arg(long)] + #[arg(long, num_args = 0..=1, default_missing_value = "true")] stream: Option, }, } @@ -104,7 +104,7 @@ pub enum BenchCommand { concurrency: usize, #[arg(long)] limit: Option, - #[arg(long)] + #[arg(long, num_args = 0..=1, default_missing_value = "true")] stream: Option, #[arg(long, default_value_t = 32768)] max_tokens: u32, @@ -120,7 +120,7 @@ pub enum BenchCommand { concurrency: usize, #[arg(long)] limit: Option, - #[arg(long)] + #[arg(long, num_args = 0..=1, default_missing_value = "true")] stream: Option, #[arg(long, default_value_t = 32768)] max_tokens: u32, @@ -212,7 +212,18 @@ async fn dispatch_bench(command: BenchCommand) -> Result<()> { limit, stream, max_tokens, - } => run_aime_benchmark(config, provider, model, concurrency, limit, stream, max_tokens).await, + } => { + run_aime_benchmark( + config, + provider, + model, + concurrency, + limit, + stream, + max_tokens, + ) + .await + } BenchCommand::GpqaDiamond { config, provider, @@ -221,7 +232,18 @@ async fn dispatch_bench(command: BenchCommand) -> Result<()> { limit, stream, max_tokens, - } => run_gpqa_benchmark(config, provider, model, concurrency, limit, stream, max_tokens).await, + } => { + run_gpqa_benchmark( + config, + provider, + model, + concurrency, + limit, + stream, + max_tokens, + ) + .await + } } } @@ -280,7 +302,11 @@ async fn run_aime_benchmark( pb.inc(1); match result { Ok(response) => { - metrics.record_success(response.status, response.elapsed_ms as u64, response.first_token_ms.map(|ms| ms as u64)); + 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); @@ -385,7 +411,11 @@ async fn run_gpqa_benchmark( pb.inc(1); match result { Ok(response) => { - metrics.record_success(response.status, response.elapsed_ms as u64, response.first_token_ms.map(|ms| ms as u64)); + 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()); @@ -490,7 +520,11 @@ async fn run_rpm(config_path: PathBuf, options: RpmCommandOptions) -> Result<()> let success = result.result.is_ok(); mode_summary.record(result.phase, result.second, success); match result.result { - Ok(response) => metrics.record_success(response.status, response.elapsed_ms as u64, response.first_token_ms.map(|ms| ms as u64)), + Ok(response) => 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)), } } @@ -1220,6 +1254,49 @@ mod tests { 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 rpm_command_parses_window_boundary_offset() { let cli = Cli::try_parse_from([ diff --git a/src/protocols/anthropic.rs b/src/protocols/anthropic.rs index 03ad004..f210dfa 100644 --- a/src/protocols/anthropic.rs +++ b/src/protocols/anthropic.rs @@ -192,14 +192,14 @@ mod tests { .await; let request = ModelRequest { - base_url: server.uri(), - api_token: "test-token".to_string(), - model: "claude-test".to_string(), - prompt: "hello".to_string(), - temperature: 0.0, - max_tokens: 1024, - stream: false -}; + base_url: server.uri(), + api_token: "test-token".to_string(), + model: "claude-test".to_string(), + prompt: "hello".to_string(), + temperature: 0.0, + max_tokens: 1024, + stream: false, + }; let response = super::send(&Client::new(), &request) .await @@ -223,14 +223,14 @@ mod tests { .await; let request = ModelRequest { - base_url: server.uri(), - api_token: "test-token".to_string(), - model: "claude-test".to_string(), - prompt: "hello".to_string(), - temperature: 0.0, - max_tokens: 1024, - stream: false -}; + base_url: server.uri(), + api_token: "test-token".to_string(), + model: "claude-test".to_string(), + prompt: "hello".to_string(), + temperature: 0.0, + max_tokens: 1024, + stream: false, + }; let error = super::send(&Client::new(), &request) .await @@ -254,14 +254,14 @@ mod tests { .await; let request = ModelRequest { - base_url: server.uri(), - api_token: "test-token".to_string(), - model: "claude-test".to_string(), - prompt: "hello".to_string(), - temperature: 0.0, - max_tokens: 1024, - stream: false -}; + base_url: server.uri(), + api_token: "test-token".to_string(), + model: "claude-test".to_string(), + prompt: "hello".to_string(), + temperature: 0.0, + max_tokens: 1024, + stream: false, + }; let message = super::send(&Client::new(), &request) .await @@ -289,14 +289,14 @@ mod tests { .await; let request = ModelRequest { - base_url: format!("{}/v1", server.uri()), - api_token: "test-token".to_string(), - model: "claude-test".to_string(), - prompt: "hello".to_string(), - temperature: 0.0, - max_tokens: 1024, - stream: false -}; + base_url: format!("{}/v1", server.uri()), + api_token: "test-token".to_string(), + model: "claude-test".to_string(), + prompt: "hello".to_string(), + temperature: 0.0, + max_tokens: 1024, + stream: false, + }; let response = super::send(&Client::new(), &request) .await @@ -304,4 +304,52 @@ mod tests { assert_eq!(response.text, "prefixed"); } + + #[tokio::test] + async fn extracts_streaming_messages_text_and_ttft() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .and(header("x-api-key", "test-token")) + .and(header("anthropic-version", "2023-06-01")) + .and(body_json(serde_json::json!({ + "model": "claude-test", + "messages": [{"role": "user", "content": "hello"}], + "temperature": 0.0, + "max_tokens": 1024, + "stream": true + }))) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string( + "event: content_block_delta\n\ + data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"hi \"}}\n\n\ + event: content_block_delta\n\ + data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"there\"}}\n\n\ + event: message_stop\n\ + data: {\"type\":\"message_stop\"}\n\n", + ), + ) + .mount(&server) + .await; + + let request = ModelRequest { + base_url: server.uri(), + api_token: "test-token".to_string(), + model: "claude-test".to_string(), + prompt: "hello".to_string(), + temperature: 0.0, + max_tokens: 1024, + stream: true, + }; + + let response = super::send_stream(&Client::new(), &request) + .await + .expect("stream response"); + + assert_eq!(response.status, 200); + assert_eq!(response.text, "hi there"); + assert!(response.first_token_ms.is_some()); + } } diff --git a/src/protocols/openai.rs b/src/protocols/openai.rs index f3a2601..1508878 100644 --- a/src/protocols/openai.rs +++ b/src/protocols/openai.rs @@ -189,14 +189,14 @@ mod tests { .await; let request = ModelRequest { - base_url: server.uri(), - api_token: "test-token".to_string(), - model: "gpt-test".to_string(), - prompt: "hello".to_string(), - temperature: 0.0, - max_tokens: 1024, - stream: false -}; + base_url: server.uri(), + api_token: "test-token".to_string(), + model: "gpt-test".to_string(), + prompt: "hello".to_string(), + temperature: 0.0, + max_tokens: 1024, + stream: false, + }; let response = super::send(&Client::new(), &request) .await @@ -220,14 +220,14 @@ mod tests { .await; let request = ModelRequest { - base_url: format!("{}/", server.uri()), - api_token: "test-token".to_string(), - model: "gpt-test".to_string(), - prompt: "hello".to_string(), - temperature: 0.0, - max_tokens: 1024, - stream: false -}; + base_url: format!("{}/", server.uri()), + api_token: "test-token".to_string(), + model: "gpt-test".to_string(), + prompt: "hello".to_string(), + temperature: 0.0, + max_tokens: 1024, + stream: false, + }; let error = super::send(&Client::new(), &request) .await @@ -250,14 +250,14 @@ mod tests { .await; let request = ModelRequest { - base_url: server.uri(), - api_token: "test-token".to_string(), - model: "gpt-test".to_string(), - prompt: "hello".to_string(), - temperature: 0.0, - max_tokens: 1024, - stream: false -}; + base_url: server.uri(), + api_token: "test-token".to_string(), + model: "gpt-test".to_string(), + prompt: "hello".to_string(), + temperature: 0.0, + max_tokens: 1024, + stream: false, + }; let message = super::send(&Client::new(), &request) .await @@ -287,14 +287,14 @@ mod tests { .await; let request = ModelRequest { - base_url: format!("{}/v1", server.uri()), - api_token: "test-token".to_string(), - model: "gpt-test".to_string(), - prompt: "hello".to_string(), - temperature: 0.0, - max_tokens: 1024, - stream: false -}; + base_url: format!("{}/v1", server.uri()), + api_token: "test-token".to_string(), + model: "gpt-test".to_string(), + prompt: "hello".to_string(), + temperature: 0.0, + max_tokens: 1024, + stream: false, + }; let response = super::send(&Client::new(), &request) .await @@ -302,4 +302,48 @@ mod tests { assert_eq!(response.text, "prefixed"); } + + #[tokio::test] + async fn extracts_streaming_chat_completion_text_and_ttft() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .and(header("authorization", "Bearer test-token")) + .and(body_json(serde_json::json!({ + "model": "gpt-test", + "messages": [{"role": "user", "content": "hello"}], + "temperature": 0.0, + "max_tokens": 1024, + "stream": true + }))) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string( + "data: {\"choices\":[{\"delta\":{\"content\":\"hi \"}}]}\n\n\ + data: {\"choices\":[{\"delta\":{\"content\":\"there\"}}]}\n\n\ + data: [DONE]\n\n", + ), + ) + .mount(&server) + .await; + + let request = ModelRequest { + base_url: server.uri(), + api_token: "test-token".to_string(), + model: "gpt-test".to_string(), + prompt: "hello".to_string(), + temperature: 0.0, + max_tokens: 1024, + stream: true, + }; + + let response = super::send_stream(&Client::new(), &request) + .await + .expect("stream response"); + + assert_eq!(response.status, 200); + assert_eq!(response.text, "hi there"); + assert!(response.first_token_ms.is_some()); + } }