diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 1b5d5aea6f..9c59d14028 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -34,8 +34,8 @@ pub use parser::{ParserSelection, validate_parser_overrides}; pub use renderer::hf::ChatTemplateContentFormatOption; pub use renderer::{ ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DeepSeekV41ChatRenderer, - DynChatRenderer, HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, RenderedPrompt, - RendererSelection, + DynChatRenderer, HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, MediaPartSource, + RenderedPrompt, RendererSelection, }; pub use request::{ ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool, @@ -112,6 +112,10 @@ impl ChatRequestProcessor { request: &ChatRequest, rendered: RenderedPrompt, ) -> Result<(Prompt, Option)> { + let media_is_empty = match &rendered.media_order { + Some(media_order) => media_order.is_empty(), + None => !request.has_multimodal(), + }; match self.model_dtype { Some(model_dtype) => { multimodal::finalize_rendered_prompt( @@ -122,7 +126,7 @@ impl ChatRequestProcessor { ) .await } - None if !request.has_multimodal() => Ok((rendered.prompt, None)), + None if media_is_empty => Ok((rendered.prompt, None)), None => Err(Error::UnsupportedMultimodalRenderer), } } diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index 674979e6ce..a5dd3ba8ea 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -16,7 +16,7 @@ use std::fs; use std::path::Path; use std::sync::{Arc, LazyLock}; -use itertools::izip; +use itertools::{Either, izip}; use llm_multimodal::{ AsyncMultiModalTracker, AudioClip, AudioPreProcessor, EncoderFieldLayouts, FieldLayout, ImageFrame, MediaConnector, MediaConnectorConfig, MediaContentPart, Modality, ModelMetadata, @@ -33,7 +33,7 @@ use vllm_text::Prompt; use vllm_text::tokenizer::{DynTokenizer, Tokenizer}; use crate::error::{Error, Result, bail_multimodal, multimodal}; -use crate::renderer::RenderedPrompt; +use crate::renderer::{MediaPartSource, RenderedPrompt}; use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest}; mod audio; @@ -549,7 +549,8 @@ pub(crate) async fn finalize_rendered_prompt( info: Option<&MultimodalModelInfo>, model_dtype: ModelDtype, ) -> Result<(Prompt, Option)> { - if !request.has_multimodal() { + let media_parts = extract_media_parts(request, rendered.media_order.as_deref())?; + if media_parts.is_empty() { return Ok((rendered.prompt, None)); } let info = info.ok_or(Error::UnsupportedMultimodalRenderer)?; @@ -561,63 +562,92 @@ pub(crate) async fn finalize_rendered_prompt( .map_err(|error| multimodal!("{error}"))?, Prompt::TokenIds(token_ids) => token_ids, }; - let media_parts = extract_media_parts(request)?; let prepared = info.prepare_multimodal(media_parts, &mut prompt_token_ids, model_dtype).await?; Ok((Prompt::TokenIds(prompt_token_ids), Some(prepared))) } -/// Extract media parts from chat messages in message/content order. -/// -/// Assistant history is skipped because generated assistant blocks are already -/// represented as text for prompt rendering in this crate. -fn extract_media_parts(request: &ChatRequest) -> Result> { - let mut all_parts = Vec::new(); - for message in &request.messages { - let content = match message { - ChatMessage::System { content } - | ChatMessage::Developer { content, .. } - | ChatMessage::User { content } - | ChatMessage::ToolResponse { content, .. } => content, - ChatMessage::Assistant { .. } => continue, - }; - let ChatContent::Parts(parts) = content else { - continue; - }; - for part in parts { - match part { - ChatContentPart::Text { .. } => {} - ChatContentPart::ImageUrl { - image_url, - detail, - uuid, - } => all_parts.push(MediaContentPart::ImageUrl { - url: image_url.clone(), - detail: *detail, - uuid: uuid.clone(), - }), - ChatContentPart::VideoUrl { video_url, uuid } => { - all_parts.push(MediaContentPart::VideoUrl { - url: video_url.clone(), - uuid: uuid.clone(), - }) - } - ChatContentPart::InputAudio { data, format, uuid } => { - all_parts.push(MediaContentPart::AudioUrl { - url: input_audio_data_url(data, format.as_deref())?, - uuid: uuid.clone(), - }) - } - ChatContentPart::AudioUrl { audio_url, uuid } => { - all_parts.push(MediaContentPart::AudioUrl { - url: audio_url.clone(), - uuid: uuid.clone(), - }) +/// Resolve media parts in the placeholder order reported by the renderer. +fn extract_media_parts( + request: &ChatRequest, + media_order: Option<&[MediaPartSource]>, +) -> Result> { + let parts = match media_order { + // Resolve exactly the sources selected by the renderer, in placeholder order. + Some(media_order) => Either::Left(media_order.iter().map(|source| -> Result<_> { + let message = request.messages.get(source.message_index).ok_or_else(|| { + multimodal!( + "renderer reported missing chat message {} for multimodal content", + source.message_index + ) + })?; + let content = match message { + ChatMessage::System { content } + | ChatMessage::Developer { content, .. } + | ChatMessage::User { content } + | ChatMessage::ToolResponse { content, .. } => content, + ChatMessage::Assistant { .. } => { + bail_multimodal!("renderer reported multimodal assistant content") } + }; + let ChatContent::Parts(parts) = content else { + bail_multimodal!("renderer reported multimodal content in a text-only chat message") + }; + let part = parts.get(source.content_part_index).ok_or_else(|| { + multimodal!( + "renderer reported missing content part {} in chat message {}", + source.content_part_index, + source.message_index + ) + })?; + Ok(part) + })), + // Fall back to all media in request message and content-part order (e.g. Jinja). + None => Either::Right(request.messages.iter().flat_map(|message| { + let content = match message { + ChatMessage::System { content } + | ChatMessage::Developer { content, .. } + | ChatMessage::User { content } + | ChatMessage::ToolResponse { content, .. } => Some(content), + ChatMessage::Assistant { .. } => None, + }; + match content { + Some(ChatContent::Parts(parts)) => parts.as_slice(), + _ => &[], } - } - } - Ok(all_parts) + .iter() + .filter(|part| part.is_multimodal()) + .map(Ok) + })), + }; + parts + .map(|part| match part? { + ChatContentPart::Text { .. } => { + bail_multimodal!("renderer reported a text part as multimodal content") + } + ChatContentPart::ImageUrl { + image_url, + detail, + uuid, + } => Ok(MediaContentPart::ImageUrl { + url: image_url.clone(), + detail: *detail, + uuid: uuid.clone(), + }), + ChatContentPart::VideoUrl { video_url, uuid } => Ok(MediaContentPart::VideoUrl { + url: video_url.clone(), + uuid: uuid.clone(), + }), + ChatContentPart::InputAudio { data, format, uuid } => Ok(MediaContentPart::AudioUrl { + url: input_audio_data_url(data, format.as_deref())?, + uuid: uuid.clone(), + }), + ChatContentPart::AudioUrl { audio_url, uuid } => Ok(MediaContentPart::AudioUrl { + url: audio_url.clone(), + uuid: uuid.clone(), + }), + }) + .collect() } /// Wrap OpenAI base64 audio in a data URL consumed by `MediaConnector`. @@ -1043,6 +1073,75 @@ mod tests { assert!(input_audio_data_url("AAAA", Some("flac")).is_err()); } + #[test] + fn extract_media_parts_uses_request_order_or_explicit_selection() { + let request = ChatRequest { + messages: vec![ + ChatMessage::user(vec![ + ChatContentPart::text("before image"), + ChatContentPart::ImageUrl { + image_url: "data:image/png;base64,image-b".to_string(), + detail: None, + uuid: Some("image-b".to_string()), + }, + ]), + ChatMessage::user(vec![ChatContentPart::ImageUrl { + image_url: "data:image/png;base64,image-a".to_string(), + detail: None, + uuid: Some("image-a".to_string()), + }]), + ChatMessage::assistant_text("answer"), + ChatMessage::user("follow-up"), + ], + ..ChatRequest::for_test() + }; + let media_order = [ + MediaPartSource { + message_index: 1, + content_part_index: 0, + }, + MediaPartSource { + message_index: 0, + content_part_index: 1, + }, + ]; + + let selections = [None, Some(&media_order[..]), Some(&[][..])]; + let uuids = selections.map(|order| { + extract_media_parts(&request, order) + .unwrap() + .into_iter() + .map(|part| match part { + MediaContentPart::ImageUrl { uuid, .. } => uuid, + _ => panic!("expected image URL"), + }) + .collect::>() + }); + + expect_test::expect![[r#" + [ + [ + Some( + "image-b", + ), + Some( + "image-a", + ), + ], + [ + Some( + "image-a", + ), + Some( + "image-b", + ), + ], + [], + ] + "#]] + .assert_debug_eq(&uuids); + } + fn image_url_part() -> MediaContentPart { MediaContentPart::ImageUrl { url: "https://example.com/image.png".to_string(), diff --git a/rust/src/chat/src/renderer/deepseek.rs b/rust/src/chat/src/renderer/deepseek.rs index 5f900ab88b..5925c5fd16 100644 --- a/rust/src/chat/src/renderer/deepseek.rs +++ b/rust/src/chat/src/renderer/deepseek.rs @@ -16,6 +16,7 @@ use serde_json_fmt::JsonFormat; use llm_multimodal::DEEPSEEK_V41_IMAGE_PLACEHOLDER; +use super::MediaPartSource; use crate::error::{Error, Result}; use crate::request::{ ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ReasoningEffort, @@ -87,6 +88,14 @@ struct RenderedToolSchema<'a> { /// Render one chat request into the final prompt string. pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Result { + render_request_with_media_order(request, dialect).map(|(prompt, _)| prompt) +} + +/// Render one chat request and record media in placeholder order. +pub(super) fn render_request_with_media_order( + request: &ChatRequest, + dialect: DsDialect, +) -> Result<(String, Vec)> { let (thinking_mode, reasoning_effort_prompt) = match dialect { DsDialect::V4 => { resolve_thinking_options(request).map(|(mode, prompt)| (mode, Cow::Borrowed(prompt)))? @@ -112,6 +121,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul dialect, ); let mut out = String::from(BOS_TOKEN); + let mut media_order = Vec::new(); if dialect == DsDialect::V41 && (thinking_mode == ThinkingMode::Thinking || synthetic_tool_system @@ -126,7 +136,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul let mut request_tools_attached = false; let mut render_index = 0isize; if synthetic_tool_system { - render_system_message(&mut out, None, request_tools, dialect)?; + render_system_message(&mut out, None, request_tools, dialect, &mut media_order)?; request_tools_attached = true; render_index += 1; } @@ -157,14 +167,22 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul } else { &[] }; - render_system_message(&mut out, Some(content), tools, dialect)?; + render_system_message( + &mut out, + Some((message_index, content)), + tools, + dialect, + &mut media_order, + )?; } ChatMessage::Developer { content, tools } => { render_developer_message( &mut out, + message_index, content, tools.as_deref().unwrap_or(&[]), dialect, + &mut media_order, )?; } ChatMessage::User { .. } | ChatMessage::ToolResponse { .. } => { @@ -174,6 +192,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul message_index, &last_tool_call_order, dialect, + &mut media_order, )?; } ChatMessage::Assistant { content } => { @@ -223,7 +242,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul } } - Ok(out) + Ok((out, media_order)) } /// Resolve DeepSeek V4's thinking controls. Unlike the Python tokenizer @@ -488,12 +507,13 @@ fn render_tool_schema(out: &mut String, tool: &ChatTool) -> Result<()> { /// Render a system turn, optionally followed by the V4 tool preamble. fn render_system_message( out: &mut String, - content: Option<&ChatContent>, + content: Option<(usize, &ChatContent)>, tools: &[ChatTool], dialect: DsDialect, + media_order: &mut Vec, ) -> Result<()> { - if let Some(content) = content { - write_chat_content(out, content, dialect)?; + if let Some((message_index, content)) = content { + write_chat_content(out, content, dialect, message_index, media_order)?; } if !tools.is_empty() { out.push_str("\n\n"); @@ -505,9 +525,11 @@ fn render_system_message( /// Developer messages are rendered as user-like turns with optional tools. fn render_developer_message( out: &mut String, + message_index: usize, content: &ChatContent, tools: &[ChatTool], dialect: DsDialect, + media_order: &mut Vec, ) -> Result<()> { if content.is_empty() { return Err(Error::ChatTemplate( @@ -516,7 +538,7 @@ fn render_developer_message( } out.push_str(USER_SP_TOKEN); - write_chat_content(out, content, dialect)?; + write_chat_content(out, content, dialect, message_index, media_order)?; if !tools.is_empty() { out.push_str("\n\n"); render_tools(out, tools, dialect)?; @@ -531,6 +553,7 @@ fn render_user_content_block( message_index: usize, tool_call_order: &HashMap, dialect: DsDialect, + media_order: &mut Vec, ) -> Result<()> { let (block_start, block_end) = user_content_block_bounds(messages, message_index); let mut sorted_tool_indices = @@ -542,7 +565,9 @@ fn render_user_content_block( out.push_str("\n\n"); } match &messages[message_index] { - ChatMessage::User { content } => write_chat_content(out, content, dialect)?, + ChatMessage::User { content } => { + write_chat_content(out, content, dialect, message_index, media_order)?; + } ChatMessage::ToolResponse { .. } => { let sorted_index = sorted_tool_indices .next() @@ -550,7 +575,7 @@ fn render_user_content_block( let ChatMessage::ToolResponse { content, .. } = &messages[sorted_index] else { unreachable!("sorted tool response index should reference a tool message"); }; - write_tool_result(out, content, dialect)?; + write_tool_result(out, content, dialect, sorted_index, media_order)?; } _ => unreachable!("user content block should only contain user content messages"), } @@ -597,9 +622,15 @@ fn sorted_tool_response_indices( } /// Render one tool response payload inside a V4 `` block. -fn write_tool_result(out: &mut String, content: &ChatContent, dialect: DsDialect) -> Result<()> { +fn write_tool_result( + out: &mut String, + content: &ChatContent, + dialect: DsDialect, + message_index: usize, + media_order: &mut Vec, +) -> Result<()> { out.push_str(""); - write_chat_content(out, content, dialect)?; + write_chat_content(out, content, dialect, message_index, media_order)?; out.push_str(""); Ok(()) } @@ -727,7 +758,13 @@ fn encode_arguments_to_dsml( /// The V4.1 dialect inlines the image placeholder at each image part's /// position unconditionally (matching the Python encoding's /// `IMAGE_PLACEHOLDER`); other dialects reject multimodal parts. -fn write_chat_content(out: &mut String, content: &ChatContent, dialect: DsDialect) -> Result<()> { +fn write_chat_content( + out: &mut String, + content: &ChatContent, + dialect: DsDialect, + message_index: usize, + media_order: &mut Vec, +) -> Result<()> { match content { ChatContent::Text(text) => out.push_str(text), ChatContent::Parts(parts) => { @@ -738,6 +775,10 @@ fn write_chat_content(out: &mut String, content: &ChatContent, dialect: DsDialec match part { ChatContentPart::ImageUrl { .. } if dialect == DsDialect::V41 => { out.push_str(DEEPSEEK_V41_IMAGE_PLACEHOLDER); + media_order.push(MediaPartSource { + message_index, + content_part_index: index, + }); } _ => out.push_str(part.as_text()?), } diff --git a/rust/src/chat/src/renderer/deepseek_v32/mod.rs b/rust/src/chat/src/renderer/deepseek_v32/mod.rs index 1a41e14f8f..03c2765d19 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/mod.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/mod.rs @@ -26,6 +26,7 @@ impl ChatRenderer for DeepSeekV32ChatRenderer { Ok(RenderedPrompt { prompt: Prompt::Text(encoding::render_request(request)?), + media_order: None, effective_template_kwargs: request_template_kwargs(request), }) } diff --git a/rust/src/chat/src/renderer/deepseek_v4/mod.rs b/rust/src/chat/src/renderer/deepseek_v4/mod.rs index 155a6ee140..54b111e505 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/mod.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/mod.rs @@ -24,6 +24,7 @@ impl ChatRenderer for DeepSeekV4ChatRenderer { Ok(RenderedPrompt { prompt: Prompt::Text(deepseek::render_request(request, DsDialect::V4)?), + media_order: None, effective_template_kwargs: request_template_kwargs(request), }) } diff --git a/rust/src/chat/src/renderer/deepseek_v41/mod.rs b/rust/src/chat/src/renderer/deepseek_v41/mod.rs index 4e5ddb3725..7737f7ebda 100644 --- a/rust/src/chat/src/renderer/deepseek_v41/mod.rs +++ b/rust/src/chat/src/renderer/deepseek_v41/mod.rs @@ -21,9 +21,12 @@ impl DeepSeekV41ChatRenderer { impl ChatRenderer for DeepSeekV41ChatRenderer { fn render(&self, request: &ChatRequest) -> Result { request.validate()?; + let (prompt, media_order) = + deepseek::render_request_with_media_order(request, DsDialect::V41)?; Ok(RenderedPrompt { - prompt: Prompt::Text(deepseek::render_request(request, DsDialect::V41)?), + prompt: Prompt::Text(prompt), + media_order: Some(media_order), effective_template_kwargs: request_template_kwargs(request), }) } diff --git a/rust/src/chat/src/renderer/deepseek_v41/tests.rs b/rust/src/chat/src/renderer/deepseek_v41/tests.rs index d7b0c538bb..3babb5c269 100644 --- a/rust/src/chat/src/renderer/deepseek_v41/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v41/tests.rs @@ -8,7 +8,7 @@ use serde_json::{Value, json}; use super::DeepSeekV41ChatRenderer; use crate::ChatRenderer; -use crate::event::AssistantContentBlock; +use crate::event::{AssistantContentBlock, AssistantToolCall}; use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ReasoningEffort}; @@ -83,6 +83,82 @@ fn inlines_image_placeholder_at_image_part_positions() { ); } +#[test] +fn media_order_follows_reordered_tool_results() { + let renderer = DeepSeekV41ChatRenderer::new(); + let request = ChatRequest { + messages: vec![ + ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::ToolCall(AssistantToolCall { + id: "call-a".to_string(), + name: "tool-a".to_string(), + arguments: "{}".to_string(), + }), + AssistantContentBlock::ToolCall(AssistantToolCall { + id: "call-b".to_string(), + name: "tool-b".to_string(), + arguments: "{}".to_string(), + }), + ]), + ChatMessage::tool_response( + vec![ChatContentPart::text("result-b"), image_part("image-b")], + "call-b", + ), + ChatMessage::tool_response( + vec![ChatContentPart::text("result-a"), image_part("image-a")], + "call-a", + ), + ], + ..ChatRequest::for_test() + }; + + let rendered = renderer.render(&request).unwrap(); + let media_order = rendered.media_order.unwrap(); + let prompt = rendered.prompt.into_text().unwrap(); + + assert!(prompt.find("result-a").unwrap() < prompt.find("result-b").unwrap()); + expect![[r#" + [ + MediaPartSource { + message_index: 2, + content_part_index: 1, + }, + MediaPartSource { + message_index: 1, + content_part_index: 1, + }, + ] + "#]] + .assert_debug_eq(&media_order); +} + +#[test] +fn media_order_excludes_dropped_historical_developer_content() { + let renderer = DeepSeekV41ChatRenderer::new(); + let request = ChatRequest { + messages: vec![ + ChatMessage::developer(vec![image_part("historical-image")], None), + ChatMessage::user("question"), + ], + ..ChatRequest::for_test() + }; + + let rendered = renderer.render(&request).unwrap(); + let media_order = rendered.media_order.unwrap(); + let prompt = rendered.prompt.into_text().unwrap(); + + assert!(!prompt.contains("<｜deepseek_image｜>")); + assert!(media_order.is_empty()); +} + +fn image_part(uuid: &str) -> ChatContentPart { + ChatContentPart::ImageUrl { + image_url: format!("data:image/png;base64,{uuid}"), + detail: None, + uuid: Some(uuid.to_string()), + } +} + #[test] fn maps_reasoning_effort_to_reference_numeric_budget() { for (effort, budget) in [ diff --git a/rust/src/chat/src/renderer/harmony/mod.rs b/rust/src/chat/src/renderer/harmony/mod.rs index b07d8458d7..71431cc92b 100644 --- a/rust/src/chat/src/renderer/harmony/mod.rs +++ b/rust/src/chat/src/renderer/harmony/mod.rs @@ -116,6 +116,7 @@ impl ChatRenderer for HarmonyChatRenderer { fn render(&self, request: &ChatRequest) -> Result { Ok(RenderedPrompt { prompt: Prompt::TokenIds(self.render_token_ids(request)?), + media_order: None, effective_template_kwargs: request_template_kwargs(request), }) } diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index 5e7ba75e0a..ba2b76d100 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -217,6 +217,7 @@ impl HfChatRenderer { Ok(RenderedPrompt { prompt: Prompt::Text(prompt), + media_order: None, effective_template_kwargs, }) } diff --git a/rust/src/chat/src/renderer/inkling/mod.rs b/rust/src/chat/src/renderer/inkling/mod.rs index a33e3414f2..e8afe66396 100644 --- a/rust/src/chat/src/renderer/inkling/mod.rs +++ b/rust/src/chat/src/renderer/inkling/mod.rs @@ -346,6 +346,7 @@ impl ChatRenderer for InklingChatRenderer { Ok(RenderedPrompt { prompt: Prompt::TokenIds(out), + media_order: None, effective_template_kwargs, }) } diff --git a/rust/src/chat/src/renderer/kimi_k3/encoding.rs b/rust/src/chat/src/renderer/kimi_k3/encoding.rs index 392b6b642d..88b4e3ae61 100644 --- a/rust/src/chat/src/renderer/kimi_k3/encoding.rs +++ b/rust/src/chat/src/renderer/kimi_k3/encoding.rs @@ -10,6 +10,7 @@ use std::collections::HashMap; use serde_json::{Map, Value, json}; use vllm_tokenizer::Tokenizer; +use super::super::MediaPartSource; use crate::error::{Error, Result}; use crate::request::{ ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, @@ -29,6 +30,7 @@ const VALID_THINKING_EFFORTS: &[&str] = &["low", "high", "max"]; pub(super) struct K3TokenWriter<'a> { tokenizer: &'a dyn Tokenizer, token_ids: Vec, + media_order: Vec, } impl<'a> K3TokenWriter<'a> { @@ -36,6 +38,7 @@ impl<'a> K3TokenWriter<'a> { Self { tokenizer, token_ids: Vec::new(), + media_order: Vec::new(), } } @@ -55,13 +58,25 @@ impl<'a> K3TokenWriter<'a> { Ok(()) } - pub(super) fn finish(self) -> Vec { - self.token_ids + fn image(&mut self, message_index: usize, content_part_index: usize) -> Result<()> { + self.control(IMAGE_PLACEHOLDER)?; + self.media_order.push(MediaPartSource { + message_index, + content_part_index, + }); + Ok(()) + } + + pub(super) fn finish(self) -> (Vec, Vec) { + (self.token_ids, self.media_order) } } -/// Render and tokenize one chat request using K3's segment-aware contract. -pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -> Result> { +/// Render and tokenize one chat request while recording media in placeholder order. +pub(super) fn render_request_with_media_order( + request: &ChatRequest, + tokenizer: &dyn Tokenizer, +) -> Result<(Vec, Vec)> { let thinking = thinking_enabled(request)?; let thinking_effort = thinking.then(|| thinking_effort(request)).transpose()?; let tools = request_tools(request); @@ -126,7 +141,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) - resolved.sort_by_key(|item| (item.0, item.1)); } - for (position, _, content, name) in resolved { + for (position, message_index, content, name) in resolved { let tool_name = name.as_deref().ok_or_else(|| { Error::ChatTemplate( "Kimi K3 tool messages need a resolvable tool name: \ @@ -135,7 +150,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) - .to_string(), ) })?; - write_tool_message(out, tool_name, position, &content)?; + write_tool_message(out, tool_name, position, message_index, &content)?; } Ok(()) }; @@ -161,14 +176,14 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) - } if !local_tools.is_empty() => { write_tool_declare(&mut out, local_tools, true)?; if !content_is_empty(content) { - write_role_message(&mut out, "system", None, content)?; + write_role_message(&mut out, "system", None, message_index, content)?; } } ChatMessage::System { content } | ChatMessage::Developer { content, .. } => { - write_role_message(&mut out, "system", None, content)?; + write_role_message(&mut out, "system", None, message_index, content)?; } ChatMessage::User { content } => { - write_role_message(&mut out, "user", None, content)?; + write_role_message(&mut out, "user", None, message_index, content)?; } ChatMessage::Assistant { content } => { tool_call_id_index.clear(); @@ -364,6 +379,7 @@ fn write_role_message( out: &mut K3TokenWriter<'_>, role: &str, name: Option<&str>, + message_index: usize, content: &ChatContent, ) -> Result<()> { let mut attrs = vec![("role", role.to_string())]; @@ -372,7 +388,7 @@ fn write_role_message( } let attr_refs: Vec<(&str, &str)> = attrs.iter().map(|(k, v)| (*k, v.as_str())).collect(); write_open_tag(out, "message", &attr_refs)?; - write_content(out, content)?; + write_content(out, message_index, content)?; write_close_tag(out, "message")?; out.control(END_OF_MSG) } @@ -381,6 +397,7 @@ fn write_tool_message( out: &mut K3TokenWriter<'_>, tool_name: &str, index: usize, + message_index: usize, content: &ChatContent, ) -> Result<()> { let index_str = index.to_string(); @@ -389,7 +406,7 @@ fn write_tool_message( "message", &[("role", "tool"), ("tool", tool_name), ("index", &index_str)], )?; - write_content(out, content)?; + write_content(out, message_index, content)?; write_close_tag(out, "message")?; out.control(END_OF_MSG) } @@ -471,14 +488,20 @@ fn write_assistant_tool_call( write_close_tag(out, "call") } -fn write_content(out: &mut K3TokenWriter<'_>, content: &ChatContent) -> Result<()> { +fn write_content( + out: &mut K3TokenWriter<'_>, + message_index: usize, + content: &ChatContent, +) -> Result<()> { match content { ChatContent::Text(text) => write_text_with_images(out, text), ChatContent::Parts(parts) => { - for part in parts { + for (content_part_index, part) in parts.iter().enumerate() { match part { ChatContentPart::Text { text } => write_text_with_images(out, text)?, - ChatContentPart::ImageUrl { .. } => out.control(IMAGE_PLACEHOLDER)?, + ChatContentPart::ImageUrl { .. } => { + out.image(message_index, content_part_index)?; + } ChatContentPart::VideoUrl { .. } => { return Err(Error::UnsupportedMultimodalContent("video_url")); } diff --git a/rust/src/chat/src/renderer/kimi_k3/mod.rs b/rust/src/chat/src/renderer/kimi_k3/mod.rs index 09bf356f5d..90750eeac4 100644 --- a/rust/src/chat/src/renderer/kimi_k3/mod.rs +++ b/rust/src/chat/src/renderer/kimi_k3/mod.rs @@ -30,9 +30,12 @@ impl KimiK3ChatRenderer { impl ChatRenderer for KimiK3ChatRenderer { fn render(&self, request: &ChatRequest) -> Result { request.validate()?; + let (token_ids, media_order) = + encoding::render_request_with_media_order(request, self.tokenizer.as_ref())?; Ok(RenderedPrompt { - prompt: Prompt::TokenIds(encoding::render_request(request, self.tokenizer.as_ref())?), + prompt: Prompt::TokenIds(token_ids), + media_order: Some(media_order), effective_template_kwargs: request_template_kwargs(request), }) } diff --git a/rust/src/chat/src/renderer/kimi_k3/tests.rs b/rust/src/chat/src/renderer/kimi_k3/tests.rs index 1d9f984dbb..d0ecbe44d3 100644 --- a/rust/src/chat/src/renderer/kimi_k3/tests.rs +++ b/rust/src/chat/src/renderer/kimi_k3/tests.rs @@ -284,6 +284,65 @@ fn partial_tool_results_keep_assistant_call_position() { ); } +#[test] +fn media_order_follows_reordered_tool_results() { + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ + ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::ToolCall(AssistantToolCall { + id: "call-a".to_string(), + name: "tool-a".to_string(), + arguments: "{}".to_string(), + }), + AssistantContentBlock::ToolCall(AssistantToolCall { + id: "call-b".to_string(), + name: "tool-b".to_string(), + arguments: "{}".to_string(), + }), + ]), + ChatMessage::tool_response( + vec![ChatContentPart::text("result-b"), image_part("image-b")], + "call-b", + ), + ChatMessage::tool_response( + vec![ChatContentPart::text("result-a"), image_part("image-a")], + "call-a", + ), + ]; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + let renderer = KimiK3ChatRenderer::new(Arc::new(test_tokenizer())); + + let rendered = renderer.render(&request).unwrap(); + let media_order = rendered.media_order.unwrap(); + let Prompt::TokenIds(token_ids) = rendered.prompt else { + panic!("kimi k3 renderer should return token IDs") + }; + let prompt = test_tokenizer().decode(&token_ids, false).unwrap(); + + assert!(prompt.find("result-a").unwrap() < prompt.find("result-b").unwrap()); + expect![[r#" + [ + MediaPartSource { + message_index: 2, + content_part_index: 1, + }, + MediaPartSource { + message_index: 1, + content_part_index: 1, + }, + ] + "#]] + .assert_debug_eq(&media_order); +} + +fn image_part(uuid: &str) -> ChatContentPart { + ChatContentPart::ImageUrl { + image_url: format!("data:image/png;base64,{uuid}"), + detail: None, + uuid: Some(uuid.to_string()), + } +} + #[test] fn defaults_thinking_effort_to_max() { let rendered = render_request(&crate::request::ChatRequest::for_test()); diff --git a/rust/src/chat/src/renderer/mod.rs b/rust/src/chat/src/renderer/mod.rs index a4649a34f7..bc1831c375 100644 --- a/rust/src/chat/src/renderer/mod.rs +++ b/rust/src/chat/src/renderer/mod.rs @@ -30,11 +30,24 @@ pub use inkling::InklingChatRenderer; pub use kimi_k3::KimiK3ChatRenderer; pub use selection::RendererSelection; +/// Location of one multimodal content part in the source chat request. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct MediaPartSource { + /// Zero-based message index in the source request. + pub message_index: usize, + /// Zero-based content-part index within the source message. + pub content_part_index: usize, +} + /// Rendered chat prompt submitted to the text backend. #[derive(Debug, Clone, PartialEq)] pub struct RenderedPrompt { /// The rendered prompt, either as text or already tokenized. pub prompt: Prompt, + /// Media sources in rendered placeholder order. + /// `None` uses the message and content-part order in the request (e.g. Jinja). + /// `Some` uses exactly the listed sources; an empty list omits all media. + pub media_order: Option>, /// Effective chat-template kwargs visible to the renderer after applying /// server defaults, request overrides, and typed reasoning controls. pub effective_template_kwargs: HashMap, diff --git a/rust/src/chat/tests/chat.rs b/rust/src/chat/tests/chat.rs index ab99b4d8ba..ed96134cdc 100644 --- a/rust/src/chat/tests/chat.rs +++ b/rust/src/chat/tests/chat.rs @@ -250,6 +250,7 @@ impl ChatRenderer for FakeChatBackend { Ok(RenderedPrompt { prompt: Prompt::Text(prompt), + media_order: None, effective_template_kwargs: request.chat_options.template_kwargs.clone(), }) } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index f13dc11b27..f5f74d557e 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -252,6 +252,7 @@ impl ChatRenderer for FakeTextBackend { fn render(&self, _request: &ChatRequest) -> vllm_chat::Result { Ok(RenderedPrompt { prompt: Prompt::Text(String::new()), + media_order: None, effective_template_kwargs: Default::default(), }) } diff --git a/rust/src/server/src/routes/http_client_tests.rs b/rust/src/server/src/routes/http_client_tests.rs index 216670b695..a0e0872e64 100644 --- a/rust/src/server/src/routes/http_client_tests.rs +++ b/rust/src/server/src/routes/http_client_tests.rs @@ -186,6 +186,7 @@ impl ChatRenderer for FakeChatBackend { } Ok(RenderedPrompt { prompt: Prompt::Text(prompt), + media_order: None, effective_template_kwargs: Default::default(), }) } diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 5ee2b51f3f..6fbe22e501 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -540,6 +540,7 @@ impl ChatRenderer for FakeChatBackend { } Ok(vllm_chat::RenderedPrompt { prompt: Prompt::Text(prompt), + media_order: None, effective_template_kwargs: Default::default(), }) }