diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index a87c52fb5578..804414a89d95 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -15,7 +15,9 @@ use axum::extract::State; use axum::http::HeaderMap; use axum::response::sse::Event; use axum::response::{IntoResponse, Response}; -use futures::{Stream, StreamExt as _, pin_mut}; +use futures::future::try_join_all; +use futures::stream::{BoxStream, select_all}; +use futures::{Stream, StreamExt as _, TryStreamExt as _, pin_mut}; use serde_json::Value; use thiserror_ext::AsReport as _; use tracing::{debug, error, info, trace}; @@ -38,6 +40,9 @@ use crate::routes::openai::chat_completions::types::{ use crate::routes::openai::utils::logprobs::{ decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps, }; +use crate::routes::openai::utils::parallel_sampling::{ + choice_request_id, choice_seed, merge_choice_usage, +}; use crate::routes::openai::utils::types::{ ChatLogProbs, FunctionCallDelta, FunctionCallResponse, StreamResponseEnvelope, ToolCall, ToolCallDelta, Usage, @@ -80,44 +85,134 @@ pub async fn chat_completions( let created = unix_timestamp(); let api_server_options = state.api_server_options; - let chat_stream = - match state.chat.chat(prepared.chat_request).instrument(request_span.clone()).await { - Ok(stream) => stream, + let mut chat_streams = Vec::with_capacity(prepared.n as usize); + for index in 0..prepared.n { + let chat_request = match choice_chat_request(&prepared.chat_request, index, prepared.n) { + Ok(chat_request) => chat_request, + Err(error) => return error.into_response(), + }; + match state.chat.chat(chat_request).instrument(request_span.clone()).await { + Ok(stream) => chat_streams.push(stream), Err(error) => { return chat_submit_error("failed to submit chat request", error).into_response(); } - }; + } + } if stream { - let chunk_stream = chat_completion_chunk_stream( - chat_stream, - prepared.request_id, - prepared.response_model, - created, - api_server_options, - prepared.options, - ); + let chunk_streams = chat_streams + .into_iter() + .zip(0..) + .map(|(chat_stream, index)| { + chat_completion_chunk_stream( + chat_stream, + prepared.request_id.clone(), + prepared.response_model.clone(), + created, + api_server_options, + prepared.options.clone(), + ) + .map_ok(move |chunk| with_choice_index(chunk, index)) + .boxed() + }) + .collect(); + let chunk_stream = merge_choice_chunk_streams(chunk_streams); let sse_stream = chat_completion_sse_stream(chunk_stream).instrument(request_span); sse_response(sse_stream, api_server_options.sse_keep_alive_interval) } else { - let response = match collect_chat_completion( - chat_stream, - prepared.request_id, - prepared.response_model, - created, - api_server_options, - prepared.options, - ) + let responses = match try_join_all(chat_streams.into_iter().map(|chat_stream| { + collect_chat_completion( + chat_stream, + prepared.request_id.clone(), + prepared.response_model.clone(), + created, + api_server_options, + prepared.options.clone(), + ) + })) .instrument(request_span.clone()) .await { - Ok(response) => response, + Ok(responses) => responses, Err(error) => return error.into_response(), }; - Json(response).into_response() + Json(merge_choice_responses(responses)).into_response() + } +} + +/// Build the child chat request for one of `n` choices. +fn choice_chat_request(base: &ChatRequest, index: u32, n: u32) -> Result { + let mut chat_request = base.clone(); + chat_request.request_id = choice_request_id(&base.request_id, index, n); + chat_request.sampling_params.seed = choice_seed(base.sampling_params.seed, index)?; + Ok(chat_request) +} + +/// Set the choice index on every choice of one streamed chunk. +fn with_choice_index( + mut chunk: ChatCompletionStreamResponse, + index: u32, +) -> ChatCompletionStreamResponse { + for choice in &mut chunk.choices { + choice.index = index; + } + chunk +} + +/// Interleave per-choice chunk streams into one response stream. +/// +/// Each choice's usage-only chunk (usage set, no choices) is held back and +/// folded into one final usage chunk, so clients see a single usage chunk +/// after all choices finish. +#[try_stream] +async fn merge_choice_chunk_streams( + streams: Vec>>, + mut y: TryYielder, +) -> Result<(), ApiError> { + let mut chunks = select_all(streams); + let mut merged_usage_chunk: Option = None; + + while let Some(chunk) = chunks.next().await { + let mut chunk = chunk?; + if !chunk.choices.is_empty() || chunk.usage.is_none() { + y.yield_ok(chunk).await; + continue; + } + if let Some(merged) = merged_usage_chunk.take() { + chunk.usage = merged + .usage + .zip(chunk.usage) + .map(|(usage, choice_usage)| merge_choice_usage(usage, choice_usage)); + } + merged_usage_chunk = Some(chunk); + } + + if let Some(usage_chunk) = merged_usage_chunk { + y.yield_ok(usage_chunk).await; + } + Ok(()) +} + +/// Fold per-choice responses into one response with indexed choices. +/// +/// Response-level fields such as prompt logprobs and `kv_transfer_params` come +/// from the first choice. +fn merge_choice_responses(responses: Vec) -> ChatCompletionResponse { + let mut responses = responses.into_iter(); + let mut merged = responses.next().expect("n must be at least 1"); + for response in responses { + merged.usage = merged + .usage + .zip(response.usage) + .map(|(usage, choice_usage)| merge_choice_usage(usage, choice_usage)); + merged.choices.extend(response.choices); + } + for (choice, index) in merged.choices.iter_mut().zip(0..) { + choice.index = index; } + merged } async fn collect_chat_completion( diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index d4df61427028..ee802e198c4c 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -30,6 +30,8 @@ pub(super) struct PreparedRequest { pub request_id: String, /// Public model ID echoed back to the client. pub response_model: String, + /// Number of choices to sample for this request. + pub n: u32, /// Public response rendering options for route-layer helpers. pub options: ResponseOptions, /// Lowered chat request for `vllm-chat`. @@ -205,6 +207,7 @@ pub(super) fn prepare_chat_request( Ok(PreparedRequest { request_id, response_model, + n: request.n.unwrap_or(1), options: ResponseOptions { include_usage, include_continuous_usage, diff --git a/rust/src/server/src/routes/openai/chat_completions/validate.rs b/rust/src/server/src/routes/openai/chat_completions/validate.rs index 1bfc1d9a091d..89778a655c9b 100644 --- a/rust/src/server/src/routes/openai/chat_completions/validate.rs +++ b/rust/src/server/src/routes/openai/chat_completions/validate.rs @@ -24,10 +24,6 @@ pub(super) fn validate_request_compat( ); } - if request.n.unwrap_or(1) > 1 { - bail_invalid_request!(param = "n", "Only n=1 is supported."); - } - if request.top_logprobs.is_some() && !request.logprobs { bail_invalid_request!( param = "top_logprobs", diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index 7074b8c6293c..10390ba23823 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -15,7 +15,9 @@ use axum::extract::State; use axum::http::HeaderMap; use axum::response::sse::Event; use axum::response::{IntoResponse, Response}; -use futures::{Stream, StreamExt as _, pin_mut}; +use futures::future::try_join_all; +use futures::stream::{BoxStream, select_all}; +use futures::{Stream, StreamExt as _, TryStreamExt as _, pin_mut}; use serde_json::Value; use thiserror_ext::AsReport as _; use tracing::{debug, error, info, trace}; @@ -41,6 +43,9 @@ use crate::routes::openai::completions::types::{ CompletionChoice, CompletionResponse, CompletionSseChunk, CompletionStreamChoice, CompletionStreamResponse, }; +use crate::routes::openai::utils::parallel_sampling::{ + choice_request_id, choice_seed, merge_choice_usage, +}; use crate::routes::openai::utils::types::LogProbs; use crate::routes::openai::utils::usage::ContinuousUsage; use crate::routes::openai::utils::validated_json::ValidatedJson; @@ -87,49 +92,133 @@ pub async fn completions( let created = unix_timestamp(); let api_server_options = state.api_server_options; - let text_stream = match state - .chat - .text() - .generate(prepared.text_request) - .instrument(request_span.clone()) - .await - { - Ok(stream) => stream, - Err(error) => { - return text_submit_error("failed to submit completion request", error).into_response(); + let mut text_streams = Vec::with_capacity(prepared.n as usize); + for index in 0..prepared.n { + let text_request = match choice_text_request(&prepared.text_request, index, prepared.n) { + Ok(text_request) => text_request, + Err(error) => return error.into_response(), + }; + match state.chat.text().generate(text_request).instrument(request_span.clone()).await { + Ok(stream) => text_streams.push(stream), + Err(error) => { + return text_submit_error("failed to submit completion request", error) + .into_response(); + } } - }; + } if stream { - let chunk_stream = completion_chunk_stream( - text_stream, - prepared.request_id, - prepared.response_model, - created, - api_server_options, - prepared.options, - ); + let chunk_streams = text_streams + .into_iter() + .zip(0..) + .map(|(text_stream, index)| { + completion_chunk_stream( + text_stream, + prepared.request_id.clone(), + prepared.response_model.clone(), + created, + api_server_options, + prepared.options.clone(), + ) + .map_ok(move |chunk| with_choice_index(chunk, index)) + .boxed() + }) + .collect(); + let chunk_stream = merge_choice_chunk_streams(chunk_streams); let sse_stream = completion_sse_stream(chunk_stream).instrument(request_span); sse_response(sse_stream, api_server_options.sse_keep_alive_interval) } else { - let response = match collect_completion( - text_stream, - prepared.request_id, - prepared.response_model, - created, - api_server_options, - prepared.options, - ) + let responses = match try_join_all(text_streams.into_iter().map(|text_stream| { + collect_completion( + text_stream, + prepared.request_id.clone(), + prepared.response_model.clone(), + created, + api_server_options, + prepared.options.clone(), + ) + })) .instrument(request_span.clone()) .await { - Ok(response) => response, + Ok(responses) => responses, Err(error) => return error.into_response(), }; - Json(response).into_response() + Json(merge_choice_responses(responses)).into_response() + } +} + +/// Build the child text request for one of `n` choices. +fn choice_text_request(base: &TextRequest, index: u32, n: u32) -> Result { + let mut text_request = base.clone(); + text_request.request_id = choice_request_id(&base.request_id, index, n); + text_request.sampling_params.seed = choice_seed(base.sampling_params.seed, index)?; + Ok(text_request) +} + +/// Set the choice index on every choice of one streamed chunk. +fn with_choice_index(mut chunk: CompletionSseChunk, index: u32) -> CompletionSseChunk { + if let CompletionSseChunk::Chunk(response) = &mut chunk { + for choice in &mut response.choices { + choice.index = index; + } + } + chunk +} + +/// Interleave per-choice chunk streams into one response stream. +/// +/// Each choice's usage chunk is held back and folded into one final usage +/// chunk, so clients see a single usage chunk after all choices finish. +#[try_stream] +async fn merge_choice_chunk_streams( + streams: Vec>>, + mut y: TryYielder, +) -> Result<(), ApiError> { + let mut chunks = select_all(streams); + let mut merged_usage_chunk: Option = None; + + while let Some(chunk) = chunks.next().await { + match chunk? { + CompletionSseChunk::Usage(mut usage_chunk) => { + if let Some(merged) = merged_usage_chunk.take() { + usage_chunk.usage = merged + .usage + .zip(usage_chunk.usage) + .map(|(usage, choice_usage)| merge_choice_usage(usage, choice_usage)); + } + merged_usage_chunk = Some(usage_chunk); + } + chunk => y.yield_ok(chunk).await, + } + } + + if let Some(usage_chunk) = merged_usage_chunk { + y.yield_ok(CompletionSseChunk::Usage(usage_chunk)).await; + } + Ok(()) +} + +/// Fold per-choice responses into one response with indexed choices. +/// +/// Response-level fields such as `kv_transfer_params` come from the first +/// choice. +fn merge_choice_responses(responses: Vec) -> CompletionResponse { + let mut responses = responses.into_iter(); + let mut merged = responses.next().expect("n must be at least 1"); + for response in responses { + merged.usage = merged + .usage + .zip(response.usage) + .map(|(usage, choice_usage)| merge_choice_usage(usage, choice_usage)); + merged.choices.extend(response.choices); + } + for (choice, index) in merged.choices.iter_mut().zip(0..) { + choice.index = index; } + merged } async fn collect_completion( diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index 53a6ff9906ce..9bb7c0cb5473 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -24,6 +24,8 @@ pub(super) struct PreparedRequest { pub request_id: String, /// Public model ID echoed back to the client. pub response_model: String, + /// Number of choices to sample for this request. + pub n: u32, /// Public response rendering options for route-layer helpers. pub options: ResponseOptions, /// Lowered text request for the shared `vllm-text` facade. @@ -165,6 +167,7 @@ pub(super) fn prepare_completion_request( Ok(PreparedRequest { request_id, response_model, + n: request.n.unwrap_or(1), options: ResponseOptions { include_usage, include_continuous_usage, diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs index 2a6901370682..12b29f5a3c65 100644 --- a/rust/src/server/src/routes/openai/completions/types.rs +++ b/rust/src/server/src/routes/openai/completions/types.rs @@ -59,6 +59,7 @@ pub struct CompletionRequest { pub max_tokens: Option, /// How many completions to generate for each prompt + #[validate(range(min = 1, max = 10))] pub n: Option, /// Number between -2.0 and 2.0. Positive values penalize new tokens based diff --git a/rust/src/server/src/routes/openai/completions/validate.rs b/rust/src/server/src/routes/openai/completions/validate.rs index e41ad7a85d4e..0bcbee387248 100644 --- a/rust/src/server/src/routes/openai/completions/validate.rs +++ b/rust/src/server/src/routes/openai/completions/validate.rs @@ -26,10 +26,6 @@ pub(super) fn validate_request_compat( ); } - if request.n.unwrap_or(1) > 1 { - bail_invalid_request!(param = "n", "Only n=1 is supported."); - } - if request.max_tokens == Some(0) && !request.echo { bail_invalid_request!( param = "max_tokens", diff --git a/rust/src/server/src/routes/openai/utils/mod.rs b/rust/src/server/src/routes/openai/utils/mod.rs index 4bbea924a38b..0ed2ce876282 100644 --- a/rust/src/server/src/routes/openai/utils/mod.rs +++ b/rust/src/server/src/routes/openai/utils/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project pub mod logprobs; +pub mod parallel_sampling; pub mod structured_outputs; pub mod types; pub mod usage; diff --git a/rust/src/server/src/routes/openai/utils/parallel_sampling.rs b/rust/src/server/src/routes/openai/utils/parallel_sampling.rs new file mode 100644 index 000000000000..3a2154bae87b --- /dev/null +++ b/rust/src/server/src/routes/openai/utils/parallel_sampling.rs @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Helpers for OpenAI `n > 1` parallel sampling. +//! +//! Like Python vLLM's `ParentRequest`, each choice is submitted as its own +//! child engine request, and the route merges the per-choice outputs back into +//! one response. + +use super::types::Usage; +use crate::error::{ApiError, invalid_request}; + +/// Request ID of one child choice request. +/// +/// Uses Python's `{index}_{parent}` scheme when `n > 1`, and keeps the parent +/// ID for `n = 1` so single-choice requests are unaffected. +pub(crate) fn choice_request_id(parent_request_id: &str, index: u32, n: u32) -> String { + if n > 1 { + format!("{index}_{parent_request_id}") + } else { + parent_request_id.to_string() + } +} + +/// Seed of one child choice request. +/// +/// Offset by the choice index, matching Python, so that seeded choices differ +/// from each other while staying reproducible. +pub(crate) fn choice_seed(seed: Option, index: u32) -> Result, ApiError> { + seed.map(|seed| { + seed.checked_add(i64::from(index)).ok_or_else(|| { + invalid_request!( + param = "seed", + "`seed + n - 1` must fit in a signed 64-bit integer." + ) + }) + }) + .transpose() +} + +/// Fold the usage of one more choice into the usage of a multi-choice response. +/// +/// All choices share one prompt, so prompt tokens (and prompt cache details) +/// are counted once while completion and reasoning tokens are summed. +pub(crate) fn merge_choice_usage(usage: Usage, choice_usage: Usage) -> Usage { + let completion_tokens = + usage.completion_tokens.unwrap_or(0) + choice_usage.completion_tokens.unwrap_or(0); + let mut completion_tokens_details = usage.completion_tokens_details; + completion_tokens_details.reasoning_tokens += + choice_usage.completion_tokens_details.reasoning_tokens; + Usage { + total_tokens: usage.prompt_tokens + completion_tokens, + completion_tokens: Some(completion_tokens), + completion_tokens_details, + ..usage + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn choice_request_id_keeps_parent_id_for_single_choice() { + assert_eq!(choice_request_id("cmpl-1", 0, 1), "cmpl-1"); + assert_eq!(choice_request_id("cmpl-1", 0, 2), "0_cmpl-1"); + assert_eq!(choice_request_id("cmpl-1", 1, 2), "1_cmpl-1"); + } + + #[test] + fn choice_seed_offsets_by_choice_index() { + assert_eq!(choice_seed(None, 1).expect("no seed"), None); + assert_eq!(choice_seed(Some(41), 1).expect("seed fits"), Some(42)); + assert!(choice_seed(Some(i64::MAX), 1).is_err()); + } + + #[test] + fn merge_choice_usage_counts_prompt_once() { + let usage = merge_choice_usage( + Usage::from_counts(5, 3, Some(4), 1), + Usage::from_counts(5, 2, Some(4), 2), + ); + + assert_eq!(usage.prompt_tokens, 5); + assert_eq!(usage.completion_tokens, Some(5)); + assert_eq!(usage.total_tokens, 10); + assert_eq!( + usage.prompt_tokens_details.map(|details| details.cached_tokens), + Some(4) + ); + assert_eq!(usage.completion_tokens_details.reasoning_tokens, 3); + } +} diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 5c1fff2cae76..18e6fa698a7e 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -1243,6 +1243,98 @@ where ) } +/// Build an app whose mock engine expects one request per `output_specs_by_choice` +/// entry, checks them together, then streams each entry to its request. +async fn test_app_with_parallel_engine_outputs( + output_specs_by_choice: Vec, Option)>>, + check_requests: F, +) -> (axum::Router, MockEngineTask) +where + F: FnOnce(&[EngineCoreRequest]) + Send + 'static, +{ + let ipc = IpcNamespace::new().expect("create ipc namespace"); + let handshake_address = ipc.handshake_endpoint(); + let engine_id = b"engine-openai-parallel".to_vec(); + + let engine_task = MockEngineTask::new(spawn_mock_engine_task( + handshake_address.clone(), + engine_id.clone(), + move |dealer, push| { + boxed_test_future(async move { + let mut requests = Vec::with_capacity(output_specs_by_choice.len()); + for _ in 0..output_specs_by_choice.len() { + let add = recv_engine_message(dealer).await; + let request: EngineCoreRequest = + rmp_serde::from_slice(&add[1]).expect("decode request"); + requests.push(request); + } + check_requests(&requests); + + for (request, output_specs) in requests.iter().zip(output_specs_by_choice) { + send_outputs( + push, + engine_outputs_for_request(&request.request_id, output_specs), + ) + .await; + } + }) + }, + )); + + let client = EngineCoreClient::connect( + EngineCoreClientConfig::new_single(handshake_address) + .with_model_name("test-model") + .with_local_input_output_addresses( + Some(ipc.input_endpoint()), + Some(ipc.output_endpoint()), + ), + ) + .await + .expect("connect client"); + + let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new())); + ( + build_router(Arc::new(AppState::new( + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], + chat, + ))), + engine_task, + ) +} + +/// Output specs for `n=2` requests: choice 0 decodes "hi", choice 1 decodes "yo". +fn two_choice_output_specs() -> Vec, Option)>> { + vec![ + default_stream_output_specs(), + vec![ + (vec![b'y' as u32], None), + (vec![b'o' as u32], None), + (vec![b'!' as u32], Some(EngineCoreFinishReason::Stop)), + ], + ] +} + +/// Assert that `n=2` fans out into child requests with Python-style +/// `{index}_{parent}` IDs and per-choice seed offsets. +fn assert_two_choice_child_requests(requests: &[EngineCoreRequest], parent_request_id: &str) { + let summary = requests + .iter() + .map(|request| { + ( + request.request_id.as_str(), + request.sampling_params.as_ref().expect("sampling params").seed, + ) + }) + .collect::>(); + assert_eq!( + summary, + [ + (format!("0_{parent_request_id}").as_str(), Some(41)), + (format!("1_{parent_request_id}").as_str(), Some(42)), + ] + ); +} + async fn test_chat_with_engine_handle() -> (ChatLlm, MockEngineTask) { test_chat_with_engine_outputs(b"engine-openai-chat", default_stream_output_specs()).await } @@ -3626,6 +3718,124 @@ async fn stream_without_include_usage_keeps_existing_shape() { assert!(text.trim_end().ends_with("data: [DONE]"), "{text}"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn non_stream_chat_completions_support_n_greater_than_one() { + let (app, engine_task) = + test_app_with_parallel_engine_outputs(two_choice_output_specs(), |requests| { + assert_two_choice_child_requests(requests, "chatcmpl-choice-parent"); + }) + .await; + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "request_id": "choice-parent", + "stream": false, + "n": 2, + "seed": 41, + "messages": [{"role": "user", "content": "hello"}] + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + engine_task.await.expect("mock engine task"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); + + assert_eq!(json["id"], "chatcmpl-choice-parent"); + assert_eq!(json["choices"][0]["index"], 0); + assert_eq!(json["choices"][0]["message"]["content"], "hi"); + assert_eq!(json["choices"][1]["index"], 1); + assert_eq!(json["choices"][1]["message"]["content"], "yo"); + // The prompt is counted once; completion tokens are summed across choices. + assert_eq!(json["usage"]["prompt_tokens"], 22); + assert_eq!(json["usage"]["completion_tokens"], 6); + assert_eq!(json["usage"]["total_tokens"], 28); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn stream_chat_completions_support_n_greater_than_one() { + let (app, engine_task) = + test_app_with_parallel_engine_outputs(two_choice_output_specs(), |_| {}).await; + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "stream": true, + "n": 2, + "stream_options": { + "include_usage": true, + "continuous_usage_stats": true + }, + "messages": [{"role": "user", "content": "hello"}] + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + engine_task.await.expect("mock engine task"); + let text = String::from_utf8(body.to_vec()).expect("utf8 body"); + let payloads = sse_json_payloads(&text); + + let content_for = |index: u64| { + payloads + .iter() + .flat_map(|payload| payload["choices"].as_array().into_iter().flatten()) + .filter(|choice| choice["index"] == index) + .filter_map(|choice| choice["delta"]["content"].as_str()) + .collect::() + }; + assert_eq!(content_for(0), "hi", "{text}"); + assert_eq!(content_for(1), "yo", "{text}"); + + // Continuous usage stays per choice on every choice chunk. + assert!( + payloads + .iter() + .filter(|payload| payload["choices"] != json!([])) + .all(|payload| { + payload["usage"]["completion_tokens"].as_u64().is_some_and(|tokens| tokens <= 3) + }), + "{text}" + ); + let usage_chunks = payloads + .iter() + .filter(|payload| payload["choices"] == json!([])) + .collect::>(); + assert_eq!(usage_chunks.len(), 1, "{text}"); + assert_eq!(usage_chunks[0]["usage"]["prompt_tokens"], 22); + assert_eq!(usage_chunks[0]["usage"]["completion_tokens"], 6); + assert_eq!(payloads.last(), Some(usage_chunks[0]), "{text}"); + assert_eq!( + sse_data_payloads(&text).iter().filter(|p| **p == "[DONE]").count(), + 1 + ); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn completions_invalid_request_returns_openai_error() { @@ -3656,6 +3866,34 @@ async fn completions_invalid_request_returns_openai_error() { assert_eq!(json["error"]["type"], "invalid_request_error"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn completions_reject_n_outside_supported_range() { + for n in [0, 11] { + let mut app = test_app().await; + let response = app + .call( + Request::builder() + .method("POST") + .uri("/v1/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "prompt": "hello", + "n": n + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST, "n={n}"); + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn completions_empty_allowed_token_ids_returns_openai_error() { @@ -4998,6 +5236,112 @@ async fn completions_stream_continuous_usage_stats_adds_usage_to_chunks() { assert_eq!(usage_chunk["usage"]["completion_tokens"], 3); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn non_stream_completions_support_n_greater_than_one() { + let (app, engine_task) = + test_app_with_parallel_engine_outputs(two_choice_output_specs(), |requests| { + assert_two_choice_child_requests(requests, "cmpl-choice-parent"); + }) + .await; + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/v1/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "request_id": "choice-parent", + "prompt": "hello", + "stream": false, + "n": 2, + "seed": 41 + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + engine_task.await.expect("mock engine task"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); + + assert_eq!(json["id"], "cmpl-choice-parent"); + assert_eq!(json["choices"][0]["index"], 0); + assert_eq!(json["choices"][0]["text"], "hi"); + assert_eq!(json["choices"][1]["index"], 1); + assert_eq!(json["choices"][1]["text"], "yo"); + // The prompt is counted once; completion tokens are summed across choices. + assert_eq!(json["usage"]["prompt_tokens"], 6); + assert_eq!(json["usage"]["completion_tokens"], 6); + assert_eq!(json["usage"]["total_tokens"], 12); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn stream_completions_support_n_greater_than_one() { + let (app, engine_task) = + test_app_with_parallel_engine_outputs(two_choice_output_specs(), |_| {}).await; + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/v1/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "prompt": "hello", + "stream": true, + "n": 2, + "stream_options": {"include_usage": true} + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + engine_task.await.expect("mock engine task"); + let text = String::from_utf8(body.to_vec()).expect("utf8 body"); + let payloads = sse_json_payloads(&text); + + let text_for = |index: u64| { + payloads + .iter() + .flat_map(|payload| payload["choices"].as_array().into_iter().flatten()) + .filter(|choice| choice["index"] == index) + .filter_map(|choice| choice["text"].as_str()) + .collect::() + }; + assert_eq!(text_for(0), "hi", "{text}"); + assert_eq!(text_for(1), "yo", "{text}"); + + let usage_chunks = payloads + .iter() + .filter(|payload| payload.get("usage").is_some()) + .collect::>(); + assert_eq!(usage_chunks.len(), 1, "{text}"); + assert_eq!(usage_chunks[0]["choices"], json!([])); + assert_eq!(usage_chunks[0]["usage"]["prompt_tokens"], 6); + assert_eq!(usage_chunks[0]["usage"]["completion_tokens"], 6); + assert_eq!(payloads.last(), Some(usage_chunks[0]), "{text}"); + assert_eq!( + sse_data_payloads(&text).iter().filter(|p| **p == "[DONE]").count(), + 1 + ); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn completions_echo_stream_emits_separate_prompt_chunk() {