diff --git a/rust/src/chat/src/error.rs b/rust/src/chat/src/error.rs index 25d8d01568..bbd9957200 100644 --- a/rust/src/chat/src/error.rs +++ b/rust/src/chat/src/error.rs @@ -66,6 +66,8 @@ pub enum Error { ToolCallStreamInvariant { message: String }, #[error(transparent)] Text(#[from] vllm_text::Error), + #[error(transparent)] + Tokenizer(#[from] vllm_tokenizer::TokenizerError), } pub type Result = std::result::Result; diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index e0bf965215..ae366a4de0 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -52,7 +52,7 @@ mod stream; use vllm_engine_core_client::EngineCoreClient; use vllm_engine_core_client::protocol::ModelDtype; use vllm_llm::Llm; -use vllm_text::{TextLlm, TextRequest}; +use vllm_text::{Prompt, TextLlm, TextRequest}; /// Validate explicit parser override names without starting request processing. pub fn validate_parser_overrides( @@ -198,6 +198,33 @@ impl ChatLlm { Ok(ChatEventStream::new(request.request_id, structured_stream)) } + /// Render through the chat template and tokenize, without submitting to the engine. + /// + /// Same render → [`multimodal::finalize_rendered_prompt`] → encode pipeline as + /// [`Self::chat`], but stops after token IDs so `/tokenize` counts match what + /// generation would see. Used by `POST /tokenize` (chat form). + pub async fn tokenize_chat(&self, request: ChatRequest) -> Result> { + request.validate()?; + + let rendered = self.backend.chat_renderer().render(&request)?; + let (prompt, _mm_features) = multimodal::finalize_rendered_prompt( + &request, + rendered, + self.backend.multimodal_model_info(), + self.model_dtype, + ) + .await?; + + let tokenizer = self.text.tokenizer(); + let token_ids = match prompt { + // Rendered string from the template (usual chat path). + Prompt::Text(text) => tokenizer.encode(&text, request.add_special_tokens)?, + // Already tokenized (e.g. multimodal path); pass through unchanged. + Prompt::TokenIds(ids) => ids, + }; + Ok(token_ids) + } + /// Shut down the underlying LLM client and its background tasks. pub async fn shutdown(self) -> Result<()> { self.text.shutdown().await?; diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index bf2218c27e..a942d0a3c9 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -9,6 +9,7 @@ pub(crate) mod openai; mod pause; mod server_info; mod sleep; +mod tokenize; mod version; use std::sync::Arc; @@ -72,7 +73,9 @@ fn build_router_with_options( .route("/v1/models", get(openai::list_models)) .route("/v1/completions", post(openai::completions)) .route("/v1/chat/completions", post(openai::chat_completions)) - // vLLM specific inference endpoints + // vLLM specific endpoints + .route("/tokenize", post(tokenize::tokenize)) + .route("/detokenize", post(tokenize::detokenize)) .route("/inference/v1/generate", post(inference::generate)); if runtime_lora_updating_enabled { 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 ba75c9c1a8..04ae791bc3 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -156,8 +156,7 @@ pub(crate) fn prepare_chat_request( return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false), }) } - -fn normalize_generation_prompt_mode( +pub(crate) fn normalize_generation_prompt_mode( add_generation_prompt: Option, continue_final_message: bool, messages: &[VllmChatMessage], @@ -204,7 +203,7 @@ fn extract_last_assistant_content(messages: &[ChatMessage]) -> Option { } /// Lower one OpenAI chat message into the `vllm-chat` message shape. -fn convert_message(message: ChatMessage) -> Result { +pub(crate) fn convert_message(message: ChatMessage) -> Result { match message { ChatMessage::System { content, .. } => { Ok(VllmChatMessage::system(convert_content(content)?)) @@ -316,7 +315,7 @@ fn convert_assistant_tool_calls( .collect() } -fn convert_tools(tools: Option>) -> Result, ApiError> { +pub(crate) fn convert_tools(tools: Option>) -> Result, ApiError> { tools .unwrap_or_default() .into_iter() diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index 00557ad53d..3efef62213 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -9,9 +9,9 @@ use vllm_chat::ReasoningEffort; use crate::routes::openai::utils::structured_outputs::ResponseFormat; use crate::routes::openai::utils::types::{ - ChatLogProbs, ChatMessage, MessageContent, Normalizable, StreamOptions, StringOrArray, Tool, - ToolCall, ToolCallDelta, ToolChoice, ToolChoiceValue, ToolReference, UNKNOWN_MODEL_ID, Usage, - default_true, validate_stop, validate_top_p_value, + ChatLogProbs, ChatMessage, Normalizable, StreamOptions, StringOrArray, Tool, ToolCall, + ToolCallDelta, ToolChoice, ToolChoiceValue, ToolReference, UNKNOWN_MODEL_ID, Usage, + default_true, validate_messages, validate_stop, validate_top_p_value, }; /// vLLM-compatible request type for the Chat Completions API. @@ -430,32 +430,6 @@ fn default_model() -> String { UNKNOWN_MODEL_ID.to_string() } -/// Validates messages array is not empty and has valid content -fn validate_messages(messages: &[ChatMessage]) -> Result<(), validator::ValidationError> { - if messages.is_empty() { - return Err(validator::ValidationError::new("messages cannot be empty")); - } - - for msg in messages { - if let ChatMessage::User { content, .. } = msg { - match content { - MessageContent::Text(text) if text.is_empty() => { - return Err(validator::ValidationError::new( - "message content cannot be empty", - )); - } - MessageContent::Parts(parts) if parts.is_empty() => { - return Err(validator::ValidationError::new( - "message content parts cannot be empty", - )); - } - _ => {} - } - } - } - Ok(()) -} - /// Schema-level validation for cross-field dependencies fn validate_chat_cross_parameters( req: &ChatCompletionRequest, diff --git a/rust/src/server/src/routes/openai/utils/types.rs b/rust/src/server/src/routes/openai/utils/types.rs index ff747a5daf..9e0acd04cc 100644 --- a/rust/src/server/src/routes/openai/utils/types.rs +++ b/rust/src/server/src/routes/openai/utils/types.rs @@ -412,6 +412,41 @@ pub struct ListModelsResponse { pub data: Vec, } +// ============================================================================ +// Shared validation helpers +// ============================================================================ + +/// Validates a messages array is non-empty and has valid user-message content. +/// +/// Used by both `POST /v1/chat/completions` and `POST /tokenize` (chat form) +/// so validation behaviour stays in lockstep. +pub(crate) fn validate_messages( + messages: &[ChatMessage], +) -> Result<(), validator::ValidationError> { + if messages.is_empty() { + return Err(validator::ValidationError::new("messages cannot be empty")); + } + + for msg in messages { + if let ChatMessage::User { content, .. } = msg { + match content { + MessageContent::Text(text) if text.is_empty() => { + return Err(validator::ValidationError::new( + "message content cannot be empty", + )); + } + MessageContent::Parts(parts) if parts.is_empty() => { + return Err(validator::ValidationError::new( + "message content parts cannot be empty", + )); + } + _ => {} + } + } + } + Ok(()) +} + // ============================================================================ // Normalizable trait // ============================================================================ diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 5a2dea7017..67aa9cccff 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -407,6 +407,9 @@ struct FakeChatBackend { multimodal_model_info: Option, } +/// Synthetic BOS id used when `add_special_tokens` is true in tests. +const FAKE_BOS_TOKEN_ID: u32 = 1; + #[derive(Debug)] struct FakeChatTokenizer; @@ -414,9 +417,12 @@ impl Tokenizer for FakeChatTokenizer { fn encode( &self, text: &str, - _add_special_tokens: bool, + add_special_tokens: bool, ) -> vllm_text::tokenizer::Result> { let mut token_ids = Vec::new(); + if add_special_tokens { + token_ids.push(FAKE_BOS_TOKEN_ID); + } let mut rest = text; while !rest.is_empty() { if let Some(stripped) = rest.strip_prefix("") { @@ -460,6 +466,7 @@ impl Tokenizer for FakeChatTokenizer { fn id_to_token(&self, id: u32) -> Option { match id { + FAKE_BOS_TOKEN_ID => Some("".to_string()), 999 => Some("".to_string()), 151655 => Some("<|image_pad|>".to_string()), 0xF001 => Some("".to_string()), @@ -468,6 +475,7 @@ impl Tokenizer for FakeChatTokenizer { 0xF004 => Some("<|END_THINKING|>".to_string()), 0xF005 => Some("◁think▷".to_string()), 0xF006 => Some("◁/think▷".to_string()), + id if id < 128 => char::from_u32(id).map(|ch| ch.to_string()), _ => None, } } @@ -2442,7 +2450,8 @@ async fn non_stream_completions_echo_prepends_prompt_text() { "model": "Qwen/Qwen1.5-0.5B-Chat", "prompt": "hello", "echo": true, - "stream": false + "stream": false, + "add_special_tokens": false }) .to_string(), )) @@ -2645,7 +2654,8 @@ async fn non_stream_completions_include_prompt_logprobs() { "prompt": "hello", "stream": false, "echo": true, - "logprobs": 1 + "logprobs": 1, + "add_special_tokens": false }) .to_string(), )) @@ -3426,7 +3436,8 @@ async fn completions_echo_stream_emits_separate_prompt_chunk() { "prompt": "hello", "echo": true, "stream": true, - "stream_options": {"include_usage": true} + "stream_options": {"include_usage": true}, + "add_special_tokens": false }) .to_string(), )) @@ -4864,3 +4875,392 @@ async fn completions_empty_stop_string_returns_validation_error() { assert_eq!(response.status(), StatusCode::BAD_REQUEST); } + +async fn post_json( + app: &mut axum::Router, + uri: &str, + body: serde_json::Value, +) -> (StatusCode, serde_json::Value) { + let response = app + .call( + Request::builder() + .method("POST") + .uri(uri) + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .expect("build request"), + ) + .await + .expect("call app"); + let status = response.status(); + let bytes = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + let json: serde_json::Value = serde_json::from_slice(&bytes) + .unwrap_or_else(|_| json!({ "raw": String::from_utf8_lossy(&bytes) })); + (status, json) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn tokenize_completion_round_trips_through_detokenize() { + let mut app = test_app().await; + let prompt = "Hello world"; + + let (_, tokenize_json) = post_json( + &mut app, + "/tokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "prompt": prompt, + "add_special_tokens": false, + }), + ) + .await; + let tokens = tokenize_json["tokens"] + .as_array() + .expect("tokens array") + .iter() + .map(|v| v.as_u64().expect("token id") as u32) + .collect::>(); + + let (status, detokenize_json) = post_json( + &mut app, + "/detokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "tokens": tokens, + }), + ) + .await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(detokenize_json["prompt"], prompt); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn tokenize_completion_add_special_tokens_changes_ids() { + let mut app = test_app().await; + + let (_, with_special) = post_json( + &mut app, + "/tokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "prompt": "hi", + "add_special_tokens": true, + }), + ) + .await; + let (_, without_special) = post_json( + &mut app, + "/tokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "prompt": "hi", + "add_special_tokens": false, + }), + ) + .await; + + let with_ids: Vec = with_special["tokens"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_u64().unwrap() as u32) + .collect(); + let without_ids: Vec = without_special["tokens"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_u64().unwrap() as u32) + .collect(); + + assert_ne!(with_ids, without_ids); + assert_eq!(with_ids.first().copied(), Some(FAKE_BOS_TOKEN_ID)); + assert_eq!(without_ids.first().copied(), Some(b'h' as u32)); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn tokenize_completion_return_token_strs_matches_tokens() { + let mut app = test_app().await; + + let (status, json) = post_json( + &mut app, + "/tokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "prompt": "hi", + "add_special_tokens": false, + "return_token_strs": true, + }), + ) + .await; + + assert_eq!(status, StatusCode::OK); + let tokens = json["tokens"].as_array().expect("tokens"); + let token_strs = json["token_strs"].as_array().expect("token_strs"); + assert_eq!(tokens.len(), token_strs.len()); + assert_eq!(token_strs.len(), json["count"].as_u64().unwrap() as usize); + assert!(!token_strs[0].as_str().unwrap().is_empty()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn tokenize_completion_count_and_max_model_len() { + let mut app = test_app().await; + + let (status, json) = post_json( + &mut app, + "/tokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "prompt": "hello", + "add_special_tokens": false, + }), + ) + .await; + + assert_eq!(status, StatusCode::OK); + assert_eq!( + json["count"].as_u64().unwrap() as usize, + json["tokens"].as_array().unwrap().len() + ); + assert!(json["max_model_len"].as_u64().unwrap() > 0); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn tokenize_chat_includes_generation_prompt_in_token_count() { + let mut app = test_app().await; + let messages = json!([{"role": "user", "content": "hi"}]); + + let (_, with_prompt) = post_json( + &mut app, + "/tokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "messages": messages, + "add_generation_prompt": true, + "add_special_tokens": false, + }), + ) + .await; + let (_, without_prompt) = post_json( + &mut app, + "/tokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "messages": messages, + "add_generation_prompt": false, + "add_special_tokens": false, + }), + ) + .await; + + let with_len = with_prompt["tokens"].as_array().unwrap().len(); + let without_len = without_prompt["tokens"].as_array().unwrap().len(); + assert!(with_len > without_len); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn tokenize_chat_conflicting_generation_flags_returns_400() { + let mut app = test_app().await; + + let (status, json) = post_json( + &mut app, + "/tokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "messages": [{"role": "user", "content": "hi"}], + "add_generation_prompt": true, + "continue_final_message": true, + }), + ) + .await; + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(json["error"]["type"], "invalid_request_error"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn tokenize_chat_empty_messages_returns_400() { + let mut app = test_app().await; + + let (status, json) = post_json( + &mut app, + "/tokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "messages": [], + }), + ) + .await; + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(json["error"]["type"], "invalid_request_error"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn tokenize_chat_empty_message_content_returns_400() { + let mut app = test_app().await; + + let (status, json) = post_json( + &mut app, + "/tokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "messages": [{"role": "user", "content": ""}], + }), + ) + .await; + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(json["error"]["type"], "invalid_request_error"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn tokenize_unknown_model_returns_404() { + let mut app = test_app().await; + + let (status, json) = post_json( + &mut app, + "/tokenize", + json!({ + "model": "does-not-exist", + "prompt": "hello", + }), + ) + .await; + + assert_eq!(status, StatusCode::NOT_FOUND); + assert_eq!(json["error"]["type"], "invalid_request_error"); + assert_eq!(json["error"]["code"], "model_not_found"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn detokenize_unknown_model_returns_404() { + let mut app = test_app().await; + + let (status, json) = post_json( + &mut app, + "/detokenize", + json!({ + "model": "does-not-exist", + "tokens": [72, 101, 108, 108, 111], + }), + ) + .await; + + assert_eq!(status, StatusCode::NOT_FOUND); + assert_eq!(json["error"]["type"], "invalid_request_error"); + assert_eq!(json["error"]["code"], "model_not_found"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn detokenize_empty_tokens_returns_empty_prompt() { + let mut app = test_app().await; + + let (status, json) = post_json( + &mut app, + "/detokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "tokens": [], + }), + ) + .await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(json["prompt"], ""); +} + +/// Decode an explicit token sequence — pins `/detokenize` independently of +/// `/tokenize` (the round-trip test alone would pass even if encode and decode +/// were both wrong in mirrored ways). +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn detokenize_decodes_known_token_ids() { + let mut app = test_app().await; + + let (status, json) = post_json( + &mut app, + "/detokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "tokens": [72, 101, 108, 108, 111], + }), + ) + .await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(json["prompt"], "Hello"); +} + +/// `continue_final_message` without a trailing assistant message must 400. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn tokenize_chat_continue_without_assistant_returns_400() { + let mut app = test_app().await; + + let (status, json) = post_json( + &mut app, + "/tokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "messages": [{"role": "user", "content": "hi"}], + "add_generation_prompt": false, + "continue_final_message": true, + }), + ) + .await; + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(json["error"]["type"], "invalid_request_error"); +} + +/// `continue_final_message` must not append a new generation suffix vs `add_generation_prompt`. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn tokenize_chat_continue_final_vs_new_assistant_differs() { + let mut app = test_app().await; + let messages = json!([ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "partial,"} + ]); + + let (_, continue_final) = post_json( + &mut app, + "/tokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "messages": messages, + "add_generation_prompt": false, + "continue_final_message": true, + "add_special_tokens": false, + }), + ) + .await; + let (_, new_assistant) = post_json( + &mut app, + "/tokenize", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "messages": messages, + "add_generation_prompt": true, + "continue_final_message": false, + "add_special_tokens": false, + }), + ) + .await; + + let continue_len = continue_final["tokens"].as_array().unwrap().len(); + let new_len = new_assistant["tokens"].as_array().unwrap().len(); + assert!(new_len > continue_len); +} diff --git a/rust/src/server/src/routes/tokenize.rs b/rust/src/server/src/routes/tokenize.rs new file mode 100644 index 0000000000..3fb4f149e6 --- /dev/null +++ b/rust/src/server/src/routes/tokenize.rs @@ -0,0 +1,148 @@ +//! `POST /tokenize` and `POST /detokenize` (root paths, matching Python). +//! +//! Encode/decode runs entirely in-process via [`DynTokenizer`]; the inference +//! engine is not involved. + +mod types; + +use std::sync::Arc; + +use axum::Json; +use axum::extract::State; +use axum::http::HeaderMap; +use axum::response::{IntoResponse, Response}; +use thiserror_ext::AsReport as _; + +use crate::error::{ApiError, server_error}; +use crate::routes::openai::utils::validated_json::ValidatedJson; +use crate::routes::tokenize::types::{ + DetokenizeRequest, DetokenizeResponse, TokenizeChatRequest, TokenizeCompletionRequest, + TokenizeRequest, TokenizeResponse, +}; +use crate::state::AppState; +use crate::utils::resolve_base_request_id; + +/// Match Python `tokenize-{base}` where base is `X-Request-Id` or a new UUID. +fn tokenize_request_id(headers: &HeaderMap) -> String { + let base = resolve_base_request_id( + headers.get("X-Request-Id").and_then(|value| value.to_str().ok()), + None, + ); + format!("tokenize-{base}") +} + +/// Reject an unknown model name, matching the other handlers. +fn check_model(state: &AppState, model: Option<&str>) -> Result<(), ApiError> { + if let Some(model) = model + && !state.served_model_names().iter().any(|n| n == model) + { + return Err(ApiError::model_not_found(model.to_string())); + } + Ok(()) +} + +/// Build the `token_strs` vector when requested, via the tokenizer vocab. +fn token_strs(tokenizer: &vllm_text::tokenizer::DynTokenizer, ids: &[u32]) -> Vec { + // Unknown IDs yield "" — intentional; matches Python's convert_ids_to_tokens behaviour. + ids.iter().map(|&id| tokenizer.id_to_token(id).unwrap_or_default()).collect() +} + +pub async fn tokenize( + State(state): State>, + headers: HeaderMap, + ValidatedJson(body): ValidatedJson, +) -> Response { + let request_id = tokenize_request_id(&headers); + let tokenizer = state.chat.text().tokenizer(); + let max_model_len = state.chat.engine_core_client().max_model_len(); + + let result = match body { + // Completion form: encode the raw `prompt` string (no chat template). + TokenizeRequest::Completion(req) => tokenize_completion(&state, &tokenizer, req), + // Chat form: render `messages` through the template, then encode (see `tokenize_chat`). + TokenizeRequest::Chat(req) => tokenize_chat(&state, &request_id, req).await, + }; + + match result { + Ok((tokens, want_strs)) => { + let token_strs = want_strs.then(|| token_strs(&tokenizer, &tokens)); + Json(TokenizeResponse { + count: tokens.len(), + max_model_len, + tokens, + token_strs, + }) + .into_response() + } + Err(error) => error.into_response(), + } +} + +fn tokenize_completion( + state: &AppState, + tokenizer: &vllm_text::tokenizer::DynTokenizer, + req: TokenizeCompletionRequest, +) -> Result<(Vec, bool), ApiError> { + check_model(state, req.model.as_deref())?; + let tokens = tokenizer + .encode(&req.prompt, req.add_special_tokens) + .map_err(|e| server_error!("tokenize failed: {}", e.to_report_string()))?; + Ok((tokens, req.return_token_strs)) +} + +/// HTTP adapter for the chat-shaped `/tokenize` body. +/// +/// Not [`vllm_chat::ChatLlm::tokenize_chat`]: this checks the model name and maps +/// errors to [`ApiError`]; the chat-crate method does render → finalize → encode. +async fn tokenize_chat( + state: &AppState, + request_id: &str, + req: TokenizeChatRequest, +) -> Result<(Vec, bool), ApiError> { + check_model(state, req.model.as_deref())?; + let return_token_strs = req.return_token_strs; + // `continue_final_message` / `add_generation_prompt` mutual exclusion is + // enforced in `normalize_generation_prompt_mode` inside `into_chat_request`. + let tokens = state + .chat + .tokenize_chat(req.into_chat_request(request_id.to_string())?) + .await + .map_err(|e| server_error!("tokenize failed: {}", e.to_report_string()))?; + Ok((tokens, return_token_strs)) +} + +pub async fn detokenize( + State(state): State>, + ValidatedJson(body): ValidatedJson, +) -> Response { + if let Err(error) = check_model(&state, body.model.as_deref()) { + return error.into_response(); + } + let tokenizer = state.chat.text().tokenizer(); + match tokenizer.decode(&body.tokens, /* skip_special_tokens = */ false) { + Ok(prompt) => Json(DetokenizeResponse { prompt }).into_response(), + Err(e) => server_error!("detokenize failed: {}", e.to_report_string()).into_response(), + } +} + +#[cfg(test)] +mod tests { + use axum::http::{HeaderMap, HeaderValue}; + + use super::tokenize_request_id; + + #[test] + fn tokenize_request_id_prefers_x_request_id_header() { + let mut headers = HeaderMap::new(); + headers.insert("X-Request-Id", HeaderValue::from_static("client-req-1")); + assert_eq!(tokenize_request_id(&headers), "tokenize-client-req-1"); + } + + #[test] + fn tokenize_request_id_generates_uuid_when_header_missing() { + let headers = HeaderMap::new(); + let id = tokenize_request_id(&headers); + assert!(id.starts_with("tokenize-")); + assert_ne!(id, "tokenize-"); + } +} diff --git a/rust/src/server/src/routes/tokenize/types.rs b/rust/src/server/src/routes/tokenize/types.rs new file mode 100644 index 0000000000..8067c09f66 --- /dev/null +++ b/rust/src/server/src/routes/tokenize/types.rs @@ -0,0 +1,209 @@ +use std::collections::HashMap; + +use itertools::Itertools as _; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use validator::{Validate, ValidationErrors}; +use vllm_chat::{ChatOptions, ChatRequest, ChatToolChoice, SamplingParams}; +use vllm_text::output::TextDecodeOptions; + +use crate::error::ApiError; +use crate::routes::openai::chat_completions::convert::{ + convert_message, convert_tools, normalize_generation_prompt_mode, +}; +use crate::routes::openai::utils::types::{ + ChatMessage, Normalizable, Tool, default_true, validate_messages, +}; + +/// `POST /tokenize` body. Untagged: a JSON object with `messages` parses as the +/// chat variant; one with `prompt` parses as the completion variant. +#[derive(Debug, Clone, Deserialize)] +#[serde(untagged)] +pub enum TokenizeRequest { + Chat(TokenizeChatRequest), + Completion(TokenizeCompletionRequest), +} + +#[derive(Debug, Clone, Deserialize)] +pub struct TokenizeCompletionRequest { + pub model: Option, + pub prompt: String, + #[serde(default = "default_true")] + pub add_special_tokens: bool, + #[serde(default)] + pub return_token_strs: bool, +} + +#[derive(Debug, Clone, Deserialize, Validate)] +pub struct TokenizeChatRequest { + pub model: Option, + #[validate(custom(function = "validate_messages"))] + pub messages: Vec, + #[serde(default = "default_true")] + pub add_generation_prompt: bool, + #[serde(default)] + pub continue_final_message: bool, + #[serde(default)] // chat default is FALSE (template adds specials) + pub add_special_tokens: bool, + #[serde(default)] + pub return_token_strs: bool, + #[serde(default)] + pub chat_template: Option, + #[serde(default)] + pub chat_template_kwargs: Option>, + #[serde(default)] + pub tools: Option>, +} + +impl TokenizeChatRequest { + /// Lower this tokenize body into a [`ChatRequest`] for template rendering. + /// + /// Reuses [`convert_message`] and [`normalize_generation_prompt_mode`] from + /// `chat_completions/convert` so message lowering and generation-prompt + /// rules match chat completions. Only fields that affect rendering are set; + /// `sampling_params`, `decode_options`, etc. stay at default because + /// tokenize never generates. + pub fn into_chat_request(self, request_id: String) -> Result { + let messages: Vec<_> = self.messages.into_iter().map(convert_message).try_collect()?; + let generation_prompt_mode = normalize_generation_prompt_mode( + Some(self.add_generation_prompt), + self.continue_final_message, + &messages, + )?; + + Ok(ChatRequest { + request_id, + messages, + sampling_params: SamplingParams::default(), + chat_options: ChatOptions { + generation_prompt_mode, + chat_template: self.chat_template, + reasoning_effort: None, + template_kwargs: self.chat_template_kwargs.unwrap_or_default(), + }, + tools: convert_tools(self.tools)?, + tool_choice: ChatToolChoice::Auto, + decode_options: TextDecodeOptions::default(), + intermediate: false, + priority: 0, + documents: None, + cache_salt: None, + add_special_tokens: self.add_special_tokens, + data_parallel_rank: None, + lora_request: None, + }) + } +} + +#[derive(Debug, Clone, Deserialize)] +pub struct DetokenizeRequest { + pub model: Option, + pub tokens: Vec, +} + +#[derive(Debug, Clone, Serialize)] +pub struct TokenizeResponse { + pub count: usize, + pub max_model_len: u32, + pub tokens: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub token_strs: Option>, +} + +#[derive(Debug, Clone, Serialize)] +pub struct DetokenizeResponse { + pub prompt: String, +} + +// ---- trait impls required by ValidatedJson ---- +impl Validate for TokenizeRequest { + fn validate(&self) -> Result<(), ValidationErrors> { + if let Self::Chat(req) = self { + req.validate()?; + } + Ok(()) + } +} +impl Validate for DetokenizeRequest { + fn validate(&self) -> Result<(), ValidationErrors> { + Ok(()) + } +} +impl Normalizable for TokenizeRequest {} // default no-op normalize() +impl Normalizable for DetokenizeRequest {} + +#[cfg(test)] +mod tests { + use super::*; + use crate::routes::openai::utils::types::{ChatMessage, MessageContent}; + use serde_json::json; + use vllm_chat::ChatTool; + + #[test] + fn tokenize_request_converts_openai_tools() { + // The untagged `TokenizeRequest` must resolve a messages+tools body to + // the chat variant and accept standard OpenAI tool objects + // (`{"type":"function",...}`), then convert them to `ChatTool`. + let request: TokenizeRequest = serde_json::from_value(json!({ + "messages": [{"role": "user", "content": "hi"}], + "tools": [{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + }], + })) + .expect("OpenAI tool JSON deserializes to the chat variant"); + + let TokenizeRequest::Chat(req) = request else { + panic!("messages+tools body should parse as the chat variant"); + }; + + let chat_request = + req.into_chat_request("tokenize-test".to_string()).expect("request is valid"); + + assert_eq!( + chat_request.tools, + vec![ChatTool { + name: "get_weather".to_string(), + description: Some("Get weather".to_string()), + parameters: json!({ + "type": "object", + "properties": {"city": {"type": "string"}}, + }), + strict: None, + }] + ); + } + + #[test] + fn into_chat_request_rejects_conflicting_generation_flags() { + let req = TokenizeChatRequest { + model: None, + messages: vec![ChatMessage::User { + content: MessageContent::Text("hi".to_string()), + name: None, + }], + add_generation_prompt: true, + continue_final_message: true, + add_special_tokens: false, + return_token_strs: false, + chat_template: None, + chat_template_kwargs: None, + tools: None, + }; + + let error = req + .into_chat_request("tokenize-test".to_string()) + .expect_err("conflicting flags"); + assert_eq!( + error.to_error_response().error.message, + "Cannot set both `continue_final_message` and `add_generation_prompt` to True." + ); + } +}