From 2cae98dfa59bc58d939f526465ff025670d002f0 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 8 Jul 2026 19:57:04 +0800 Subject: [PATCH] [Rust Frontend] Handle `continue_final_message` with renderer sentinel (#47844) Signed-off-by: Bugen Zhao --- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/rust_frontend.yaml | 2 +- rust/src/chat/src/renderer/hf/mod.rs | 213 +++++++++++++++++++++-- 3 files changed, 200 insertions(+), 17 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index a0846135ea..0b72c852e7 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2026,7 +2026,7 @@ steps: - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load" - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" - - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info and not test_tokenize_chat" + - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info" - label: Rust Frontend Core Correctness # TBD timeout_in_minutes: 180 diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index 6b2aa4cb0b..1dfe912aa8 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -64,7 +64,7 @@ steps: - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" # - pytest -v -s entrypoints/serve/dev/test_sleep.py # /tokenizer_info is not implemented in the Rust frontend (the CLI flag is accepted as a no-op). - - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info and not test_tokenize_chat" + - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info" - label: Rust Frontend Core Correctness timeout_in_minutes: 30 diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index 3ad6a5c7c7..d9031d73a4 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -155,11 +155,23 @@ impl HfChatRenderer { effective_template: &CompiledChatTemplate, request: &ChatRequest, ) -> Result { - let messages = to_template_messages( + let mut messages = to_template_messages( &request.messages, effective_template.content_format(), self.multimodal.as_ref(), )?; + + // Handling of `continue_final_message`: + // Append a sentinel tag to the final message content, render as usual, then + // truncate the rendered prompt at the tag so any template suffix after the + // final message content (e.g. the end-of-turn marker) is dropped. + let final_message_text = if request.chat_options.continue_final_message() { + let final_message = messages.last_mut().ok_or(Error::EmptyMessages)?; + Some(append_continue_final_message_tag(final_message)?) + } else { + None + }; + let tools = request.tool_parsing_enabled().then(|| to_template_tools(&request.tools)); trace!( message_count = messages.len(), @@ -183,6 +195,13 @@ impl HfChatRenderer { }) .map_err(|error| Error::ChatTemplate(error.to_report_string()))?; + let prompt = match &final_message_text { + Some(final_message_text) => { + truncate_prompt_at_continue_final_message_tag(prompt, final_message_text)? + } + None => prompt, + }; + trace!( prompt_len = prompt.len(), prompt, "rendered chat template prompt" @@ -429,6 +448,74 @@ fn to_template_string_content( } } +/// Sentinel appended to the final message content when `continue_final_message` +/// is requested, used to locate the truncation point in the rendered prompt. +/// +/// Same literal as `transformers`. Occurrences of this string earlier in the +/// prompt are harmless because truncation uses the rightmost match, and the +/// appended sentinel ends up last as long as the template renders messages in +/// order. +const CONTINUE_FINAL_MESSAGE_TAG: &str = "CONTINUE_FINAL_MESSAGE_TAG "; + +/// Append [`CONTINUE_FINAL_MESSAGE_TAG`] to the trailing text of the final +/// message, returning the original text for post-render validation. +// TODO: transformers v5 also allows continuing a non-`content` field (e.g. +// `reasoning_content`) by passing a field name; only the boolean form is +// supported here. +fn append_continue_final_message_tag(message: &mut TemplateMessage) -> Result { + let text = match &mut message.content { + TemplateContent::String(text) => Some(text), + // Pick the last text part in the message. + TemplateContent::OpenAi(parts) => parts.iter_mut().rev().find_map(|part| match part { + TemplateContentPart::Text { text } => Some(text), + TemplateContentPart::Image => None, + }), + }; + let text = text.ok_or_else(|| { + Error::ChatTemplate( + "continue_final_message is set but there is no text to continue \ + in the final message" + .to_string(), + ) + })?; + + let original = text.clone(); + text.push_str(CONTINUE_FINAL_MESSAGE_TAG); + Ok(original) +} + +/// Truncate the rendered prompt at [`CONTINUE_FINAL_MESSAGE_TAG`] so that it +/// ends exactly with the final message content, dropping any template suffix +/// such as end-of-turn markers. +fn truncate_prompt_at_continue_final_message_tag( + mut rendered: String, + final_message_text: &str, +) -> Result { + let tag_loc = rendered + .rfind(CONTINUE_FINAL_MESSAGE_TAG.trim_end()) + .filter(|_| rendered.contains(final_message_text.trim())); + let Some(tag_loc) = tag_loc else { + return Err(Error::ChatTemplate(format!( + "continue_final_message is set but the final message does not appear \ + in the prompt after applying the chat template! This can happen if \ + the chat template deletes portions of the final message. Final \ + message to continue: {}", + final_message_text.trim(), + ))); + }; + + if rendered[tag_loc..].starts_with(CONTINUE_FINAL_MESSAGE_TAG) { + // The template preserved spacing, so a plain cut at the tag suffices. + rendered.truncate(tag_loc); + } else { + // The template trimmed the trailing spacing of the message content, so + // apply the same trimming to the retained prefix. + rendered.truncate(tag_loc); + rendered.truncate(rendered.trim_end().len()); + } + Ok(rendered) +} + fn to_template_tools(tools: &[ChatTool]) -> Vec { tools .iter() @@ -548,28 +635,124 @@ mod tests { ChatRole::Assistant, "The capital of", )]); + let template = + "{% if continue_final_message %}continue:{% endif %}{{ messages[0].content }}"; - assert_eq!( - render( - Some("{% if continue_final_message %}continue{% else %}new{% endif %}"), - &request, - ) - .unwrap(), - "new" - ); + assert_eq!(render(Some(template), &request).unwrap(), "The capital of"); request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant; assert_eq!( - render( - Some("{% if continue_final_message %}continue{% else %}new{% endif %}"), - &request, - ) - .unwrap(), - "continue" + render(Some(template), &request).unwrap(), + "continue:The capital of" ); } + #[test] + fn continue_final_message_truncates_template_suffix() { + let mut request = sample_request(vec![ + ChatMessage::text(ChatRole::User, "What is the capital of France?"), + ChatMessage::text(ChatRole::Assistant, "The capital of"), + ]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant; + + // The Qwen3 template is unaware of `continue_final_message`; the + // end-of-turn marker it appends must still be stripped. + let rendered = render(Some(QWEN3_0_6B_TEMPLATE), &request).unwrap(); + + expect![[r#" + <|im_start|>user + What is the capital of France?<|im_end|> + <|im_start|>assistant + + + + + The capital of"#]] + .assert_eq(&rendered); + } + + #[test] + fn continue_final_message_trims_like_the_template_does() { + let mut request = sample_request(vec![ChatMessage::text(ChatRole::Assistant, "Sure, ")]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant; + + // The template trims the trailing spacing of the message content, so + // the truncated prompt must be trimmed the same way. + let rendered = render( + Some("{{ messages[0].content.strip() }}<|im_end|>"), + &request, + ) + .unwrap(); + + assert_eq!(rendered, "Sure,"); + } + + #[test] + fn continue_final_message_appends_to_last_text_part() { + // The renderer itself is role-agnostic like transformers (the + // assistant-final restriction is enforced by request validation + // upstream), so a multimodal user message exercises the part + // selection: the sentinel must attach to the last *text* part, + // skipping the trailing image. + let mut request = sample_request(vec![ChatMessage::user(vec![ + ChatContentPart::text("Sure,"), + ChatContentPart::image_url("data:image/png;base64,test"), + ])]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant; + + let rendered = render_mm( + "{% for item in messages[0].content %}{% if item.type == 'image' %}{% else %}{{ item.text }}{% endif %}{% endfor %}<|im_end|>", + &request, + ChatTemplateContentFormatOption::OpenAi, + ) + .unwrap() + .prompt; + + // Anything rendered after the continued text (here the image + // placeholder and the end marker) is truncated away, matching + // transformers. + assert_eq!(rendered, Prompt::Text("Sure,".to_string())); + } + + #[test] + fn continue_final_message_composes_with_aware_templates() { + // A template that reads `continue_final_message` and skips its own + // end-of-turn marker must produce the same prompt as an unaware one: + // the sentinel truncation degenerates to a cut at the very end. + let mut request = sample_request(vec![ + ChatMessage::text(ChatRole::User, "hi"), + ChatMessage::text(ChatRole::Assistant, "Sure,"), + ]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant; + + let aware = "{% for m in messages %}<|im_start|>{{ m.role }}\n{{ m.content }}{% if not (loop.last and continue_final_message) %}<|im_end|>\n{% endif %}{% endfor %}"; + let unaware = "{% for m in messages %}<|im_start|>{{ m.role }}\n{{ m.content }}<|im_end|>\n{% endfor %}"; + + let expected = "<|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\nSure,"; + assert_eq!(render(Some(aware), &request).unwrap(), expected); + assert_eq!(render(Some(unaware), &request).unwrap(), expected); + } + + #[test] + fn continue_final_message_errors_when_template_drops_final_message() { + let mut request = sample_request(vec![ + ChatMessage::text(ChatRole::User, "hi"), + ChatMessage::text(ChatRole::Assistant, "Sure,"), + ]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant; + + let error = render( + Some( + "{% for m in messages %}{% if m.role == 'user' %}{{ m.content }}{% endif %}{% endfor %}", + ), + &request, + ) + .unwrap_err(); + + assert!(matches!(error, Error::ChatTemplate(_))); + } + #[test] fn chat_template_flattens_text_parts_for_string_templates() { let request = sample_request(vec![ChatMessage::user(vec![