diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index 543a7e806c4b..3fdea298a57c 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -3,6 +3,7 @@ mod types; mod validate; use std::convert::Infallible; +use std::pin::Pin; use std::result::Result; use std::sync::Arc; @@ -12,6 +13,8 @@ use axum::extract::State; use axum::http::HeaderMap; use axum::response::sse::{Event, Sse}; use axum::response::{IntoResponse, Response}; +use futures::future::try_join_all; +use futures::stream::SelectAll; use futures::{Stream, StreamExt as _, pin_mut}; use serde_json::Value; use thiserror_ext::AsReport as _; @@ -40,6 +43,30 @@ use crate::routes::openai::utils::validated_json::ValidatedJson; use crate::state::AppState; use crate::utils::{resolve_request_context, unix_timestamp}; +type ChatCompletionEventStream = + Pin> + Send>>; + +#[derive(Debug, Clone, Copy)] +struct ChoiceUsage { + prompt_tokens: u32, + completion_tokens: u32, +} + +#[derive(Debug)] +struct ChatChoiceOutput { + choice: ChatCompletionChoice, + usage: ChoiceUsage, + prompt_logprobs: Option>>>, + prompt_token_ids: Option>, + kv_transfer_params: Option, +} + +#[derive(Debug)] +enum ChatCompletionStreamEvent { + Chunk(ChatCompletionStreamResponse), + Finished(ChoiceUsage), +} + /// Validate one chat completion request and proxy it into the shared /// `vllm-chat` stack. pub async fn chat_completions( @@ -64,52 +91,98 @@ pub async fn chat_completions( let created = unix_timestamp(); let log_request = state.enable_log_requests; - let chat_stream = - match state.chat.chat(prepared.chat_request).instrument(request_span.clone()).await { - Ok(stream) => stream, - Err(error) => { - return server_error!( - "failed to submit chat request: {}", - error.to_report_string() - ) - .into_response(); - } - }; - if stream { - let chunk_stream = chat_completion_chunk_stream( - chat_stream, + let mut event_streams: SelectAll = SelectAll::new(); + for choice_index in 0..prepared.n { + let chat_request = match child_chat_request( + &prepared.chat_request, + &prepared.request_id, + choice_index, + prepared.n, + ) { + Ok(request) => request, + Err(error) => return error.into_response(), + }; + let chat_stream = + match state.chat.chat(chat_request).instrument(request_span.clone()).await { + Ok(stream) => stream, + Err(error) => { + return server_error!( + "failed to submit chat request: {}", + error.to_report_string() + ) + .into_response(); + } + }; + event_streams.push(Box::pin(chat_completion_choice_event_stream( + chat_stream, + prepared.request_id.clone(), + prepared.response_model.clone(), + created, + choice_index, + log_request, + prepared.requested_logprobs, + prepared.echo.clone(), + prepared.return_token_ids, + prepared.return_tokens_as_token_ids, + ))); + } + + let chunk_stream = chat_completion_parallel_chunk_stream( + event_streams, prepared.request_id, prepared.response_model, created, - log_request, prepared.include_usage, - prepared.requested_logprobs, - prepared.echo, - prepared.return_token_ids, - prepared.return_tokens_as_token_ids, ); let sse_stream = chat_completion_sse_stream(chunk_stream).instrument(request_span); Sse::new(sse_stream).into_response() } else { - let response = match collect_chat_completion( - chat_stream, + let mut choice_futures = Vec::with_capacity(prepared.n as usize); + for choice_index in 0..prepared.n { + let chat_request = match child_chat_request( + &prepared.chat_request, + &prepared.request_id, + choice_index, + prepared.n, + ) { + Ok(request) => request, + Err(error) => return error.into_response(), + }; + let chat_stream = + match state.chat.chat(chat_request).instrument(request_span.clone()).await { + Ok(stream) => stream, + Err(error) => { + return server_error!( + "failed to submit chat request: {}", + error.to_report_string() + ) + .into_response(); + } + }; + choice_futures.push(collect_chat_completion_choice( + chat_stream, + choice_index, + prepared.requested_logprobs, + prepared.include_prompt_logprobs, + prepared.echo.clone(), + prepared.return_token_ids, + prepared.return_tokens_as_token_ids, + )); + } + + let choices = match try_join_all(choice_futures).instrument(request_span.clone()).await { + Ok(choices) => choices, + Err(error) => return error.into_response(), + }; + + let response = chat_completion_response_from_choices( prepared.request_id, prepared.response_model, created, - prepared.requested_logprobs, - prepared.include_prompt_logprobs, - prepared.echo, - prepared.return_token_ids, - prepared.return_tokens_as_token_ids, - ) - .instrument(request_span.clone()) - .await - { - Ok(response) => response, - Err(error) => return error.into_response(), - }; + choices, + ); if log_request { let usage = response.usage.as_ref(); @@ -127,17 +200,39 @@ pub async fn chat_completions( } } -async fn collect_chat_completion( +fn child_chat_request( + base: &vllm_chat::ChatRequest, + parent_request_id: &str, + choice_index: u32, + choice_count: u32, +) -> Result { + let mut request = base.clone(); + if choice_count > 1 { + request.request_id = format!("{choice_index}_{parent_request_id}"); + } else { + request.request_id = parent_request_id.to_string(); + } + if let Some(seed) = request.sampling_params.seed { + request.sampling_params.seed = + Some(seed.checked_add(i64::from(choice_index)).ok_or_else(|| { + ApiError::invalid_request( + "`seed + choice index` must fit within a signed 64-bit integer.".to_string(), + Some("seed"), + ) + })?); + } + Ok(request) +} + +async fn collect_chat_completion_choice( stream: ChatEventStream, - request_id: String, - response_model: String, - created: u64, + choice_index: u32, requested_logprobs: bool, include_prompt_logprobs: bool, echo: Option, return_token_ids: bool, return_tokens_as_token_ids: bool, -) -> Result { +) -> Result { let collected = stream.collect_message().await.map_err(|error| { server_error!( "failed to collect chat completion response: {}", @@ -191,15 +286,9 @@ async fn collect_chat_completion( } else { None }; - let usage = Usage::from_counts(prompt_token_count as u32, output_token_count as u32); - - Ok(ChatCompletionResponse { - id: request_id, - object: "chat.completion".to_string(), - created, - model: response_model, - choices: vec![ChatCompletionChoice { - index: 0, + Ok(ChatChoiceOutput { + choice: ChatCompletionChoice { + index: choice_index, message: ChatCompletionMessage { role: AssistantRole, content: match &echo { @@ -213,22 +302,51 @@ async fn collect_chat_completion( finish_reason: Some(finish_reason), stop_reason, token_ids: return_token_ids.then_some(token_ids), - }], - usage: Some(usage), - system_fingerprint: None, + }, + usage: ChoiceUsage { + prompt_tokens: prompt_token_count as u32, + completion_tokens: output_token_count as u32, + }, prompt_logprobs, prompt_token_ids: return_token_ids.then(|| prompt_token_ids.to_vec()), kv_transfer_params, }) } +fn chat_completion_response_from_choices( + request_id: String, + response_model: String, + created: u64, + outputs: Vec, +) -> ChatCompletionResponse { + let prompt_tokens = outputs.first().map_or(0, |output| output.usage.prompt_tokens); + let completion_tokens = outputs.iter().map(|output| output.usage.completion_tokens).sum(); + let prompt_logprobs = outputs.iter().find_map(|output| output.prompt_logprobs.clone()); + let prompt_token_ids = outputs.iter().find_map(|output| output.prompt_token_ids.clone()); + let kv_transfer_params = outputs.iter().find_map(|output| output.kv_transfer_params.clone()); + + ChatCompletionResponse { + id: request_id, + object: "chat.completion".to_string(), + created, + model: response_model, + choices: outputs.into_iter().map(|output| output.choice).collect(), + usage: Some(Usage::from_counts(prompt_tokens, completion_tokens)), + system_fingerprint: None, + prompt_logprobs, + prompt_token_ids, + kv_transfer_params, + } +} + /// Convert one internal chat event stream into OpenAI chat-completion chunks. #[try_stream] async fn chat_completion_chunk_stream( - mut stream: impl ChatEventStreamTrait + Unpin, + stream: impl ChatEventStreamTrait + Unpin, request_id: String, response_model: String, created: u64, + choice_index: u32, log_request: bool, include_usage: bool, requested_logprobs: bool, @@ -236,6 +354,91 @@ async fn chat_completion_chunk_stream( return_token_ids: bool, return_tokens_as_token_ids: bool, mut y: TryYielder, +) -> Result<(), ApiError> { + let events = chat_completion_choice_event_stream( + stream, + request_id.clone(), + response_model.clone(), + created, + choice_index, + log_request, + requested_logprobs, + echo, + return_token_ids, + return_tokens_as_token_ids, + ); + pin_mut!(events); + + while let Some(next) = events.next().await { + match next? { + ChatCompletionStreamEvent::Chunk(chunk) => y.yield_ok(chunk).await, + ChatCompletionStreamEvent::Finished(usage) => { + if include_usage { + y.yield_ok(usage_chunk( + &request_id, + &response_model, + created, + Usage::from_counts(usage.prompt_tokens, usage.completion_tokens), + )) + .await; + } + } + } + } + + Ok(()) +} + +#[try_stream] +async fn chat_completion_parallel_chunk_stream( + stream: impl Stream>, + request_id: String, + response_model: String, + created: u64, + include_usage: bool, + mut y: TryYielder, +) -> Result<(), ApiError> { + pin_mut!(stream); + let mut prompt_tokens = None; + let mut completion_tokens = 0_u32; + + while let Some(next) = stream.next().await { + match next? { + ChatCompletionStreamEvent::Chunk(chunk) => y.yield_ok(chunk).await, + ChatCompletionStreamEvent::Finished(usage) => { + prompt_tokens.get_or_insert(usage.prompt_tokens); + completion_tokens = completion_tokens.saturating_add(usage.completion_tokens); + } + } + } + + if include_usage { + y.yield_ok(usage_chunk( + &request_id, + &response_model, + created, + Usage::from_counts(prompt_tokens.unwrap_or(0), completion_tokens), + )) + .await; + } + + Ok(()) +} + +/// Convert one internal chat event stream into indexed OpenAI chat-completion events. +#[try_stream] +async fn chat_completion_choice_event_stream( + mut stream: impl ChatEventStreamTrait + Unpin, + request_id: String, + response_model: String, + created: u64, + choice_index: u32, + log_request: bool, + requested_logprobs: bool, + echo: Option, + return_token_ids: bool, + return_tokens_as_token_ids: bool, + mut y: TryYielder, ) -> Result<(), ApiError> { let mut saw_tool_calls = false; @@ -250,20 +453,21 @@ async fn chat_completion_chunk_stream( Ok(ChatEvent::Start { prompt_token_ids, .. }) => { - let mut chunk = start_chunk(&request_id, &response_model, created); + let mut chunk = start_chunk(&request_id, &response_model, created, choice_index); if return_token_ids { chunk.prompt_token_ids = Some(prompt_token_ids.to_vec()); } - y.yield_ok(chunk).await; + y.yield_ok(ChatCompletionStreamEvent::Chunk(chunk)).await; // When echo=true, emit the last assistant message content as a delta chunk. if let Some(echo_text) = &echo { - y.yield_ok(block_delta_chunk( + y.yield_ok(ChatCompletionStreamEvent::Chunk(block_delta_chunk( &request_id, &response_model, created, + choice_index, AssistantBlockKind::Text, echo_text.clone(), - )) + ))) .await; } } @@ -271,13 +475,14 @@ async fn chat_completion_chunk_stream( if let Some(pending_chunk) = pending_chunk.as_mut() { pending_chunk.push_block_delta(kind, delta); } else { - y.yield_ok(block_delta_chunk( + y.yield_ok(ChatCompletionStreamEvent::Chunk(block_delta_chunk( &request_id, &response_model, created, + choice_index, kind, delta, - )) + ))) .await; } } @@ -294,18 +499,22 @@ async fn chat_completion_chunk_stream( if let Some(pending_chunk) = pending_chunk.as_mut() { pending_chunk.logprobs = openai_logprobs; pending_chunk.token_ids = openai_token_ids; - if let Some(chunk) = - pending_chunk.take_chunk(&request_id, &response_model, created) - { - y.yield_ok(chunk).await; + if let Some(chunk) = pending_chunk.take_chunk( + &request_id, + &response_model, + created, + choice_index, + ) { + y.yield_ok(ChatCompletionStreamEvent::Chunk(chunk)).await; } } else if let Some(logprobs) = openai_logprobs { - y.yield_ok(logprobs_only_chunk( + y.yield_ok(ChatCompletionStreamEvent::Chunk(logprobs_only_chunk( &request_id, &response_model, created, + choice_index, logprobs, - )) + ))) .await; } } @@ -326,14 +535,15 @@ async fn chat_completion_chunk_stream( if let Some(pending_chunk) = pending_chunk.as_mut() { pending_chunk.push_tool_call_start(tool_index, id, name); } else { - y.yield_ok(tool_call_start_chunk( + y.yield_ok(ChatCompletionStreamEvent::Chunk(tool_call_start_chunk( &request_id, &response_model, created, + choice_index, tool_index, id, name, - )) + ))) .await; } } @@ -342,13 +552,14 @@ async fn chat_completion_chunk_stream( if let Some(pending_chunk) = pending_chunk.as_mut() { pending_chunk.push_tool_call_arguments(tool_index, delta); } else { - y.yield_ok(tool_call_arguments_chunk( + y.yield_ok(ChatCompletionStreamEvent::Chunk(tool_call_arguments_chunk( &request_id, &response_model, created, + choice_index, tool_index, delta, - )) + ))) .await; } } @@ -373,20 +584,25 @@ async fn chat_completion_chunk_stream( } if let Some(pending_chunk) = pending_chunk.as_mut() - && let Some(chunk) = - pending_chunk.take_chunk(&request_id, &response_model, created) + && let Some(chunk) = pending_chunk.take_chunk( + &request_id, + &response_model, + created, + choice_index, + ) { - y.yield_ok(chunk).await; + y.yield_ok(ChatCompletionStreamEvent::Chunk(chunk)).await; } match final_chunk( &request_id, &response_model, created, + choice_index, finish_reason, saw_tool_calls, ) { - Ok(chunk) => y.yield_ok(chunk).await, + Ok(chunk) => y.yield_ok(ChatCompletionStreamEvent::Chunk(chunk)).await, Err(error) => { error!( error = %error.to_error_response().error.message, @@ -396,15 +612,11 @@ async fn chat_completion_chunk_stream( } } - if include_usage { - y.yield_ok(usage_chunk( - &request_id, - &response_model, - created, - Usage::from_counts(prompt_token_count as u32, output_token_count as u32), - )) - .await; - } + y.yield_ok(ChatCompletionStreamEvent::Finished(ChoiceUsage { + prompt_tokens: prompt_token_count as u32, + completion_tokens: output_token_count as u32, + })) + .await; return Ok(()); } @@ -508,6 +720,7 @@ impl PendingChatChunk { request_id: &str, response_model: &str, created: u64, + choice_index: u32, ) -> Option { let has_delta = self.delta.content.is_some() || self.delta.reasoning.is_some() @@ -520,6 +733,7 @@ impl PendingChatChunk { let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created); chunk.choices.push(ChatCompletionStreamChoice { + index: choice_index, delta: self.take_delta(), logprobs, token_ids, @@ -602,9 +816,11 @@ fn start_chunk( request_id: &str, response_model: &str, created: u64, + choice_index: u32, ) -> ChatCompletionStreamResponse { let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created); chunk.choices.push(ChatCompletionStreamChoice { + index: choice_index, delta: ChatMessageDelta { role: Some(AssistantRole), ..Default::default() @@ -619,6 +835,7 @@ fn block_delta_chunk( request_id: &str, response_model: &str, created: u64, + choice_index: u32, kind: AssistantBlockKind, delta: String, ) -> ChatCompletionStreamResponse { @@ -638,6 +855,7 @@ fn block_delta_chunk( let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created); chunk.choices.push(ChatCompletionStreamChoice { + index: choice_index, delta, ..Default::default() }); @@ -648,12 +866,14 @@ fn tool_call_start_chunk( request_id: &str, response_model: &str, created: u64, + choice_index: u32, tool_index: u32, id: String, name: String, ) -> ChatCompletionStreamResponse { let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created); chunk.choices.push(ChatCompletionStreamChoice { + index: choice_index, delta: ChatMessageDelta { tool_calls: Some(vec![ToolCallDelta { index: tool_index, @@ -675,11 +895,13 @@ fn tool_call_arguments_chunk( request_id: &str, response_model: &str, created: u64, + choice_index: u32, tool_index: u32, delta: String, ) -> ChatCompletionStreamResponse { let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created); chunk.choices.push(ChatCompletionStreamChoice { + index: choice_index, delta: ChatMessageDelta { tool_calls: Some(vec![ToolCallDelta { index: tool_index, @@ -701,10 +923,12 @@ fn logprobs_only_chunk( request_id: &str, response_model: &str, created: u64, + choice_index: u32, logprobs: ChatLogProbs, ) -> ChatCompletionStreamResponse { let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created); chunk.choices.push(ChatCompletionStreamChoice { + index: choice_index, logprobs: Some(logprobs), ..Default::default() }); @@ -716,6 +940,7 @@ fn final_chunk( request_id: &str, response_model: &str, created: u64, + choice_index: u32, finish_reason: FinishReason, saw_tool_calls: bool, ) -> Result { @@ -730,6 +955,7 @@ fn final_chunk( let mut chunk = ChatCompletionStreamResponse::new(request_id, response_model, created); chunk.choices.push(ChatCompletionStreamChoice { + index: choice_index, finish_reason: Some(finish_reason.to_string()), stop_reason, ..Default::default() @@ -775,6 +1001,7 @@ mod tests { "chatcmpl-1", "model", 1, + 0, AssistantBlockKind::Text, "hello".to_string(), ); @@ -789,6 +1016,7 @@ mod tests { "chatcmpl-1", "model", 1, + 0, AssistantBlockKind::Reasoning, "thinking".to_string(), ); @@ -806,6 +1034,7 @@ mod tests { "chatcmpl-1", "model", 1, + 0, FinishReason::Stop(Some(StopReason::Text("stop".to_string()))), false, ) @@ -817,7 +1046,7 @@ mod tests { #[test] fn final_chunk_maps_length_finish_reason() { - let chunk = final_chunk("chatcmpl-1", "model", 1, FinishReason::Length, false) + let chunk = final_chunk("chatcmpl-1", "model", 1, 0, FinishReason::Length, false) .expect("finish reason is valid"); assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("length")); @@ -826,7 +1055,7 @@ mod tests { #[test] fn final_chunk_maps_abort_finish_reason() { - let chunk = final_chunk("chatcmpl-1", "model", 1, FinishReason::Abort, false) + let chunk = final_chunk("chatcmpl-1", "model", 1, 0, FinishReason::Abort, false) .expect("abort is a valid finish reason"); assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("abort")); @@ -835,12 +1064,12 @@ mod tests { #[test] fn final_chunk_rejects_error_finish_reason() { - assert!(final_chunk("chatcmpl-1", "model", 1, FinishReason::Error, false).is_err()); + assert!(final_chunk("chatcmpl-1", "model", 1, 0, FinishReason::Error, false).is_err()); } #[test] fn final_chunk_maps_stop_to_tool_calls_when_tool_calls_were_streamed() { - let chunk = final_chunk("chatcmpl-1", "model", 1, FinishReason::stop_eos(), true) + let chunk = final_chunk("chatcmpl-1", "model", 1, 0, FinishReason::stop_eos(), true) .expect("finish reason is valid"); assert_eq!( @@ -892,6 +1121,7 @@ mod tests { "chatcmpl-1".to_string(), "model".to_string(), 1, + 0, false, false, true, @@ -955,6 +1185,7 @@ mod tests { "chatcmpl-1".to_string(), "model".to_string(), 1, + 0, false, false, true, @@ -1014,6 +1245,7 @@ mod tests { "chatcmpl-1".to_string(), "model".to_string(), 1, + 0, false, false, false, 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 2701bef809c8..fa0e5b323790 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -25,6 +25,8 @@ pub struct PreparedRequest { pub response_model: String, /// Whether the caller asked for the final streamed usage chunk. pub include_usage: bool, + /// Number of chat completion choices requested by the caller. + pub n: u32, /// Whether the caller requested output logprobs on chat choices. pub requested_logprobs: bool, /// Whether the caller requested top-level prompt logprobs. @@ -73,6 +75,7 @@ pub(crate) fn prepare_chat_request( let include_usage = (request.stream_options.as_ref()) .and_then(|options| options.include_usage) .unwrap_or(false); + let n = request.n.unwrap_or(1); let requested_logprobs = request.logprobs; // Auto-enable prompt logprobs for non-streaming echo, matching Python vLLM's @@ -144,6 +147,7 @@ pub(crate) fn prepare_chat_request( request_id, response_model, include_usage, + n, requested_logprobs, include_prompt_logprobs, chat_request, 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 fbd10eea0cb0..13f5c4dc19f6 100644 --- a/rust/src/server/src/routes/openai/chat_completions/validate.rs +++ b/rust/src/server/src/routes/openai/chat_completions/validate.rs @@ -18,10 +18,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", @@ -217,6 +213,17 @@ mod tests { .expect("stop strings should be accepted"); } + #[test] + fn validate_request_compat_accepts_n_greater_than_one() { + let request = ChatCompletionRequest { + n: Some(2), + ..base_request() + }; + + validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])) + .expect("n>1 should be accepted"); + } + #[test] fn validate_request_compat_accepts_non_zero_penalties_and_function_tools() { let request = ChatCompletionRequest { diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index 9eda8b9d2a51..2ead5d267336 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -3,6 +3,7 @@ mod types; mod validate; use std::convert::Infallible; +use std::pin::Pin; use std::result::Result; use std::sync::Arc; @@ -12,7 +13,10 @@ use axum::extract::State; use axum::http::HeaderMap; use axum::response::sse::{Event, Sse}; use axum::response::{IntoResponse, Response}; +use futures::future::try_join_all; +use futures::stream::SelectAll; use futures::{Stream, StreamExt as _, pin_mut}; +use serde_json::Value; use thiserror_ext::AsReport as _; use tracing::{debug, error, info, trace}; use tracing_futures::Instrument as _; @@ -34,6 +38,28 @@ use crate::routes::openai::utils::validated_json::ValidatedJson; use crate::state::AppState; use crate::utils::{resolve_request_context, unix_timestamp}; +type CompletionEventStream = + Pin> + Send>>; + +#[derive(Debug, Clone, Copy)] +struct ChoiceUsage { + prompt_tokens: u32, + completion_tokens: u32, +} + +#[derive(Debug)] +struct CompletionChoiceOutput { + choice: CompletionChoice, + usage: ChoiceUsage, + kv_transfer_params: Option, +} + +#[derive(Debug)] +enum CompletionStreamEvent { + Chunk(CompletionStreamResponse), + Finished(ChoiceUsage), +} + /// Validate one completions request and proxy it into the shared `vllm-text` /// stack. pub async fn completions( @@ -60,57 +86,108 @@ pub async fn completions( let include_prompt_logprobs = prepared.text_request.sampling_params.prompt_logprobs.is_some(); let log_request = state.enable_log_requests; - let text_stream = match state - .chat - .text() - .generate(prepared.text_request) - .instrument(request_span.clone()) - .await - { - Ok(stream) => stream, - Err(error) => { - return server_error!( - "failed to submit completion request: {}", - error.to_report_string() - ) - .into_response(); + if stream { + let mut event_streams: SelectAll = SelectAll::new(); + for choice_index in 0..prepared.n { + let text_request = match child_text_request( + &prepared.text_request, + &prepared.request_id, + choice_index, + prepared.n, + ) { + Ok(request) => request, + Err(error) => return error.into_response(), + }; + let text_stream = match state + .chat + .text() + .generate(text_request) + .instrument(request_span.clone()) + .await + { + Ok(stream) => stream, + Err(error) => { + return server_error!( + "failed to submit completion request: {}", + error.to_report_string() + ) + .into_response(); + } + }; + event_streams.push(Box::pin(completion_choice_event_stream( + text_stream, + prepared.request_id.clone(), + prepared.response_model.clone(), + created, + choice_index, + log_request, + prepared.echo.clone(), + logprobs, + prepared.return_token_ids, + prepared.return_tokens_as_token_ids, + ))); } - }; - if stream { - let chunk_stream = completion_chunk_stream( - text_stream, + let chunk_stream = completion_parallel_chunk_stream( + event_streams, prepared.request_id, prepared.response_model, created, - log_request, prepared.include_usage, - prepared.echo, - logprobs, - prepared.return_token_ids, - prepared.return_tokens_as_token_ids, ); let sse_stream = completion_sse_stream(chunk_stream).instrument(request_span); Sse::new(sse_stream).into_response() } else { - let response = match collect_completion( - text_stream, + let mut choice_futures = Vec::with_capacity(prepared.n as usize); + for choice_index in 0..prepared.n { + let text_request = match child_text_request( + &prepared.text_request, + &prepared.request_id, + choice_index, + prepared.n, + ) { + Ok(request) => request, + Err(error) => return error.into_response(), + }; + let text_stream = match state + .chat + .text() + .generate(text_request) + .instrument(request_span.clone()) + .await + { + Ok(stream) => stream, + Err(error) => { + return server_error!( + "failed to submit completion request: {}", + error.to_report_string() + ) + .into_response(); + } + }; + choice_futures.push(collect_completion_choice( + text_stream, + choice_index, + prepared.echo.clone(), + logprobs, + include_prompt_logprobs, + prepared.return_token_ids, + prepared.return_tokens_as_token_ids, + )); + } + + let choices = match try_join_all(choice_futures).instrument(request_span.clone()).await { + Ok(choices) => choices, + Err(error) => return error.into_response(), + }; + + let response = completion_response_from_choices( prepared.request_id, prepared.response_model, created, - prepared.echo, - logprobs, - include_prompt_logprobs, - prepared.return_token_ids, - prepared.return_tokens_as_token_ids, - ) - .instrument(request_span.clone()) - .await - { - Ok(response) => response, - Err(error) => return error.into_response(), - }; + choices, + ); if log_request { let usage = response.usage.as_ref(); @@ -128,17 +205,39 @@ pub async fn completions( } } -async fn collect_completion( +fn child_text_request( + base: &vllm_text::TextRequest, + parent_request_id: &str, + choice_index: u32, + choice_count: u32, +) -> Result { + let mut request = base.clone(); + if choice_count > 1 { + request.request_id = format!("{choice_index}_{parent_request_id}"); + } else { + request.request_id = parent_request_id.to_string(); + } + if let Some(seed) = request.sampling_params.seed { + request.sampling_params.seed = + Some(seed.checked_add(i64::from(choice_index)).ok_or_else(|| { + ApiError::invalid_request( + "`seed + choice index` must fit within a signed 64-bit integer.".to_string(), + Some("seed"), + ) + })?); + } + Ok(request) +} + +async fn collect_completion_choice( stream: impl TextOutputStream, - request_id: String, - response_model: String, - created: u64, + choice_index: u32, echo: Option, requested_logprobs: Option, include_prompt_logprobs: bool, return_token_ids: bool, return_tokens_as_token_ids: bool, -) -> Result { +) -> Result { let collected = stream .collect_output() .await @@ -176,13 +275,9 @@ async fn collect_completion( Some(prompt) => format!("{prompt}{}", collected.text), }; - Ok(CompletionResponse { - id: request_id, - object: "text_completion".to_string(), - created, - model: response_model, - choices: vec![CompletionChoice { - index: 0, + Ok(CompletionChoiceOutput { + choice: CompletionChoice { + index: choice_index, text, logprobs, finish_reason: Some(completion_finish_reason_to_openai(finish_reason)?.into()), @@ -190,16 +285,37 @@ async fn collect_completion( prompt_logprobs, token_ids: return_token_ids.then(|| collected.token_ids.clone()), prompt_token_ids: return_token_ids.then(|| collected.prompt_token_ids.to_vec()), - }], - usage: Some(Usage::from_counts( - collected.prompt_token_ids.len() as u32, - collected.token_ids.len() as u32, - )), - system_fingerprint: None, + }, + usage: ChoiceUsage { + prompt_tokens: collected.prompt_token_ids.len() as u32, + completion_tokens: collected.token_ids.len() as u32, + }, kv_transfer_params: collected.kv_transfer_params, }) } +fn completion_response_from_choices( + request_id: String, + response_model: String, + created: u64, + outputs: Vec, +) -> CompletionResponse { + let prompt_tokens = outputs.first().map_or(0, |output| output.usage.prompt_tokens); + let completion_tokens = outputs.iter().map(|output| output.usage.completion_tokens).sum(); + let kv_transfer_params = outputs.iter().find_map(|output| output.kv_transfer_params.clone()); + + CompletionResponse { + id: request_id, + object: "text_completion".to_string(), + created, + model: response_model, + choices: outputs.into_iter().map(|output| output.choice).collect(), + usage: Some(Usage::from_counts(prompt_tokens, completion_tokens)), + system_fingerprint: None, + kv_transfer_params, + } +} + /// Convert one internal decoded-text stream into OpenAI completions chunks. #[try_stream] async fn completion_chunk_stream( @@ -207,6 +323,7 @@ async fn completion_chunk_stream( request_id: String, response_model: String, created: u64, + choice_index: u32, log_request: bool, include_usage: bool, echo: Option, @@ -214,6 +331,94 @@ async fn completion_chunk_stream( return_token_ids: bool, return_tokens_as_token_ids: bool, mut y: TryYielder, +) -> Result<(), ApiError> { + let events = completion_choice_event_stream( + stream, + request_id.clone(), + response_model.clone(), + created, + choice_index, + log_request, + echo, + requested_logprobs, + return_token_ids, + return_tokens_as_token_ids, + ); + pin_mut!(events); + + while let Some(next) = events.next().await { + match next? { + CompletionStreamEvent::Chunk(chunk) => { + y.yield_ok(CompletionSseChunk::Chunk(chunk)).await; + } + CompletionStreamEvent::Finished(usage) => { + if include_usage { + y.yield_ok(CompletionSseChunk::Usage(usage_chunk( + &request_id, + &response_model, + created, + Usage::from_counts(usage.prompt_tokens, usage.completion_tokens), + ))) + .await; + } + } + } + } + + Ok(()) +} + +#[try_stream] +async fn completion_parallel_chunk_stream( + stream: impl Stream>, + request_id: String, + response_model: String, + created: u64, + include_usage: bool, + mut y: TryYielder, +) -> Result<(), ApiError> { + pin_mut!(stream); + let mut prompt_tokens = None; + let mut completion_tokens = 0_u32; + + while let Some(next) = stream.next().await { + match next? { + CompletionStreamEvent::Chunk(chunk) => { + y.yield_ok(CompletionSseChunk::Chunk(chunk)).await; + } + CompletionStreamEvent::Finished(usage) => { + prompt_tokens.get_or_insert(usage.prompt_tokens); + completion_tokens = completion_tokens.saturating_add(usage.completion_tokens); + } + } + } + + if include_usage { + y.yield_ok(CompletionSseChunk::Usage(usage_chunk( + &request_id, + &response_model, + created, + Usage::from_counts(prompt_tokens.unwrap_or(0), completion_tokens), + ))) + .await; + } + + Ok(()) +} + +#[try_stream] +async fn completion_choice_event_stream( + stream: impl TextOutputStream, + request_id: String, + response_model: String, + created: u64, + choice_index: u32, + log_request: bool, + echo: Option, + requested_logprobs: Option, + return_token_ids: bool, + return_tokens_as_token_ids: bool, + mut y: TryYielder, ) -> Result<(), ApiError> { pin_mut!(stream); let mut visible_text_len = 0_u32; @@ -227,24 +432,36 @@ async fn completion_chunk_stream( debug!("completion stream started"); if let Some(prompt) = echo.as_ref() { visible_text_len = text_len(prompt); - let mut chunk = - delta_chunk(&request_id, &response_model, created, prompt.clone(), None); + let mut chunk = delta_chunk( + &request_id, + &response_model, + created, + choice_index, + prompt.clone(), + None, + ); if return_token_ids && first_chunk { if let Some(choice) = chunk.choices.first_mut() { choice.prompt_token_ids = Some(prompt_token_ids.to_vec()); } first_chunk = false; } - y.yield_ok(CompletionSseChunk::Chunk(chunk)).await; + y.yield_ok(CompletionStreamEvent::Chunk(chunk)).await; } else if return_token_ids { // Emit a chunk with prompt_token_ids in the first streaming response - let mut chunk = - delta_chunk(&request_id, &response_model, created, String::new(), None); + let mut chunk = delta_chunk( + &request_id, + &response_model, + created, + choice_index, + String::new(), + None, + ); if let Some(choice) = chunk.choices.first_mut() { choice.prompt_token_ids = Some(prompt_token_ids.to_vec()); } first_chunk = false; - y.yield_ok(CompletionSseChunk::Chunk(chunk)).await; + y.yield_ok(CompletionStreamEvent::Chunk(chunk)).await; } } Ok(DecodedTextEvent::TextDelta { @@ -268,44 +485,45 @@ async fn completion_chunk_stream( } else { None }; - let mut chunk = delta_chunk(&request_id, &response_model, created, delta, logprobs); + let mut chunk = delta_chunk( + &request_id, + &response_model, + created, + choice_index, + delta, + logprobs, + ); if return_token_ids && let Some(choice) = chunk.choices.first_mut() { choice.token_ids = Some(token_ids); } - y.yield_ok(CompletionSseChunk::Chunk(chunk)).await; + y.yield_ok(CompletionStreamEvent::Chunk(chunk)).await; visible_text_len = visible_text_len.saturating_add(delta_text_len); if let Some(finished) = finished { + let finish_reason = finished.finish_reason.clone(); if log_request { info!( stream = true, model = %response_model, prompt_tokens = finished.prompt_token_count, output_tokens = finished.output_token_count, - finish_reason = finished.finish_reason.as_str(), + finish_reason = finish_reason.as_str(), "completion finished" ); } - y.yield_ok(CompletionSseChunk::Chunk(final_chunk( + y.yield_ok(CompletionStreamEvent::Chunk(final_chunk( &request_id, &response_model, created, - finished.finish_reason, + choice_index, + finish_reason, )?)) .await; - - if include_usage { - y.yield_ok(CompletionSseChunk::Usage(usage_chunk( - &request_id, - &response_model, - created, - Usage::from_counts( - finished.prompt_token_count as u32, - finished.output_token_count as u32, - ), - ))) - .await; - } + y.yield_ok(CompletionStreamEvent::Finished(ChoiceUsage { + prompt_tokens: finished.prompt_token_count as u32, + completion_tokens: finished.output_token_count as u32, + })) + .await; } } Err(error) => { @@ -324,11 +542,13 @@ fn delta_chunk( request_id: &str, response_model: &str, created: u64, + choice_index: u32, text: String, logprobs: Option, ) -> CompletionStreamResponse { let mut chunk = CompletionStreamResponse::new(request_id, response_model, created); chunk.choices.push(CompletionStreamChoice { + index: choice_index, text, logprobs, ..Default::default() @@ -340,12 +560,14 @@ fn final_chunk( request_id: &str, response_model: &str, created: u64, + choice_index: u32, finish_reason: FinishReason, ) -> Result { let finish_reason = completion_finish_reason_to_openai(finish_reason)?; let mut chunk = CompletionStreamResponse::new(request_id, response_model, created); chunk.choices.push(CompletionStreamChoice { + index: choice_index, finish_reason: Some(finish_reason.to_string()), ..Default::default() }); @@ -436,7 +658,7 @@ mod tests { #[test] fn final_chunk_maps_stop_finish_reason() { - let chunk = final_chunk("cmpl-1", "model", 1, FinishReason::stop_eos()) + let chunk = final_chunk("cmpl-1", "model", 1, 0, FinishReason::stop_eos()) .expect("finish reason valid"); assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("stop")); assert_eq!(chunk.choices[0].text, ""); @@ -444,21 +666,21 @@ mod tests { #[test] fn final_chunk_maps_length_finish_reason() { - let chunk = - final_chunk("cmpl-1", "model", 1, FinishReason::Length).expect("finish reason valid"); + let chunk = final_chunk("cmpl-1", "model", 1, 0, FinishReason::Length) + .expect("finish reason valid"); assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("length")); } #[test] fn final_chunk_maps_abort_finish_reason() { let chunk = - final_chunk("cmpl-1", "model", 1, FinishReason::Abort).expect("finish reason valid"); + final_chunk("cmpl-1", "model", 1, 0, FinishReason::Abort).expect("finish reason valid"); assert_eq!(chunk.choices[0].finish_reason.as_deref(), Some("abort")); } #[test] fn final_chunk_rejects_error_finish_reason() { - assert!(final_chunk("cmpl-1", "model", 1, FinishReason::Error).is_err()); + assert!(final_chunk("cmpl-1", "model", 1, 0, FinishReason::Error).is_err()); } #[tokio::test] @@ -526,6 +748,7 @@ mod tests { "cmpl-1".to_string(), "model".to_string(), 1, + 0, false, false, None, diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index 2d4ff089397c..e8416bc51c5b 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -17,6 +17,8 @@ pub struct PreparedRequest { pub response_model: String, /// Whether the caller asked for the final streamed usage chunk. pub include_usage: bool, + /// Number of completion choices requested by the caller. + pub n: u32, /// Lowered text request for the shared `vllm-text` facade. pub text_request: TextRequest, /// Original text prompt that should be echoed back northbound when @@ -64,6 +66,7 @@ pub(crate) fn prepare_completion_request( let include_usage = (request.stream_options.as_ref()) .and_then(|options| options.include_usage) .unwrap_or(false); + let n = request.n.unwrap_or(1); let echo = request.echo.then(|| request.prompt.as_text().cloned()).flatten(); let structured_outputs = @@ -117,6 +120,7 @@ pub(crate) fn prepare_completion_request( request_id, response_model, include_usage, + n, text_request, echo, return_token_ids: request.return_token_ids.unwrap_or(false), diff --git a/rust/src/server/src/routes/openai/completions/validate.rs b/rust/src/server/src/routes/openai/completions/validate.rs index a53609234b60..0fb26e955740 100644 --- a/rust/src/server/src/routes/openai/completions/validate.rs +++ b/rust/src/server/src/routes/openai/completions/validate.rs @@ -22,8 +22,8 @@ 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.n == Some(0) { + bail_invalid_request!(param = "n", "n must be at least 1."); } if request.max_tokens == Some(0) { @@ -135,6 +135,28 @@ mod tests { ); } + #[test] + fn validate_request_compat_accepts_n_greater_than_one() { + let request = CompletionRequest { + n: Some(2), + ..base_request() + }; + assert!( + validate_request_compat(&request, &served_names(&["Qwen/Qwen1.5-0.5B-Chat"])).is_ok() + ); + } + + #[test] + fn validate_request_compat_rejects_zero_n() { + let request = CompletionRequest { + n: Some(0), + ..base_request() + }; + assert!( + validate_request_compat(&request, &served_names(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err() + ); + } + #[test] fn validate_request_compat_accepts_any_served_name() { let request = base_request(); diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index a3e437e04802..33b2603f6e57 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -149,6 +149,18 @@ fn assert_adapter_a_lora_request(request: &EngineCoreRequest) { assert_eq!(lora.lora_path, "org/adapter-a"); } +fn stream_output_specs_from_text(text: &str) -> Vec<(Vec, Option)> { + let token_ids = bytes_to_token_ids(text.as_bytes()); + token_ids + .into_iter() + .enumerate() + .map(|(index, token_id)| { + let finish_reason = (index + 1 == text.len()).then_some(EngineCoreFinishReason::Stop); + (vec![token_id], finish_reason) + }) + .collect() +} + fn sse_data_payloads(text: &str) -> Vec<&str> { text.lines().filter_map(|line| line.strip_prefix("data: ")).collect() } @@ -871,6 +883,71 @@ async fn test_app_with_stream_output_specs( ) } +async fn test_app_with_parallel_stream_output_specs( + output_specs_by_choice: Vec, Option)>>, +) -> (axum::Router, MockEngineTask) { + test_app_with_parallel_stream_output_specs_and_request_check(output_specs_by_choice, |_| {}) + .await +} + +async fn test_app_with_parallel_stream_output_specs_and_request_check( + 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.into_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, + ) +} + async fn test_app_with_backend_and_stream_output_specs( backend: Arc, output_specs: Vec<(Vec, Option)>, @@ -1640,6 +1717,105 @@ async fn non_stream_chat_returns_json_response() { assert_eq!(json["usage"]["total_tokens"], 25); } +#[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_stream_output_specs(vec![ + stream_output_specs_from_text("hi!"), + stream_output_specs_from_text("yo!"), + ]) + .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": false, + "n": 2, + "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["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"); + 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 non_stream_chat_completions_fans_out_choice_request_ids_and_seeds() { + let (app, engine_task) = test_app_with_parallel_stream_output_specs_and_request_check( + vec![ + stream_output_specs_from_text("hi!"), + stream_output_specs_from_text("yo!"), + ], + |requests| { + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].request_id, "0_chatcmpl-choice-parent"); + assert_eq!(requests[1].request_id, "1_chatcmpl-choice-parent"); + assert_eq!( + requests[0].sampling_params.as_ref().expect("sampling params").seed, + Some(41) + ); + assert_eq!( + requests[1].sampling_params.as_ref().expect("sampling params").seed, + Some(42) + ); + }, + ) + .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"][1]["index"], 1); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn non_stream_chat_image_url_reaches_engine_mm_features() { @@ -2210,6 +2386,71 @@ async fn include_usage_adds_final_usage_chunk_before_done() { assert_eq!(usage_chunk["usage"]["total_tokens"], 25); } +#[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_stream_output_specs(vec![ + stream_output_specs_from_text("hi!"), + stream_output_specs_from_text("yo!"), + ]) + .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}, + "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_data_payloads(&text); + + assert!( + payloads.iter().any(|payload| { + payload.contains("\"index\":0") && payload.contains("\"content\":\"h\"") + }), + "{text}" + ); + assert!( + payloads.iter().any(|payload| { + payload.contains("\"index\":1") && payload.contains("\"content\":\"y\"") + }), + "{text}" + ); + assert_eq!( + payloads.iter().filter(|payload| **payload == "[DONE]").count(), + 1 + ); + + let usage_chunk: serde_json::Value = serde_json::from_str( + payloads + .iter() + .find(|payload| payload.contains("\"usage\":")) + .expect("usage chunk"), + ) + .expect("usage chunk json"); + assert_eq!(usage_chunk["choices"], json!([])); + assert_eq!(usage_chunk["usage"]["prompt_tokens"], 22); + assert_eq!(usage_chunk["usage"]["completion_tokens"], 6); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn stream_without_include_usage_keeps_existing_shape() { @@ -2318,6 +2559,105 @@ async fn non_stream_completions_return_json_response() { assert_eq!(json["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_stream_output_specs(vec![ + stream_output_specs_from_text("hi!"), + stream_output_specs_from_text("yo!"), + ]) + .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": false, + "n": 2 + }) + .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["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"); + assert_eq!(json["usage"]["prompt_tokens"], 5); + assert_eq!(json["usage"]["completion_tokens"], 6); + assert_eq!(json["usage"]["total_tokens"], 11); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn non_stream_completions_fans_out_choice_request_ids_and_seeds() { + let (app, engine_task) = test_app_with_parallel_stream_output_specs_and_request_check( + vec![ + stream_output_specs_from_text("hi!"), + stream_output_specs_from_text("yo!"), + ], + |requests| { + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].request_id, "0_cmpl-choice-parent"); + assert_eq!(requests[1].request_id, "1_cmpl-choice-parent"); + assert_eq!( + requests[0].sampling_params.as_ref().expect("sampling params").seed, + Some(41) + ); + assert_eq!( + requests[1].sampling_params.as_ref().expect("sampling params").seed, + Some(42) + ); + }, + ) + .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"][1]["index"], 1); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn non_stream_completions_echo_prepends_prompt_text() { @@ -3301,6 +3641,71 @@ async fn completions_happy_path_returns_sse_stream() { assert_eq!(usage_chunk["usage"]["completion_tokens"], 3); } +#[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_stream_output_specs(vec![ + stream_output_specs_from_text("hi!"), + stream_output_specs_from_text("yo!"), + ]) + .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_data_payloads(&text); + + assert!( + payloads + .iter() + .any(|payload| payload.contains("\"index\":0") && payload.contains("\"text\":\"h\"")), + "{text}" + ); + assert!( + payloads + .iter() + .any(|payload| payload.contains("\"index\":1") && payload.contains("\"text\":\"y\"")), + "{text}" + ); + assert_eq!( + payloads.iter().filter(|payload| **payload == "[DONE]").count(), + 1 + ); + + let usage_chunk: serde_json::Value = serde_json::from_str( + payloads + .iter() + .find(|payload| payload.contains("\"usage\":")) + .expect("usage chunk"), + ) + .expect("usage chunk json"); + assert_eq!(usage_chunk["choices"], json!([])); + assert_eq!(usage_chunk["usage"]["prompt_tokens"], 5); + assert_eq!(usage_chunk["usage"]["completion_tokens"], 6); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn completions_echo_stream_emits_separate_prompt_chunk() {