[Rust Frontend] Add /tokenize and /detokenize endpoints (#44222)
Signed-off-by: Tan Ngoc Do <darkknightkhtn2008@gmail.com> Signed-off-by: TanNgocDo <darkknightkhtn2008@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Bugen Zhao <i@bugenzhao.com>
This commit is contained in:
TanNgocDo
co-authored by
Claude Opus 4.8
Bugen Zhao
Claude Opus 4.8
Bugen Zhao
parent
80e2c4462d
commit
69fdaffbcd
@@ -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<T> = std::result::Result<T, Error>;
|
||||
|
||||
@@ -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<Vec<u32>> {
|
||||
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?;
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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<bool>,
|
||||
continue_final_message: bool,
|
||||
messages: &[VllmChatMessage],
|
||||
@@ -204,7 +203,7 @@ fn extract_last_assistant_content(messages: &[ChatMessage]) -> Option<String> {
|
||||
}
|
||||
|
||||
/// Lower one OpenAI chat message into the `vllm-chat` message shape.
|
||||
fn convert_message(message: ChatMessage) -> Result<VllmChatMessage, ApiError> {
|
||||
pub(crate) fn convert_message(message: ChatMessage) -> Result<VllmChatMessage, ApiError> {
|
||||
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<Vec<Tool>>) -> Result<Vec<ChatTool>, ApiError> {
|
||||
pub(crate) fn convert_tools(tools: Option<Vec<Tool>>) -> Result<Vec<ChatTool>, ApiError> {
|
||||
tools
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -412,6 +412,41 @@ pub struct ListModelsResponse {
|
||||
pub data: Vec<ModelObject>,
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 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
|
||||
// ============================================================================
|
||||
|
||||
@@ -407,6 +407,9 @@ struct FakeChatBackend {
|
||||
multimodal_model_info: Option<vllm_chat::multimodal::MultimodalModelInfo>,
|
||||
}
|
||||
|
||||
/// 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<Vec<u32>> {
|
||||
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("<image>") {
|
||||
@@ -460,6 +466,7 @@ impl Tokenizer for FakeChatTokenizer {
|
||||
|
||||
fn id_to_token(&self, id: u32) -> Option<String> {
|
||||
match id {
|
||||
FAKE_BOS_TOKEN_ID => Some("<bos>".to_string()),
|
||||
999 => Some("<image>".to_string()),
|
||||
151655 => Some("<|image_pad|>".to_string()),
|
||||
0xF001 => Some("<think>".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::<Vec<_>>();
|
||||
|
||||
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<u32> = with_special["tokens"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|v| v.as_u64().unwrap() as u32)
|
||||
.collect();
|
||||
let without_ids: Vec<u32> = 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);
|
||||
}
|
||||
|
||||
@@ -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<String> {
|
||||
// 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<Arc<AppState>>,
|
||||
headers: HeaderMap,
|
||||
ValidatedJson(body): ValidatedJson<TokenizeRequest>,
|
||||
) -> 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<u32>, 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<u32>, 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<Arc<AppState>>,
|
||||
ValidatedJson(body): ValidatedJson<DetokenizeRequest>,
|
||||
) -> 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-");
|
||||
}
|
||||
}
|
||||
@@ -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<String>,
|
||||
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<String>,
|
||||
#[validate(custom(function = "validate_messages"))]
|
||||
pub messages: Vec<ChatMessage>,
|
||||
#[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<String>,
|
||||
#[serde(default)]
|
||||
pub chat_template_kwargs: Option<HashMap<String, Value>>,
|
||||
#[serde(default)]
|
||||
pub tools: Option<Vec<Tool>>,
|
||||
}
|
||||
|
||||
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<ChatRequest, ApiError> {
|
||||
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<String>,
|
||||
pub tokens: Vec<u32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct TokenizeResponse {
|
||||
pub count: usize,
|
||||
pub max_model_len: u32,
|
||||
pub tokens: Vec<u32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub token_strs: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[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."
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user