Просмотр исходного кода

fix: support bare stream flag and cover SSE parsing

main
orangels 1 месяц назад
Родитель
Сommit
60cc009c83
3 измененных файлов: 246 добавлений и 77 удалений
  1. +90
    -13
      src/cli.rs
  2. +80
    -32
      src/protocols/anthropic.rs
  3. +76
    -32
      src/protocols/openai.rs

+ 90
- 13
src/cli.rs Просмотреть файл

@@ -4,10 +4,10 @@ use crate::config::AppConfig;
use crate::metrics::{LatencySummary, Metrics, MetricsSummary}; use crate::metrics::{LatencySummary, Metrics, MetricsSummary};
use crate::report::{ use crate::report::{
BenchmarkParamsReport, BenchmarkReport, BenchmarkSummaryReport, CorrectCaseReport, 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::{ use crate::rpm_modes::{
ProbePhase, RpmMode, ScheduledProbe, burst_schedule, sliding_window_schedule, ProbePhase, RpmMode, ScheduledProbe, burst_schedule, sliding_window_schedule,
@@ -47,7 +47,7 @@ pub enum Command {
model: Option<String>, model: Option<String>,
#[arg(long)] #[arg(long)]
prompt: String, prompt: String,
#[arg(long)]
#[arg(long, num_args = 0..=1, default_missing_value = "true")]
stream: Option<bool>, stream: Option<bool>,
}, },
Dataset { Dataset {
@@ -81,7 +81,7 @@ pub enum Command {
concurrency: Option<usize>, concurrency: Option<usize>,
#[arg(long)] #[arg(long)]
prompt: String, prompt: String,
#[arg(long)]
#[arg(long, num_args = 0..=1, default_missing_value = "true")]
stream: Option<bool>, stream: Option<bool>,
}, },
} }
@@ -104,7 +104,7 @@ pub enum BenchCommand {
concurrency: usize, concurrency: usize,
#[arg(long)] #[arg(long)]
limit: Option<usize>, limit: Option<usize>,
#[arg(long)]
#[arg(long, num_args = 0..=1, default_missing_value = "true")]
stream: Option<bool>, stream: Option<bool>,
#[arg(long, default_value_t = 32768)] #[arg(long, default_value_t = 32768)]
max_tokens: u32, max_tokens: u32,
@@ -120,7 +120,7 @@ pub enum BenchCommand {
concurrency: usize, concurrency: usize,
#[arg(long)] #[arg(long)]
limit: Option<usize>, limit: Option<usize>,
#[arg(long)]
#[arg(long, num_args = 0..=1, default_missing_value = "true")]
stream: Option<bool>, stream: Option<bool>,
#[arg(long, default_value_t = 32768)] #[arg(long, default_value_t = 32768)]
max_tokens: u32, max_tokens: u32,
@@ -212,7 +212,18 @@ async fn dispatch_bench(command: BenchCommand) -> Result<()> {
limit, limit,
stream, stream,
max_tokens, 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 { BenchCommand::GpqaDiamond {
config, config,
provider, provider,
@@ -221,7 +232,18 @@ async fn dispatch_bench(command: BenchCommand) -> Result<()> {
limit, limit,
stream, stream,
max_tokens, 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); pb.inc(1);
match result { match result {
Ok(response) => { 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) let actual = judge::extract_final_integer(&response.text)
.unwrap_or_else(|| "no_answer".to_string()); .unwrap_or_else(|| "no_answer".to_string());
let correct = judge::judge_integer(&response.text, &case.answer); let correct = judge::judge_integer(&response.text, &case.answer);
@@ -385,7 +411,11 @@ async fn run_gpqa_benchmark(
pb.inc(1); pb.inc(1);
match result { match result {
Ok(response) => { 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) let actual = judge::extract_choice(&response.text)
.map(|choice| choice.to_string()) .map(|choice| choice.to_string())
.unwrap_or_else(|| "no_answer".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(); let success = result.result.is_ok();
mode_summary.record(result.phase, result.second, success); mode_summary.record(result.phase, result.second, success);
match result.result { 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)), Err(error) => metrics.record_failure(error_code(&error)),
} }
} }
@@ -1220,6 +1254,49 @@ mod tests {
assert_eq!(probe_seconds, Some(30)); 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] #[test]
fn rpm_command_parses_window_boundary_offset() { fn rpm_command_parses_window_boundary_offset() {
let cli = Cli::try_parse_from([ let cli = Cli::try_parse_from([


+ 80
- 32
src/protocols/anthropic.rs Просмотреть файл

@@ -192,14 +192,14 @@ mod tests {
.await; .await;


let request = ModelRequest { 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) let response = super::send(&Client::new(), &request)
.await .await
@@ -223,14 +223,14 @@ mod tests {
.await; .await;


let request = ModelRequest { 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) let error = super::send(&Client::new(), &request)
.await .await
@@ -254,14 +254,14 @@ mod tests {
.await; .await;


let request = ModelRequest { 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) let message = super::send(&Client::new(), &request)
.await .await
@@ -289,14 +289,14 @@ mod tests {
.await; .await;


let request = ModelRequest { 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) let response = super::send(&Client::new(), &request)
.await .await
@@ -304,4 +304,52 @@ mod tests {


assert_eq!(response.text, "prefixed"); 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());
}
} }

+ 76
- 32
src/protocols/openai.rs Просмотреть файл

@@ -189,14 +189,14 @@ mod tests {
.await; .await;


let request = ModelRequest { 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) let response = super::send(&Client::new(), &request)
.await .await
@@ -220,14 +220,14 @@ mod tests {
.await; .await;


let request = ModelRequest { 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) let error = super::send(&Client::new(), &request)
.await .await
@@ -250,14 +250,14 @@ mod tests {
.await; .await;


let request = ModelRequest { 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) let message = super::send(&Client::new(), &request)
.await .await
@@ -287,14 +287,14 @@ mod tests {
.await; .await;


let request = ModelRequest { 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) let response = super::send(&Client::new(), &request)
.await .await
@@ -302,4 +302,48 @@ mod tests {


assert_eq!(response.text, "prefixed"); 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());
}
} }

Загрузка…
Отмена
Сохранить