pub mod media;
#[cfg(feature = "mm-routing")]
pub mod mm_routing;
pub mod prompt;
pub mod speculative_prefill;
pub(crate) mod structural_tag;
pub(crate) mod tool_choice;
pub mod tools;
use anyhow::Context;
use anyhow::{Result, bail};
use dynamo_protocols::types::{
ChatCompletionMessageContent, ChatCompletionRequestMessage,
ChatCompletionRequestToolMessageContent, ChatCompletionRequestToolMessageContentPart,
ChatCompletionRequestUserMessageContent, ChatCompletionRequestUserMessageContentPart,
ChatCompletionToolChoiceOption, EncodingFormat,
};
use dynamo_renderer::{OAIPromptFormatter, PromptRenderError, RenderedPrompt};
use dynamo_runtime::config::{
env_is_falsey, environment_names::llm as env_llm, is_truthy, parse_bool_opt,
};
use dynamo_runtime::error::{DynamoError, ErrorType, PublicDetails};
use dynamo_runtime::telemetry::{LIFECYCLE_TRACE_CONTEXT_KEY, LifecycleStage, LifecycleTrace};
use either::Either;
use futures::Stream;
use futures::stream::{self, StreamExt};
use std::borrow::Cow;
use std::time::Instant;
use dynamo_runtime::dynamo_nvtx_range;
use dynamo_runtime::metrics::frontend_perf::{
DETOKENIZE_TOKEN_COUNT, DETOKENIZE_TOTAL_US, STAGE_DURATION_SECONDS, STAGE_PREPROCESS,
StageGuard, TEMPLATE_SECONDS, TOKENIZE_SECONDS,
};
use std::{
any::Any,
collections::{HashMap, HashSet},
pin::Pin,
sync::{Arc, Mutex, OnceLock},
};
use tokio_util::sync::CancellationToken;
use tracing::{self, Instrument};
#[cfg(all(feature = "mm-routing", feature = "media-ffmpeg"))]
use crate::local_model::runtime_config::{
ModelRuntimeConfig, SGLANG_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY,
VLLM_NEMOTRON_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY,
VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY,
};
use crate::local_model::runtime_config::{
SGLANG_GENERATE_CAPABILITY, TOKEN_BUDGET_RUNTIME_KEY, TokenBudget,
};
#[cfg(feature = "mm-routing")]
use crate::model_card::ModelInfoType;
use crate::model_card::{ModelDeploymentCard, ModelInfo, PromptFormatterArtifact};
#[cfg(feature = "mm-routing")]
use crate::preprocessor::media::MediaFetcher;
use crate::preprocessor::media::{MediaDecoder, MediaLoader, max_data_url_bytes};
use crate::protocols::common::preprocessor::{
MultimodalData, MultimodalDataMap, MultimodalUuidMap, PreprocessedRequestBuilder, RoutingHints,
};
use crate::protocols::common::timing::RequestTracker;
use crate::tokenizers::Encoding;
use dynamo_parsers::{
ReasoningParser, ReasoningParserType, reasoning::ParserResult,
tool_calling::parsers::get_tool_parser_map,
};
use dynamo_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream};
use dynamo_runtime::pipeline::{
AsyncEngineContext, Context as PipelineContext, Error, ManyOut, Operator, SingleIn, async_trait,
};
use dynamo_runtime::protocols::annotated::{Annotated, AnnotationsProvider};
use crate::protocols::{
TokenIdType,
common::{
OutputOptionsProvider, SamplingOptionsProvider, StopConditionsProvider,
extensions::{
AgentHints, NvExtProvider, merge_response_nvext, request_cache_salt,
routing_constraints_to_kv,
},
},
openai::{
DeltaGeneratorExt, ParsingOptions,
chat_completions::{
NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse,
scrub_synthetic_chunk_metadata,
},
completions::{NvCreateCompletionRequest, NvCreateCompletionResponse},
embeddings::{NvCreateEmbeddingRequest, NvCreateEmbeddingResponse},
},
};
use crate::tokenizers::traits::Tokenizer;
use crate::preprocessor::prompt::{
MediaRequestExt, apply_continue_final_message, prompt_formatter_from_mdc,
};
use crate::protocols::openai::common_ext::CommonExtProvider;
use dynamo_renderer::{OAIChatLikeRequest, PromptFormatter, PromptInput, TextInput, TokenInput};
pub use crate::protocols::common::llm_backend::{BackendOutput, PreprocessedRequest};
pub use crate::protocols::common::metrics::{
ANNOTATION_LLM_METRICS, ANNOTATION_PAYLOAD_USAGE, LLMMetricAnnotation,
};
pub use crate::protocols::common::preprocessor::PreprocessedEmbeddingRequest;
use crate::protocols::common::llm_backend::EmbeddingsEngineOutput;
fn routing_priorities(hints: Option<&AgentHints>) -> (Option<f64>, Option<u32>, Option<i32>) {
let priority_jump = hints.and_then(|h| {
h.priority
.map(|priority| priority.max(0) as f64)
.or(h.latency_sensitivity)
});
let strict_priority = hints.and_then(|h| h.strict_priority);
let priority = hints.and_then(|h| h.priority);
(priority_jump, strict_priority, priority)
}
pub(crate) fn invalid_argument_error(message: impl Into<String>) -> anyhow::Error {
DynamoError::builder()
.error_type(ErrorType::InvalidArgument)
.message(message.into())
.build()
.into()
}
#[cfg(feature = "mm-routing")]
enum ImageDimFetchFailure {
InvalidArgument(String),
Recoverable(String),
}
#[cfg(feature = "mm-routing")]
impl ImageDimFetchFailure {
fn from_error(error: anyhow::Error) -> Self {
if MediaFetcher::is_policy_rejection(&error) {
Self::InvalidArgument(error.to_string())
} else {
Self::Recoverable(error.to_string())
}
}
fn to_error(&self) -> anyhow::Error {
match self {
Self::InvalidArgument(message) => invalid_argument_error(message.clone()),
Self::Recoverable(message) => anyhow::anyhow!("fetch_image_dims failed: {message}"),
}
}
}
fn validate_legacy_jail_nvext_choice_count(
n: u8,
extra_fields: Option<&[String]>,
is_legacy_jail: bool,
) -> Result<()> {
if n <= 1 || !is_legacy_jail {
return Ok(());
}
const CHOICE_SPECIFIC_FIELDS: [&str; 3] = ["engine_data", "routed_experts", "stop_reason"];
if let Some(field) = extra_fields.and_then(|fields| {
fields
.iter()
.find(|field| CHOICE_SPECIFIC_FIELDS.contains(&field.as_str()))
}) {
return Err(invalid_argument_error(format!(
"legacy tool-call parsing requires n = 1 when nvext.extra_fields requests choice-specific field `{field}`"
)));
}
Ok(())
}
#[derive(Clone, Copy, Debug)]
pub struct PromptReasoningPrefill {
legacy: bool,
unified: bool,
prefix: Option<&'static str>,
}
impl From<bool> for PromptReasoningPrefill {
fn from(value: bool) -> Self {
Self {
legacy: value,
unified: value,
prefix: None,
}
}
}
impl PromptReasoningPrefill {
fn for_route(self, route: &ToolProcessingRoute) -> bool {
match route {
ToolProcessingRoute::MuseUnified(_) | ToolProcessingRoute::Unified(_) => self.unified,
_ => self.legacy,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
enum ToolProcessingRoute {
MuseUnified(&'static str),
Unified(&'static str),
ParserV2(String),
LegacyJail(Option<String>),
PassThrough,
}
impl ToolProcessingRoute {
fn uses_legacy_jail(&self) -> bool {
matches!(self, Self::LegacyJail(_))
}
}
fn tool_content_part_as_user(
part: &ChatCompletionRequestToolMessageContentPart,
) -> Cow<'_, ChatCompletionRequestUserMessageContentPart> {
Cow::Owned(match part {
ChatCompletionRequestToolMessageContentPart::Text(part) => {
ChatCompletionRequestUserMessageContentPart::Text(part.clone())
}
ChatCompletionRequestToolMessageContentPart::ImageUrl(part) => {
ChatCompletionRequestUserMessageContentPart::ImageUrl(part.clone())
}
ChatCompletionRequestToolMessageContentPart::VideoUrl(part) => {
ChatCompletionRequestUserMessageContentPart::VideoUrl(part.clone())
}
ChatCompletionRequestToolMessageContentPart::AudioUrl(part) => {
ChatCompletionRequestUserMessageContentPart::AudioUrl(part.clone())
}
})
}
enum MultimodalContentPart<'a> {
User(&'a ChatCompletionRequestUserMessageContentPart),
Tool(&'a ChatCompletionRequestToolMessageContentPart),
}
impl<'a> MultimodalContentPart<'a> {
fn as_user(&self) -> Cow<'a, ChatCompletionRequestUserMessageContentPart> {
match self {
Self::User(part) => Cow::Borrowed(part),
Self::Tool(part) => tool_content_part_as_user(part),
}
}
fn media_info(&self) -> Option<(&'static str, Option<url::Url>, Option<String>)> {
match self {
Self::User(part) => match *part {
ChatCompletionRequestUserMessageContentPart::ImageUrl(part) => Some((
"image_url",
part.image_url.as_ref().map(|media| media.url.clone()),
part.uuid.clone(),
)),
ChatCompletionRequestUserMessageContentPart::VideoUrl(part) => Some((
"video_url",
part.video_url.as_ref().map(|media| media.url.clone()),
part.uuid.clone(),
)),
ChatCompletionRequestUserMessageContentPart::AudioUrl(part) => Some((
"audio_url",
part.audio_url.as_ref().map(|media| media.url.clone()),
part.uuid.clone(),
)),
_ => None,
},
Self::Tool(part) => match *part {
ChatCompletionRequestToolMessageContentPart::ImageUrl(part) => Some((
"image_url",
part.image_url.as_ref().map(|media| media.url.clone()),
part.uuid.clone(),
)),
ChatCompletionRequestToolMessageContentPart::VideoUrl(part) => Some((
"video_url",
part.video_url.as_ref().map(|media| media.url.clone()),
part.uuid.clone(),
)),
ChatCompletionRequestToolMessageContentPart::AudioUrl(part) => Some((
"audio_url",
part.audio_url.as_ref().map(|media| media.url.clone()),
part.uuid.clone(),
)),
_ => None,
},
}
}
}
fn multimodal_content_parts(
message: &ChatCompletionRequestMessage,
) -> Option<impl Iterator<Item = MultimodalContentPart<'_>>> {
match message {
ChatCompletionRequestMessage::User(user) => match &user.content {
ChatCompletionRequestUserMessageContent::Array(parts) => {
Some(Either::Left(parts.iter().map(MultimodalContentPart::User)))
}
ChatCompletionRequestUserMessageContent::Text(_) => None,
},
ChatCompletionRequestMessage::Tool(tool) => match &tool.content {
ChatCompletionRequestToolMessageContent::Array(parts) => {
Some(Either::Right(parts.iter().map(MultimodalContentPart::Tool)))
}
ChatCompletionRequestToolMessageContent::Text(_) => None,
},
_ => None,
}
}
#[cfg(feature = "mm-routing")]
fn image_content_part_url(
content_part: &ChatCompletionRequestUserMessageContentPart,
) -> Option<&str> {
let ChatCompletionRequestUserMessageContentPart::ImageUrl(part) = content_part else {
return None;
};
part.image_url.as_ref().map(|image| image.url.as_str())
}
pub(crate) fn decode_base64_to_floats(encoded: &str) -> Result<Vec<f32>, String> {
use base64::{Engine as _, engine::general_purpose::STANDARD};
let bytes = STANDARD
.decode(encoded)
.map_err(|error| error.to_string())?;
if bytes.len() % std::mem::size_of::<f32>() != 0 {
return Err(format!(
"base64-decoded embedding byte length {} is not a multiple of 4",
bytes.len()
));
}
Ok(bytes
.chunks_exact(std::mem::size_of::<f32>())
.map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
.collect())
}
#[cfg(test)]
mod embedding_base64_transport_tests {
use super::{OpenAIPreprocessor, decode_base64_to_floats};
use crate::protocols::common::llm_backend::EmbeddingsEngineOutput;
use crate::protocols::openai::embeddings::NvCreateEmbeddingRequest;
use base64::{Engine as _, engine::general_purpose::STANDARD};
use dynamo_runtime::protocols::annotated::Annotated;
use futures::StreamExt;
use futures::stream;
#[test]
fn portable_embedding_bytes_decode() {
let expected = vec![0.0, 1.0, -1.0, 2.5, -42.5, 3.25, f32::MIN, f32::MAX];
let bytes = expected
.iter()
.flat_map(|value: &f32| value.to_le_bytes())
.collect::<Vec<_>>();
let encoded = STANDARD.encode(bytes);
assert_eq!(decode_base64_to_floats(&encoded).unwrap(), expected);
}
#[test]
fn portable_embedding_bytes_reject_invalid_base64() {
assert!(decode_base64_to_floats("not!valid!base64").is_err());
}
#[test]
fn portable_embedding_bytes_reject_partial_float() {
let encoded = STANDARD.encode([0_u8; 5]);
let error = decode_base64_to_floats(&encoded).unwrap_err();
assert!(error.contains("not a multiple of 4"));
}
#[test]
fn empty_embeddings_engine_output_surfaces_as_error() {
let request: NvCreateEmbeddingRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"input": "hello"
}))
.unwrap();
let output = Annotated::from_data(EmbeddingsEngineOutput {
embeddings: vec![],
prompt_tokens: 0,
total_tokens: 0,
});
let items = futures::executor::block_on(
OpenAIPreprocessor::transform_embedding_postprocessor_stream(
stream::iter(vec![output]),
request,
)
.collect::<Vec<_>>(),
);
assert_eq!(items.len(), 1);
let item = items.into_iter().next().unwrap();
assert!(item.is_error(), "empty embeddings must be a stream error");
let err = item.into_result().unwrap_err();
assert!(
err.to_string().contains("empty `embeddings` field"),
"error should name the empty embeddings field, got: {err}"
);
}
}
pub const ANNOTATION_FORMATTED_PROMPT: &str = "formatted_prompt";
pub const ANNOTATION_TOKEN_IDS: &str = "token_ids";
pub(crate) const DEFAULT_THINKING_MODE_RUNTIME_KEY: &str = "default_thinking_mode";
fn drain_router_routing_data(
data: &mut Option<BackendOutput>,
tracker: Option<&crate::protocols::common::timing::RequestTracker>,
) {
let Some(routing_data) = data.as_mut().and_then(|d| d.routing_data.take()) else {
return;
};
let Some(tracker) = tracker else {
return;
};
if let Some(timing) = routing_data.timing {
tracker.set_external_timing(timing);
}
if let Some(worker_id) = routing_data.worker_id {
tracker.set_external_worker_info(worker_id);
}
if let Some(token_ids) = routing_data.token_ids {
tracker.set_external_query_token_ids(token_ids);
}
}
fn attach_metrics_annotation<Resp>(response: &mut Annotated<Resp>, metrics: &LLMMetricAnnotation) {
if let Ok(metrics_annotated) = metrics.to_annotation::<()>() {
response.event = metrics_annotated.event;
response.comment = metrics_annotated.comment;
} else {
tracing::warn!("Failed to serialize LLM metrics annotation");
}
}
fn attach_llm_metrics<Resp>(response: &mut Annotated<Resp>, metrics: LLMMetricAnnotation)
where
Resp: 'static,
{
if response.event.is_some() {
return;
}
if let Some(data) = response.data.as_mut()
&& let Some(chat_response) =
(data as &mut dyn Any).downcast_mut::<NvCreateChatCompletionStreamResponse>()
{
chat_response.llm_metrics = Some(metrics);
return;
}
attach_metrics_annotation(response, &metrics);
}
#[inline]
fn build_llm_metric_annotation(
tracker: Option<&RequestTracker>,
input_tokens: usize,
output_tokens: usize,
chunk_tokens: usize,
cached_tokens: Option<usize>,
mm_counts: MultimodalCounts,
image_tokens: Option<usize>,
) -> LLMMetricAnnotation {
LLMMetricAnnotation {
input_tokens,
output_tokens,
chunk_tokens,
cached_tokens,
image_count: mm_counts.image,
video_count: mm_counts.video,
audio_count: mm_counts.audio,
image_tokens,
prefill_worker_id: tracker.and_then(|t| t.prefill_worker_id()),
prefill_dp_rank: tracker.and_then(|t| t.prefill_dp_rank()),
prefill_worker_type: tracker
.and_then(|t| t.prefill_worker_type())
.map(String::from),
decode_worker_id: tracker.and_then(|t| t.decode_worker_id()),
decode_dp_rank: tracker.and_then(|t| t.decode_dp_rank()),
decode_worker_type: tracker
.and_then(|t| t.decode_worker_type())
.map(String::from),
tokenize_latency: tracker.and_then(|t| t.tokenize_latency()),
detokenize_total_latency: tracker.and_then(|t| t.detokenize_total_latency()),
detokenize_count: tracker.map(|t| t.detokenize_count()),
}
}
struct ChoiceReasoningState {
parser: Box<dyn ReasoningParser>,
guided_json_bypass_decision: Option<bool>,
pending_reasoning: String,
pending_content: String,
left_reasoning: bool,
drained: bool,
parser_finished: bool,
}
#[derive(Debug, Default)]
struct ReasoningUsageEstimator {
total: u32,
active_choices: HashSet<u32>,
}
impl ReasoningUsageEstimator {
fn observe(&mut self, chunk: &NvCreateChatCompletionStreamResponse) {
let token_count = chunk
.llm_metrics
.as_ref()
.map_or(0, |metrics| metrics.chunk_tokens)
.try_into()
.unwrap_or(u32::MAX);
let has_reasoning = chunk
.inner
.choices
.iter()
.any(|choice| choice.delta.reasoning_content.is_some());
let active_without_visible_output = chunk.inner.choices.iter().any(|choice| {
self.active_choices.contains(&choice.index) && !choice_has_visible_output(choice)
});
let usage_chunk_while_reasoning =
chunk.inner.choices.is_empty() && !self.active_choices.is_empty();
if token_count > 0
&& (has_reasoning || active_without_visible_output || usage_chunk_while_reasoning)
{
self.total = self.total.saturating_add(token_count);
}
for choice in &chunk.inner.choices {
if choice.delta.reasoning_content.is_some() {
self.active_choices.insert(choice.index);
} else if choice_has_visible_output(choice) {
self.active_choices.remove(&choice.index);
}
}
}
fn annotate(&self, usage: &mut dynamo_protocols::types::CompletionUsage) {
let details = usage.completion_tokens_details.get_or_insert_default();
if details.reasoning_tokens.unwrap_or(0) == 0 {
details.reasoning_tokens = Some(self.total);
}
}
}
fn choice_has_visible_output(choice: &dynamo_protocols::types::ChatChoiceStream) -> bool {
choice.delta.content.is_some()
|| choice
.delta
.tool_calls
.as_ref()
.is_some_and(|calls| !calls.is_empty())
}
fn annotate_reasoning_usage<S>(
stream: S,
) -> impl Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send
where
S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
{
let mut estimator = ReasoningUsageEstimator::default();
stream.map(move |mut response| {
if let Some(chunk) = response.data.as_mut() {
estimator.observe(chunk);
if let Some(usage) = chunk.inner.usage.as_mut() {
estimator.annotate(usage);
}
}
response
})
}
fn drain_deferred_reasoning(state: &mut ChoiceReasoningState) -> (Option<String>, Option<String>) {
if state.drained {
return (None, None);
}
state.drained = true;
let pending = std::mem::take(&mut state.pending_reasoning);
let pending_content = std::mem::take(&mut state.pending_content);
let result = if state.parser_finished {
ParserResult::default()
} else {
state.parser_finished = true;
state.parser.finish_reasoning_stream()
};
if state.left_reasoning {
let mut content = pending_content;
content.push_str(&result.normal_text);
(
(!content.is_empty()).then_some(content),
(!result.reasoning_text.is_empty()).then_some(result.reasoning_text),
)
} else {
let mut text = pending;
text.push_str(&result.reasoning_text);
text.push_str(&result.normal_text);
if text.is_empty() {
text = pending_content;
}
((!text.is_empty()).then_some(text), None)
}
}
struct ReasoningState {
stream: Pin<Box<dyn Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send>>,
parser_name: String,
prompt_injected_reasoning: bool,
bypass_bare_guided_json: bool,
choices: HashMap<u32, ChoiceReasoningState>,
last_response: Option<Annotated<NvCreateChatCompletionStreamResponse>>,
defer_reasoning_for_nonempty_content: bool,
saw_terminal_error: bool,
}
#[derive(Default)]
struct DeferredUnifiedChoice {
pending_reasoning: String,
pending_content: String,
saw_visible_output: bool,
drained: bool,
}
impl DeferredUnifiedChoice {
fn release_split(&mut self) -> (Option<String>, Option<String>) {
let content = std::mem::take(&mut self.pending_content);
let reasoning = std::mem::take(&mut self.pending_reasoning);
(
(!content.is_empty()).then_some(content),
(!reasoning.is_empty()).then_some(reasoning),
)
}
fn drain(&mut self) -> (Option<String>, Option<String>) {
if self.drained {
return (None, None);
}
self.drained = true;
if self.saw_visible_output {
return self.release_split();
}
let reasoning = std::mem::take(&mut self.pending_reasoning);
let content = std::mem::take(&mut self.pending_content);
if !reasoning.trim().is_empty() {
(Some(reasoning), None)
} else {
((!content.is_empty()).then_some(content), None)
}
}
}
#[derive(Debug, Clone, Copy)]
pub struct MmImageEntry {
pub mm_hash: u64,
pub width: u32,
pub height: u32,
}
#[cfg(feature = "mm-routing")]
#[derive(Debug, Clone, PartialEq, Eq)]
struct TrackedMmRoutingReplacement {
mm_hash: u64,
target_tokens: Vec<TokenIdType>,
worker_tokens: Vec<TokenIdType>,
routing_tokens: Vec<TokenIdType>,
}
#[cfg_attr(
not(feature = "mm-routing"),
allow(
dead_code,
reason = "the default build keeps the shared media-gathering return type but emits no routing entries"
)
)]
#[derive(Debug, Clone)]
enum MmRoutingEntry {
Image {
mm_hash: u64,
width: u32,
height: u32,
},
#[cfg_attr(
not(feature = "media-ffmpeg"),
allow(dead_code, reason = "video entries require the FFmpeg media decoder")
)]
Video {
mm_hash: u64,
placeholder_token_id: TokenIdType,
event_video_token_id: Option<TokenIdType>,
target_tokens: Vec<TokenIdType>,
replacement_tokens: Vec<TokenIdType>,
},
}
#[cfg(feature = "mm-routing")]
impl MmRoutingEntry {
fn mm_hash(&self) -> u64 {
match self {
Self::Image { mm_hash, .. } | Self::Video { mm_hash, .. } => *mm_hash,
}
}
}
#[cfg(feature = "mm-routing")]
fn exact_mm_routing_layout_accepts_next_entry(
previous_entry_was_video: &mut bool,
current_entry_is_video: bool,
) -> bool {
let accepted = !(*previous_entry_was_video && current_entry_is_video);
*previous_entry_was_video = current_entry_is_video;
accepted
}
#[cfg(feature = "mm-routing")]
fn exact_mm_routing_entries_are_unambiguous(entries: &[MmRoutingEntry]) -> bool {
let mut previous_entry_was_video = false;
entries.iter().all(|entry| {
exact_mm_routing_layout_accepts_next_entry(
&mut previous_entry_was_video,
matches!(entry, MmRoutingEntry::Video { .. }),
)
})
}
#[cfg(feature = "mm-routing")]
fn routing_bos_to_prepend(
configured_bos: Option<TokenIdType>,
has_image_entry: bool,
) -> Option<TokenIdType> {
has_image_entry.then_some(configured_bos).flatten()
}
#[cfg(feature = "mm-routing")]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RoutingImagePromptLayout {
RepeatedPad,
WrappedRepeatedPad {
image_start: TokenIdType,
image_end: TokenIdType,
},
KimiK3 {
media_begin: TokenIdType,
media_content: TokenIdType,
media_end: TokenIdType,
},
}
#[cfg(feature = "mm-routing")]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum RoutingImageDimensionPolicy {
Encoded,
ExifTransposed,
}
#[cfg(feature = "mm-routing")]
fn runtime_has_vllm_generate_capability(
runtime_config: &crate::local_model::runtime_config::ModelRuntimeConfig,
) -> bool {
use crate::local_model::runtime_config::VLLM_INFERENCE_V1_GENERATE_CAPABILITY;
runtime_config
.get_engine_specific::<bool>(VLLM_INFERENCE_V1_GENERATE_CAPABILITY)
.ok()
.flatten()
.unwrap_or(false)
}
#[cfg(feature = "mm-routing")]
fn routing_image_dimension_policy(
runtime_config: &crate::local_model::runtime_config::ModelRuntimeConfig,
frontend_decoding: bool,
prompt_layout: Option<RoutingImagePromptLayout>,
) -> RoutingImageDimensionPolicy {
if runtime_has_vllm_generate_capability(runtime_config)
&& !frontend_decoding
&& matches!(prompt_layout, Some(RoutingImagePromptLayout::KimiK3 { .. }))
{
RoutingImageDimensionPolicy::ExifTransposed
} else {
RoutingImageDimensionPolicy::Encoded
}
}
#[cfg(feature = "mm-routing")]
fn encode_routing_segment(
tokenizer: &dyn Tokenizer,
text: &str,
allow_special: bool,
) -> Result<Vec<TokenIdType>> {
let segment = crate::tokenizers::EncodeSegment::new(text, allow_special);
Ok(tokenizer.encode_segments(&[segment])?.token_ids().to_vec())
}
#[cfg(feature = "mm-routing")]
fn resolve_routing_image_prompt_layout(
tokenizer: &dyn Tokenizer,
kind: mm_routing::image::ImagePromptKind,
expected_image_token_id: Option<TokenIdType>,
) -> Result<RoutingImagePromptLayout> {
if kind == mm_routing::image::ImagePromptKind::RepeatedPad {
return Ok(RoutingImagePromptLayout::RepeatedPad);
}
let resolve_segmented_control = |label: &str, token: &str| -> Result<TokenIdType> {
let ids = encode_routing_segment(tokenizer, token, true)?;
if ids.len() != 1 {
bail!(
"{label} routing control token {token:?} encoded to {} ids ({ids:?}); expected exactly one",
ids.len()
);
}
Ok(ids[0])
};
if kind == mm_routing::image::ImagePromptKind::Nemotron {
let resolve_atomic_control = |label: &str, token: &str| -> Result<TokenIdType> {
let ids = tokenizer.encode(token)?.token_ids().to_vec();
if ids.len() != 1 {
bail!(
"{label} routing control token {token:?} encoded to {} ids ({ids:?}); expected exactly one",
ids.len()
);
}
Ok(ids[0])
};
let image_context = resolve_atomic_control(
"Nemotron image-context",
mm_routing::nemotron::IMAGE_CONTEXT,
)?;
let expected_image_token_id =
expected_image_token_id.context("Nemotron image context token id is unavailable")?;
anyhow::ensure!(
image_context == expected_image_token_id,
"Nemotron tokenizer image-context id {image_context} does not match config id {expected_image_token_id}"
);
return Ok(RoutingImagePromptLayout::WrappedRepeatedPad {
image_start: resolve_atomic_control(
"Nemotron image-start",
mm_routing::nemotron::IMAGE_START,
)?,
image_end: resolve_atomic_control(
"Nemotron image-end",
mm_routing::nemotron::IMAGE_END,
)?,
});
}
Ok(RoutingImagePromptLayout::KimiK3 {
media_begin: resolve_segmented_control("Kimi-K3 media-begin", "<|media_begin|>")?,
media_content: resolve_segmented_control("Kimi-K3 media-content", "<|media_content|>")?,
media_end: resolve_segmented_control("Kimi-K3 media-end", "<|media_end|>")?,
})
}
#[cfg(feature = "mm-routing")]
fn append_mm_routing_replacement(
expanded: &mut Vec<TokenIdType>,
tokenizer: &dyn Tokenizer,
layout: RoutingImagePromptLayout,
image: MmImageEntry,
num_image_tokens: usize,
) -> Result<()> {
let fill_token = dynamo_kv_router::protocols::pad_value_for_mm_hash(image.mm_hash);
append_mm_routing_replacement_with_fill(
expanded,
tokenizer,
layout,
image,
num_image_tokens,
fill_token,
)
}
#[cfg(feature = "mm-routing")]
fn append_mm_routing_replacement_with_fill(
expanded: &mut Vec<TokenIdType>,
tokenizer: &dyn Tokenizer,
layout: RoutingImagePromptLayout,
image: MmImageEntry,
num_image_tokens: usize,
fill_token: TokenIdType,
) -> Result<()> {
match layout {
RoutingImagePromptLayout::RepeatedPad => {
expanded.extend(std::iter::repeat_n(fill_token, num_image_tokens));
}
RoutingImagePromptLayout::WrappedRepeatedPad {
image_start,
image_end,
} => {
expanded.push(image_start);
expanded.extend(std::iter::repeat_n(fill_token, num_image_tokens));
expanded.push(image_end);
}
RoutingImagePromptLayout::KimiK3 {
media_begin,
media_content,
media_end,
} => {
expanded.push(media_begin);
let dimensions = format!("image {}x{}", image.width, image.height);
expanded.extend(encode_routing_segment(tokenizer, &dimensions, false)?);
expanded.push(media_content);
expanded.extend(std::iter::repeat_n(fill_token, num_image_tokens));
expanded.push(media_end);
}
}
Ok(())
}
#[cfg(feature = "mm-routing")]
fn apply_tracked_mm_replacements(
routing_prepend_bos: Option<TokenIdType>,
replacements: &[TrackedMmRoutingReplacement],
token_ids: &[TokenIdType],
block_size: usize,
image_token_id: Option<TokenIdType>,
video_token_id: Option<TokenIdType>,
kv_event_mm_identity: mm_routing::KvEventMmIdentity,
) -> Result<(
Vec<TokenIdType>,
usize,
Vec<Option<dynamo_kv_router::protocols::BlockExtraInfo>>,
)> {
use dynamo_kv_router::protocols::{BlockExtraInfo, BlockMmObjectInfo};
use dynamo_kv_router::zmq_wire::{normalize_mm_placeholder_runs, normalize_mm_token_runs};
anyhow::ensure!(block_size > 0, "MM routing block size must be positive");
anyhow::ensure!(
replacements.iter().all(|replacement| {
!replacement.target_tokens.is_empty()
&& replacement.worker_tokens.len() == replacement.routing_tokens.len()
}),
"tracked MM routing replacements must have a non-empty target and equal token lengths"
);
let replacement_tokens = replacements.iter().try_fold(0usize, |total, replacement| {
total
.checked_add(replacement.routing_tokens.len())
.context("MM routing replacement capacity overflow")
})?;
let capacity = token_ids
.len()
.checked_add(replacement_tokens)
.and_then(|value| value.checked_add(routing_prepend_bos.is_some() as usize))
.context("MM routing token capacity overflow")?;
let mut worker_tokens = Vec::with_capacity(capacity);
let mut routing_tokens = Vec::with_capacity(capacity);
if let Some(bos) = routing_prepend_bos {
worker_tokens.push(bos);
routing_tokens.push(bos);
}
let mut distinct_targets: Vec<&[TokenIdType]> = Vec::new();
for replacement in replacements {
let target = replacement.target_tokens.as_slice();
if !distinct_targets.contains(&target) {
distinct_targets.push(target);
}
}
let target_matches_at = |target: &[TokenIdType], token_index: usize| {
token_ids.get(token_index..token_index.saturating_add(target.len())) == Some(target)
};
let mut spans = Vec::with_capacity(replacements.len());
let mut replacement_index = 0usize;
let mut token_index = 0usize;
while token_index < token_ids.len() {
if let Some(replacement) = replacements.get(replacement_index)
&& target_matches_at(&replacement.target_tokens, token_index)
{
let start = worker_tokens.len();
worker_tokens.extend_from_slice(&replacement.worker_tokens);
routing_tokens.extend_from_slice(&replacement.routing_tokens);
spans.push((start, worker_tokens.len(), replacement.mm_hash));
token_index += replacement.target_tokens.len();
replacement_index += 1;
continue;
}
anyhow::ensure!(
!distinct_targets
.iter()
.any(|target| target_matches_at(target, token_index)),
"multimodal placeholders do not match request order"
);
worker_tokens.push(token_ids[token_index]);
routing_tokens.push(token_ids[token_index]);
token_index += 1;
}
anyhow::ensure!(
replacement_index == replacements.len(),
"tokenized prompt is missing multimodal placeholders"
);
let expanded_prompt_len = routing_tokens.len();
let padded_len = expanded_prompt_len.div_ceil(block_size) * block_size;
worker_tokens.resize(padded_len, 0);
routing_tokens.resize(padded_len, 0);
let mut block_mm_infos = vec![None; padded_len / block_size];
for (block_index, block_start) in (0..padded_len).step_by(block_size).enumerate() {
let block_end = block_start + block_size;
let mm_hashes: Vec<u64> = spans
.iter()
.filter(|(start, end, _)| *start < block_end && *end > block_start)
.map(|(_, _, mm_hash)| *mm_hash)
.collect();
if mm_hashes.is_empty() {
continue;
}
let worker_block = &worker_tokens[block_start..block_end];
let routing_block = &mut routing_tokens[block_start..block_end];
let normalized = if video_token_id.is_some() {
normalize_mm_placeholder_runs(worker_block, image_token_id, video_token_id, &mm_hashes)
.map(|(tokens, _)| tokens)
} else if let Some(image_token_id) = image_token_id
&& worker_block.contains(&image_token_id)
{
normalize_mm_token_runs(worker_block, image_token_id, &mm_hashes)
.map(|(tokens, _)| tokens)
} else {
None
};
match normalized {
Some(normalized) => {
anyhow::ensure!(
normalized == routing_block,
"frontend MM replacement differs from KV-event normalization"
);
}
None if kv_event_mm_identity == mm_routing::KvEventMmIdentity::PadValueTokens => {
}
None => {
routing_block.copy_from_slice(worker_block);
block_mm_infos[block_index] = Some(BlockExtraInfo {
mm_objects: mm_hashes
.into_iter()
.map(|mm_hash| BlockMmObjectInfo {
mm_hash,
offsets: Vec::new(),
})
.collect(),
});
}
}
}
Ok((routing_tokens, expanded_prompt_len, block_mm_infos))
}
#[cfg(feature = "mm-routing")]
fn expand_mm_routing_tokens(
tokenizer: &dyn Tokenizer,
prompt_layout: RoutingImagePromptLayout,
routing_prepend_bos: Option<TokenIdType>,
find_token_id: TokenIdType,
mm_image_entries: &[MmImageEntry],
n_tokens: &[usize],
token_ids: &[TokenIdType],
) -> Result<(Vec<TokenIdType>, usize)> {
debug_assert_eq!(mm_image_entries.len(), n_tokens.len());
let n_total: usize = n_tokens.iter().sum();
let bos_extra = routing_prepend_bos.is_some() as usize;
let mut expanded = Vec::with_capacity(token_ids.len() + n_total + bos_extra);
if let Some(bos) = routing_prepend_bos {
expanded.push(bos);
}
let mut image_idx = 0usize;
for &token_id in token_ids {
if token_id == find_token_id && image_idx < mm_image_entries.len() {
append_mm_routing_replacement(
&mut expanded,
tokenizer,
prompt_layout,
mm_image_entries[image_idx],
n_tokens[image_idx],
)?;
image_idx += 1;
} else {
expanded.push(token_id);
}
}
let expanded_prompt_len = expanded.len();
Ok((expanded, expanded_prompt_len))
}
#[cfg(feature = "mm-routing")]
#[allow(clippy::too_many_arguments)]
fn try_expand_mm_routing_tokens(
tokenizer: &dyn Tokenizer,
prompt_layout: RoutingImagePromptLayout,
routing_prepend_bos: Option<TokenIdType>,
find_token_id: TokenIdType,
mm_image_entries: &[MmImageEntry],
n_tokens: &[usize],
token_ids: &[TokenIdType],
model_id: &str,
) -> Option<(Vec<TokenIdType>, usize)> {
match expand_mm_routing_tokens(
tokenizer,
prompt_layout,
routing_prepend_bos,
find_token_id,
mm_image_entries,
n_tokens,
token_ids,
) {
Ok(expanded) => Some(expanded),
Err(error) => {
tracing::warn!(
target: "mm_routing",
model = model_id,
%error,
"routing-only image prompt expansion failed; skipping MM routing info"
);
None
}
}
}
struct MediaFetchTask<'a> {
modality: &'static str,
slot_idx: usize,
content_part: Cow<'a, ChatCompletionRequestUserMessageContentPart>,
}
#[derive(Debug, Clone, Copy, Default)]
pub struct MultimodalCounts {
pub image: usize,
pub video: usize,
pub audio: usize,
}
impl MultimodalCounts {
fn from_preprocessed(request: &PreprocessedRequest) -> Self {
let count = |key: &str| {
request
.multi_modal_data
.as_ref()
.and_then(|m| m.get(key))
.map_or(0, |v| v.len())
};
Self {
image: count("image_url"),
video: count("video_url"),
audio: count("audio_url"),
}
}
}
#[cfg(feature = "mm-routing")]
fn checked_add_image_tokens(total: Option<usize>, next: usize) -> Option<usize> {
total?.checked_add(next)
}
#[cfg(feature = "mm-routing")]
fn has_mm_processor_override(value: Option<&serde_json::Value>) -> bool {
value.is_some_and(|value| match value {
serde_json::Value::Null => false,
serde_json::Value::Object(map) => !map.is_empty(),
_ => true,
})
}
#[cfg(feature = "mm-routing")]
fn exact_mm_routing_preconditions_met(
has_user_uuid: bool,
resolved_image_count: usize,
total_image_count: usize,
has_processor_override: bool,
) -> bool {
!has_user_uuid && !has_processor_override && resolved_image_count == total_image_count
}
#[cfg(all(feature = "mm-routing", feature = "media-ffmpeg"))]
fn should_hash_decoded_video(
exact_mm_routing_eligible: bool,
has_user_uuid: bool,
has_processor_override: bool,
has_video_routing_processor: bool,
) -> bool {
exact_mm_routing_eligible
&& !has_user_uuid
&& !has_processor_override
&& has_video_routing_processor
}
#[cfg(feature = "mm-routing")]
fn exact_mm_routing_supports_modality(modality: &str, frontend_decoding: bool) -> bool {
modality == "image_url" || (modality == "video_url" && frontend_decoding)
}
#[cfg(feature = "mm-routing")]
fn aggregate_image_tokens(
image_tokens: Option<usize>,
resolved_image_count: usize,
total_image_count: usize,
has_processor_override: bool,
) -> Option<usize> {
if total_image_count == 0 || has_processor_override || resolved_image_count != total_image_count
{
return None;
}
image_tokens
}
#[cfg(feature = "mm-routing")]
fn mdc_model_dir(mdc: &ModelDeploymentCard) -> Option<std::path::PathBuf> {
let ModelInfoType::HfConfigJson(cf) = mdc.model_info.as_ref()?;
cf.path()?.parent().map(std::path::PathBuf::from)
}
#[cfg(feature = "mm-routing")]
static DIM_FETCH_MEDIA_FETCHER: std::sync::LazyLock<crate::preprocessor::media::MediaFetcher> =
std::sync::LazyLock::new(crate::preprocessor::media::MediaFetcher::from_env);
#[cfg(feature = "mm-routing")]
static DIM_FETCH_HTTP_CLIENT: std::sync::LazyLock<reqwest::Client> =
std::sync::LazyLock::new(|| {
DIM_FETCH_MEDIA_FETCHER
.build_http_client()
.expect("dim-fetch http client construction failed")
});
pub(crate) const PRESERVE_OMITTED_MAX_TOKENS_CONTEXT_KEY: &str =
"dynamo.llm.preserve_omitted_max_tokens";
pub(crate) const REQUEST_PARSING_OPTIONS_CONTEXT_KEY: &str = "dynamo.llm.request_parsing_options";
const EMBEDDING_ADD_SPECIAL_TOKENS_ENV: &str = "DYN_EMBEDDING_TOKENIZATION_ADD_SPECIAL_TOKENS";
fn parse_embedding_add_special_tokens(value: &str) -> Option<bool> {
parse_bool_opt(value)
}
fn off_runtime<T>(f: impl FnOnce() -> T) -> T {
match tokio::runtime::Handle::try_current().map(|h| h.runtime_flavor()) {
Ok(tokio::runtime::RuntimeFlavor::MultiThread) => tokio::task::block_in_place(f),
_ => f(),
}
}
fn embedding_add_special_tokens_env() -> Result<Option<bool>> {
match std::env::var(EMBEDDING_ADD_SPECIAL_TOKENS_ENV) {
Ok(value) => match parse_embedding_add_special_tokens(&value) {
Some(value) => Ok(Some(value)),
None => bail!(
"invalid value {value:?} for {EMBEDDING_ADD_SPECIAL_TOKENS_ENV}; \
expected true/false/on/off/yes/no/1/0"
),
},
Err(std::env::VarError::NotPresent) => Ok(None),
Err(std::env::VarError::NotUnicode(_)) => bail!(
"{EMBEDDING_ADD_SPECIAL_TOKENS_ENV} must be valid Unicode and one of \
true/false/on/off/yes/no/1/0"
),
}
}
#[cfg(test)]
mod embedding_add_special_tokens_env_tests {
use super::parse_embedding_add_special_tokens;
#[test]
fn parser_uses_the_documented_truth_table() {
for value in ["1", "true", "on", "yes", " TRUE ", "On", "Yes"] {
assert_eq!(parse_embedding_add_special_tokens(value), Some(true));
}
for value in ["0", "false", "off", "no", " FALSE ", "Off", "No"] {
assert_eq!(parse_embedding_add_special_tokens(value), Some(false));
}
for value in ["", "yes-please"] {
assert_eq!(parse_embedding_add_special_tokens(value), None);
}
}
}
fn embedding_chat_template_present(mdc: &ModelDeploymentCard) -> Result<bool> {
if mdc.chat_template_file.is_some() {
return Ok(true);
}
let Some(artifact) = mdc.prompt_formatter.as_ref() else {
return Ok(false);
};
let PromptFormatterArtifact::HfTokenizerConfigJson(checked_file) = artifact else {
return Ok(true);
};
let Some(path) = checked_file.path() else {
return Ok(true);
};
let contents = std::fs::read_to_string(path)
.with_context(|| format!("fs:read_to_string '{}'", path.display()))?;
let config: serde_json::Value = serde_json::from_str(&contents)
.with_context(|| format!("Failed to parse '{}'", path.display()))?;
Ok(config
.get("chat_template")
.is_some_and(|template| !template.is_null()))
}
fn embedding_prompt_formatter(mdc: &ModelDeploymentCard) -> Result<PromptFormatter> {
match prompt_formatter_from_mdc(mdc) {
Ok(formatter) => Ok(formatter),
Err(_) if !embedding_chat_template_present(mdc)? => Ok(PromptFormatter::no_op()),
Err(err) => Err(err),
}
}
fn attach_agent_context_from_context(
request: &mut PreprocessedRequest,
context: &PipelineContext<()>,
) {
if let Ok(agent_context) = context.get::<crate::protocols::common::extensions::AgentContext>(
crate::protocols::common::extensions::AGENT_CONTEXT_CONTEXT_KEY,
) {
request.agent_context = Some(agent_context.as_ref().clone());
}
}
fn attach_image_cache_scope_from_context(
request: &mut PreprocessedRequest,
context: &PipelineContext<()>,
) {
use crate::protocols::common::extensions::{SESSION_AFFINITY_CONTEXT_KEY, SessionAffinityId};
if let Ok(session_affinity) = context.get::<SessionAffinityId>(SESSION_AFFINITY_CONTEXT_KEY) {
request.image_cache_scope = Some(session_affinity.as_str().to_owned());
}
}
fn attach_request_context_metadata(
request: &mut PreprocessedRequest,
context: &PipelineContext<()>,
) {
attach_agent_context_from_context(request, context);
attach_image_cache_scope_from_context(request, context);
}
struct NormalizedArgsRequest<'a, R> {
inner: &'a R,
normalize_tool_call_args: bool,
continue_final_message: bool,
}
impl<R: OAIChatLikeRequest> OAIChatLikeRequest for NormalizedArgsRequest<'_, R> {
fn model(&self) -> String {
self.inner.model()
}
fn messages(&self) -> minijinja::value::Value {
let mut json = serde_json::to_value(self.inner.typed_messages().unwrap_or_default())
.unwrap_or_default();
if self.normalize_tool_call_args
&& let Err(e) = crate::preprocessor::prompt::normalize_tool_call_arguments(&mut json)
{
tracing::error!(
error = %e,
"tool_call arguments normalization failed; template rendering may fail \
if it calls .items() on a string"
);
}
if self.continue_final_message
&& let Err(e) =
crate::preprocessor::prompt::append_continue_final_message_tag(&mut json)
{
tracing::debug!(
error = %e,
"continue_final_message marker not appended; truncation will reject the request"
);
}
minijinja::value::Value::from_serialize(&json)
}
fn tools(&self) -> Option<minijinja::value::Value> {
self.inner.tools()
}
fn tool_choice(&self) -> Option<minijinja::value::Value> {
self.inner.tool_choice()
}
fn response_format(&self) -> Option<minijinja::value::Value> {
self.inner.response_format()
}
fn should_add_generation_prompt(&self) -> bool {
self.inner.should_add_generation_prompt()
}
fn extract_text(&self) -> Option<TextInput> {
self.inner.extract_text()
}
fn chat_template_args(&self) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
self.inner.chat_template_args()
}
fn mm_processor_kwargs(&self) -> Option<&serde_json::Value> {
self.inner.mm_processor_kwargs()
}
}
impl<R: AnnotationsProvider> AnnotationsProvider for NormalizedArgsRequest<'_, R> {
fn annotations(&self) -> Option<Vec<String>> {
self.inner.annotations()
}
}
impl<R: SamplingOptionsProvider> SamplingOptionsProvider for NormalizedArgsRequest<'_, R> {
fn extract_sampling_options(
&self,
) -> anyhow::Result<crate::protocols::common::SamplingOptions> {
self.inner.extract_sampling_options()
}
}
impl<R: StopConditionsProvider> StopConditionsProvider for NormalizedArgsRequest<'_, R> {
fn extract_stop_conditions(&self) -> anyhow::Result<crate::protocols::common::StopConditions> {
self.inner.extract_stop_conditions()
}
}
impl<R: OutputOptionsProvider> OutputOptionsProvider for NormalizedArgsRequest<'_, R> {
fn extract_output_options(&self) -> anyhow::Result<crate::protocols::common::OutputOptions> {
self.inner.extract_output_options()
}
}
impl<R: NvExtProvider> NvExtProvider for NormalizedArgsRequest<'_, R> {
fn nvext(&self) -> Option<&crate::protocols::common::extensions::NvExt> {
self.inner.nvext()
}
fn raw_prompt(&self) -> Option<String> {
self.inner.raw_prompt()
}
fn unsupported_fields(&self) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
self.inner.unsupported_fields()
}
}
#[derive(Clone, Copy, Debug, Default)]
struct PreprocessRequestOptions {
preserve_omitted_max_tokens: bool,
}
struct EmbeddingTokenizerState {
model_card: ModelDeploymentCard,
with_special_tokens: OnceLock<Arc<dyn Tokenizer>>,
without_special_tokens: OnceLock<Arc<dyn Tokenizer>>,
initialization_lock: Mutex<()>,
add_special_tokens_default: Option<bool>,
}
impl EmbeddingTokenizerState {
fn new(model_card: &ModelDeploymentCard) -> Result<Self> {
Ok(Self {
model_card: model_card.clone(),
with_special_tokens: OnceLock::new(),
without_special_tokens: OnceLock::new(),
initialization_lock: Mutex::new(()),
add_special_tokens_default: embedding_add_special_tokens_env()?,
})
}
fn tokenizer(&self, add_special_tokens: bool) -> Result<Arc<dyn Tokenizer>> {
let tokenizer = if add_special_tokens {
&self.with_special_tokens
} else {
&self.without_special_tokens
};
if let Some(tokenizer) = tokenizer.get() {
return Ok(tokenizer.clone());
}
let _guard = self
.initialization_lock
.lock()
.map_err(|_| anyhow::anyhow!("embedding tokenizer initialization lock was poisoned"))?;
if let Some(tokenizer) = tokenizer.get() {
return Ok(tokenizer.clone());
}
let initialized = self.model_card.embedding_tokenizer_with_options(
crate::tokenizers::TokenizerOptions { add_special_tokens },
)?;
let initialized: Arc<dyn Tokenizer> = (*initialized).clone();
Ok(tokenizer.get_or_init(|| initialized).clone())
}
}
pub struct OpenAIPreprocessor {
mdcsum: String,
formatter: Arc<dyn OAIPromptFormatter>,
tokenizer: Arc<dyn Tokenizer>,
embedding_tokenizers: Option<EmbeddingTokenizerState>,
model_info: Arc<dyn ModelInfo>,
token_id_bound: Option<usize>,
lora_name: Option<String>,
runtime_config: crate::local_model::runtime_config::ModelRuntimeConfig,
structural_tag_reasoning_boundary: structural_tag::ResolvedReasoningBoundary,
kv_cache_block_size: usize,
tool_call_parser: Option<String>,
parser_requires_special_tokens: bool,
normalize_tool_call_args: bool,
media_loader: Option<MediaLoader>,
token_budget: Option<TokenBudget>,
context_length: u32,
speculative_prefill_tasks: speculative_prefill::PrefillTasks,
#[cfg(feature = "mm-routing")]
image_token_counter: Option<mm_routing::image::ImageRoutingProcessor>,
#[cfg(all(feature = "mm-routing", feature = "media-ffmpeg"))]
video_routing_processor: Option<mm_routing::VideoRoutingProcessor>,
#[cfg(feature = "mm-routing")]
routing_image_token_id: Option<crate::protocols::TokenIdType>,
#[cfg(feature = "mm-routing")]
routing_image_prompt_layout: Option<RoutingImagePromptLayout>,
#[cfg(feature = "mm-routing")]
routing_image_dimension_policy: RoutingImageDimensionPolicy,
#[cfg(feature = "mm-routing")]
routing_prepend_bos: Option<crate::protocols::TokenIdType>,
}
pub(crate) const LORA_NAME_CONTEXT_KEY: &str = "discovery.lora_name";
fn token_id_bound(model_vocab: Option<usize>, tokenizer_bound: Option<usize>) -> Option<usize> {
model_vocab
.into_iter()
.chain(tokenizer_bound)
.filter(|&size| size > 0)
.max()
}
fn ensure_token_ids_in_vocab(
field: &str,
tokens: &[crate::protocols::TokenIdType],
bound: Option<usize>,
) -> anyhow::Result<()> {
if let Some(bound) = bound
&& let Some(&bad) = tokens.iter().find(|&&t| t as usize >= bound)
{
return Err(invalid_argument_error(format!(
"{field} token id {bad} is out of range (must be < {bound})"
)));
}
Ok(())
}
#[cfg(all(feature = "mm-routing", feature = "media-ffmpeg"))]
fn resolve_qwen_video_processor_contract(
runtime_config: &ModelRuntimeConfig,
) -> Result<Option<mm_routing::QwenVideoProcessorContract>> {
let vllm_contract = runtime_config
.get_engine_specific::<mm_routing::VllmQwenVideoProcessorContract>(
VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY,
)
.with_context(|| {
format!(
"invalid Qwen video processor runtime metadata under {VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY}"
)
})?
.map(mm_routing::QwenVideoProcessorContract::try_from)
.transpose()?;
let sglang_contract = runtime_config
.get_engine_specific::<mm_routing::SglangQwenVideoProcessorContract>(
SGLANG_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY,
)
.with_context(|| {
format!(
"invalid Qwen video processor runtime metadata under {SGLANG_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY}"
)
})?
.map(mm_routing::QwenVideoProcessorContract::try_from)
.transpose()?;
anyhow::ensure!(
vllm_contract.is_none() || sglang_contract.is_none(),
"multiple Qwen video processor contracts were published"
);
Ok(vllm_contract.or(sglang_contract))
}
impl OpenAIPreprocessor {
fn omitted_max_tokens_default(
prompt_len: usize,
token_limit: Option<u32>,
options: PreprocessRequestOptions,
) -> Option<u32> {
if options.preserve_omitted_max_tokens {
return None;
}
token_limit.map(|limit| limit.saturating_sub(prompt_len as u32))
}
fn exact_prompt_len(
expanded_image_prompt_len: Option<usize>,
multi_modal_data: Option<&MultimodalDataMap>,
token_ids_len: usize,
) -> Option<usize> {
let has_images = multi_modal_data
.and_then(|media| media.get("image_url"))
.is_some_and(|items| !items.is_empty());
let has_other_media = multi_modal_data.is_some_and(|media| {
media
.iter()
.any(|(kind, items)| kind != "image_url" && !items.is_empty())
});
if has_other_media {
None
} else if let Some(expanded_prompt_len) =
expanded_image_prompt_len.filter(|length| *length > 0)
{
Some(expanded_prompt_len)
} else if has_images {
None
} else {
Some(token_ids_len)
}
}
fn validate_requested_token_budget(
prompt_len: usize,
max_tokens: Option<u32>,
token_budget: Option<&TokenBudget>,
) -> Result<()> {
let Some(token_budget) = token_budget else {
return Ok(());
};
let combined_limit = token_budget.combined_limit as usize;
if prompt_len >= combined_limit {
if !token_budget.reject_prompt_overflow {
return Ok(());
}
return Err(Self::prompt_overflow_error(prompt_len, combined_limit).into());
}
let requested_tokens = prompt_len.saturating_add(max_tokens.unwrap_or(1) as usize);
if requested_tokens > combined_limit && token_budget.reject_total_overflow {
let request_description = match max_tokens {
Some(max_tokens) => format!(
"your request has {prompt_len} input tokens and asks for {max_tokens} output \
tokens ({requested_tokens} tokens total)"
),
None => format!(
"your request has {prompt_len} input tokens and requires room for at least one \
output token ({requested_tokens} tokens minimum)"
),
};
return Err(DynamoError::builder()
.error_type(ErrorType::InvalidArgument)
.message(format!(
"This model configuration accepts at most {} combined input and output \
tokens. However, {}. Please reduce the input length or requested output \
length.",
combined_limit, request_description,
))
.public_details(PublicDetails::ContextLength {
limit: combined_limit as u64,
actual: Some(requested_tokens as u64),
})
.build()
.into());
}
Ok(())
}
fn validate_preprocessed_token_budget(
request: &PreprocessedRequest,
token_budget: Option<&TokenBudget>,
) -> Result<Option<usize>> {
if request.prompt_embeds.is_some() {
return Ok(None);
}
let exact_prompt_len = Self::exact_prompt_len(
request
.mm_routing_info
.as_ref()
.map(|mm| mm.expanded_prompt_len),
request.multi_modal_data.as_ref(),
request.token_ids.len(),
);
if let Some(prompt_len) = exact_prompt_len {
Self::validate_requested_token_budget(
prompt_len,
request.stop_conditions.max_tokens,
token_budget,
)?;
}
Ok(exact_prompt_len)
}
fn nvext_passthrough_args<R: NvExtProvider>(
request: &R,
) -> Option<serde_json::Map<String, serde_json::Value>> {
let mut nvext_passthrough = serde_json::Map::new();
if let Some(nvext) = request.nvext() {
if let Some(ref fields) = nvext.extra_fields {
nvext_passthrough.insert("extra_fields".to_string(), serde_json::json!(fields));
}
if let Some(ref metadata_upload) = nvext.metadata_upload {
nvext_passthrough.insert(
"metadata_upload".to_string(),
serde_json::json!(metadata_upload),
);
}
if nvext.token_data.is_some() {
nvext_passthrough.insert("token_in".to_string(), serde_json::Value::Bool(true));
}
}
if let Some(salt) = request_cache_salt(request) {
nvext_passthrough.insert("cache_salt".to_string(), serde_json::json!(salt));
}
if nvext_passthrough.is_empty() {
None
} else {
Some(nvext_passthrough)
}
}
fn sampling_passthrough_args<R: NvExtProvider>(
request: &R,
) -> Option<serde_json::Map<String, serde_json::Value>> {
let mut sampling_passthrough = serde_json::Map::new();
if let Some(fields) = request.unsupported_fields() {
for key in [
"detokenize",
"allowed_token_ids",
"bad_words_token_ids",
"logprob_token_ids",
] {
if let Some(value) = fields.get(key) {
sampling_passthrough.insert(key.to_string(), value.clone());
}
}
}
if sampling_passthrough.is_empty() {
None
} else {
Some(sampling_passthrough)
}
}
fn backend_extra_args<R: OAIChatLikeRequest + NvExtProvider>(
request: &R,
reasoning_parser_configured: bool,
reasoning_ended: Option<bool>,
) -> Option<serde_json::Value> {
let mut extra_args = serde_json::Map::new();
if let Some(nvext_passthrough) = Self::nvext_passthrough_args(request) {
extra_args.insert(
"nvext".to_string(),
serde_json::Value::Object(nvext_passthrough),
);
}
if let Some(sampling_passthrough) = Self::sampling_passthrough_args(request) {
extra_args.insert(
"sampling_options".to_string(),
serde_json::Value::Object(sampling_passthrough),
);
}
if reasoning_parser_configured
&& let Some(chat_template_args) = request.chat_template_args()
{
extra_args.insert(
"reasoning_parser_kwargs".to_string(),
serde_json::json!({ "chat_template_kwargs": chat_template_args }),
);
}
if reasoning_parser_configured && let Some(reasoning_ended) = reasoning_ended {
extra_args.insert("reasoning_ended".to_string(), reasoning_ended.into());
}
if extra_args.is_empty() {
None
} else {
Some(serde_json::Value::Object(extra_args))
}
}
fn has_request_thinking_control(
chat_template_args: Option<&HashMap<String, serde_json::Value>>,
) -> bool {
chat_template_args.is_some_and(|args| {
[
"thinking",
"enable_thinking",
"thinking_mode",
"reasoning_effort",
]
.iter()
.any(|key| args.contains_key(*key))
})
}
fn request_has_client_thinking_control(request: &NvCreateChatCompletionRequest) -> bool {
request.thinking.is_some()
|| Self::has_request_thinking_control(request.chat_template_args.as_ref())
}
fn apply_default_thinking_mode_from_runtime_config(
runtime_config: &crate::local_model::runtime_config::ModelRuntimeConfig,
request: &mut NvCreateChatCompletionRequest,
) {
let default_mode = runtime_config
.runtime_data
.get(DEFAULT_THINKING_MODE_RUNTIME_KEY)
.and_then(serde_json::Value::as_str);
let from_client = Self::request_has_client_thinking_control(request);
Self::apply_default_thinking_mode_value(
default_mode,
&mut request.chat_template_args,
from_client,
);
}
fn apply_default_thinking_mode_value(
default_mode: Option<&str>,
chat_template_args: &mut Option<HashMap<String, serde_json::Value>>,
thinking_control_from_client: bool,
) {
if thinking_control_from_client {
return;
}
let Some(default_mode) = default_mode else {
return;
};
let enabled = match default_mode {
"enabled" => true,
"disabled" => false,
other => {
tracing::warn!(
default_thinking_mode = other,
"Ignoring invalid runtime_config default_thinking_mode; expected 'enabled' or 'disabled'"
);
return;
}
};
let args = chat_template_args.get_or_insert_with(HashMap::new);
args.insert("thinking".to_string(), serde_json::Value::Bool(enabled));
args.insert(
"enable_thinking".to_string(),
serde_json::Value::Bool(enabled),
);
args.insert(
"thinking_mode".to_string(),
serde_json::Value::String(if enabled { "enabled" } else { "disabled" }.to_string()),
);
}
fn apply_default_thinking_mode(&self, request: &mut NvCreateChatCompletionRequest) {
Self::apply_default_thinking_mode_from_runtime_config(&self.runtime_config, request);
}
fn guided_output_requires_reasoning<R: OAIChatLikeRequest>(
request: &R,
reasoning_parser: Option<&str>,
) -> bool {
if reasoning_parser.is_none() {
return false;
}
let is_guided_tool_choice = Self::has_guided_tool_choice(request);
let is_structured_response = Self::has_structured_response_format(request);
let structured_response_requires_reasoning = is_structured_response
&& Self::structured_response_supports_sglang_reasoning_gate(reasoning_parser);
(is_guided_tool_choice || structured_response_requires_reasoning)
&& Self::sglang_effective_reasoning_enabled(
reasoning_parser,
request.chat_template_args(),
)
}
fn request_requires_reasoning<R: OAIChatLikeRequest>(
request: &R,
reasoning_parser: Option<&str>,
has_thinking_budget: bool,
) -> bool {
Self::guided_output_requires_reasoning(request, reasoning_parser)
|| (has_thinking_budget
&& Self::sglang_effective_reasoning_enabled(
reasoning_parser,
request.chat_template_args(),
))
}
fn structured_response_supports_sglang_reasoning_gate(reasoning_parser: Option<&str>) -> bool {
!matches!(reasoning_parser, Some("gpt_oss"))
}
pub(crate) fn has_structured_response_format<R: OAIChatLikeRequest>(request: &R) -> bool {
request.response_format().is_some_and(|format| {
format
.get_attr("type")
.ok()
.is_some_and(|kind| kind.as_str().is_some_and(|kind| kind != "text"))
})
}
fn sglang_effective_reasoning_enabled(
reasoning_parser: Option<&str>,
chat_template_args: Option<&std::collections::HashMap<String, serde_json::Value>>,
) -> bool {
let thinking_enabled = dynamo_renderer::thinking_bool_from_args(chat_template_args);
match reasoning_parser {
Some("qwen3" | "glm45" | "nemotron_nano" | "nemotron3" | "nemotron_v3") => {
thinking_enabled != Some(false)
}
Some("kimi_k25" | "kimi_k3" | "kimi-k3") => thinking_enabled != Some(false),
Some("minimax_m2") => {
Self::deepseek_renderer_reasoning_enabled(chat_template_args, true)
}
Some("deepseek_v3" | "deepseek_v3_1") => {
Self::deepseek_renderer_reasoning_enabled(chat_template_args, false)
}
Some(
"deepseek_v3_2" | "deepseek_v4" | "deepseek-v4" | "deepseekv4" | "deepseek_v41",
) => Self::deepseek_renderer_reasoning_enabled(chat_template_args, true),
Some("gemma4" | "gemma-4") => thinking_enabled == Some(true),
Some("mistral") => Self::mistral_reasoning_enabled(chat_template_args),
Some("minimax_m3" | "minimax-m3") => chat_template_args
.and_then(|args| args.get("thinking_mode"))
.and_then(serde_json::Value::as_str)
.is_none_or(|mode| mode != "disabled"),
Some("deepseek_r1" | "step3" | "gpt_oss" | "kimi") => true,
_ => false,
}
}
fn deepseek_renderer_reasoning_enabled(
chat_template_args: Option<&std::collections::HashMap<String, serde_json::Value>>,
default_enabled: bool,
) -> bool {
if let Some(enabled) = dynamo_renderer::thinking_bool_from_args(chat_template_args) {
return enabled;
}
if let Some(mode) = chat_template_args
.and_then(|args| args.get("thinking_mode"))
.and_then(serde_json::Value::as_str)
{
return match mode {
"chat" => false,
"thinking" => true,
_ => default_enabled,
};
}
default_enabled
}
#[cfg(test)]
fn normalize_thinking_arg(
request: &mut NvCreateChatCompletionRequest,
reasoning_parser: Option<&str>,
tool_call_parser: Option<&str>,
) {
let thinking_control_from_client = Self::request_has_client_thinking_control(request);
Self::normalize_thinking_arg_with_source(
request,
reasoning_parser,
tool_call_parser,
thinking_control_from_client,
);
}
fn normalize_thinking_arg_with_source(
request: &mut NvCreateChatCompletionRequest,
reasoning_parser: Option<&str>,
tool_call_parser: Option<&str>,
thinking_control_from_client: bool,
) {
let constrained_generation = !thinking_control_from_client
&& Self::is_minimax_m3_family(reasoning_parser, tool_call_parser)
&& Self::has_minimax_m3_constrained_generation(request);
Self::normalize_thinking_args_with_source(
&mut request.chat_template_args,
reasoning_parser,
tool_call_parser,
thinking_control_from_client,
constrained_generation,
);
}
fn normalize_thinking_args_with_source(
chat_template_args: &mut Option<HashMap<String, serde_json::Value>>,
reasoning_parser: Option<&str>,
tool_call_parser: Option<&str>,
thinking_control_from_client: bool,
constrained_generation: bool,
) {
let normalized =
Self::normalize_thinking_aliases(chat_template_args.as_ref(), reasoning_parser);
if Self::is_minimax_m3_family(reasoning_parser, tool_call_parser) {
Self::normalize_minimax_m3_thinking_mode(
chat_template_args,
normalized,
thinking_control_from_client,
constrained_generation,
);
return;
}
Self::apply_normalized_thinking_aliases(chat_template_args, normalized);
}
fn normalize_thinking_aliases(
chat_template_args: Option<&HashMap<String, serde_json::Value>>,
reasoning_parser: Option<&str>,
) -> Option<bool> {
chat_template_args
.and_then(|args| {
for key in ["thinking", "enable_thinking", "thinking_mode"] {
match args.get(key) {
Some(serde_json::Value::Bool(b)) => {
if key == "thinking_mode" {
return Some(*b);
}
return dynamo_renderer::thinking_bool_from_args(Some(args));
}
Some(serde_json::Value::String(value)) => {
if key == "thinking_mode" {
if value.eq_ignore_ascii_case("enabled") {
return Some(true);
}
if value.eq_ignore_ascii_case("disabled") {
return Some(false);
}
if let Some(b) = parse_bool_opt(value) {
return Some(b);
}
return None;
}
return Some(is_truthy(value));
}
Some(serde_json::Value::Number(value)) => {
if let Some(value) = value.as_f64() {
return Some(value != 0.0);
}
}
_ => {}
}
}
None
})
.or_else(|| {
matches!(reasoning_parser, Some("kimi_k25" | "kimi_k3" | "kimi-k3")).then_some(true)
})
}
fn normalize_minimax_m3_thinking_mode(
chat_template_args: &mut Option<HashMap<String, serde_json::Value>>,
normalized: Option<bool>,
thinking_control_from_client: bool,
constrained_generation: bool,
) {
let explicit_thinking_mode_is_adaptive = chat_template_args
.as_ref()
.and_then(|args| args.get("thinking_mode"))
.and_then(|v| v.as_str())
.is_some_and(|s| s.eq_ignore_ascii_case("adaptive"));
if !thinking_control_from_client && constrained_generation {
let args = chat_template_args.get_or_insert_default();
args.insert(
"thinking_mode".to_string(),
serde_json::Value::String("disabled".to_string()),
);
args.insert("thinking".to_string(), serde_json::Value::Bool(false));
args.insert(
"enable_thinking".to_string(),
serde_json::Value::Bool(false),
);
return;
}
Self::apply_normalized_thinking_aliases(chat_template_args, normalized);
let Some(normalized) = normalized else {
return;
};
if !explicit_thinking_mode_is_adaptive {
let args = chat_template_args.get_or_insert_default();
args.insert(
"thinking_mode".to_string(),
serde_json::Value::String(
if normalized { "enabled" } else { "disabled" }.to_string(),
),
);
}
}
fn apply_normalized_thinking_aliases(
chat_template_args: &mut Option<HashMap<String, serde_json::Value>>,
normalized: Option<bool>,
) {
let Some(normalized) = normalized else {
return;
};
let args = chat_template_args.get_or_insert_default();
args.insert("thinking".to_string(), serde_json::Value::Bool(normalized));
args.insert(
"enable_thinking".to_string(),
serde_json::Value::Bool(normalized),
);
}
fn is_minimax_m3_family(
reasoning_parser: Option<&str>,
tool_call_parser: Option<&str>,
) -> bool {
matches!(reasoning_parser, Some("minimax_m3") | Some("minimax-m3"))
|| matches!(
tool_call_parser,
Some("minimax_m3")
| Some("minimax-m3")
| Some("minimax_m3_nom")
| Some("minimax-m3-nom")
)
}
fn has_minimax_m3_constrained_generation<R: OAIChatLikeRequest>(request: &R) -> bool {
Self::has_structured_response_format(request) || Self::has_guided_tool_choice(request)
}
fn has_guided_tool_choice<R: OAIChatLikeRequest>(request: &R) -> bool {
request.tool_choice().is_some_and(|tool_choice| {
match tool_choice.as_str() {
Some("required") => true,
Some(_) => false,
None => true,
}
})
}
fn normalize_kimi_k3_named_tool_choice(
request: &mut NvCreateChatCompletionRequest,
tool_call_parser: Option<&str>,
) {
let uses_kimi_k3_parser =
tool_call_parser.is_some_and(|parser| matches!(parser, "kimi_k3" | "kimi-k3"));
let is_named = matches!(
request.inner.tool_choice.as_ref(),
Some(ChatCompletionToolChoiceOption::Named(_))
);
if !uses_kimi_k3_parser || !is_named {
return;
}
let args = request.chat_template_args.get_or_insert_default();
args.insert("thinking".to_string(), serde_json::Value::Bool(false));
args.insert(
"enable_thinking".to_string(),
serde_json::Value::Bool(false),
);
}
fn mistral_reasoning_enabled(
chat_template_args: Option<&std::collections::HashMap<String, serde_json::Value>>,
) -> bool {
chat_template_args
.and_then(|args| args.get("reasoning_effort"))
.and_then(serde_json::Value::as_str)
.is_some_and(|effort| effort != "none")
}
pub fn new(mdc: ModelDeploymentCard) -> Result<Arc<Self>> {
let formatter = prompt_formatter_from_mdc(&mdc)?;
let tokenizer = mdc.tokenizer()?;
match formatter {
PromptFormatter::OAI(formatter) => Self::new_with_parts(mdc, formatter, tokenizer),
}
}
pub fn new_for_embeddings(mdc: ModelDeploymentCard) -> Result<Arc<Self>> {
if !mdc.model_type.supports_embedding() {
anyhow::bail!("embedding preprocessor requires an embedding-capable model");
}
let tokenizer = mdc.tokenizer()?;
let PromptFormatter::OAI(formatter) = embedding_prompt_formatter(&mdc)?;
let embedding_tokenizers = EmbeddingTokenizerState::new(&mdc)?;
Self::new_with_parts_inner(mdc, formatter, tokenizer, Some(embedding_tokenizers), None)
}
pub fn new_with_parts(
mdc: ModelDeploymentCard,
formatter: Arc<dyn OAIPromptFormatter>,
tokenizer: crate::tokenizers::Tokenizer,
) -> Result<Arc<Self>> {
Self::new_with_parts_and_cancel(mdc, formatter, tokenizer, None)
}
pub fn new_with_parts_and_cancel(
mdc: ModelDeploymentCard,
formatter: Arc<dyn OAIPromptFormatter>,
tokenizer: crate::tokenizers::Tokenizer,
speculative_prefill_cancel: Option<CancellationToken>,
) -> Result<Arc<Self>> {
Self::new_with_parts_inner(mdc, formatter, tokenizer, None, speculative_prefill_cancel)
}
fn new_with_parts_inner(
mdc: ModelDeploymentCard,
formatter: Arc<dyn OAIPromptFormatter>,
tokenizer: crate::tokenizers::Tokenizer,
embedding_tokenizers: Option<EmbeddingTokenizerState>,
speculative_prefill_cancel: Option<CancellationToken>,
) -> Result<Arc<Self>> {
let mdcsum = mdc.mdcsum().to_string();
let tokenizer: Arc<dyn Tokenizer> = (*tokenizer).clone();
let lora_name = mdc.lora.as_ref().map(|l| l.name.clone());
let Some(ref model_info) = mdc.model_info else {
anyhow::bail!(
"Blank ModelDeploymentCard cannot be used for pre-processing, no model_info"
);
};
let model_info = model_info.get_model_info()?;
let tokenizer_id_bound = mdc.tokenizer_id_bound().unwrap_or_else(|| {
tracing::warn!(
model = %mdc.display_name,
"Tokenizer not loaded from this model card: only vocab_size bounds client token ids"
);
None
});
let token_id_bound = token_id_bound(model_info.vocab_size(), tokenizer_id_bound);
let tool_call_parser = mdc.runtime_config.tool_call_parser.clone();
if mdc.model_type.supports_chat() {
crate::protocols::openai::chat_completions::tool_parser_v2::validate_parser_version(
tool_call_parser.as_deref(),
mdc.runtime_config.reasoning_parser.as_deref(),
)?;
tracing::info!(
model = %mdc.display_name,
parser_version = ?crate::protocols::openai::chat_completions::tool_parser_v2::selected_version()?,
unified_family = ?crate::protocols::openai::chat_completions::unified_parser::selected_family(
tool_call_parser.as_deref(),
mdc.runtime_config.reasoning_parser.as_deref(),
),
tool_parser = ?tool_call_parser,
reasoning_parser = ?mdc.runtime_config.reasoning_parser,
"Selected model parser route"
);
}
let parser_requires_special_tokens = mdc.model_type.supports_chat()
&& Self::parser_requires_special_tokens(
tool_call_parser.as_deref(),
mdc.runtime_config.reasoning_parser.as_deref(),
)?;
let normalize_tool_call_args = mdc.runtime_config.tool_call_arguments_format
== crate::local_model::runtime_config::ToolCallArgumentsFormat::JsonObject
|| mdc.runtime_config.tool_call_parser.as_deref() == Some("glm47");
if let Some(ref lora_name) = lora_name {
tracing::info!(model = %mdc.display_name, lora_name, "LoRA adapter detected in MDC");
}
let runtime_config = mdc.runtime_config.clone();
let structural_tag_reasoning_boundary =
structural_tag::validate_runtime_config(&runtime_config)?;
let token_budget = match runtime_config
.get_engine_specific::<TokenBudget>(TOKEN_BUDGET_RUNTIME_KEY)
{
Ok(token_budget) => token_budget,
Err(error) => {
tracing::warn!(
%error,
key = TOKEN_BUDGET_RUNTIME_KEY,
"Ignoring invalid runtime metadata; token overflow handling will be delegated to the backend"
);
None
}
};
let kv_cache_block_size = mdc.kv_cache_block_size as usize;
let context_length = mdc.effective_context_length();
#[cfg(feature = "mm-routing")]
let model_dir_for_routing: Option<std::path::PathBuf> = mdc_model_dir(&mdc);
#[cfg(feature = "mm-routing")]
let fastokens_active = runtime_config.effective_tokenizer_backend().is_fastokens();
#[cfg(feature = "mm-routing")]
if fastokens_active && model_dir_for_routing.is_some() {
tracing::warn!(
target: "mm_routing",
"fastokens tokenizer backend is active; MM-aware KV routing disabled. \
Image-token metrics remain enabled when SMG supports the model."
);
}
#[cfg(feature = "mm-routing")]
let image_token_inputs: Option<(String, String, std::path::PathBuf)> =
model_dir_for_routing.as_ref().map(|p| {
(
mdc.source_path().to_string(),
model_info.model_type(),
p.clone(),
)
});
let media_loader = match mdc.media_decoder {
Some(media_decoder) => Some(MediaLoader::new(media_decoder, mdc.media_fetcher)?),
None => None,
};
#[cfg(feature = "mm-routing")]
let (
image_token_counter,
routing_image_token_id,
routing_image_prompt_layout,
bos_token_string,
) = match image_token_inputs.as_ref() {
Some((model_id, model_type, model_dir)) => {
let (counter, counter_err): (
Option<mm_routing::image::ImageRoutingProcessor>,
Option<String>,
) = match mm_routing::image::ImageRoutingProcessor::try_new_for_runtime(
model_id,
Some(model_type),
model_dir,
if runtime_has_vllm_generate_capability(&runtime_config) {
mm_routing::image::ImageRoutingRuntime::VllmNativeGenerate
} else {
mm_routing::image::ImageRoutingRuntime::BackendNeutral
},
) {
Ok(c) => (Some(c), None),
Err(e) => (None, Some(e.to_string())),
};
let (img_tok, prompt_layout, bos_tok_string) = if fastokens_active {
(None, None, None)
} else {
let routing_tokens = mm_routing::image::resolve_routing_tokens(
model_id,
model_dir,
counter.as_ref(),
);
let prompt_layout =
routing_tokens.image_prompt_kind.and_then(|kind| {
match resolve_routing_image_prompt_layout(
tokenizer.as_ref(),
kind,
routing_tokens.chat_placeholder_token_id,
) {
Ok(layout) => Some(layout),
Err(e) => {
tracing::warn!(
target: "mm_routing",
model = %model_id,
error = %e,
"model-specific routing image prompt could not be resolved; exact MM routing disabled"
);
None
}
}
});
let exact_mm_routing_ready = counter.is_some() && prompt_layout.is_some();
(
routing_tokens.exact_routing_image_token_id(exact_mm_routing_ready),
prompt_layout.filter(|_| counter.is_some()),
routing_tokens.bos_token_string,
)
};
match (counter.is_some(), img_tok.is_some()) {
(true, true) => tracing::info!(
target: "mm_routing",
model = %model_id,
model_dir = %model_dir.display(),
"MM-aware KV routing enabled"
),
_ if fastokens_active => {}
(counter_ok, img_ok) => {
let mut reasons: Vec<String> = Vec::new();
if !counter_ok {
reasons.push(format!(
"image routing processor unavailable ({})",
counter_err.as_deref().unwrap_or("unknown error")
));
}
if !img_ok {
reasons.push(
"image-placeholder token or model-specific prompt layout \
unresolvable from \
config.json / processor_config.json / \
tokenizer_config.json / vocab probe"
.to_string(),
);
}
tracing::warn!(
target: "mm_routing",
model = %model_id,
reasons = %reasons.join("; "),
"{} is not supported for MM-aware KV routing ({}). \
Falling back to KV routing without MM awareness — \
text-prefix overlap still works but the router \
cannot distinguish requests by image content.",
model_id,
reasons.join("; ")
);
}
}
(counter, img_tok, prompt_layout, bos_tok_string)
}
None => {
tracing::debug!(
target: "mm_routing",
"model directory not derivable from MDC; MM-aware routing disabled"
);
(None, None, None, None)
}
};
#[cfg(all(feature = "mm-routing", feature = "media-ffmpeg"))]
let video_routing_processor = {
let qwen_contract = match resolve_qwen_video_processor_contract(&runtime_config) {
Ok(contract) => contract,
Err(error) => {
tracing::warn!(
target: "mm_routing",
%error,
"invalid Qwen video processor runtime metadata; exact video routing disabled"
);
None
}
};
let nemotron_contract = match runtime_config
.get_engine_specific::<mm_routing::NemotronVideoProcessorContract>(
VLLM_NEMOTRON_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY,
) {
Ok(target) => target,
Err(error) => {
tracing::warn!(
target: "mm_routing",
%error,
key = VLLM_NEMOTRON_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY,
"invalid Nemotron video processor runtime metadata; exact video routing disabled"
);
None
}
};
let contracts = mm_routing::VideoProcessorContracts {
qwen: qwen_contract,
nemotron: nemotron_contract,
};
if fastokens_active {
None
} else {
match (
image_token_inputs.as_ref(),
contracts.qwen.is_some() || contracts.nemotron.is_some(),
) {
(Some((model_id, model_type, model_dir)), true) => {
match mm_routing::VideoRoutingProcessor::try_new(
model_id,
model_type,
model_dir,
tokenizer.clone(),
contracts,
) {
Ok(processor) => processor,
Err(error) => {
tracing::warn!(
target: "mm_routing",
model = %model_id,
%error,
"exact video-aware KV routing disabled for this model"
);
None
}
}
}
_ => None,
}
}
};
#[cfg(all(feature = "mm-routing", feature = "media-ffmpeg"))]
if video_routing_processor.is_some()
&& let Some((model_id, model_type, _)) = image_token_inputs.as_ref()
{
tracing::info!(
target: "mm_routing",
model = %model_id,
model_type = %model_type,
"exact video-aware KV routing enabled"
);
}
#[cfg(feature = "mm-routing")]
let routing_image_dimension_policy = routing_image_dimension_policy(
&runtime_config,
media_loader.is_some(),
routing_image_prompt_layout,
);
#[cfg(feature = "mm-routing")]
if image_token_counter.is_some() || routing_image_token_id.is_some() {
std::sync::LazyLock::force(&DIM_FETCH_MEDIA_FETCHER);
std::sync::LazyLock::force(&DIM_FETCH_HTTP_CLIENT);
}
#[cfg(feature = "mm-routing")]
let routing_prepend_bos = match bos_token_string {
Some(bos_text) => match tokenizer.encode(&bos_text) {
Ok(enc) if enc.token_ids().len() == 1 => {
let id = enc.token_ids()[0];
tracing::debug!(
target: "mm_routing",
bos_token = %bos_text,
bos_token_id = id,
"routing-side BOS prepend enabled (tokenizer_config.json add_bos_token=true)"
);
Some(id)
}
Ok(enc) => {
tracing::debug!(
target: "mm_routing",
bos_token = %bos_text,
round_trip_ids = ?enc.token_ids(),
"BOS token does not round-trip to a single id; routing-side prepend disabled"
);
None
}
Err(e) => {
tracing::debug!(
target: "mm_routing",
bos_token = %bos_text,
error = %e,
"BOS token failed to re-encode; routing-side prepend disabled"
);
None
}
},
None => None,
};
Ok(Arc::new(Self {
formatter,
tokenizer,
embedding_tokenizers,
model_info,
token_id_bound,
mdcsum,
lora_name,
runtime_config,
structural_tag_reasoning_boundary,
kv_cache_block_size,
tool_call_parser,
parser_requires_special_tokens,
normalize_tool_call_args,
media_loader,
token_budget,
context_length,
speculative_prefill_tasks: speculative_prefill::PrefillTasks::new(
speculative_prefill_cancel.as_ref(),
),
#[cfg(feature = "mm-routing")]
image_token_counter,
#[cfg(all(feature = "mm-routing", feature = "media-ffmpeg"))]
video_routing_processor,
#[cfg(feature = "mm-routing")]
routing_image_token_id,
#[cfg(feature = "mm-routing")]
routing_image_prompt_layout,
#[cfg(feature = "mm-routing")]
routing_image_dimension_policy,
#[cfg(feature = "mm-routing")]
routing_prepend_bos,
}))
}
pub fn tokenize(&self, s: &str) -> anyhow::Result<Encoding> {
self.tokenizer.encode(s)
}
pub fn tokenize_rendered_prompt(&self, prompt: &RenderedPrompt) -> anyhow::Result<Encoding> {
match prompt.encode_segments() {
Some(segments) => self.tokenizer.encode_segments(&segments),
None => self.tokenize(prompt.as_str()),
}
}
pub async fn preprocess_request<
R: OAIChatLikeRequest
+ MediaRequestExt
+ AnnotationsProvider
+ SamplingOptionsProvider
+ StopConditionsProvider
+ OutputOptionsProvider
+ NvExtProvider
+ CommonExtProvider,
>(
&self,
request: &R,
tracker: Option<&RequestTracker>,
) -> Result<(
PreprocessedRequest,
HashMap<String, String>,
PromptReasoningPrefill,
)> {
let (request, annotations, prompt_injected_reasoning, _image_tokens) = self
.preprocess_request_with_options(
request,
tracker,
PreprocessRequestOptions::default(),
None,
)
.await?;
Ok((request, annotations, prompt_injected_reasoning))
}
async fn preprocess_request_with_options<
R: OAIChatLikeRequest
+ MediaRequestExt
+ AnnotationsProvider
+ SamplingOptionsProvider
+ StopConditionsProvider
+ OutputOptionsProvider
+ NvExtProvider
+ CommonExtProvider,
>(
&self,
request: &R,
tracker: Option<&RequestTracker>,
options: PreprocessRequestOptions,
lora_name: Option<String>,
) -> Result<(
PreprocessedRequest,
HashMap<String, String>,
PromptReasoningPrefill,
Option<usize>,
)> {
let _stage_guard = StageGuard::new(STAGE_PREPROCESS, "");
let preprocess_start = Instant::now();
let mut builder = self.builder_with_lora(request, lora_name)?;
let template_start = Instant::now();
let formatted_prompt = {
let _nvtx = dynamo_nvtx_range!("preprocess.template");
off_runtime(|| self.apply_template(request))
.with_context(|| "Failed to apply prompt template")?
};
TEMPLATE_SECONDS.observe(template_start.elapsed().as_secs_f64());
let prompt_injected_reasoning = Self::prompt_reasoning_prefill_for_parsers(
self.tool_call_parser.as_deref(),
self.runtime_config.reasoning_parser.as_deref(),
formatted_prompt.as_ref().map(RenderedPrompt::as_str),
);
let tokenize_start = Instant::now();
let (token_ids, annotations) = {
let _nvtx = dynamo_nvtx_range!("preprocess.tokenize");
self.gather_tokens(request, formatted_prompt.as_ref(), tracker)
.await
.with_context(|| "Failed to gather tokens")?
};
TOKENIZE_SECONDS.observe(tokenize_start.elapsed().as_secs_f64());
let (_mm_routing_entries, image_tokens) = self
.gather_multi_modal_data_with_image_tokens(
request,
&mut builder,
formatted_prompt.as_ref().map(RenderedPrompt::as_str),
&token_ids,
)
.await
.with_context(|| "Failed to gather multimodal data")?;
builder.token_ids(token_ids);
STAGE_DURATION_SECONDS
.with_label_values(&[STAGE_PREPROCESS])
.observe(preprocess_start.elapsed().as_secs_f64());
if let Some(nvext) = request.nvext()
&& let Some(router_params) = &nvext.router
{
builder.router(Some(router_params.clone()));
}
let mut preprocessed = builder.build()?;
if let Some(reasoning_ended) = Self::prompt_injected_reasoning_ended_arg(
self.runtime_config.reasoning_parser.as_deref(),
formatted_prompt.as_ref().map(RenderedPrompt::as_str),
) {
let extra_args = preprocessed
.extra_args
.get_or_insert_with(|| serde_json::json!({}));
let extra_args = extra_args
.as_object_mut()
.context("preprocessed extra_args must be an object")?;
extra_args.insert("reasoning_ended".to_string(), reasoning_ended.into());
}
let exact_prompt_len =
Self::validate_preprocessed_token_budget(&preprocessed, self.token_budget.as_ref())?;
if preprocessed.stop_conditions.max_tokens.is_none()
&& let Some(prompt_len) = exact_prompt_len
&& !self.token_budget.is_some_and(|token_budget| {
prompt_len >= token_budget.combined_limit as usize
&& !token_budget.reject_prompt_overflow
})
&& let Some(max_tokens) = Self::omitted_max_tokens_default(
prompt_len,
self.token_budget
.map(|token_budget| token_budget.combined_limit),
options,
)
{
preprocessed.stop_conditions.max_tokens = Some(max_tokens);
Self::validate_preprocessed_token_budget(&preprocessed, self.token_budget.as_ref())?;
}
Ok((
preprocessed,
annotations,
prompt_injected_reasoning,
image_tokens,
))
}
pub fn builder<
R: OAIChatLikeRequest
+ MediaRequestExt
+ AnnotationsProvider
+ SamplingOptionsProvider
+ StopConditionsProvider
+ OutputOptionsProvider
+ NvExtProvider,
>(
&self,
request: &R,
) -> Result<PreprocessedRequestBuilder> {
self.builder_with_lora(request, None)
}
fn builder_with_lora<
R: OAIChatLikeRequest
+ MediaRequestExt
+ AnnotationsProvider
+ SamplingOptionsProvider
+ StopConditionsProvider
+ OutputOptionsProvider
+ NvExtProvider,
>(
&self,
request: &R,
lora_name_override: Option<String>,
) -> Result<PreprocessedRequestBuilder> {
let mut builder = PreprocessedRequest::builder();
builder.model(request.model());
let mut stop_conditions = request.extract_stop_conditions()?;
let eos_token_ids = self.model_info.eos_token_ids();
let hidden_eos_token_ids = eos_token_ids.clone();
if let Some(stop_tokens) = &mut stop_conditions.stop_token_ids_hidden {
for eos_token_id in hidden_eos_token_ids {
if !stop_tokens.contains(&eos_token_id) {
stop_tokens.push(eos_token_id);
}
}
} else {
stop_conditions.stop_token_ids_hidden = Some(hidden_eos_token_ids);
}
let mut visible_tool_parser_end_token_ids = Vec::new();
if let Some(stop_tokens) = &mut stop_conditions.stop_token_ids_hidden {
visible_tool_parser_end_token_ids =
self.remove_tool_parser_end_tokens_from_hidden_stops(request, stop_tokens)?;
}
if !visible_tool_parser_end_token_ids.is_empty() {
let visible_stops = stop_conditions
.stop_token_ids_visible
.get_or_insert_with(Vec::new);
for token_id in visible_tool_parser_end_token_ids {
if !visible_stops.contains(&token_id) {
visible_stops.push(token_id);
}
}
}
stop_conditions.apply_ignore_eos();
if !stop_conditions.ignore_eos.unwrap_or(false) {
builder.eos_token_ids(eos_token_ids);
}
let has_thinking_budget = stop_conditions.max_thinking_tokens.is_some();
if has_thinking_budget
&& self
.runtime_config
.supports_runtime_capability(SGLANG_GENERATE_CAPABILITY)
&& Self::has_structured_response_format(request)
&& !Self::structured_response_supports_sglang_reasoning_gate(
self.runtime_config.reasoning_parser.as_deref(),
)
{
return Err(invalid_argument_error(
"thinking_token_budget is not supported with GPT-OSS structured output on SGLang",
));
}
builder.stop_conditions(stop_conditions);
builder.sampling_options(request.extract_sampling_options()?);
let mut output_options = request.extract_output_options()?;
if output_options.skip_special_tokens.is_none() && self.parser_requires_special_tokens {
output_options.skip_special_tokens = Some(false);
} else if Self::special_tokens_will_be_stripped(
output_options.skip_special_tokens,
self.parser_requires_special_tokens,
) {
static WARNED: std::sync::Once = std::sync::Once::new();
WARNED.call_once(|| {
tracing::warn!(
tool_call_parser = ?self.tool_call_parser,
reasoning_parser = ?self.runtime_config.reasoning_parser,
"skip_special_tokens=true requested while a special-token-dependent parser is active; the engine's special tokens will be stripped before parsing, so tool_calls/reasoning_content will be empty and the markup will leak into content. Unset skip_special_tokens (Dynamo defaults it to false for these parsers) or set it to false."
);
});
}
builder.output_options(output_options);
builder.annotations(request.annotations().unwrap_or_default());
builder.mdc_sum(Some(self.mdcsum.clone()));
let lora_name = self.lora_name.clone().or(lora_name_override);
let cache_namespace = request_cache_salt(request).map(str::to_owned);
if let Some(nvext) = request.nvext() {
let hints = nvext.agent_hints.as_ref();
let (priority_jump, strict_priority, priority) = routing_priorities(hints);
builder.request_timestamp_ms(nvext.request_timestamp_ms);
let routing = RoutingHints {
backend_instance_id: nvext.backend_instance_id,
prefill_worker_id: nvext.prefill_worker_id,
decode_worker_id: nvext.decode_worker_id,
dp_rank: nvext.dp_rank,
prefill_dp_rank: nvext.prefill_dp_rank,
expected_output_tokens: hints.and_then(|h| h.osl),
priority_jump,
strict_priority,
priority,
lora_name,
cache_namespace: cache_namespace.clone(),
allowed_worker_ids: None,
routing_constraints: nvext
.routing_constraints
.clone()
.map(routing_constraints_to_kv),
};
builder.routing(Some(routing));
} else if lora_name.is_some() || cache_namespace.is_some() {
builder.routing(Some(RoutingHints {
lora_name,
cache_namespace,
..Default::default()
}));
}
if let Some(extra_args) = Self::backend_extra_args(
request,
self.runtime_config.reasoning_parser.is_some(),
None,
) {
builder.extra_args(Some(extra_args));
}
builder.require_reasoning(Self::request_requires_reasoning(
request,
self.runtime_config.reasoning_parser.as_deref(),
has_thinking_budget,
));
builder.mm_processor_kwargs(request.mm_processor_kwargs().cloned());
if self.media_loader.is_none() {
builder.media_io_kwargs(request.media_io_kwargs().cloned());
}
Ok(builder)
}
fn remove_tool_parser_end_tokens_from_hidden_stops<R: OAIChatLikeRequest>(
&self,
request: &R,
hidden_stop_token_ids: &mut Vec<TokenIdType>,
) -> Result<Vec<TokenIdType>> {
let has_tools = Self::request_has_effective_tools(request);
let tool_choice_none = request
.tool_choice()
.as_ref()
.and_then(|tool_choice| tool_choice.as_str())
== Some("none");
if !Self::should_keep_tool_parser_end_tokens_visible(has_tools, tool_choice_none) {
return Ok(Vec::new());
}
let Some(tool_call_parser) = self.tool_call_parser.as_deref().filter(|p| !p.is_empty())
else {
return Ok(Vec::new());
};
let Some(tool_call_config) = get_tool_parser_map().get(tool_call_parser) else {
return Ok(Vec::new());
};
let mut visible_stop_token_ids = Vec::new();
for end_token in tool_call_config.parser_config.tool_call_end_tokens() {
if end_token.is_empty() {
continue;
}
let encoded = self.tokenizer.encode(&end_token).with_context(|| {
format!(
"Failed to encode tool-call end token {end_token:?} for parser {tool_call_parser:?}"
)
})?;
let was_hidden_eos =
Self::remove_single_token_marker(hidden_stop_token_ids, encoded.token_ids());
if !was_hidden_eos {
tracing::debug!(
token_ids = ?encoded.token_ids(),
end_token,
parser = tool_call_parser,
"Tool-call end token was not a single hidden EOS token"
);
continue;
}
if let [token_id] = encoded.token_ids()
&& !visible_stop_token_ids.contains(token_id)
{
visible_stop_token_ids.push(*token_id);
}
}
Ok(visible_stop_token_ids)
}
fn request_has_effective_tools<R: OAIChatLikeRequest>(request: &R) -> bool {
request
.tools()
.as_ref()
.and_then(|tools| tools.len())
.is_some_and(|len| len > 0)
|| request.typed_messages().is_some_and(|messages| {
messages.iter().any(|message| {
matches!(
message,
dynamo_protocols::types::ChatCompletionRequestMessage::System(system)
if system.tools.as_ref().is_some_and(|tools| !tools.is_empty())
)
})
})
}
fn should_keep_tool_parser_end_tokens_visible(has_tools: bool, tool_choice_none: bool) -> bool {
has_tools && !tool_choice_none
}
fn remove_single_token_marker(
hidden_eos_token_ids: &mut Vec<TokenIdType>,
marker_token_ids: &[TokenIdType],
) -> bool {
let [marker_token_id] = marker_token_ids else {
return false;
};
let before = hidden_eos_token_ids.len();
hidden_eos_token_ids.retain(|token_id| token_id != marker_token_id);
hidden_eos_token_ids.len() != before
}
fn map_prompt_render_error(error: anyhow::Error) -> anyhow::Error {
tracing::debug!(?error, "Chat template rendering failed");
let message = match error.downcast_ref::<PromptRenderError>() {
Some(PromptRenderError::InvalidRequest(message)) => message.clone(),
None => format!("{error:#}"),
};
invalid_argument_error(message)
}
pub fn apply_template<
R: OAIChatLikeRequest
+ AnnotationsProvider
+ SamplingOptionsProvider
+ StopConditionsProvider
+ OutputOptionsProvider
+ NvExtProvider
+ CommonExtProvider,
>(
&self,
request: &R,
) -> Result<Option<RenderedPrompt>> {
let continue_final = request.get_continue_final_message() == Some(true);
if continue_final
&& self
.formatter
.media_message_order(request)
.is_some_and(|order| {
order
.last()
.is_some_and(|&source| source != order.len() - 1)
})
{
return Err(invalid_argument_error(
"Cannot continue the final message because the prompt formatter moves it before other messages",
));
}
let formatted_prompt = if self.normalize_tool_call_args || continue_final {
self.apply_template_inner(&NormalizedArgsRequest {
inner: request,
normalize_tool_call_args: self.normalize_tool_call_args,
continue_final_message: continue_final,
})?
} else {
self.apply_template_inner(request)?
};
let Some(prompt) = formatted_prompt else {
return Ok(None);
};
if !continue_final {
return Ok(Some(prompt));
}
apply_continue_final_message(prompt)
.map_err(|error| invalid_argument_error(format!("{error:#}")))
.map(Some)
}
fn apply_template_inner<
R: OAIChatLikeRequest
+ AnnotationsProvider
+ SamplingOptionsProvider
+ StopConditionsProvider
+ OutputOptionsProvider
+ NvExtProvider,
>(
&self,
request: &R,
) -> Result<Option<RenderedPrompt>> {
if let PromptInput::Text(_) = request.prompt_input_type()
&& let Some(TextInput::Single(_)) = request.extract_text()
{
let use_raw_prompt = request
.nvext()
.is_some_and(|ext| ext.use_raw_prompt.unwrap_or(false));
let formatted_prompt = if use_raw_prompt {
match request.raw_prompt() {
Some(prompt) => RenderedPrompt::text(prompt),
None => {
tracing::warn!("Raw prompt requested but not available");
self.formatter
.render_prompt(request)
.map_err(Self::map_prompt_render_error)?
}
}
} else {
self.formatter
.render_prompt(request)
.map_err(Self::map_prompt_render_error)?
};
Ok(Some(formatted_prompt))
} else {
Ok(None)
}
}
fn strip_inline_data_urls(messages: &mut serde_json::Value) {
let Some(arr) = messages.as_array_mut() else {
return;
};
for msg in arr {
let Some(content) = msg.get_mut("content") else {
continue;
};
let Some(parts) = content.as_array_mut() else {
continue;
};
for part in parts {
for key in ["image_url", "video_url", "audio_url"] {
if let Some(media) = part.get_mut(key)
&& let Some(url) = media.get_mut("url")
&& url.as_str().is_some_and(|s| s.starts_with("data:"))
{
*url = serde_json::Value::String(String::new());
}
}
}
}
}
fn replace_reserved_media_slot(
media_map: &mut MultimodalDataMap,
modality: &str,
slot_idx: usize,
value: MultimodalData,
) -> Result<()> {
let slot = media_map
.get_mut(modality)
.and_then(|slots| slots.get_mut(slot_idx))
.with_context(|| {
format!(
"missing reserved multimodal slot {modality}[{slot_idx}] during media decode"
)
})?;
*slot = value;
Ok(())
}
pub async fn gather_multi_modal_data<
R: OAIChatLikeRequest + MediaRequestExt + NvExtProvider,
>(
&self,
request: &R,
builder: &mut PreprocessedRequestBuilder,
formatted_prompt: Option<&str>,
token_ids: &[crate::protocols::TokenIdType],
) -> Result<Vec<MmImageEntry>> {
let (entries, _image_tokens) = self
.gather_multi_modal_data_with_image_tokens(
request,
builder,
formatted_prompt,
token_ids,
)
.await?;
Ok(entries
.into_iter()
.filter_map(|entry| match entry {
MmRoutingEntry::Image {
mm_hash,
width,
height,
} => Some(MmImageEntry {
mm_hash,
width,
height,
}),
MmRoutingEntry::Video { .. } => None,
})
.collect())
}
async fn gather_multi_modal_data_with_image_tokens<
R: OAIChatLikeRequest + MediaRequestExt + NvExtProvider,
>(
&self,
request: &R,
builder: &mut PreprocessedRequestBuilder,
formatted_prompt: Option<&str>,
token_ids: &[crate::protocols::TokenIdType],
) -> Result<(Vec<MmRoutingEntry>, Option<usize>)> {
#[cfg(not(feature = "mm-routing"))]
let _ = token_ids;
let mut media_map: MultimodalDataMap = HashMap::new();
let mut uuid_map: MultimodalUuidMap = HashMap::new();
let mut has_user_uuid = false;
let mut fetch_tasks: Vec<MediaFetchTask<'_>> = Vec::new();
#[cfg(feature = "mm-routing")]
let mut mm_routing_entries: Vec<MmRoutingEntry> = Vec::new();
#[cfg(feature = "mm-routing")]
let mut image_tokens = self.image_token_counter.as_ref().and_then(|counter| {
(!counter.uses_request_context_budget()).then_some(0usize)
});
#[cfg(feature = "mm-routing")]
let mut exact_mm_routing_eligible = true;
#[cfg(feature = "mm-routing")]
let mut previous_routing_entry_was_video = false;
#[cfg(feature = "mm-routing")]
let mut total_image_count: usize = 0;
#[cfg(feature = "mm-routing")]
let mut url_passthrough_images: Vec<(u64, String)> = Vec::new();
let Some(messages) = request.typed_messages() else {
return Ok((Vec::new(), None));
};
let has_media_loader = self.media_loader.is_some();
let message_order = self.formatter.media_message_order(request);
for index in 0..messages.len() {
let message = &messages[message_order.as_ref().map_or(index, |order| order[index])];
let Some(content_parts) = multimodal_content_parts(message) else {
continue;
};
for content_part in content_parts {
let Some((type_str, url, uuid)) = content_part.media_info() else {
continue;
};
#[cfg(feature = "mm-routing")]
if type_str == "image_url" {
total_image_count += 1;
}
#[cfg(feature = "mm-routing")]
if !exact_mm_routing_supports_modality(type_str, has_media_loader)
|| !exact_mm_routing_layout_accepts_next_entry(
&mut previous_routing_entry_was_video,
type_str == "video_url",
)
{
exact_mm_routing_eligible = false;
}
if uuid.as_deref().is_some_and(str::is_empty) {
return Err(invalid_argument_error(format!(
"{type_str} uuid must be a non-empty string"
)));
}
let slots = media_map.entry(type_str.to_string()).or_default();
let slot_idx = slots.len();
has_user_uuid |= uuid.is_some();
uuid_map
.entry(type_str.to_string())
.or_default()
.push(uuid.clone());
match (url, uuid) {
(Some(url), _) => {
if url.scheme() == "data" {
let size = url.as_str().len();
let limit = max_data_url_bytes();
if size > limit {
let message = format!(
"{type_str} data: URL is {size} bytes, exceeds the {limit}-byte limit. \
To raise the limit, set DYN_MM_MAX_DATA_URL_MB (in megabytes) on \
both the frontend and the workers."
);
return Err(crate::protocols::common::invalid_argument_error(
message,
));
}
}
if has_media_loader {
fetch_tasks.push(MediaFetchTask {
modality: type_str,
slot_idx,
content_part: content_part.as_user(),
});
} else {
#[cfg(feature = "mm-routing")]
if type_str == "image_url" {
let mm_hash = Self::hash_image_url(url.as_str());
url_passthrough_images.push((mm_hash, url.as_str().to_string()));
}
}
slots.push(MultimodalData::Url(url));
}
(None, Some(uuid)) if type_str == "image_url" => {
slots.push(MultimodalData::UuidOnly(uuid));
}
(None, Some(_)) => {
return Err(invalid_argument_error(format!(
"UUID-only cache reuse is not supported for media modality `{type_str}`; provide a media URL"
)));
}
(None, None) => {
return Err(invalid_argument_error(format!(
"{type_str} part has neither `url` nor `uuid`; at least one is required"
)));
}
}
}
}
#[cfg(feature = "mm-routing")]
let has_processor_override = has_mm_processor_override(request.mm_processor_kwargs());
#[cfg(all(feature = "mm-routing", feature = "media-ffmpeg"))]
let hash_decoded_video = should_hash_decoded_video(
exact_mm_routing_eligible,
has_user_uuid,
has_processor_override,
self.video_routing_processor.is_some(),
);
#[cfg(not(all(feature = "mm-routing", feature = "media-ffmpeg")))]
let hash_decoded_video = false;
if !fetch_tasks.is_empty() {
let loader = self.media_loader.as_ref().unwrap();
let media_io_kwargs = request
.media_io_kwargs()
.map(<MediaDecoder as serde::Deserialize>::deserialize)
.transpose()
.map_err(|e| invalid_argument_error(format!("invalid media_io_kwargs: {e}")))?;
let results = futures::future::join_all(fetch_tasks.iter().map(|task| {
loader.fetch_and_decode_media_part_with_video_hash(
task.content_part.as_ref(),
media_io_kwargs.as_ref(),
hash_decoded_video,
)
}))
.await;
for (task, result) in fetch_tasks.into_iter().zip(results) {
let rdma_descriptor = result?;
#[cfg(feature = "mm-routing")]
if task.modality == "image_url" {
let shape = &rdma_descriptor.tensor_info.shape;
if shape.len() >= 2 {
let h = shape[0] as u32;
let w = shape[1] as u32;
let (mm_hash, hash_source) = match rdma_descriptor.content_hash() {
Some(h) => (h, "decoded_bytes"),
None => {
let source_url = image_content_part_url(task.content_part.as_ref())
.ok_or_else(|| {
invalid_argument_error(
"image_url task must contain an image URL",
)
})?;
(Self::hash_image_url(source_url), "url_fallback")
}
};
if let Some(counter) = self
.image_token_counter
.as_ref()
.filter(|counter| !counter.uses_request_context_budget())
{
let n = counter.count_tokens(w, h);
tracing::debug!(
target: "mm_routing",
model = counter.model_id(),
width = w,
height = h,
tokens = n,
mm_hash = mm_hash,
source = hash_source,
"image-token count"
);
image_tokens = checked_add_image_tokens(image_tokens, n);
}
mm_routing_entries.push(MmRoutingEntry::Image {
mm_hash,
width: w,
height: h,
});
}
} else if task.modality == "video_url" {
#[cfg(feature = "media-ffmpeg")]
{
let video_entry = (|| -> Result<MmRoutingEntry> {
let processor =
self.video_routing_processor.as_ref().ok_or_else(|| {
anyhow::anyhow!("model has no exact video routing processor")
})?;
anyhow::ensure!(
!has_processor_override,
"request mm_processor_kwargs can change the video token layout"
);
let (frame_count, width, height) =
rdma_descriptor.video_dimensions()?;
let metadata = rdma_descriptor.video_metadata()?;
let mm_hash = rdma_descriptor.video_content_hash()?;
let routing =
processor.build_replacement(&mm_routing::VideoRoutingInput {
frame_count,
width,
height,
source_fps: metadata.source_fps,
sampled_timestamps: &metadata.sampled_timestamps,
})?;
tracing::debug!(
target: "mm_routing",
n_frames = frame_count,
width,
height,
replacement_tokens = routing.replacement_tokens.len(),
mm_hash,
"video routing metadata resolved"
);
Ok(MmRoutingEntry::Video {
mm_hash,
placeholder_token_id: routing.placeholder_token_id,
event_video_token_id: routing.event_video_token_id,
target_tokens: routing.target_tokens,
replacement_tokens: routing.replacement_tokens,
})
})();
match video_entry {
Ok(entry) => mm_routing_entries.push(entry),
Err(error) => {
exact_mm_routing_eligible = false;
tracing::debug!(
target: "mm_routing",
%error,
"mm-routing: decoded video is not eligible for exact routing; falling back to text-prefix routing"
);
}
}
}
#[cfg(not(feature = "media-ffmpeg"))]
{
exact_mm_routing_eligible = false;
}
}
Self::replace_reserved_media_slot(
&mut media_map,
task.modality,
task.slot_idx,
MultimodalData::Decoded(rdma_descriptor),
)?;
}
}
#[cfg(feature = "mm-routing")]
if !has_user_uuid && !url_passthrough_images.is_empty() {
let dim_results =
futures::future::join_all(url_passthrough_images.iter().map(|(mm_hash, url)| {
Self::fetch_image_dims(*mm_hash, url, self.routing_image_dimension_policy)
}))
.await;
for ((mm_hash, url), dim_res) in url_passthrough_images.into_iter().zip(dim_results) {
match dim_res {
Ok((w, h)) => {
if let Some(counter) = self
.image_token_counter
.as_ref()
.filter(|counter| !counter.uses_request_context_budget())
{
let n = counter.count_tokens(w, h);
tracing::debug!(
target: "mm_routing",
model = counter.model_id(),
width = w,
height = h,
tokens = n,
mm_hash = mm_hash,
source = "url_passthrough_header_fetch",
"image-token count"
);
image_tokens = checked_add_image_tokens(image_tokens, n);
}
mm_routing_entries.push(MmRoutingEntry::Image {
mm_hash,
width: w,
height: h,
});
}
Err(e) => {
if MediaFetcher::is_policy_rejection(&e) {
return Err(e);
}
let url_for_log = if url.starts_with("data:") {
url.split_once(',')
.map(|(p, _)| format!("{p},<redacted>"))
.unwrap_or_else(|| "data:<redacted>".to_string())
} else {
url.to_string()
};
tracing::warn!(
target: "mm_routing",
url = %url_for_log,
error = %e,
"mm-routing: failed to fetch image dims; MM routing entry skipped"
);
}
}
}
}
if !media_map.is_empty() {
builder.multi_modal_data(Some(media_map));
if has_user_uuid {
builder.multi_modal_uuids(Some(uuid_map));
}
#[cfg(feature = "mm-routing")]
if has_user_uuid {
mm_routing_entries.clear();
}
let messages_json = serde_json::to_value(request.messages())?;
let mut extra_args = serde_json::json!({
"messages": messages_json
});
Self::strip_inline_data_urls(&mut extra_args["messages"]);
if let Some(prompt) = formatted_prompt {
extra_args["formatted_prompt"] = serde_json::Value::String(prompt.to_string());
}
if let Some(serde_json::Value::Object(backend_extra_args)) = Self::backend_extra_args(
request,
self.runtime_config.reasoning_parser.is_some(),
Self::prompt_injected_reasoning_ended_arg(
self.runtime_config.reasoning_parser.as_deref(),
formatted_prompt,
),
) {
let extra_args_obj = extra_args
.as_object_mut()
.expect("multimodal extra_args must be an object");
extra_args_obj.extend(backend_extra_args);
}
#[cfg(feature = "mm-routing")]
let resolved_image_count = mm_routing_entries
.iter()
.filter(|entry| matches!(entry, MmRoutingEntry::Image { .. }))
.count();
#[cfg(feature = "mm-routing")]
let exact_mm_routing_ready = exact_mm_routing_eligible
&& exact_mm_routing_preconditions_met(
has_user_uuid,
resolved_image_count,
total_image_count,
has_processor_override,
);
#[cfg(feature = "mm-routing")]
let request_budgeted_image_counter = self
.image_token_counter
.as_ref()
.filter(|counter| counter.uses_request_context_budget());
#[cfg(feature = "mm-routing")]
let image_text_prompt_len = if exact_mm_routing_ready {
match request_budgeted_image_counter {
Some(counter) => {
counter
.context_budget_text_len(
self.tokenizer.clone(),
formatted_prompt,
resolved_image_count,
)
.await
}
None => None,
}
} else {
None
};
#[cfg(feature = "mm-routing")]
if let Some(counter) = request_budgeted_image_counter {
image_tokens = image_text_prompt_len.and_then(|text_prompt_len| {
let dimensions: Vec<_> = mm_routing_entries
.iter()
.filter_map(|entry| match entry {
MmRoutingEntry::Image { width, height, .. } => Some((*width, *height)),
MmRoutingEntry::Video { .. } => None,
})
.collect();
counter
.count_tokens_for_images(
&dimensions,
self.context_length as usize,
text_prompt_len,
)
.ok()
.and_then(|counts| counts.into_iter().try_fold(0usize, usize::checked_add))
});
}
#[cfg(feature = "mm-routing")]
let mm_routing_info = if exact_mm_routing_ready {
self.build_mm_exact_routing_info(
&mm_routing_entries,
token_ids,
image_text_prompt_len,
)
} else {
None
};
#[cfg(feature = "mm-routing")]
if let Some(mm_routing_info) = mm_routing_info {
let hex = |entry: &MmRoutingEntry| {
serde_json::Value::String(format!("{:016x}", entry.mm_hash()))
};
if mm_routing_entries
.iter()
.all(|entry| matches!(entry, MmRoutingEntry::Image { .. }))
{
extra_args["mm_hashes"] =
serde_json::Value::Array(mm_routing_entries.iter().map(hex).collect());
} else {
let mut grouped = serde_json::Map::new();
let image_hashes: Vec<_> = mm_routing_entries
.iter()
.filter(|entry| matches!(entry, MmRoutingEntry::Image { .. }))
.map(hex)
.collect();
let video_hashes: Vec<_> = mm_routing_entries
.iter()
.filter(|entry| matches!(entry, MmRoutingEntry::Video { .. }))
.map(hex)
.collect();
if !image_hashes.is_empty() {
grouped.insert("image".to_string(), serde_json::Value::Array(image_hashes));
}
if !video_hashes.is_empty() {
grouped.insert("video".to_string(), serde_json::Value::Array(video_hashes));
}
extra_args["mm_hashes_by_modality"] = serde_json::Value::Object(grouped);
}
builder.mm_routing_info(Some(mm_routing_info));
} else if has_processor_override && total_image_count > 0 {
tracing::debug!(
target: "mm_routing",
"mm-routing: exact MM routing disabled because mm_processor_kwargs is non-empty"
);
} else if !mm_routing_entries.is_empty() {
tracing::warn!(
target: "mm_routing",
resolved = mm_routing_entries.len(),
expected_images = total_image_count,
"mm-routing: exact MM routing info not built (media resolution or placeholder-order mismatch); skipping mm_hashes forwarding"
);
}
builder.extra_args(Some(extra_args));
}
#[cfg(feature = "mm-routing")]
let image_tokens = aggregate_image_tokens(
image_tokens,
mm_routing_entries
.iter()
.filter(|entry| matches!(entry, MmRoutingEntry::Image { .. }))
.count(),
total_image_count,
has_processor_override,
);
#[cfg(feature = "mm-routing")]
return Ok((mm_routing_entries, image_tokens));
#[cfg(not(feature = "mm-routing"))]
Ok((Vec::new(), None))
}
#[cfg(feature = "mm-routing")]
pub fn gather_mm_exact_routing_info(
&self,
builder: &mut PreprocessedRequestBuilder,
image_entries: &[MmImageEntry],
token_ids: &[crate::protocols::TokenIdType],
) -> Result<()> {
let entries: Vec<_> = image_entries
.iter()
.map(|entry| MmRoutingEntry::Image {
mm_hash: entry.mm_hash,
width: entry.width,
height: entry.height,
})
.collect();
if let Some(info) = self.build_mm_exact_routing_info(&entries, token_ids, None) {
builder.mm_routing_info(Some(info));
}
Ok(())
}
#[cfg(feature = "mm-routing")]
fn build_mm_exact_routing_info(
&self,
entries: &[MmRoutingEntry],
token_ids: &[crate::protocols::TokenIdType],
image_text_prompt_len: Option<usize>,
) -> Option<crate::protocols::common::preprocessor::MmRoutingInfo> {
use crate::protocols::common::preprocessor::MmRoutingInfo;
if entries.is_empty() {
return None;
}
if !exact_mm_routing_entries_are_unambiguous(entries) {
tracing::debug!(
target: "mm_routing",
"consecutive video objects cannot be mapped exactly in vLLM KV events; skipping MM routing info"
);
return None;
}
let image_token_id = self.routing_image_token_id;
let image_counter_required = entries
.iter()
.any(|entry| matches!(entry, MmRoutingEntry::Image { .. }));
if image_counter_required && self.image_token_counter.is_none() {
tracing::debug!(
target: "mm_routing",
"image_token_counter unavailable; skipping MM routing info"
);
return None;
}
if image_counter_required && image_token_id.is_none() {
tracing::debug!(
target: "mm_routing",
"routing_image_token_id unresolved; skipping MM routing info"
);
return None;
}
if image_counter_required && self.routing_image_prompt_layout.is_none() {
tracing::debug!(
target: "mm_routing",
"routing_image_prompt_layout unresolved; skipping MM routing info"
);
return None;
}
let block_size = self.kv_cache_block_size;
if block_size == 0 {
tracing::debug!(
target: "mm_routing",
"kv_cache_block_size is 0; skipping MM routing info"
);
return None;
}
let (mut expanded, expanded_prompt_len, block_mm_infos) = if entries
.iter()
.all(|entry| matches!(entry, MmRoutingEntry::Image { .. }))
{
let counter = self
.image_token_counter
.as_ref()
.expect("image counter requirement checked above");
let images: Vec<MmImageEntry> = entries
.iter()
.map(|entry| match entry {
MmRoutingEntry::Image {
mm_hash,
width,
height,
} => MmImageEntry {
mm_hash: *mm_hash,
width: *width,
height: *height,
},
MmRoutingEntry::Video { .. } => unreachable!("all entries checked as images"),
})
.collect();
let image_token_id = image_token_id.expect("image token requirement checked above");
let placeholder_count = token_ids
.iter()
.filter(|&&token_id| token_id == image_token_id)
.count();
if placeholder_count != images.len() {
tracing::warn!(
target: "mm_routing",
placeholder_count,
image_count = images.len(),
routing_image_token_id = image_token_id,
"placeholder token count in tokenized prompt does not match image count; \
skipping MM routing info (text-prefix routing only)"
);
return None;
}
let dimensions: Vec<(u32, u32)> = images
.iter()
.map(|image| (image.width, image.height))
.collect();
let text_prompt_len = if counter.uses_request_context_budget() {
match image_text_prompt_len {
Some(text_prompt_len) => text_prompt_len,
None => {
tracing::debug!(
target: "mm_routing",
model = counter.model_id(),
"request text length unavailable for batch image budgeting; skipping MM routing info"
);
return None;
}
}
} else {
token_ids.len()
};
let n_tokens = match counter.count_tokens_for_images(
&dimensions,
self.context_length as usize,
text_prompt_len,
) {
Ok(counts) => counts,
Err(error) => {
tracing::warn!(
target: "mm_routing",
model = counter.model_id(),
%error,
"request-wide image-token accounting failed; skipping MM routing info"
);
return None;
}
};
let (expanded, expanded_prompt_len) = try_expand_mm_routing_tokens(
self.tokenizer.as_ref(),
self.routing_image_prompt_layout
.expect("image prompt layout requirement checked above"),
self.routing_prepend_bos,
image_token_id,
&images,
&n_tokens,
token_ids,
counter.model_id(),
)?;
(expanded, expanded_prompt_len, Vec::new())
} else {
if image_counter_required
&& self.image_token_counter.as_ref().is_some_and(
mm_routing::image::ImageRoutingProcessor::uses_request_context_budget,
)
{
tracing::debug!(
target: "mm_routing",
"request-budgeted image expansion is unsupported for mixed modalities; skipping MM routing info"
);
return None;
}
let mut replacements = Vec::with_capacity(entries.len());
let video_token_id = entries.iter().find_map(|entry| match entry {
MmRoutingEntry::Video {
event_video_token_id,
..
} => *event_video_token_id,
MmRoutingEntry::Image { .. } => None,
});
if video_token_id.is_none() && entries.len() != 1 {
tracing::debug!(
target: "mm_routing",
"shared image/video placeholders support one video-only object; skipping exact routing for this media layout"
);
return None;
}
for entry in entries {
let replacement = match entry {
MmRoutingEntry::Image {
mm_hash,
width,
height,
} => {
let counter = self
.image_token_counter
.as_ref()
.expect("image counter requirement checked above");
let token_count = counter.count_tokens(*width, *height);
let mut routing_tokens = Vec::with_capacity(token_count);
let mut worker_tokens = Vec::with_capacity(token_count);
let image = MmImageEntry {
mm_hash: *mm_hash,
width: *width,
height: *height,
};
let layout = self
.routing_image_prompt_layout
.expect("image prompt layout requirement checked above");
let image_token_id =
image_token_id.expect("image token requirement checked above");
let result = append_mm_routing_replacement(
&mut routing_tokens,
self.tokenizer.as_ref(),
layout,
image,
token_count,
)
.and_then(|()| {
append_mm_routing_replacement_with_fill(
&mut worker_tokens,
self.tokenizer.as_ref(),
layout,
image,
token_count,
image_token_id,
)
});
if let Err(error) = result {
tracing::warn!(
target: "mm_routing",
model = counter.model_id(),
%error,
"routing-only image prompt expansion failed; skipping MM routing info"
);
return None;
}
TrackedMmRoutingReplacement {
mm_hash: *mm_hash,
target_tokens: vec![image_token_id],
worker_tokens,
routing_tokens,
}
}
MmRoutingEntry::Video {
mm_hash,
placeholder_token_id,
event_video_token_id: _,
target_tokens,
replacement_tokens,
} => {
let fill_token =
dynamo_kv_router::protocols::pad_value_for_mm_hash(*mm_hash);
TrackedMmRoutingReplacement {
mm_hash: *mm_hash,
target_tokens: target_tokens.clone(),
worker_tokens: replacement_tokens.clone(),
routing_tokens: replacement_tokens
.iter()
.map(|replacement_id| {
if *replacement_id == *placeholder_token_id {
fill_token
} else {
*replacement_id
}
})
.collect(),
}
}
};
replacements.push(replacement);
}
let routing_bos =
routing_bos_to_prepend(self.routing_prepend_bos, image_counter_required);
#[cfg(feature = "media-ffmpeg")]
let kv_event_mm_identity = self
.video_routing_processor
.as_ref()
.map_or(mm_routing::KvEventMmIdentity::MmMetadata, |processor| {
processor.kv_event_mm_identity()
});
#[cfg(not(feature = "media-ffmpeg"))]
let kv_event_mm_identity = mm_routing::KvEventMmIdentity::MmMetadata;
match apply_tracked_mm_replacements(
routing_bos,
&replacements,
token_ids,
block_size,
image_token_id,
video_token_id,
kv_event_mm_identity,
) {
Ok(expanded) => expanded,
Err(error) => {
tracing::warn!(
target: "mm_routing",
media_count = entries.len(),
%error,
"placeholder token sequence does not match multimodal content order; \
skipping MM routing info (text-prefix routing only)"
);
return None;
}
}
};
let total_tokens = expanded.len().div_ceil(block_size) * block_size;
if expanded.len() < total_tokens {
expanded.resize(total_tokens, 0);
}
tracing::debug!(
target: "mm_routing",
n_media = entries.len(),
block_size,
total_tokens,
"MmRoutingInfo built (exact, pad_value)"
);
Some(MmRoutingInfo {
routing_token_ids: expanded,
block_mm_infos,
expanded_prompt_len,
})
}
#[cfg(feature = "mm-routing")]
fn hash_image_url(url: &str) -> u64 {
xxhash_rust::xxh3::xxh3_64(url.as_bytes())
}
#[cfg(feature = "mm-routing")]
async fn fetch_image_dims(
mm_hash: u64,
url: &str,
dimension_policy: RoutingImageDimensionPolicy,
) -> Result<(u32, u32)> {
use moka::future::Cache;
use std::sync::LazyLock;
static DIM_CACHE: LazyLock<Cache<(u64, RoutingImageDimensionPolicy), (u32, u32)>> =
LazyLock::new(|| {
Cache::builder()
.max_capacity(100_000)
.time_to_live(std::time::Duration::from_secs(24 * 60 * 60))
.build()
});
let cache_key = (mm_hash, dimension_policy);
if let Some(dims) = DIM_CACHE.get(&cache_key).await {
return Ok(dims);
}
let url_owned = url.to_string();
DIM_CACHE
.try_get_with(cache_key, async move {
Self::fetch_image_dims_uncached(&url_owned, dimension_policy)
.await
.map_err(ImageDimFetchFailure::from_error)
})
.await
.map_err(|error| error.to_error())
}
#[cfg(feature = "mm-routing")]
async fn fetch_image_dims_uncached(
url: &str,
dimension_policy: RoutingImageDimensionPolicy,
) -> Result<(u32, u32)> {
const SMALL_RANGE: usize = 4 * 1024 - 1;
const LARGE_RANGE: usize = 64 * 1024 - 1;
const DIM_FETCH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
if let Some(rest) = url.strip_prefix("data:") {
let comma = rest
.find(',')
.ok_or_else(|| anyhow::anyhow!("malformed data URI: no comma"))?;
let prefix = &rest[..comma];
let payload = &rest[comma + 1..];
let bytes: Vec<u8> = if prefix.contains(";base64") {
use base64::Engine;
base64::engine::general_purpose::STANDARD
.decode(payload)
.map_err(|e| anyhow::anyhow!("data URI base64 decode: {}", e))?
} else {
payload.as_bytes().to_vec()
};
return Self::dimensions_from_image_bytes(&bytes, dimension_policy);
}
if !(url.starts_with("http://") || url.starts_with("https://")) {
anyhow::bail!("unsupported url scheme for dim fetch: {}", url);
}
let parsed = url::Url::parse(url)?;
DIM_FETCH_MEDIA_FETCHER
.check_if_url_allowed_with_dns(&parsed)
.await?;
let mut range_end = SMALL_RANGE;
loop {
let resp = DIM_FETCH_HTTP_CLIENT
.get(url)
.header("Range", format!("bytes=0-{}", range_end))
.timeout(DIM_FETCH_TIMEOUT)
.send()
.await
.map_err(|error| {
crate::preprocessor::media::MediaFetcher::map_fetch_error(error.into())
})?;
let status = resp.status();
if status != reqwest::StatusCode::PARTIAL_CONTENT {
anyhow::bail!(
"image dim fetch expected 206 Partial Content, got HTTP {}",
status
);
}
let bytes = resp.bytes().await?;
match Self::dimensions_from_image_bytes(&bytes, dimension_policy) {
Ok((w, h)) => return Ok((w, h)),
Err(_) if range_end < LARGE_RANGE => {
range_end = LARGE_RANGE;
continue;
}
Err(e) => anyhow::bail!("image header parse failed after 64KB: {}", e),
}
}
}
#[cfg(feature = "mm-routing")]
fn dimensions_from_image_bytes(
bytes: &[u8],
dimension_policy: RoutingImageDimensionPolicy,
) -> Result<(u32, u32)> {
use image::{ImageDecoder, ImageReader, metadata::Orientation};
use std::io::Cursor;
let reader = ImageReader::new(Cursor::new(bytes)).with_guessed_format()?;
if dimension_policy == RoutingImageDimensionPolicy::Encoded {
return Ok(reader.into_dimensions()?);
}
let mut decoder = reader.into_decoder()?;
let (width, height) = decoder.dimensions();
let swaps_axes = matches!(
decoder.orientation()?,
Orientation::Rotate90
| Orientation::Rotate270
| Orientation::Rotate90FlipH
| Orientation::Rotate270FlipH
);
Ok(if swaps_axes {
(height, width)
} else {
(width, height)
})
}
pub async fn gather_tokens<
R: OAIChatLikeRequest
+ AnnotationsProvider
+ SamplingOptionsProvider
+ StopConditionsProvider
+ OutputOptionsProvider
+ NvExtProvider,
>(
&self,
request: &R,
formatted_prompt: Option<&RenderedPrompt>,
tracker: Option<&RequestTracker>,
) -> Result<(Vec<crate::protocols::TokenIdType>, HashMap<String, String>)> {
let mut annotations = HashMap::new();
let mut tokens_out: Vec<crate::protocols::TokenIdType> = Vec::new();
match request.prompt_input_type() {
PromptInput::Tokens(_) => {
if let Some(token_input) = request.extract_tokens() {
match token_input {
TokenInput::Single(tokens) => {
tokens_out = tokens;
}
TokenInput::Batch(token_batches) => {
if token_batches.len() == 1 {
tokens_out = token_batches[0].clone();
} else {
bail!(
"Batch token input not supported for more than one token in requests (got {})",
token_batches.len()
);
}
}
}
ensure_token_ids_in_vocab("prompt", &tokens_out, self.token_id_bound)?;
}
}
PromptInput::Text(_) => {
if let Some(text_input) = request.extract_text() {
match text_input {
TextInput::Single(raw_prompt) => {
if let Some(f) = formatted_prompt
&& request.has_annotation(ANNOTATION_FORMATTED_PROMPT)
{
annotations.insert(
ANNOTATION_FORMATTED_PROMPT.to_string(),
f.as_str().to_string(),
);
}
let has_backend_instance_id = request
.nvext()
.and_then(|ext| ext.backend_instance_id)
.is_some();
let token_data =
request.nvext().and_then(|ext| ext.token_data.as_ref());
let (tokens_vec, skip_token_annotation) = if let Some(tokens) =
token_data
{
ensure_token_ids_in_vocab(
"nvext.token_data",
tokens,
self.token_id_bound,
)?;
tracing::info!(
token_count = tokens.len(),
first_tokens = ?&tokens[..std::cmp::min(5, tokens.len())],
"[SIDECAR-SKIP-TOKENIZE] Found nvext.token_data — using pre-computed tokens, SKIPPING tokenization"
);
(tokens.clone(), true)
} else if has_backend_instance_id {
tracing::warn!(
"backend_instance_id provided but no token_data; tokenizing prompt"
);
let encoding = self
.encode_prompt_with_timing(
formatted_prompt,
raw_prompt.as_str(),
tracker,
)
.await?;
(encoding.token_ids().to_vec(), false)
} else {
let encoding = self
.encode_prompt_with_timing(
formatted_prompt,
raw_prompt.as_str(),
tracker,
)
.await?;
(encoding.token_ids().to_vec(), false)
};
if request.has_annotation(ANNOTATION_TOKEN_IDS)
&& !skip_token_annotation
{
annotations.insert(
ANNOTATION_TOKEN_IDS.to_string(),
serde_json::to_string(&tokens_vec)?,
);
}
tokens_out = tokens_vec;
}
TextInput::Batch(texts) => {
if texts.len() == 1 {
let encoding = self.encode_with_timing(&texts[0], tracker).await?;
let tokens = encoding.token_ids().to_vec();
tokens_out = tokens;
} else {
bail!(
"Batch text input not supported for more than one text in requests (got {})",
texts.len()
);
}
}
}
}
}
}
Self::capture_prompt_token_ids(request, tracker, &tokens_out);
Ok((tokens_out, annotations))
}
fn capture_prompt_token_ids<R: NvExtProvider>(
request: &R,
tracker: Option<&RequestTracker>,
token_ids: &[TokenIdType],
) {
let requested = request
.nvext()
.and_then(|nvext| nvext.extra_fields.as_ref())
.is_some_and(|fields| fields.iter().any(|field| field == "prompt_token_ids"));
if requested && let Some(tracker) = tracker {
tracker.set_prompt_token_ids(token_ids.to_vec());
}
}
fn prompt_overflow_error(token_count: usize, combined_limit: usize) -> DynamoError {
DynamoError::builder()
.error_type(ErrorType::InvalidArgument)
.message(format!(
"This model's maximum context length is {} tokens. \
However, your messages resulted in {} tokens. \
Please reduce the length of the messages.",
combined_limit, token_count,
))
.public_details(PublicDetails::ContextLength {
limit: combined_limit as u64,
actual: Some(token_count as u64),
})
.build()
}
async fn encode_with_timing(
&self,
prompt: &str,
tracker: Option<&RequestTracker>,
) -> anyhow::Result<Encoding> {
let encode_start = Instant::now();
let owned = if prompt.contains('\0') {
tracing::debug!("Prompt contains null bytes; stripping to avoid tokenizer divergence");
prompt.replace('\0', "")
} else {
prompt.to_string()
};
let tokenizer = self.tokenizer.clone();
let encoding = tokio::task::spawn_blocking(move || tokenizer.encode(&owned)).await??;
if let Some(t) = tracker {
t.record_tokenize_latency(encode_start.elapsed());
}
Ok(encoding)
}
async fn encode_prompt_with_timing(
&self,
formatted_prompt: Option<&RenderedPrompt>,
raw_prompt: &str,
tracker: Option<&RequestTracker>,
) -> anyhow::Result<Encoding> {
let Some(prompt) = formatted_prompt
.filter(|prompt| prompt.segments().is_some())
.cloned()
else {
return self
.encode_with_timing(
formatted_prompt
.map(RenderedPrompt::as_str)
.unwrap_or(raw_prompt),
tracker,
)
.await;
};
let encode_start = Instant::now();
let tokenizer = self.tokenizer.clone();
let encoding = tokio::task::spawn_blocking(move || {
let segments = prompt
.encode_segments()
.expect("prompt was checked for rendered segments");
tokenizer.encode_segments(&segments)
})
.await??;
if let Some(t) = tracker {
t.record_tokenize_latency(encode_start.elapsed());
}
Ok(encoding)
}
pub async fn preprocess_embedding_request(
&self,
request: &NvCreateEmbeddingRequest,
) -> Result<(PreprocessedEmbeddingRequest, HashMap<String, String>)> {
let _stage_guard = StageGuard::new(STAGE_PREPROCESS, "");
let mut annotations = HashMap::new();
let mut builder = PreprocessedEmbeddingRequest::builder();
let embedding_tokenizers = self.embedding_tokenizers.as_ref().ok_or_else(|| {
anyhow::anyhow!(
"embedding tokenization is unavailable; construct the preprocessor with \
OpenAIPreprocessor::new_for_embeddings"
)
})?;
let effective_add_special_tokens = request
.add_special_tokens
.or(embedding_tokenizers.add_special_tokens_default)
.unwrap_or(true);
let is_text_input = matches!(
&request.inner.input,
dynamo_protocols::types::EmbeddingInput::String(_)
| dynamo_protocols::types::EmbeddingInput::StringArray(_)
);
let truncation_limit =
self.embedding_truncation_limit(request.truncate_prompt_tokens, is_text_input)?;
let mut all_token_ids = match &request.inner.input {
dynamo_protocols::types::EmbeddingInput::String(s) => {
let tokenizer = embedding_tokenizers.tokenizer(effective_add_special_tokens)?;
let encoding = tokenizer.encode(s)?;
vec![encoding.token_ids().to_vec()]
}
dynamo_protocols::types::EmbeddingInput::StringArray(arr) => {
let tokenizer = embedding_tokenizers.tokenizer(effective_add_special_tokens)?;
let input_strs: Vec<String> = arr.to_vec();
let encodings = tokio::task::spawn_blocking({
let tokenizer = tokenizer.clone();
let strs = input_strs.clone();
move || {
tokenizer.encode_batch(&strs.iter().map(|s| s.as_str()).collect::<Vec<_>>())
}
})
.await??;
encodings
.into_iter()
.map(|encoding| encoding.token_ids().to_vec())
.collect()
}
dynamo_protocols::types::EmbeddingInput::IntegerArray(token_ids) => {
ensure_token_ids_in_vocab("input", token_ids, self.token_id_bound)?;
vec![token_ids.clone()]
}
dynamo_protocols::types::EmbeddingInput::ArrayOfIntegerArray(token_arrays) => {
for token_ids in token_arrays {
ensure_token_ids_in_vocab("input", token_ids, self.token_id_bound)?;
}
token_arrays.clone()
}
};
if let Some(limit) = truncation_limit {
for token_ids in &mut all_token_ids {
token_ids.truncate(limit);
}
}
if request.has_annotation(ANNOTATION_TOKEN_IDS) {
annotations.insert(
ANNOTATION_TOKEN_IDS.to_string(),
serde_json::to_string(&all_token_ids)?,
);
}
builder.token_ids(all_token_ids);
builder.model(request.inner.model.clone());
builder.encoding_format(request.inner.encoding_format.as_ref().map(|f| match f {
EncodingFormat::Float => "float".to_string(),
EncodingFormat::Base64 => "base64".to_string(),
}));
builder.truncate_prompt_tokens(request.truncate_prompt_tokens);
builder.dimensions(request.inner.dimensions);
builder.annotations(request.annotations().unwrap_or_default());
builder.mdc_sum(Some(self.mdcsum.clone()));
Ok((builder.build()?, annotations))
}
fn embedding_truncation_limit(
&self,
requested: Option<i64>,
is_text_input: bool,
) -> Result<Option<usize>> {
let Some(requested) = requested else {
return Ok(None);
};
if requested < -1 {
return Err(invalid_argument_error(format!(
"truncate_prompt_tokens must be >= -1, got {requested}"
)));
}
if !is_text_input {
return Ok(None);
}
let model_limit = self.context_length as usize;
if requested == -1 {
if model_limit == 0 {
return Err(invalid_argument_error(
"truncate_prompt_tokens=-1 requires a configured model context length",
));
}
return Ok(Some(model_limit));
}
let requested = usize::try_from(requested).map_err(|_| {
invalid_argument_error("truncate_prompt_tokens is too large for this platform")
})?;
if model_limit > 0 && requested > model_limit {
return Err(invalid_argument_error(format!(
"truncate_prompt_tokens={requested} cannot be greater than \
max_model_len={model_limit}. Please request a smaller truncation size."
)));
}
Ok(Some(requested))
}
fn apply_unified_response_policies<S>(
stream: S,
emit_tool_calls: bool,
defer_reasoning_for_nonempty_content: bool,
reasoning_disabled: bool,
) -> Pin<Box<dyn Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send>>
where
S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
{
let stream: Pin<Box<dyn Stream<Item = _> + Send>> = Box::pin(
Self::apply_tool_call_response_policy(stream, emit_tool_calls),
);
let stream = stream.map(move |mut response| {
if reasoning_disabled && let Some(data) = response.data.as_mut() {
for choice in &mut data.inner.choices {
if let Some(reasoning) = choice.delta.reasoning_content.take() {
let text = match choice.delta.content.take() {
Some(ChatCompletionMessageContent::Text(text)) => text,
None => String::new(),
Some(other) => {
choice.delta.content = Some(other);
choice.delta.reasoning_content = Some(reasoning);
continue;
}
};
choice.delta.content = Some(ChatCompletionMessageContent::Text(format!(
"{reasoning}{text}"
)));
}
}
}
response
});
let stream: Pin<Box<dyn Stream<Item = _> + Send>> =
Box::pin(annotate_reasoning_usage(stream));
let stream: Pin<Box<dyn Stream<Item = _> + Send>> = if defer_reasoning_for_nonempty_content
{
Box::pin(Self::defer_unified_reasoning_for_nonempty_content(stream))
} else {
stream
};
let stream: Pin<Box<dyn Stream<Item = _> + Send>> = if defer_reasoning_for_nonempty_content
{
Box::pin(Self::hold_usage_until_stream_end(stream))
} else {
stream
};
stream
}
fn tool_processing_route(
&self,
request: &NvCreateChatCompletionRequest,
guided_tool_constraint: &crate::protocols::openai::GuidedToolConstraint,
) -> anyhow::Result<ToolProcessingRoute> {
use crate::protocols::openai::chat_completions::{tool_parser_v2, unified_parser};
let family = unified_parser::selected_request_family(
self.tool_call_parser.as_deref(),
self.runtime_config.reasoning_parser.as_deref(),
request.inner.tool_choice.as_ref(),
guided_tool_constraint,
);
tracing::debug!(
target: "dynamo_unified",
tool_call_parser = ?self.tool_call_parser,
reasoning_parser = ?self.runtime_config.reasoning_parser,
?family,
version = ?tool_parser_v2::selected_version()?,
"unified parser path decision"
);
if let Some(family) = family {
return Ok(ToolProcessingRoute::Unified(family));
}
let uses_tool_call_structural_tag = guided_tool_constraint.uses_structural_tag();
let selected_version = tool_parser_v2::selected_version()?;
if selected_version == tool_parser_v2::ParserVersion::Auto
&& let Some(family) = unified_parser::configured_family(
self.tool_call_parser.as_deref(),
self.runtime_config.reasoning_parser.as_deref(),
)
&& family == "muse_glimmer"
&& !uses_tool_call_structural_tag
&& matches!(
request.inner.tool_choice.as_ref(),
None | Some(ChatCompletionToolChoiceOption::Auto)
| Some(ChatCompletionToolChoiceOption::None)
)
{
return Ok(ToolProcessingRoute::MuseUnified(family));
}
let effective_tool_call_parser = self.tool_call_parser.clone().or_else(|| {
self.runtime_config
.reasoning_parser
.as_deref()
.filter(|parser| matches!(*parser, "kimi_k3" | "kimi-k3"))
.map(str::to_string)
});
let parser_unwraps_all_kimi_k3_responses = effective_tool_call_parser
.as_deref()
.is_some_and(|parser| matches!(parser, "kimi_k3" | "kimi-k3"));
let tool_call_parsing_enabled = Self::tool_call_parsing_enabled(request);
let has_tools = request.inner.has_effective_tools();
let should_jail = if tool_call_parsing_enabled || parser_unwraps_all_kimi_k3_responses {
Self::should_apply_tool_jail(
effective_tool_call_parser.as_ref(),
request.inner.tool_choice.as_ref(),
has_tools,
)?
} else {
false
};
if !should_jail {
return Ok(ToolProcessingRoute::PassThrough);
}
tool_parser_v2::validate_tool_request_mode(
selected_version,
effective_tool_call_parser.as_deref(),
self.runtime_config.reasoning_parser.as_deref(),
should_jail,
)
.map_err(|error| invalid_argument_error(format!("{error:#}")))?;
if let Some(parser_name) = effective_tool_call_parser.as_deref()
&& tool_parser_v2::enabled()
&& selected_version == tool_parser_v2::ParserVersion::V2
&& tool_parser_v2::supports_family(parser_name)
&& (uses_tool_call_structural_tag
|| matches!(
request.inner.tool_choice.as_ref(),
None | Some(ChatCompletionToolChoiceOption::Auto)
))
{
let parser_name = match parser_name {
"deepseek-v4" | "deepseekv4" => "deepseek_v4",
parser_name => parser_name,
};
return Ok(ToolProcessingRoute::ParserV2(parser_name.to_string()));
}
if selected_version == tool_parser_v2::ParserVersion::V2 {
return Err(invalid_argument_error(format!(
"{}=2 was requested, but this tool choice requires the v1 tool-call jail",
env_llm::DYN_PARSER_VERSION
)));
}
Ok(ToolProcessingRoute::LegacyJail(effective_tool_call_parser))
}
pub fn postprocessor_parsing_stream<S>(
&self,
stream: S,
request: &NvCreateChatCompletionRequest,
prompt_injected_reasoning: impl Into<PromptReasoningPrefill>,
uses_tool_call_structural_tag: bool,
) -> anyhow::Result<
impl Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
>
where
S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
{
let guided_tool_constraint = crate::preprocessor::tool_choice::guided_tool_constraint(
request,
self.tool_call_parser.as_deref(),
self.runtime_config.reasoning_parser.as_deref(),
uses_tool_call_structural_tag,
)?;
let tool_processing_route = self.tool_processing_route(request, &guided_tool_constraint)?;
self.postprocessor_parsing_stream_with_constraint(
stream,
request,
prompt_injected_reasoning.into(),
guided_tool_constraint,
tool_processing_route,
)
}
fn postprocessor_parsing_stream_with_constraint<S>(
&self,
stream: S,
request: &NvCreateChatCompletionRequest,
prompt_injected_reasoning: PromptReasoningPrefill,
guided_tool_constraint: crate::protocols::openai::GuidedToolConstraint,
tool_processing_route: ToolProcessingRoute,
) -> anyhow::Result<
impl Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
>
where
S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
{
use crate::protocols::openai::chat_completions::{tool_parser_v2, unified_parser};
let prompt_prefill = prompt_injected_reasoning;
let prompt_injected_reasoning = prompt_prefill.for_route(&tool_processing_route);
let uses_tool_call_structural_tag = guided_tool_constraint.uses_structural_tag();
let defer_reasoning_for_nonempty_content =
Self::wants_reasoning_as_content_when_empty(request.chat_template_args.as_ref());
let guided_tool_streaming = Self::guided_tool_streaming_release(
guided_tool_constraint.installs_guided_json(),
env_is_falsey(env_llm::DYN_ENABLE_GUIDED_TOOL_STREAMING),
);
if let ToolProcessingRoute::MuseUnified(family) = &tool_processing_route {
let tool_definitions =
crate::preprocessor::tool_choice::effective_tool_definitions(&request.inner)?;
let tool_definitions = (!tool_definitions.is_empty()).then_some(tool_definitions);
let reasoning_disabled = Self::request_disables_reasoning_with_family(
request,
self.tool_call_parser.as_deref(),
self.runtime_config.reasoning_parser.as_deref(),
None,
Some(family),
);
let unified: Pin<Box<dyn Stream<Item = _> + Send>> =
Box::pin(unified_parser::apply_stream_with_policy(
stream,
tool_definitions,
crate::protocols::openai::GuidedToolConstraint::None,
unified_parser::stream_prefill(family, prompt_injected_reasoning),
family,
guided_tool_streaming,
unified_parser::UnifiedRequestPolicy {
reasoning_disabled,
structured_response: Self::has_structured_response_format(request),
},
));
return Ok(Self::apply_unified_response_policies(
unified,
Self::tool_call_parsing_enabled(request),
defer_reasoning_for_nonempty_content,
reasoning_disabled,
));
}
if let ToolProcessingRoute::Unified(family) = &tool_processing_route {
let tool_definitions =
crate::preprocessor::tool_choice::effective_tool_definitions(&request.inner)?;
let tool_definitions = (!tool_definitions.is_empty()).then_some(tool_definitions);
let reasoning_disabled = Self::request_disables_reasoning_with_family(
request,
self.tool_call_parser.as_deref(),
self.runtime_config.reasoning_parser.as_deref(),
None,
Some(family),
);
let unified: Pin<Box<dyn Stream<Item = _> + Send>> =
Box::pin(unified_parser::apply_stream_with_policy(
stream,
tool_definitions,
guided_tool_constraint,
unified_parser::StreamPrefill {
starting_state: unified_parser::stream_prefill(
family,
prompt_injected_reasoning,
),
prefix: prompt_prefill.prefix,
},
family,
guided_tool_streaming,
unified_parser::UnifiedRequestPolicy {
reasoning_disabled,
structured_response: Self::has_structured_response_format(request),
},
));
return Ok(Self::apply_unified_response_policies(
unified,
Self::tool_call_parsing_enabled(request),
defer_reasoning_for_nonempty_content,
reasoning_disabled,
));
}
let is_guided_tool_choice = matches!(
request.inner.tool_choice,
Some(ChatCompletionToolChoiceOption::Required)
| Some(ChatCompletionToolChoiceOption::Named(_))
);
let is_structured_response = Self::has_structured_response_format(request);
let is_guided_output = is_guided_tool_choice || is_structured_response;
let reasoning_parser = self.runtime_config.reasoning_parser.as_deref();
let inspect_force_reasoning_guided_output = is_guided_output
&& !uses_tool_call_structural_tag
&& Self::supports_reasoning_before_guided_json(reasoning_parser);
let inspect_prompt_injected_guided_output = is_guided_output
&& prompt_injected_reasoning
&& !uses_tool_call_structural_tag
&& (Self::skips_guided_json_when_prompt_injected(reasoning_parser)
|| (is_structured_response
&& Self::skips_structured_response_when_prompt_injected(reasoning_parser)));
let inspect_unsupported_structured_response_reasoning_gate = is_structured_response
&& !uses_tool_call_structural_tag
&& !Self::structured_response_supports_sglang_reasoning_gate(reasoning_parser);
let bypass_reasoning_for_bare_guided_json = inspect_force_reasoning_guided_output
|| inspect_prompt_injected_guided_output
|| inspect_unsupported_structured_response_reasoning_gate;
let skip_reasoning_for_guided_json = is_guided_output
&& !uses_tool_call_structural_tag
&& Self::is_force_reasoning_parser(reasoning_parser)
&& !inspect_force_reasoning_guided_output;
let reasoning_disabled_by_request = Self::is_reasoning_disabled_by_request(
self.runtime_config.reasoning_parser.as_deref(),
request.chat_template_args.as_ref(),
);
let should_parse_reasoning = self.runtime_config.reasoning_parser.is_some()
&& !reasoning_disabled_by_request
&& !skip_reasoning_for_guided_json;
let should_strip_disabled_reasoning_start = reasoning_disabled_by_request
&& Self::is_nemotron_force_reasoning(self.runtime_config.reasoning_parser.as_deref())
&& !skip_reasoning_for_guided_json;
let guided_reasoning_start_token =
if should_parse_reasoning && bypass_reasoning_for_bare_guided_json {
Self::guided_json_reasoning_start_token(reasoning_parser)
} else {
None
};
let stream: Pin<Box<dyn Stream<Item = _> + Send>> =
if let Some(start_token) = guided_reasoning_start_token {
Box::pin(Self::strip_leading_reasoning_start_from_stream(
stream,
start_token,
))
} else {
Box::pin(stream)
};
let stream: Pin<Box<dyn Stream<Item = _> + Send>> = if should_parse_reasoning {
Box::pin(Self::parse_reasoning_content_from_stream_inner(
stream,
self.runtime_config.reasoning_parser.clone().unwrap(), prompt_injected_reasoning,
bypass_reasoning_for_bare_guided_json,
defer_reasoning_for_nonempty_content,
))
} else if should_strip_disabled_reasoning_start {
Box::pin(Self::strip_leading_reasoning_start_from_stream(
stream, "<think>",
))
} else {
Box::pin(stream)
};
let stream: Pin<Box<dyn Stream<Item = _> + Send>> =
if defer_reasoning_for_nonempty_content || should_strip_disabled_reasoning_start {
Box::pin(Self::hold_usage_until_stream_end(stream))
} else {
stream
};
let stream: Pin<Box<dyn Stream<Item = _> + Send>> = if should_parse_reasoning {
Box::pin(annotate_reasoning_usage(stream))
} else {
stream
};
let tool_call_parsing_enabled = Self::tool_call_parsing_enabled(request);
let tool_definitions =
crate::preprocessor::tool_choice::effective_tool_definitions(&request.inner)?;
let tool_definitions = (!tool_definitions.is_empty()).then_some(tool_definitions);
let transformed_stream: Pin<Box<dyn Stream<Item = _> + Send>> =
match tool_processing_route {
ToolProcessingRoute::ParserV2(parser_name) => Box::pin(
tool_parser_v2::apply_stream(stream, tool_definitions, parser_name),
),
ToolProcessingRoute::LegacyJail(effective_tool_call_parser) => {
Box::pin(Self::apply_tool_calling_jail(
effective_tool_call_parser,
request.inner.tool_choice.clone(),
tool_definitions,
uses_tool_call_structural_tag,
guided_tool_streaming,
stream,
))
}
ToolProcessingRoute::PassThrough => Box::pin(stream),
ToolProcessingRoute::MuseUnified(_) | ToolProcessingRoute::Unified(_) => {
unreachable!("unified routes return before legacy response processing")
}
};
Ok(Self::apply_tool_call_response_policy(
transformed_stream,
tool_call_parsing_enabled,
))
}
fn apply_tool_call_response_policy<S>(
stream: S,
tool_call_parsing_enabled: bool,
) -> Pin<Box<dyn Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static>>
where
S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
{
if tool_call_parsing_enabled {
return Box::pin(stream);
}
Box::pin(stream.map(|mut response| {
if let Some(data) = response.data.as_mut() {
for choice in &mut data.inner.choices {
choice.delta.tool_calls = None;
if choice.finish_reason
== Some(dynamo_protocols::types::FinishReason::ToolCalls)
{
choice.finish_reason = Some(dynamo_protocols::types::FinishReason::Stop);
}
}
}
response
}))
}
fn normalize_chat_stream_roles<S>(
stream: S,
) -> impl Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static
where
S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
{
let mut role_emitted_choices = HashSet::new();
stream.map(move |mut response| {
if let Some(data) = response.data.as_mut() {
for choice in &mut data.inner.choices {
choice.delta.role = role_emitted_choices
.insert(choice.index)
.then_some(dynamo_protocols::types::Role::Assistant);
}
}
response
})
}
pub fn transform_postprocessor_stream<S, Resp>(
stream: S,
generator: Box<dyn DeltaGeneratorExt<Resp>>,
context: Arc<dyn AsyncEngineContext>,
emit_payload_usage_chunk: bool,
trace_tokens_enabled: bool,
trace_finish_reason_metadata: Option<crate::request_trace::SharedFinishReasonMetadata>,
mm_counts: MultimodalCounts,
) -> impl Stream<Item = Annotated<Resp>> + Send
where
S: Stream<Item = Annotated<BackendOutput>> + Send + 'static,
Resp: Send + Sync + Clone + 'static + std::fmt::Debug,
{
Self::transform_postprocessor_stream_with_image_tokens(
stream,
generator,
context,
emit_payload_usage_chunk,
trace_tokens_enabled,
trace_finish_reason_metadata,
mm_counts,
None,
)
}
#[allow(clippy::too_many_arguments)]
fn transform_postprocessor_stream_with_image_tokens<S, Resp>(
stream: S,
generator: Box<dyn DeltaGeneratorExt<Resp>>,
context: Arc<dyn AsyncEngineContext>,
emit_payload_usage_chunk: bool,
trace_tokens_enabled: bool,
trace_finish_reason_metadata: Option<crate::request_trace::SharedFinishReasonMetadata>,
mm_counts: MultimodalCounts,
image_tokens: Option<usize>,
) -> impl Stream<Item = Annotated<Resp>> + Send
where
S: Stream<Item = Annotated<BackendOutput>> + Send + 'static,
Resp: Send + Sync + Clone + 'static + std::fmt::Debug,
{
struct DetokenizeMetricsGuard {
tracker: Option<Arc<RequestTracker>>,
}
impl DetokenizeMetricsGuard {
#[inline]
fn tracker(&self) -> Option<&RequestTracker> {
self.tracker.as_deref()
}
}
impl Drop for DetokenizeMetricsGuard {
fn drop(&mut self) {
let Some(tracker) = self.tracker() else {
return;
};
if let Some(total) = tracker.detokenize_total_latency() {
DETOKENIZE_TOTAL_US.inc_by(total.as_micros() as f64);
}
DETOKENIZE_TOKEN_COUNT.inc_by(tracker.detokenize_count() as f64);
}
}
struct State<Resp>
where
Resp: Send + Sync + Clone + 'static + std::fmt::Debug,
{
response_stream: Pin<Box<dyn Stream<Item = Annotated<BackendOutput>> + Send>>,
response_generator: Box<dyn DeltaGeneratorExt<Resp>>,
detokenize_metrics: DetokenizeMetricsGuard,
context: Arc<dyn AsyncEngineContext>,
cancelled: bool,
cumulative_output_tokens: usize,
finish_reason_sent: bool,
usage_chunk_sent: bool,
pending_client_usage: Option<Annotated<Resp>>,
finished: bool,
emit_payload_usage_chunk: bool,
trace_tokens_enabled: bool,
trace_finish_reason_metadata: Option<crate::request_trace::SharedFinishReasonMetadata>,
mm_counts: MultimodalCounts,
image_tokens: Option<usize>,
}
let tracker = generator.tracker();
let state = State {
response_stream: Box::pin(stream),
response_generator: generator,
detokenize_metrics: DetokenizeMetricsGuard { tracker },
context: context.clone(),
cancelled: false,
cumulative_output_tokens: 0,
finish_reason_sent: false,
usage_chunk_sent: false,
pending_client_usage: None,
finished: false,
emit_payload_usage_chunk,
trace_tokens_enabled,
trace_finish_reason_metadata,
mm_counts,
image_tokens,
};
stream::unfold(state, |mut inner| {
async move {
if let Some(client_chunk) = inner.pending_client_usage.take() {
inner.finished = true;
return Some((client_chunk, inner));
}
if inner.finished {
return None;
}
if let Some(mut response) = inner.response_stream.next().await {
drain_router_routing_data(
&mut response.data,
inner.detokenize_metrics.tracker(),
);
if inner.cancelled {
tracing::debug!(
request_id = inner.context.id(),
"Cancellation issued last message; closing stream"
);
return None;
}
tracing::trace!(
request_id = inner.context.id(),
"Processing common response: {:?}",
response
);
let has_finish_reason = response
.data
.as_ref()
.map(|d| d.finish_reason.is_some())
.unwrap_or(false);
let (chunk_tokens, isl) = if let Some(ref backend_output) = response.data {
let chunk_tokens = backend_output.token_ids.len();
inner.cumulative_output_tokens += chunk_tokens;
let isl = inner.response_generator.get_isl().map(|isl| isl as usize);
crate::request_trace::record_backend_finish_reason_metadata(
inner.trace_finish_reason_metadata.as_ref(),
backend_output.index,
backend_output.finish_reason.as_ref(),
backend_output.stop_reason.as_ref(),
);
(chunk_tokens, isl)
} else {
(0, None)
};
let current_osl = inner.cumulative_output_tokens;
let mut response = response.map_data(|data| {
inner
.response_generator
.choice_from_postprocessor(data)
.inspect_err(|e| {
tracing::error!(
request_id = inner.context.id(),
"Error processing common response: {:?}",
e
);
inner.cancelled = true;
inner.context.stop_generating();
})
.map_err(|e| e.to_string())
});
let tracker = inner.detokenize_metrics.tracker();
let llm_metrics = build_llm_metric_annotation(
tracker,
isl.unwrap_or(0),
current_osl,
chunk_tokens,
None,
inner.mm_counts,
inner.image_tokens,
);
if inner.trace_tokens_enabled {
crate::request_trace::record_llm_metric_tokens(
tracker,
isl,
current_osl,
None,
);
}
attach_llm_metrics(&mut response, llm_metrics);
if has_finish_reason {
inner.finish_reason_sent = true;
}
tracing::trace!(
request_id = inner.context.id(),
"OpenAI NvCreateChatCompletionStreamResponse: {:?}",
response
);
Some((response, inner))
} else {
inner.finished = true;
if inner.finish_reason_sent && !inner.usage_chunk_sent {
inner.usage_chunk_sent = true;
let usage_chunk = inner.response_generator.create_usage_chunk();
let usage = inner.response_generator.get_usage();
let tracker = inner.detokenize_metrics.tracker();
let cached_tokens = usage
.prompt_tokens_details
.as_ref()
.and_then(|d| d.cached_tokens.map(|c| c as usize));
let llm_metrics = build_llm_metric_annotation(
tracker,
usage.prompt_tokens as usize,
usage.completion_tokens as usize,
0,
cached_tokens,
inner.mm_counts,
inner.image_tokens,
);
if inner.trace_tokens_enabled {
crate::request_trace::record_llm_metric_tokens(
tracker,
Some(usage.prompt_tokens as usize),
usage.completion_tokens as usize,
cached_tokens,
);
}
let usage_requested = inner.response_generator.is_usage_enabled();
if inner.emit_payload_usage_chunk {
if usage_requested {
inner.pending_client_usage = Some(Annotated::<Resp> {
id: None,
data: Some(usage_chunk.clone()),
event: None,
comment: None,
error: None,
});
}
let annotation =
llm_metrics.to_annotation::<()>().unwrap_or_else(|e| {
tracing::warn!("Failed to serialize metrics: {}", e);
Annotated::<()>::from_data(())
});
let payload_usage = Annotated::<Resp> {
id: None,
data: Some(usage_chunk),
event: Some(ANNOTATION_PAYLOAD_USAGE.to_string()),
comment: annotation.comment,
error: None,
};
Some((payload_usage, inner))
} else {
let data = if usage_requested {
Some(usage_chunk)
} else {
None
};
let mut annotated_usage = Annotated::<Resp> {
id: None,
data,
event: None,
comment: None,
error: None,
};
attach_llm_metrics(&mut annotated_usage, llm_metrics);
Some((annotated_usage, inner))
}
} else {
None
}
}
}
})
.fuse()
}
pub fn transform_embedding_postprocessor_stream<S>(
stream: S,
original_request: NvCreateEmbeddingRequest,
) -> impl Stream<Item = Annotated<NvCreateEmbeddingResponse>> + Send
where
S: Stream<Item = Annotated<EmbeddingsEngineOutput>> + Send + 'static,
{
let encode_base64 = matches!(
original_request.inner.encoding_format,
Some(dynamo_protocols::types::EncodingFormat::Base64)
);
stream.map(move |output| {
output.map_data(|engine_output| {
if engine_output.embeddings.is_empty() {
return Err("embedding worker returned an empty `embeddings` field".to_string());
}
let embeddings: Vec<dynamo_protocols::types::Embedding> = engine_output
.embeddings
.into_iter()
.enumerate()
.map(|(index, encoded)| {
let value = if encode_base64 {
dynamo_protocols::types::EmbeddingVector::Base64(encoded)
} else {
dynamo_protocols::types::EmbeddingVector::Float(
decode_base64_to_floats(&encoded)?,
)
};
Ok::<_, String>(dynamo_protocols::types::Embedding {
index: index as u32,
object: "embedding".to_string(),
embedding: value,
})
})
.collect::<Result<Vec<_>, String>>()?;
let response = NvCreateEmbeddingResponse {
inner: dynamo_protocols::types::CreateEmbeddingResponse {
object: "list".to_string(),
model: original_request.inner.model.clone(),
data: embeddings,
usage: dynamo_protocols::types::EmbeddingUsage {
prompt_tokens: engine_output.prompt_tokens,
total_tokens: engine_output.total_tokens,
},
},
};
Ok(response)
})
})
}
pub fn should_apply_tool_jail(
tool_call_parser: Option<&String>,
tool_choice: Option<&ChatCompletionToolChoiceOption>,
has_tools: bool,
) -> std::result::Result<bool, Error> {
if tool_call_parser.is_some_and(|parser| matches!(parser.as_str(), "kimi_k3" | "kimi-k3")) {
return Ok(true);
}
match (tool_call_parser, tool_choice, has_tools) {
(None, Some(ChatCompletionToolChoiceOption::Required), true) => Ok(true),
(None, Some(ChatCompletionToolChoiceOption::Named(_)), true) => Ok(true),
(None, Some(ChatCompletionToolChoiceOption::Auto), true) => {
tracing::warn!(
"Tool choice 'auto' specified but no tool parser configured; proceeding without jailing"
);
Ok(false)
}
(Some(_), Some(ChatCompletionToolChoiceOption::None), _) => {
Ok(false) }
(Some(_), Some(_), true) => Ok(true), (Some(_), None, true) => Ok(true),
_ => Ok(false),
}
}
fn guided_tool_streaming_release(installs_guided_json: bool, rollback_disabled: bool) -> bool {
installs_guided_json && !rollback_disabled
}
pub fn apply_tool_calling_jail<S>(
tool_call_parser: Option<String>,
tool_choice: Option<dynamo_protocols::types::ChatCompletionToolChoiceOption>,
tool_definitions: Option<Vec<dynamo_parsers::tool_calling::ToolDefinition>>,
uses_tool_call_structural_tag: bool,
guided_streaming: bool,
stream: S,
) -> impl Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send
where
S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
{
use dynamo_parsers::tool_calling::jail::{
Annotated as JailAnnotated, apply_tool_calling_jail,
apply_tool_calling_jail_with_guided_streaming,
};
use std::sync::{Arc, Mutex};
#[derive(Default)]
struct PendingDynamoMetadata {
metrics_template: Option<LLMMetricAnnotation>,
chunk_tokens: usize,
nvext: Option<serde_json::Value>,
prompt_logprobs: Option<Arc<crate::protocols::common::llm_backend::PromptLogprobs>>,
response_template: Option<dynamo_protocols::types::CreateChatCompletionStreamResponse>,
}
let pending = Arc::new(Mutex::new(PendingDynamoMetadata::default()));
let pending_in = Arc::clone(&pending);
#[derive(Default)]
struct ChoiceRecovery {
input_text: String,
emitted_text: String,
}
let is_glm47 = tool_call_parser.as_deref() == Some("glm47");
let glm47_config = dynamo_parsers::tool_calling::config::Glm47ParserConfig::default();
let glm47_start = glm47_config.tool_call_start;
let glm47_end = glm47_config.tool_call_end;
let choice_recovery: Arc<Mutex<std::collections::HashMap<u32, ChoiceRecovery>>> =
Arc::new(Mutex::new(std::collections::HashMap::new()));
let choice_recovery_in = Arc::clone(&choice_recovery);
let glm47_start_in = glm47_start.clone();
let glm47_end_in = glm47_end.clone();
let terminal_error: Arc<Mutex<Option<Annotated<NvCreateChatCompletionStreamResponse>>>> =
Arc::new(Mutex::new(None));
let terminal_error_in = Arc::clone(&terminal_error);
let stream = stream.take_while(move |a| {
let is_error = a.is_error();
if is_error {
*terminal_error_in
.lock()
.expect("jail terminal error poisoned") = Some(a.clone());
}
std::future::ready(!is_error)
});
let jail_input = stream.map(move |mut a| {
let has_metadata = a.data.as_ref().is_some_and(|nv| {
nv.llm_metrics.is_some() || nv.nvext.is_some() || nv.prompt_logprobs.is_some()
});
if has_metadata {
let mut p = pending_in
.lock()
.expect("jail Dynamo metadata buffer poisoned");
if let Some(nv) = a.data.as_mut() {
if p.response_template.is_none() {
p.response_template = Some(
dynamo_protocols::types::CreateChatCompletionStreamResponse {
id: nv.inner.id.clone(),
object: nv.inner.object.clone(),
created: nv.inner.created,
model: nv.inner.model.clone(),
choices: Vec::new(),
usage: None,
service_tier: nv.inner.service_tier.clone(),
system_fingerprint: nv.inner.system_fingerprint.clone(),
},
);
}
if let Some(metrics) = nv.llm_metrics.take() {
p.chunk_tokens = p.chunk_tokens.saturating_add(metrics.chunk_tokens);
p.metrics_template = Some(metrics);
}
merge_response_nvext(&mut p.nvext, nv.nvext.take());
if let Some(prompt_logprobs) = nv.prompt_logprobs.take() {
p.prompt_logprobs = Some(prompt_logprobs);
}
}
}
if is_glm47 && let Some(data) = &a.data {
let mut recovery = choice_recovery_in
.lock()
.expect("choice recovery buffer poisoned");
for choice in &data.inner.choices {
if let Some(ChatCompletionMessageContent::Text(content)) = &choice.delta.content
{
let state = recovery.entry(choice.index).or_default();
state.input_text.push_str(content);
while let Some(marker_start) = crate::protocols::openai::chat_completions::unified_parser::first_unquoted_native_tool_call_marker(&state.input_text, "glm47") {
let after_marker =
&state.input_text[marker_start + glm47_start_in.len()..];
let Some(end) = after_marker.find(&glm47_end_in) else {
break;
};
state.input_text.drain(
..marker_start
+ glm47_start_in.len()
+ end
+ glm47_end_in.len(),
);
state.emitted_text.clear();
}
if !state.emitted_text.is_empty()
&& !state.input_text.contains(['"', '\'', '`'])
&& crate::protocols::openai::chat_completions::unified_parser::unquoted_native_tool_call_marker_or_prefix_start(&state.input_text, "glm47").is_none()
&& let Some(unemitted) =
state.input_text.strip_prefix(state.emitted_text.as_str())
{
state.input_text = unemitted.to_string();
state.emitted_text.clear();
}
}
}
}
debug_assert!(a.error.is_none(), "terminal errors must bypass the jail");
JailAnnotated {
data: a.data.map(|nv| nv.inner),
id: a.id,
event: a.event,
comment: a.comment,
error: None,
}
});
let jailed: Pin<
Box<
dyn Stream<
Item = JailAnnotated<
dynamo_protocols::types::CreateChatCompletionStreamResponse,
>,
> + Send,
>,
> = if guided_streaming {
Box::pin(apply_tool_calling_jail_with_guided_streaming(
tool_call_parser,
tool_choice,
tool_definitions,
uses_tool_call_structural_tag,
jail_input,
))
} else {
Box::pin(apply_tool_calling_jail(
tool_call_parser,
tool_choice,
tool_definitions,
uses_tool_call_structural_tag,
jail_input,
))
};
let pending_out = Arc::clone(&pending);
let pending_eof = Arc::clone(&pending);
let jailed_output = jailed.flat_map(move |a| {
debug_assert!(
a.error.is_none(),
"dynamo-parsers must not construct errors"
);
let has_choices = a.data.as_ref().is_some_and(|data| !data.choices.is_empty());
let is_payload_usage = a.event.as_deref() == Some(ANNOTATION_PAYLOAD_USAGE);
let mut prompt_logprobs = None;
let (llm_metrics, nvext) = a.data.as_ref().map_or((None, None), |_| {
let mut p = pending_out
.lock()
.expect("jail Dynamo metadata buffer poisoned");
let chunk_tokens = p.chunk_tokens;
p.chunk_tokens = 0;
let metrics = p.metrics_template.take().map(|mut metrics| {
metrics.chunk_tokens = chunk_tokens;
metrics
});
let nvext = if has_choices && !is_payload_usage {
prompt_logprobs = p.prompt_logprobs.take();
p.nvext.take()
} else {
None
};
(metrics, nvext)
});
let mut nv_chunk = Annotated {
data: a.data.map(|inner| NvCreateChatCompletionStreamResponse {
inner,
nvext,
prompt_logprobs,
llm_metrics,
tool_call_completion: Vec::new(),
}),
id: a.id,
event: a.event,
comment: a.comment,
error: None,
};
if is_glm47 && let Some(data) = &mut nv_chunk.data {
let mut recovery = choice_recovery
.lock()
.expect("choice recovery buffer poisoned");
for choice in &mut data.inner.choices {
let state = recovery.entry(choice.index).or_default();
if let Some(marker_start) = crate::protocols::openai::chat_completions::unified_parser::unquoted_native_tool_call_marker_or_prefix_start(&state.input_text, "glm47") {
let desired_content = &state.input_text[..marker_start];
let dropped_call_reported_as_length = choice.finish_reason
== Some(dynamo_protocols::types::FinishReason::Stop)
&& choice.delta.tool_calls.is_none()
&& crate::protocols::openai::chat_completions::unified_parser::first_unquoted_native_tool_call_marker(&state.input_text, "glm47").is_some();
if dropped_call_reported_as_length {
tracing::warn!(
choice_index = choice.index,
why = "dropped_native_tool_call_reported_as_length",
dropped_bytes = state.input_text.len() - desired_content.len(),
"glm47 streaming: reporting length instead of stop for a tool call dropped at end of stream"
);
choice.finish_reason =
Some(dynamo_protocols::types::FinishReason::Length);
}
if choice.finish_reason
== Some(dynamo_protocols::types::FinishReason::Length)
&& crate::protocols::openai::chat_completions::unified_parser::first_unquoted_native_tool_call_marker(&state.input_text, "glm47").is_some()
{
if !dropped_call_reported_as_length {
tracing::warn!(
choice_index = choice.index,
why = "truncated_native_tool_call_suppressed",
suppressed_bytes = state.input_text.len() - desired_content.len(),
"glm47 streaming: suppressing incomplete native tool output on length finish"
);
}
let replacement = desired_content
.strip_prefix(&state.emitted_text)
.unwrap_or_default();
choice.delta.content = (!replacement.is_empty()).then(|| {
ChatCompletionMessageContent::Text(replacement.to_string())
});
}
} else if choice.finish_reason
== Some(dynamo_protocols::types::FinishReason::Length)
&& choice.delta.tool_calls.is_none()
{
let replacement = state
.input_text
.strip_prefix(&state.emitted_text)
.unwrap_or_default();
choice.delta.content = (!replacement.is_empty())
.then(|| ChatCompletionMessageContent::Text(replacement.to_string()));
}
if let Some(ChatCompletionMessageContent::Text(content)) = &choice.delta.content {
state.emitted_text.push_str(content);
}
}
}
futures::stream::iter(std::iter::once(nv_chunk))
});
let terminal_error_out = Arc::clone(&terminal_error);
let jailed_output = jailed_output.take_while(move |_| {
let stop = terminal_error_out
.lock()
.expect("jail terminal error poisoned")
.is_some();
std::future::ready(!stop)
});
let with_eof_metadata = async_stream::stream! {
tokio::pin!(jailed_output);
while let Some(response) = jailed_output.next().await {
yield response;
}
let terminal_error = {
terminal_error
.lock()
.expect("jail terminal error poisoned")
.take()
};
if let Some(error) = terminal_error {
{
let mut p = pending_eof
.lock()
.expect("jail Dynamo metadata buffer poisoned");
p.metrics_template = None;
p.chunk_tokens = 0;
p.nvext = None;
p.prompt_logprobs = None;
p.response_template = None;
}
yield error;
return;
}
let eof_metadata = {
let mut p = pending_eof
.lock()
.expect("jail Dynamo metadata buffer poisoned");
let chunk_tokens = p.chunk_tokens;
p.chunk_tokens = 0;
let llm_metrics = p.metrics_template.take().map(|mut metrics| {
metrics.chunk_tokens = chunk_tokens;
metrics
});
let nvext = p.nvext.take();
let prompt_logprobs = p.prompt_logprobs.take();
if llm_metrics.is_none() && nvext.is_none() && prompt_logprobs.is_none() {
None
} else {
p.response_template.take().map(|inner| Annotated {
data: Some(NvCreateChatCompletionStreamResponse {
inner,
nvext,
prompt_logprobs,
llm_metrics,
tool_call_completion: Vec::new(),
}),
id: None,
event: None,
comment: None,
error: None,
})
}
};
if let Some(response) = eof_metadata {
yield response;
}
};
Self::hold_usage_until_stream_end(with_eof_metadata)
}
fn parser_requires_special_tokens(
tool_call_parser: Option<&str>,
reasoning_parser: Option<&str>,
) -> Result<bool> {
if let Some(family) =
crate::protocols::openai::chat_completions::unified_parser::selected_family(
tool_call_parser,
reasoning_parser,
)
{
return dynamo_parsers_v2::create_unified_parser_for_family(family, &[])
.with_context(|| {
format!("Failed to resolve special-token requirements for {family}")
})
.map(|parser| parser.preserve_special_tokens());
}
Ok(matches!(
tool_call_parser,
Some("gemma4")
| Some("gemma-4")
| Some("harmony")
| Some("kimi_k2")
| Some("kimi_k3")
| Some("kimi-k3")
| Some("minimax_m3")
| Some("minimax-m3")
| Some("minimax_m3_nom")
| Some("minimax-m3-nom")
| Some("inkling")
| Some("muse_glimmer")
| Some("muse")
| Some("deepseek_v41")
) || matches!(
reasoning_parser,
Some("gemma4")
| Some("gemma-4")
| Some("gpt_oss")
| Some("kimi_k25")
| Some("kimi_k3")
| Some("kimi-k3")
| Some("mistral")
| Some("minimax_m3")
| Some("minimax-m3")
| Some("inkling")
| Some("muse_glimmer")
| Some("muse")
| Some("deepseek_v41")
))
}
fn special_tokens_will_be_stripped(
skip_special_tokens: Option<bool>,
parser_requires_special_tokens: bool,
) -> bool {
skip_special_tokens == Some(true) && parser_requires_special_tokens
}
fn is_nemotron_force_reasoning(reasoning_parser: Option<&str>) -> bool {
matches!(
reasoning_parser,
Some("nemotron_nano" | "nemotron3" | "nemotron_v3")
)
}
pub(crate) fn wants_reasoning_as_content_when_empty(
chat_template_args: Option<&std::collections::HashMap<String, serde_json::Value>>,
) -> bool {
chat_template_args.is_some_and(|args| {
args.get("force_nonempty_content") == Some(&serde_json::Value::Bool(true))
})
}
pub(crate) fn stream_can_defer_all_output(
tool_call_parser: Option<&str>,
reasoning_parser: Option<&str>,
chat_template_args: Option<&std::collections::HashMap<String, serde_json::Value>>,
) -> bool {
let has_reasoning_decoder = reasoning_parser.is_some()
|| crate::protocols::openai::chat_completions::unified_parser::selected_content_decoder_family(
tool_call_parser,
reasoning_parser,
)
.is_some();
has_reasoning_decoder
&& Self::wants_reasoning_as_content_when_empty(chat_template_args)
&& !Self::is_reasoning_disabled_by_request(reasoning_parser, chat_template_args)
}
fn is_force_reasoning_parser(reasoning_parser: Option<&str>) -> bool {
matches!(
reasoning_parser,
Some(
"deepseek_r1"
| "deepseek_v3"
| "deepseek_v3_1"
| "deepseek_v3_2"
| "step3"
| "kimi_k25"
| "mistral"
| "minimax_m2"
| "minimax_append_think"
| "nemotron_nano"
| "nemotron3"
| "nemotron_v3"
)
)
}
fn supports_reasoning_before_guided_json(reasoning_parser: Option<&str>) -> bool {
matches!(
reasoning_parser,
Some(
"deepseek_r1"
| "deepseek_v3"
| "deepseek_v3_1"
| "deepseek_v3_2"
| "step3"
| "kimi_k25"
| "mistral"
| "minimax_m2"
| "nemotron_nano"
| "nemotron3"
| "nemotron_v3"
)
)
}
fn guided_json_reasoning_start_token(reasoning_parser: Option<&str>) -> Option<&'static str> {
match reasoning_parser {
Some("mistral") => Some("[THINK]"),
_ => None,
}
}
fn skips_guided_json_when_prompt_injected(reasoning_parser: Option<&str>) -> bool {
matches!(
reasoning_parser,
Some(
"deepseek_v4"
| "deepseek_v41"
| "deepseek-v4"
| "deepseekv4"
| "glm45"
| "minimax_m3"
| "minimax-m3"
)
)
}
fn skips_structured_response_when_prompt_injected(reasoning_parser: Option<&str>) -> bool {
matches!(reasoning_parser, Some("qwen3" | "kimi_k3" | "kimi-k3"))
|| Self::skips_guided_json_when_prompt_injected(reasoning_parser)
}
#[cfg(test)]
fn configured_prompt_prefill(&self, formatted_prompt: Option<&str>) -> PromptReasoningPrefill {
Self::prompt_reasoning_prefill_for_parsers(
self.tool_call_parser.as_deref(),
self.runtime_config.reasoning_parser.as_deref(),
formatted_prompt,
)
}
fn prompt_injected_reasoning_start(
reasoning_parser: Option<&str>,
formatted_prompt: Option<&str>,
) -> bool {
let Some(prompt) = formatted_prompt.map(str::trim_end) else {
return false;
};
match reasoning_parser {
Some("minimax_m3") | Some("minimax-m3") => prompt.ends_with("<mm:think>"),
Some("kimi_k3") | Some("kimi-k3") => {
let mut suffix = prompt.chars().rev().filter(|c| !c.is_whitespace());
"<|open|>think<|sep|>"
.chars()
.rev()
.all(|c| suffix.next() == Some(c))
}
Some("gemma4") => {
prompt.ends_with("<|channel>thought") || prompt.ends_with("<|channel>")
}
Some("muse" | "muse_glimmer") => prompt.ends_with("assistant to=self<|message|>"),
_ => prompt.ends_with("<think>"),
}
}
fn prompt_injected_reasoning_start_for_parsers(
tool_call_parser: Option<&str>,
reasoning_parser: Option<&str>,
formatted_prompt: Option<&str>,
) -> bool {
let parser = crate::protocols::openai::chat_completions::unified_parser::configured_family(
tool_call_parser,
reasoning_parser,
)
.or(reasoning_parser);
Self::prompt_injected_reasoning_start(parser, formatted_prompt)
}
fn prompt_reasoning_prefill_for_parsers(
tool_call_parser: Option<&str>,
reasoning_parser: Option<&str>,
formatted_prompt: Option<&str>,
) -> PromptReasoningPrefill {
let unified = Self::prompt_injected_reasoning_start_for_parsers(
tool_call_parser,
reasoning_parser,
formatted_prompt,
);
let legacy = if matches!(reasoning_parser, Some("muse" | "muse_glimmer")) {
Self::prompt_injected_reasoning_start(None, formatted_prompt)
} else {
unified
};
let prefix =
if crate::protocols::openai::chat_completions::unified_parser::configured_family(
tool_call_parser,
reasoning_parser,
) == Some("gemma4")
{
const HEADER: &str = "<|channel>thought\n";
formatted_prompt.and_then(|raw| {
("<|channel>".len()..HEADER.len())
.rev()
.map(|end| &HEADER[..end])
.find(|prefix| raw.ends_with(prefix))
})
} else {
None
};
PromptReasoningPrefill {
legacy,
unified: unified || prefix.is_some(),
prefix,
}
}
fn prompt_injected_reasoning_ended_arg(
reasoning_parser: Option<&str>,
formatted_prompt: Option<&str>,
) -> Option<bool> {
let should_forward = matches!(
reasoning_parser,
Some(
"minimax_m2" | "minimax_m3" | "minimax-m3" | "kimi_k3" | "kimi-k3" | "deepseek_v41"
)
);
if should_forward
&& Self::prompt_injected_reasoning_start(reasoning_parser, formatted_prompt)
{
Some(false)
} else {
None
}
}
pub(crate) fn request_disables_reasoning_with_family(
request: &NvCreateChatCompletionRequest,
tool_parser: Option<&str>,
reasoning_parser: Option<&str>,
default_thinking_mode: Option<&str>,
unified_family: Option<&str>,
) -> bool {
let mut chat_template_args = request.chat_template_args.clone();
let from_client = Self::request_has_client_thinking_control(request);
let constrained_generation = !from_client
&& Self::is_minimax_m3_family(reasoning_parser, tool_parser)
&& Self::has_minimax_m3_constrained_generation(request);
Self::apply_default_thinking_mode_value(
default_thinking_mode,
&mut chat_template_args,
from_client,
);
Self::normalize_thinking_args_with_source(
&mut chat_template_args,
reasoning_parser,
tool_parser,
from_client,
constrained_generation,
);
if let Some(family) = unified_family {
let thinking = dynamo_renderer::thinking_bool_from_args(chat_template_args.as_ref());
return match family {
"gemma4" => thinking != Some(true),
"deepseek_v4" | "deepseek_v41" => {
!Self::deepseek_renderer_reasoning_enabled(chat_template_args.as_ref(), true)
}
_ => thinking == Some(false),
};
}
Self::is_reasoning_disabled_by_request(reasoning_parser, chat_template_args.as_ref())
}
fn is_reasoning_disabled_by_request(
reasoning_parser: Option<&str>,
chat_template_args: Option<&std::collections::HashMap<String, serde_json::Value>>,
) -> bool {
match reasoning_parser {
Some("kimi_k25" | "kimi_k3" | "kimi-k3") => {
dynamo_renderer::thinking_bool_from_args(chat_template_args) == Some(false)
}
parser if Self::is_nemotron_force_reasoning(parser) => {
dynamo_renderer::thinking_bool_from_args(chat_template_args) == Some(false)
}
Some("deepseek_v3" | "deepseek_v3_1") => {
!Self::deepseek_renderer_reasoning_enabled(chat_template_args, false)
}
Some(
"deepseek_r1" | "deepseek_v3_2" | "deepseek_v4" | "deepseek-v4" | "deepseekv4"
| "deepseek_v41" | "minimax_m2",
) => !Self::deepseek_renderer_reasoning_enabled(chat_template_args, true),
Some("gemma4") | Some("gemma-4") => {
dynamo_renderer::thinking_bool_from_args(chat_template_args) != Some(true)
}
Some("mistral") => !Self::mistral_reasoning_enabled(chat_template_args),
Some("minimax_m3") | Some("minimax-m3") => {
if let Some(args) = chat_template_args
&& let Some(mode) = args.get("thinking_mode").and_then(|v| v.as_str())
{
return mode == "disabled";
}
false
}
_ => false,
}
}
pub fn parse_reasoning_content_from_stream<S>(
stream: S,
parser_name: String,
prompt_injected_reasoning: bool,
) -> impl Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send
where
S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
{
Self::parse_reasoning_content_from_stream_inner(
stream,
parser_name,
prompt_injected_reasoning,
false,
false,
)
}
fn parse_reasoning_content_from_stream_inner<S>(
stream: S,
parser_name: String,
prompt_injected_reasoning: bool,
bypass_bare_guided_json: bool,
defer_reasoning_for_nonempty_content: bool,
) -> impl Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send
where
S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
{
let state = ReasoningState {
stream: Box::pin(stream),
parser_name,
prompt_injected_reasoning,
bypass_bare_guided_json,
choices: HashMap::new(),
last_response: None,
defer_reasoning_for_nonempty_content,
saw_terminal_error: false,
};
stream::unfold(state, |mut state| async move {
if let Some(response) = state.stream.next().await {
if response.is_error() {
state.saw_terminal_error = true;
}
let processed_response = {
let ReasoningState {
parser_name,
prompt_injected_reasoning,
bypass_bare_guided_json,
choices,
defer_reasoning_for_nonempty_content,
..
} = &mut state;
let parser_name = &*parser_name;
let prompt_injected_reasoning = *prompt_injected_reasoning;
let bypass_bare_guided_json = *bypass_bare_guided_json;
let defer_reasoning = *defer_reasoning_for_nonempty_content;
response.map_data(|mut data| {
for choice in data.inner.choices.iter_mut() {
let choice_state = choices.entry(choice.index).or_insert_with(|| {
let mut parser =
Box::new(ReasoningParserType::get_reasoning_parser_from_name(
parser_name,
))
as Box<dyn ReasoningParser>;
if prompt_injected_reasoning {
parser.set_in_reasoning(true);
}
ChoiceReasoningState {
parser,
guided_json_bypass_decision: (!bypass_bare_guided_json)
.then_some(false),
pending_reasoning: String::new(),
pending_content: String::new(),
left_reasoning: false,
drained: false,
parser_finished: false,
}
});
let bypass_decision = if bypass_bare_guided_json {
match choice_state.guided_json_bypass_decision {
Some(decision) => Some(decision),
None => {
let decision = match choice.delta.content.as_ref() {
Some(ChatCompletionMessageContent::Text(text)) => {
let text = text.trim_start();
if text.is_empty() {
None
} else {
Some(matches!(text.as_bytes()[0], b'[' | b'{'))
}
}
_ => None,
};
if let Some(decision) = decision {
choice_state.guided_json_bypass_decision =
Some(decision);
}
decision
}
}
} else {
Some(false)
};
if bypass_decision == Some(false)
&& let Some(ChatCompletionMessageContent::Text(text)) =
choice.delta.content.as_ref()
{
let parser_result = choice_state
.parser
.parse_reasoning_streaming_incremental(text, &[]);
choice_state.drained = false;
if defer_reasoning && !choice_state.left_reasoning {
choice_state
.pending_reasoning
.push_str(&parser_result.reasoning_text);
if parser_result.normal_text.trim().is_empty() {
choice_state
.pending_content
.push_str(&parser_result.normal_text);
choice.delta.content = None;
choice.delta.reasoning_content = None;
} else {
choice_state.left_reasoning = true;
choice.delta.reasoning_content = (!choice_state
.pending_reasoning
.is_empty())
.then(|| {
std::mem::take(&mut choice_state.pending_reasoning)
});
let mut text =
std::mem::take(&mut choice_state.pending_content);
text.push_str(&parser_result.normal_text);
choice.delta.content =
Some(ChatCompletionMessageContent::Text(text));
}
} else {
choice.delta.content = parser_result
.get_some_normal_text()
.map(ChatCompletionMessageContent::Text);
choice.delta.reasoning_content =
parser_result.get_some_reasoning();
}
}
let terminal_carries_parts = matches!(
choice.delta.content,
Some(ChatCompletionMessageContent::Parts(_))
);
if defer_reasoning
&& choice.finish_reason.is_some()
&& !terminal_carries_parts
&& bypass_decision == Some(false)
{
let (content, reasoning) = drain_deferred_reasoning(choice_state);
if let Some(content) = content {
let merged = match choice.delta.content.take() {
Some(ChatCompletionMessageContent::Text(existing)) => {
existing + &content
}
Some(other) => {
choice.delta.content = Some(other);
content
}
None => content,
};
if choice.delta.content.is_none() {
choice.delta.content =
Some(ChatCompletionMessageContent::Text(merged));
}
}
if let Some(reasoning) = reasoning {
choice.delta.reasoning_content = Some(
choice.delta.reasoning_content.take().unwrap_or_default()
+ &reasoning,
);
}
}
}
Ok(data)
})
};
if state.defer_reasoning_for_nonempty_content
&& processed_response
.data
.as_ref()
.is_some_and(|data| !data.inner.choices.is_empty())
{
state.last_response = Some(processed_response.clone());
}
Some((processed_response, state))
} else if !state.defer_reasoning_for_nonempty_content || state.saw_terminal_error {
None
} else {
let mut indices: Vec<u32> = state.choices.keys().copied().collect();
indices.sort_unstable();
#[allow(clippy::type_complexity)]
let flushed: Vec<(u32, Option<String>, Option<String>)> = indices
.into_iter()
.filter_map(|index| {
let choice_state = state.choices.get_mut(&index)?;
if choice_state.guided_json_bypass_decision != Some(false) {
return None;
}
let (content, reasoning) = drain_deferred_reasoning(choice_state);
(content.is_some() || reasoning.is_some())
.then_some((index, content, reasoning))
})
.collect();
if flushed.is_empty() {
return None;
}
let mut response = state.last_response.take()?;
scrub_synthetic_chunk_metadata(&mut response);
let data = response.data.as_mut()?;
let template = data.inner.choices.first()?.clone();
data.inner.choices = flushed
.into_iter()
.map(|(index, content, reasoning)| {
let mut choice = template.clone();
choice.index = index;
choice.delta.role = None;
choice.delta.tool_calls = None;
choice.delta.function_call = None;
choice.delta.refusal = None;
choice.finish_reason = None;
choice.logprobs = None;
choice.delta.content = content.map(ChatCompletionMessageContent::Text);
choice.delta.reasoning_content = reasoning;
choice
})
.collect();
Some((response, state))
}
})
.fuse()
}
fn defer_unified_reasoning_for_nonempty_content<S>(
stream_in: S,
) -> impl Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send
where
S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
{
async_stream::stream! {
let mut states: HashMap<u32, DeferredUnifiedChoice> = HashMap::new();
let mut last_response: Option<Annotated<NvCreateChatCompletionStreamResponse>> = None;
tokio::pin!(stream_in);
while let Some(mut response) = stream_in.next().await {
if response.is_error() {
yield response;
return;
}
let Some(data) = response.data.as_mut() else {
yield response;
continue;
};
let mut prefix_choices = Vec::new();
for choice in &mut data.inner.choices {
let state = states.entry(choice.index).or_default();
if choice.delta.content.is_some()
|| choice.delta.reasoning_content.is_some()
|| choice.delta.tool_calls.is_some()
{
state.drained = false;
}
if !state.saw_visible_output
&& let Some(reasoning) = choice.delta.reasoning_content.take()
{
state.pending_reasoning.push_str(&reasoning);
}
let carries_parts = matches!(
choice.delta.content,
Some(ChatCompletionMessageContent::Parts(_))
);
match choice.delta.content.take() {
Some(ChatCompletionMessageContent::Text(text)) => {
if !state.saw_visible_output && text.trim().is_empty() {
state.pending_content.push_str(&text);
} else {
if !text.trim().is_empty() {
state.saw_visible_output = true;
}
let mut content = std::mem::take(&mut state.pending_content);
content.push_str(&text);
choice.delta.content =
Some(ChatCompletionMessageContent::Text(content));
}
}
Some(parts @ ChatCompletionMessageContent::Parts(_)) => {
choice.delta.content = Some(parts);
}
None => {}
}
let has_tool_calls = choice
.delta
.tool_calls
.as_ref()
.is_some_and(|calls| !calls.is_empty());
if carries_parts || has_tool_calls {
state.saw_visible_output = true;
}
if carries_parts {
state.pending_content.clear();
}
if state.saw_visible_output {
let (content, reasoning) = state.release_split();
if let Some(reasoning) = reasoning {
let mut prefix = choice.clone();
prefix.delta.role = None;
prefix.delta.content = None;
prefix.delta.tool_calls = None;
prefix.delta.function_call = None;
prefix.delta.refusal = None;
prefix.delta.reasoning_content = Some(reasoning);
prefix.finish_reason = None;
prefix.logprobs = None;
prefix_choices.push(prefix);
}
if let Some(content) = content {
match choice.delta.content.as_mut() {
Some(ChatCompletionMessageContent::Text(existing)) => {
existing.insert_str(0, &content);
}
Some(ChatCompletionMessageContent::Parts(_)) | None => {
let mut prefix = choice.clone();
prefix.delta.role = None;
prefix.delta.content =
Some(ChatCompletionMessageContent::Text(content));
prefix.delta.tool_calls = None;
prefix.delta.function_call = None;
prefix.delta.refusal = None;
prefix.delta.reasoning_content = None;
prefix.finish_reason = None;
prefix.logprobs = None;
prefix_choices.push(prefix);
}
}
}
}
if choice.finish_reason.is_some() {
let (content, reasoning) = state.drain();
if let Some(content) = content {
match choice.delta.content.as_mut() {
Some(ChatCompletionMessageContent::Text(existing)) => {
existing.push_str(&content);
}
Some(ChatCompletionMessageContent::Parts(_)) => {}
None => {
choice.delta.content =
Some(ChatCompletionMessageContent::Text(content));
}
}
}
if let Some(reasoning) = reasoning {
choice
.delta
.reasoning_content
.get_or_insert_default()
.push_str(&reasoning);
}
}
}
if !data.inner.choices.is_empty() {
last_response = Some(response.clone());
}
for prefix_choice in prefix_choices {
let mut prefix_response = response.clone();
scrub_synthetic_chunk_metadata(&mut prefix_response);
if let Some(prefix_data) = prefix_response.data.as_mut() {
prefix_data.inner.choices = vec![prefix_choice];
}
prefix_response.id = None;
yield prefix_response;
}
yield response;
}
let mut indices: Vec<_> = states.keys().copied().collect();
indices.sort_unstable();
let flushed: Vec<_> = indices
.into_iter()
.filter_map(|index| {
let (content, reasoning) = states.get_mut(&index)?.drain();
(content.is_some() || reasoning.is_some())
.then_some((index, content, reasoning))
})
.collect();
if !flushed.is_empty()
&& let Some(mut response) = last_response
&& scrub_synthetic_chunk_metadata(&mut response).is_some()
&& let Some(data) = response.data.as_mut()
&& let Some(template) = data.inner.choices.first().cloned()
{
data.inner.choices = flushed
.into_iter()
.map(|(index, content, reasoning)| {
let mut choice = template.clone();
choice.index = index;
choice.delta.role = None;
choice.delta.content = content.map(ChatCompletionMessageContent::Text);
choice.delta.tool_calls = None;
choice.delta.function_call = None;
choice.delta.refusal = None;
choice.delta.reasoning_content = reasoning;
choice.finish_reason = None;
choice.logprobs = None;
choice
})
.collect();
yield response;
}
}
}
fn hold_usage_until_stream_end<S>(
stream: S,
) -> impl Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send
where
S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
{
async_stream::stream! {
tokio::pin!(stream);
let mut pending_usage = None;
let mut transport_failed = false;
while let Some(response) = stream.next().await {
if response.error.is_some() {
transport_failed = true;
pending_usage = None;
yield response;
continue;
}
let is_usage_only = response.data.as_ref().is_some_and(|data| {
data.inner.choices.is_empty() && data.inner.usage.is_some()
});
if is_usage_only {
if let Some(previous) = pending_usage.replace(response) {
yield previous;
}
} else {
yield response;
}
}
if !transport_failed && let Some(usage) = pending_usage {
yield usage;
}
}
}
fn strip_leading_reasoning_start_from_stream<S>(
stream: S,
think_start_token: &'static str,
) -> impl Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send
where
S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
{
struct StripReasoningStartState {
stream:
Pin<Box<dyn Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send>>,
think_start_token: &'static str,
choices: HashMap<u32, StripChoiceState>,
last_response: Option<Annotated<NvCreateChatCompletionStreamResponse>>,
eof_flushed: bool,
}
#[derive(Default)]
struct StripChoiceState {
buffer: String,
decided: bool,
}
fn take_undecided_buffer(choice_state: &mut StripChoiceState) -> Option<String> {
if choice_state.decided || choice_state.buffer.is_empty() {
return None;
}
choice_state.decided = true;
Some(std::mem::take(&mut choice_state.buffer))
}
fn drain_undecided_buffers(
choices: &mut HashMap<u32, StripChoiceState>,
) -> HashMap<u32, String> {
choices
.iter_mut()
.filter_map(|(index, choice_state)| {
take_undecided_buffer(choice_state).map(|buffer| (*index, buffer))
})
.collect()
}
let state = StripReasoningStartState {
stream: Box::pin(stream),
think_start_token,
choices: HashMap::new(),
last_response: None,
eof_flushed: false,
};
stream::unfold(state, |mut state| async move {
if let Some(mut response) = state.stream.next().await {
let Some(mut data) = response.data.take() else {
return Some((response, state));
};
for choice in data.inner.choices.iter_mut() {
let choice_state = state.choices.entry(choice.index).or_default();
let text = match choice.delta.content.take() {
Some(ChatCompletionMessageContent::Text(text)) => text,
other => {
if let Some(buffer) = take_undecided_buffer(choice_state) {
choice.delta.content =
Some(ChatCompletionMessageContent::Text(buffer));
} else {
choice.delta.content = other;
}
continue;
}
};
let output = if choice_state.decided {
text
} else {
choice_state.buffer.push_str(&text);
let trimmed = choice_state.buffer.trim_start();
if trimmed.is_empty()
|| (state.think_start_token.starts_with(trimmed)
&& trimmed.len() < state.think_start_token.len())
{
choice.delta.content = None;
continue;
}
choice_state.decided = true;
if let Some(remainder) = trimmed.strip_prefix(state.think_start_token) {
remainder.to_string()
} else {
choice_state.buffer.clone()
}
};
choice_state.buffer.clear();
choice.delta.content = if output.is_empty() {
None
} else {
Some(ChatCompletionMessageContent::Text(output))
};
}
response.data = Some(data);
if response
.data
.as_ref()
.is_some_and(|data| !data.inner.choices.is_empty())
{
state.last_response = Some(response.clone());
}
Some((response, state))
} else if state.eof_flushed {
None
} else {
state.eof_flushed = true;
let flushed = drain_undecided_buffers(&mut state.choices);
if flushed.is_empty() {
None
} else {
let mut response = state.last_response.clone()?;
scrub_synthetic_chunk_metadata(&mut response);
let data = response.data.as_mut()?;
let mut template = data.inner.choices.first()?.clone();
template.delta.role = None;
template.delta.tool_calls = None;
template.delta.function_call = None;
template.delta.refusal = None;
template.delta.reasoning_content = None;
template.finish_reason = None;
template.logprobs = None;
let mut flushed: Vec<_> = flushed.into_iter().collect();
flushed.sort_unstable_by_key(|(index, _)| *index);
data.inner.choices = flushed
.into_iter()
.map(|(index, buffer)| {
let mut choice = template.clone();
choice.index = index;
choice.delta.content = Some(ChatCompletionMessageContent::Text(buffer));
choice
})
.collect();
if data.inner.choices.is_empty() {
None
} else {
Some((response, state))
}
}
}
})
.fuse()
}
}
fn preprocessing_lifecycle(context: &PipelineContext<()>) -> LifecycleTrace {
context
.get_optional::<LifecycleTrace>(LIFECYCLE_TRACE_CONTEXT_KEY)
.ok()
.flatten()
.map(|trace| trace.as_ref().clone())
.unwrap_or_else(|| LifecycleTrace::new(false))
}
#[async_trait]
impl
Operator<
SingleIn<NvCreateChatCompletionRequest>,
ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>,
SingleIn<PreprocessedRequest>,
ManyOut<Annotated<BackendOutput>>,
> for OpenAIPreprocessor
{
async fn generate(
&self,
request: SingleIn<NvCreateChatCompletionRequest>,
next: Arc<
dyn AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<BackendOutput>>, Error>,
>,
) -> Result<ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>, Error> {
let (mut request, context) = request.into_parts();
let lifecycle = preprocessing_lifecycle(&context);
let preprocessing = lifecycle.start(LifecycleStage::RequestPreprocessing);
let request_id = context.id().to_string();
let original_stream_flag = request.inner.stream.unwrap_or(false);
let payload_http_headers = if crate::request_trace::payload::http_header_capture_active() {
context
.get_optional::<std::collections::BTreeMap<String, String>>(
crate::request_trace::payload::HTTP_HEADERS_CONTEXT_KEY,
)
.ok()
.flatten()
} else {
None
};
let payload_handle = crate::request_trace::payload::create_handle(
&request,
&request_id,
payload_http_headers,
);
request.enable_usage_for_nonstreaming(original_stream_flag);
request.inner.stream = Some(true);
let thinking_control_from_client = Self::request_has_client_thinking_control(&request);
self.apply_default_thinking_mode(&mut request);
Self::normalize_thinking_arg_with_source(
&mut request,
self.runtime_config.reasoning_parser.as_deref(),
self.tool_call_parser.as_deref(),
thinking_control_from_client,
);
Self::normalize_kimi_k3_named_tool_choice(&mut request, self.tool_call_parser.as_deref());
let response_generator = request.response_generator(context.id().to_string());
let tracker = Some(response_generator.tracker());
let preprocess_options = PreprocessRequestOptions {
preserve_omitted_max_tokens: context
.get::<bool>(PRESERVE_OMITTED_MAX_TOKENS_CONTEXT_KEY)
.ok()
.is_some_and(|flag| *flag),
};
let (mut common_request, annotations, prompt_injected_reasoning, image_tokens) = self
.preprocess_request_with_options(
&request,
tracker.as_deref(),
preprocess_options,
context
.get_optional::<String>(LORA_NAME_CONTEXT_KEY)
.ok()
.flatten()
.map(|name| name.as_ref().clone()),
)
.instrument(preprocessing.clone())
.await?;
attach_request_context_metadata(&mut common_request, &context);
let guided_tool_constraint = self.apply_tool_choice_guided_decoding(
&request,
&mut common_request,
prompt_injected_reasoning.unified,
)?;
let tool_processing_route =
self.tool_processing_route(&request, &guided_tool_constraint)?;
let payload_parsing_options = payload_handle.as_ref().map(|_| {
context
.get_optional::<ParsingOptions>(REQUEST_PARSING_OPTIONS_CONTEXT_KEY)
.ok()
.flatten()
.map(|options| options.as_ref().clone())
.unwrap_or_else(|| {
let mut options = ParsingOptions::new(
self.tool_call_parser.clone(),
self.runtime_config.reasoning_parser.clone(),
);
options.tool_choice = request.inner.tool_choice.clone();
options.guided_tool_constraint = guided_tool_constraint.clone();
options.parallel_tool_calls = request.inner.parallel_tool_calls;
options
.with_tool_call_parsing_enabled(Self::tool_call_parsing_enabled(&request))
.with_move_reasoning_to_content_when_empty(
Self::wants_reasoning_as_content_when_empty(
request.chat_template_args.as_ref(),
),
)
})
});
validate_legacy_jail_nvext_choice_count(
request.inner.n.unwrap_or(1),
request
.nvext
.as_ref()
.and_then(|nvext| nvext.extra_fields.as_deref()),
tool_processing_route.uses_legacy_jail(),
)?;
tracing::trace!(request = ?common_request, ?prompt_injected_reasoning, "Pre-processed request");
let trace_state = crate::request_trace::build_request_end_trace_state(
&common_request,
&tracker,
&context,
self.kv_cache_block_size,
);
let trace_tokens_enabled = trace_state.is_some();
let trace_finish_reason_metadata =
crate::request_trace::finish_reason_metadata_handle(&trace_state);
common_request.tracker = tracker;
let mm_counts = MultimodalCounts::from_preprocessed(&common_request);
let mut response_generator = Box::new(response_generator);
if common_request.prompt_embeds.is_none() {
let isl = common_request.token_ids.len() as u32;
response_generator.update_isl(isl);
}
let common_request = context.map(|_| common_request);
let annotations: Vec<Annotated<NvCreateChatCompletionStreamResponse>> = annotations
.into_iter()
.flat_map(|(k, v)| Annotated::from_annotation(k, &v))
.collect();
let annotations_stream = stream::iter(annotations);
drop(preprocessing);
let response_stream = next.generate(common_request).await?;
let context = response_stream.context();
let stream = Self::transform_postprocessor_stream_with_image_tokens(
response_stream,
response_generator,
context.clone(),
payload_handle.is_some(),
trace_tokens_enabled,
trace_finish_reason_metadata,
mm_counts,
image_tokens,
);
let transformed_stream = self.postprocessor_parsing_stream_with_constraint(
stream,
&request,
prompt_injected_reasoning,
guided_tool_constraint,
tool_processing_route,
)?;
let transformed_stream = Self::normalize_chat_stream_roles(transformed_stream);
let final_stream = if let Some(payload) = payload_handle {
let parsing_options = payload_parsing_options
.expect("payload capture has matching request parsing options");
let (stream, agg_fut) =
crate::request_trace::payload_stream::scan_aggregate_with_future(
Box::pin(transformed_stream),
parsing_options,
);
tokio::spawn(async move {
let outcome = agg_fut.await;
if let Some(reason) = outcome.drop_reason.as_deref() {
tracing::debug!(
request_id = %payload.request_id(),
drop_reason = %reason,
partial_response = outcome.response.is_some(),
"request payload: response aggregation incomplete; emitting record with drop reason"
);
}
payload.emit(outcome.response.map(Arc::new), outcome.drop_reason);
});
stream
} else {
Box::pin(transformed_stream)
};
let final_stream = speculative_prefill::maybe_wrap_stream(
final_stream,
&request,
&request_id,
&next,
&self.formatter,
&self.tokenizer,
&self.speculative_prefill_tasks,
);
let final_stream =
crate::request_trace::wrap_chat_request_end_stream(final_stream, trace_state);
let stream = annotations_stream.chain(final_stream);
Ok(ResponseStream::new(Box::pin(stream), context))
}
}
#[async_trait]
impl
Operator<
SingleIn<NvCreateCompletionRequest>,
ManyOut<Annotated<NvCreateCompletionResponse>>,
SingleIn<PreprocessedRequest>,
ManyOut<Annotated<BackendOutput>>,
> for OpenAIPreprocessor
{
async fn generate(
&self,
request: SingleIn<NvCreateCompletionRequest>,
next: Arc<
dyn AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<BackendOutput>>, Error>,
>,
) -> Result<ManyOut<Annotated<NvCreateCompletionResponse>>, Error> {
let _stage_guard = StageGuard::new(STAGE_PREPROCESS, "");
let (mut request, context) = request.into_parts();
let request_id = context.id().to_string();
let original_stream_flag = request.inner.stream.unwrap_or(false);
request.enable_usage_for_nonstreaming(original_stream_flag);
request.inner.stream = Some(true);
let response_generator = request.response_generator(request_id.clone());
let mut response_generator = Box::new(response_generator);
let tracker = Some(response_generator.tracker());
let mut builder = self.builder_with_lora(
&request,
context
.get_optional::<String>(LORA_NAME_CONTEXT_KEY)
.ok()
.flatten()
.map(|name| name.as_ref().clone()),
)?;
let annotations = if let Some(ref prompt_embeds) = request.inner.prompt_embeds {
builder.token_ids(vec![]); builder.prompt_embeds(Some(prompt_embeds.clone()));
HashMap::new()
} else {
let (token_ids, ann) = self
.gather_tokens(&request, None, tracker.as_deref())
.await?;
builder.token_ids(token_ids);
ann
};
let _ = self
.gather_multi_modal_data(&request, &mut builder, None, &[])
.await?;
let mut common_request = builder.build()?;
Self::validate_preprocessed_token_budget(&common_request, self.token_budget.as_ref())?;
attach_request_context_metadata(&mut common_request, &context);
let trace_state = crate::request_trace::build_request_end_trace_state(
&common_request,
&tracker,
&context,
self.kv_cache_block_size,
);
let trace_tokens_enabled = trace_state.is_some();
let trace_finish_reason_metadata =
crate::request_trace::finish_reason_metadata_handle(&trace_state);
common_request.tracker = tracker;
if common_request.prompt_embeds.is_none() {
let isl = common_request.token_ids.len() as u32;
response_generator.update_isl(isl);
}
let common_request = context.map(|_| common_request);
let annotations: Vec<Annotated<NvCreateCompletionResponse>> = annotations
.into_iter()
.flat_map(|(k, v)| Annotated::from_annotation(k, &v))
.collect();
let annotations_stream = stream::iter(annotations);
drop(_stage_guard);
let response_stream = next.generate(common_request).await?;
let context = response_stream.context();
let stream = Self::transform_postprocessor_stream(
response_stream,
response_generator,
context.clone(),
false,
trace_tokens_enabled,
trace_finish_reason_metadata,
MultimodalCounts::default(),
);
let stream =
crate::request_trace::wrap_completion_request_end_stream(Box::pin(stream), trace_state);
let stream = annotations_stream.chain(stream);
Ok(ResponseStream::new(Box::pin(stream), context))
}
}
#[async_trait]
impl
Operator<
SingleIn<NvCreateEmbeddingRequest>,
ManyOut<Annotated<NvCreateEmbeddingResponse>>,
SingleIn<PreprocessedEmbeddingRequest>,
ManyOut<Annotated<EmbeddingsEngineOutput>>,
> for OpenAIPreprocessor
{
async fn generate(
&self,
request: SingleIn<NvCreateEmbeddingRequest>,
next: Arc<
dyn AsyncEngine<
SingleIn<PreprocessedEmbeddingRequest>,
ManyOut<Annotated<EmbeddingsEngineOutput>>,
Error,
>,
>,
) -> Result<ManyOut<Annotated<NvCreateEmbeddingResponse>>, Error> {
let (request, context) = request.into_parts();
let (preprocessed_request, annotations) =
self.preprocess_embedding_request(&request).await?;
let preprocessed_request = context.map(|_| preprocessed_request);
let response_stream = next.generate(preprocessed_request).await?;
let context = response_stream.context();
let stream = Self::transform_embedding_postprocessor_stream(response_stream, request);
let annotations_stream = stream::iter(
annotations
.into_iter()
.flat_map(|(k, v)| Annotated::from_annotation(k, &v))
.collect::<Vec<_>>(),
);
let combined_stream = annotations_stream.chain(stream);
Ok(ResponseStream::new(Box::pin(combined_stream), context))
}
}
#[cfg(test)]
mod token_data_tests {
use super::*;
use crate::common::checked_file::CheckedFile;
use crate::model_card::{ModelDeploymentCard, ModelInfoType, TokenizerKind};
const LLAMA_DIR: &str = "tests/data/sample-models/mock-llama-3.1-8b-instruct";
const TINYLLAMA_DIR: &str = "tests/data/sample-models/TinyLlama_v1.1";
const NO_VOCAB_SIZE: &str = r#"{"architectures":[],"model_type":"","eos_token_id":2}"#;
const GAP_TOKENIZER: &str = r#"{"version":"1.0","added_tokens":[],
"model":{"type":"WordLevel","vocab":{"<unk>":0,"b":2},"unk_token":"<unk>"}}"#;
const ADDED_TOKEN_TOKENIZER: &str = r#"{"version":"1.0","added_tokens":[{"id":2,
"content":"<x>","special":true,"single_word":false,"lstrip":false,"rstrip":false,
"normalized":false}],
"model":{"type":"WordLevel","vocab":{"<unk>":0,"a":1},"unk_token":"<unk>"}}"#;
fn tinyllama_card(
config: &str,
tokenizer: Option<&str>,
) -> (ModelDeploymentCard, tempfile::TempDir) {
let mut mdc = ModelDeploymentCard::load_from_disk(TINYLLAMA_DIR, None).unwrap();
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("config.json");
std::fs::write(&path, config).unwrap();
mdc.model_info = Some(ModelInfoType::HfConfigJson(
CheckedFile::from_disk(&path).unwrap(),
));
if let Some(tokenizer) = tokenizer {
let path = dir.path().join("tokenizer.json");
std::fs::write(&path, tokenizer).unwrap();
mdc.tokenizer = Some(TokenizerKind::HfTokenizerJson(
CheckedFile::from_disk(&path).unwrap(),
));
}
(mdc, dir)
}
fn tinyllama_with_config(config: &str) -> (Arc<OpenAIPreprocessor>, tempfile::TempDir) {
let (mdc, dir) = tinyllama_card(config, None);
(OpenAIPreprocessor::new(mdc).unwrap(), dir)
}
fn request(token_data: &[u32]) -> NvCreateChatCompletionRequest {
serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hi"}],
"nvext": {"token_data": token_data}
}))
.unwrap()
}
async fn assert_rejected(preprocessor: &OpenAIPreprocessor, token_data: &[u32]) {
let error = preprocessor
.preprocess_request(&request(token_data), None)
.await
.unwrap_err();
assert_eq!(
error
.downcast_ref::<DynamoError>()
.map(DynamoError::error_type),
Some(ErrorType::InvalidArgument),
"{error:#}"
);
}
async fn assert_forwarded(preprocessor: &OpenAIPreprocessor, token_data: &[u32]) {
let (preprocessed, _, _) = preprocessor
.preprocess_request(&request(token_data), None)
.await
.unwrap();
assert_eq!(preprocessed.token_ids.as_slice(), token_data);
}
#[tokio::test]
async fn token_data_is_bounded_by_model_vocab() {
let mdc = ModelDeploymentCard::load_from_disk(LLAMA_DIR, None).unwrap();
let preprocessor = OpenAIPreprocessor::new(mdc).unwrap();
assert_rejected(&preprocessor, &[1, 128256]).await;
assert_rejected(&preprocessor, &[u32::MAX]).await;
assert_forwarded(&preprocessor, &[1, 128255]).await;
}
#[tokio::test]
async fn token_data_is_bounded_by_tokenizer_vocab() {
let (preprocessor, _dir) =
tinyllama_with_config(r#"{"architectures":[],"model_type":"","eos_token_id":2}"#);
assert_rejected(&preprocessor, &[1, 32000]).await;
assert_forwarded(&preprocessor, &[1, 31999]).await;
let (preprocessor, _dir) = tinyllama_with_config(
r#"{"architectures":[],"model_type":"","eos_token_id":2,"vocab_size":31999}"#,
);
assert_forwarded(&preprocessor, &[1, 31999]).await;
assert_rejected(&preprocessor, &[1, 32000]).await;
}
#[tokio::test]
async fn token_data_is_bounded_by_largest_tokenizer_id() {
let (mdc, _dir) = tinyllama_card(NO_VOCAB_SIZE, Some(GAP_TOKENIZER));
let encoding = mdc.tokenizer().unwrap().encode("b").unwrap();
assert_eq!(encoding.token_ids(), [2]);
let preprocessor = OpenAIPreprocessor::new(mdc).unwrap();
assert_forwarded(&preprocessor, &[0, 2]).await;
assert_rejected(&preprocessor, &[0, 3]).await;
}
#[tokio::test]
async fn token_data_is_bounded_by_added_tokens() {
let (mdc, _dir) = tinyllama_card(NO_VOCAB_SIZE, Some(ADDED_TOKEN_TOKENIZER));
let encoding = mdc.tokenizer().unwrap().encode("<x>").unwrap();
assert_eq!(encoding.token_ids(), [2]);
let preprocessor = OpenAIPreprocessor::new(mdc).unwrap();
assert_forwarded(&preprocessor, &[0, 2]).await;
assert_rejected(&preprocessor, &[0, 3]).await;
}
fn completion(prompt: serde_json::Value) -> NvCreateCompletionRequest {
serde_json::from_value(serde_json::json!({"model": "test-model", "prompt": prompt}))
.unwrap()
}
#[tokio::test]
async fn token_prompt_is_bounded_by_model_vocab() {
let mdc = ModelDeploymentCard::load_from_disk(LLAMA_DIR, None).unwrap();
let preprocessor = OpenAIPreprocessor::new(mdc).unwrap();
for prompt in [
serde_json::json!([1, 128256]),
serde_json::json!([[1, 128256]]),
serde_json::json!([u32::MAX]),
] {
let error = preprocessor
.preprocess_request(&completion(prompt), None)
.await
.unwrap_err();
assert_eq!(
error
.downcast_ref::<DynamoError>()
.map(DynamoError::error_type),
Some(ErrorType::InvalidArgument),
"{error:#}"
);
}
let in_range: [u32; 2] = [1, 128255];
for prompt in [serde_json::json!(in_range), serde_json::json!([in_range])] {
let (preprocessed, _, _) = preprocessor
.preprocess_request(&completion(prompt), None)
.await
.unwrap();
assert_eq!(preprocessed.token_ids.as_slice(), in_range);
}
}
#[tokio::test]
async fn token_prompt_is_bounded_by_largest_tokenizer_id() {
let (mdc, _dir) = tinyllama_card(NO_VOCAB_SIZE, Some(GAP_TOKENIZER));
let preprocessor = OpenAIPreprocessor::new(mdc).unwrap();
let (preprocessed, _, _) = preprocessor
.preprocess_request(&completion(serde_json::json!([0, 2])), None)
.await
.unwrap();
assert_eq!(preprocessed.token_ids.as_slice(), [0, 2]);
let error = preprocessor
.preprocess_request(&completion(serde_json::json!([0, 3])), None)
.await
.unwrap_err();
assert_eq!(
error
.downcast_ref::<DynamoError>()
.map(DynamoError::error_type),
Some(ErrorType::InvalidArgument),
"{error:#}"
);
}
#[derive(Default)]
struct RecordingEmbeddingBackend(Mutex<Vec<Vec<Vec<u32>>>>);
#[async_trait]
impl
AsyncEngine<
SingleIn<PreprocessedEmbeddingRequest>,
ManyOut<Annotated<EmbeddingsEngineOutput>>,
Error,
> for RecordingEmbeddingBackend
{
async fn generate(
&self,
request: SingleIn<PreprocessedEmbeddingRequest>,
) -> Result<ManyOut<Annotated<EmbeddingsEngineOutput>>, Error> {
let (request, context) = request.transfer(());
self.0.lock().unwrap().push(request.token_ids);
Ok(ResponseStream::new(
Box::pin(stream::empty()),
context.context(),
))
}
}
#[tokio::test]
async fn embedding_token_input_is_bounded_by_model_vocab() {
let mut mdc = ModelDeploymentCard::load_from_disk(LLAMA_DIR, None).unwrap();
mdc.model_type = crate::model_type::ModelType::Embedding;
let preprocessor = OpenAIPreprocessor::new_for_embeddings(mdc).unwrap();
let backend = Arc::new(RecordingEmbeddingBackend::default());
let next: Arc<
dyn AsyncEngine<
SingleIn<PreprocessedEmbeddingRequest>,
ManyOut<Annotated<EmbeddingsEngineOutput>>,
Error,
>,
> = backend.clone();
let embed = |input: serde_json::Value| {
let request: NvCreateEmbeddingRequest =
serde_json::from_value(serde_json::json!({"model": "test-model", "input": input}))
.unwrap();
PipelineContext::new(request)
};
for input in [
serde_json::json!([1, 128256]),
serde_json::json!([[1, 2], [1, 128256]]),
serde_json::json!([u32::MAX]),
] {
let Err(error) =
Operator::generate(preprocessor.as_ref(), embed(input), next.clone()).await
else {
panic!("an out-of-range token input must fail");
};
assert_eq!(
error
.downcast_ref::<DynamoError>()
.map(DynamoError::error_type),
Some(ErrorType::InvalidArgument),
"{error:#}"
);
}
assert!(
backend.0.lock().unwrap().is_empty(),
"no request reached the backend"
);
for input in [
serde_json::json!([1, 128255]),
serde_json::json!([[1, 2], [3, 128255]]),
] {
Operator::generate(preprocessor.as_ref(), embed(input), next.clone())
.await
.unwrap();
}
assert_eq!(
*backend.0.lock().unwrap(),
vec![vec![vec![1, 128255]], vec![vec![1, 2], vec![3, 128255]]]
);
}
#[tokio::test]
async fn embedding_token_input_is_bounded_by_largest_tokenizer_id() {
let (mut mdc, _dir) = tinyllama_card(NO_VOCAB_SIZE, Some(GAP_TOKENIZER));
mdc.model_type = crate::model_type::ModelType::Embedding;
let preprocessor = OpenAIPreprocessor::new_for_embeddings(mdc).unwrap();
let backend = Arc::new(RecordingEmbeddingBackend::default());
let next: Arc<
dyn AsyncEngine<
SingleIn<PreprocessedEmbeddingRequest>,
ManyOut<Annotated<EmbeddingsEngineOutput>>,
Error,
>,
> = backend.clone();
let embed = |input: serde_json::Value| {
let request: NvCreateEmbeddingRequest =
serde_json::from_value(serde_json::json!({"model": "test-model", "input": input}))
.unwrap();
PipelineContext::new(request)
};
Operator::generate(
preprocessor.as_ref(),
embed(serde_json::json!([0, 2])),
next.clone(),
)
.await
.unwrap();
let Err(error) = Operator::generate(
preprocessor.as_ref(),
embed(serde_json::json!([0, 3])),
next.clone(),
)
.await
else {
panic!("an out-of-range token input must fail");
};
assert_eq!(
error
.downcast_ref::<DynamoError>()
.map(DynamoError::error_type),
Some(ErrorType::InvalidArgument),
"{error:#}"
);
assert_eq!(*backend.0.lock().unwrap(), vec![vec![vec![0, 2]]]);
}
#[test]
fn token_id_bound_ignores_unknown_and_zero_sizes() {
assert_eq!(token_id_bound(Some(151936), Some(151669)), Some(151936));
assert_eq!(token_id_bound(Some(128256), Some(128257)), Some(128257));
assert_eq!(token_id_bound(None, Some(32000)), Some(32000));
assert_eq!(token_id_bound(Some(0), Some(0)), None);
assert_eq!(token_id_bound(None, None), None);
assert!(ensure_token_ids_in_vocab("prompt", &[u32::MAX], None).is_ok());
}
}
#[cfg(test)]
mod strip_tests {
use super::OpenAIPreprocessor;
#[test]
fn test_strip_inline_data_urls_replaces_data_urls() {
let mut messages = serde_json::json!([{
"role": "user",
"content": [
{"type": "text", "text": "What is this?"},
{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBOR...longdata..."}},
{"type": "image_url", "image_url": {"url": "https://example.com/img.png"}}
]
}]);
OpenAIPreprocessor::strip_inline_data_urls(&mut messages);
let parts = messages[0]["content"].as_array().unwrap();
assert_eq!(parts[0]["text"], "What is this?");
assert_eq!(parts[1]["image_url"]["url"], "");
assert_eq!(parts[2]["image_url"]["url"], "https://example.com/img.png");
}
#[test]
fn test_strip_inline_data_urls_handles_video_audio() {
let mut messages = serde_json::json!([{
"role": "user",
"content": [
{"type": "video_url", "video_url": {"url": "data:video/mp4;base64,AAAA..."}},
{"type": "audio_url", "audio_url": {"url": "https://example.com/audio.wav"}}
]
}]);
OpenAIPreprocessor::strip_inline_data_urls(&mut messages);
let parts = messages[0]["content"].as_array().unwrap();
assert_eq!(parts[0]["video_url"]["url"], "");
assert_eq!(
parts[1]["audio_url"]["url"],
"https://example.com/audio.wav"
);
}
#[test]
fn test_strip_inline_data_urls_preserves_text_only() {
let mut messages = serde_json::json!([{
"role": "user",
"content": "plain text message"
}]);
let original = messages.clone();
OpenAIPreprocessor::strip_inline_data_urls(&mut messages);
assert_eq!(messages, original);
}
#[test]
fn test_strip_inline_data_urls_empty_messages() {
let mut messages = serde_json::json!([]);
OpenAIPreprocessor::strip_inline_data_urls(&mut messages);
assert_eq!(messages, serde_json::json!([]));
}
#[test]
fn test_strip_inline_data_urls_malformed_does_not_panic() {
for mut messages in [
serde_json::json!(null),
serde_json::json!({"role": "user"}),
serde_json::json!([{"role": "user"}]),
serde_json::json!([{"role": "user", "content": {"text": "x"}}]),
serde_json::json!([{
"role": "user",
"content": [
"plain",
{"type": "image_url"},
{"type": "image_url", "image_url": "https://example.com/x.png"},
{"type": "image_url", "image_url": {"url": 1}},
{"type": "image_url", "image_url": {"url": "data:image/png;base64,QQ"}}
]
}]),
] {
OpenAIPreprocessor::strip_inline_data_urls(&mut messages);
}
let mut two_inline = serde_json::json!([{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA"}},
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,BBB"}}
]
}]);
OpenAIPreprocessor::strip_inline_data_urls(&mut two_inline);
assert_eq!(two_inline[0]["content"][0]["image_url"]["url"], "");
assert_eq!(two_inline[0]["content"][1]["image_url"]["url"], "");
}
}
#[cfg(test)]
mod extra_args_media_copy_tests {
use super::*;
use crate::model_card::ModelDeploymentCard;
use crate::protocols::common::preprocessor::MultimodalData;
fn test_preprocessor() -> OpenAIPreprocessor {
let mdc = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
let preprocessor = OpenAIPreprocessor::new(mdc).unwrap();
assert!(
preprocessor.media_loader.is_none(),
"default frontend path must not decode media"
);
match Arc::try_unwrap(preprocessor) {
Ok(preprocessor) => preprocessor,
Err(_) => panic!("test preprocessor unexpectedly shared"),
}
}
fn inline_data_url() -> String {
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
.to_string()
}
#[tokio::test]
async fn deepseek_v41_display_name_fallback_orders_tool_media() {
use crate::common::checked_file::CheckedFile;
use crate::model_card::ModelInfoType;
let mut mdc = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
let dir = tempfile::tempdir().unwrap();
let config = dir.path().join("config.json");
std::fs::write(
&config,
r#"{"architectures":[],"model_type":"","eos_token_id":128009}"#,
)
.unwrap();
mdc.model_info = Some(ModelInfoType::HfConfigJson(
CheckedFile::from_disk(&config).unwrap(),
));
mdc.display_name = "DeepSeek-V4.1-Flash".into();
let preprocessor = OpenAIPreprocessor::new(mdc).unwrap();
let first = inline_data_url().replacen("image/png", "image/png;name=a", 1);
let second = inline_data_url().replacen("image/png", "image/png;name=b", 1);
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model":"alias", "messages":[
{"role":"assistant","tool_calls":[
{"id":"a","type":"function","function":{"name":"image","arguments":"{}"}},
{"id":"b","type":"function","function":{"name":"image","arguments":"{}"}}
]},
{"role":"tool","tool_call_id":"b","content":[
{"type":"image_url","image_url":{"url":second}}
]},
{"role":"tool","tool_call_id":"a","content":[
{"type":"image_url","image_url":{"url":first}}
]}
]
}))
.unwrap();
let mut builder = PreprocessedRequestBuilder::default();
builder
.model("alias".into())
.token_ids(Vec::new())
.stop_conditions(Default::default())
.sampling_options(Default::default())
.output_options(Default::default());
preprocessor
.gather_multi_modal_data_with_image_tokens(&request, &mut builder, None, &[])
.await
.unwrap();
let result = builder.build().unwrap();
let media = &result.multi_modal_data.unwrap()["image_url"];
let urls: Vec<_> = media
.iter()
.map(|image| match image {
MultimodalData::Url(url) => url.as_str(),
other => panic!("expected URL, got {other:?}"),
})
.collect();
assert_eq!(urls, [first.as_str(), second.as_str()]);
}
#[tokio::test]
async fn continue_final_message_rejects_reordered_tool_results() {
use crate::engines::ValidateRequest;
use dynamo_renderer::deepseek::v41::DeepSeekV41Formatter;
let mut preprocessor = test_preprocessor();
preprocessor.formatter = Arc::new(DeepSeekV41Formatter);
let image = inline_data_url();
let mut request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model":"alias", "continue_final_message":true, "add_generation_prompt":false,
"reasoning_effort":"none", "messages":[
{"role":"assistant","tool_calls":[
{"id":"a","type":"function","function":{"name":"image","arguments":"{}"}},
{"id":"b","type":"function","function":{"name":"image","arguments":"{}"}}
]},
{"role":"tool","tool_call_id":"b","content":[
{"type":"image_url","image_url":{"url":image}}, {"type":"text","text":"Result B"}
]},
{"role":"tool","tool_call_id":"a","content":[
{"type":"image_url","image_url":{"url":image}}, {"type":"text","text":"Result A"}
]}
]
})).unwrap();
ValidateRequest::validate(&request).unwrap();
let error = preprocessor
.preprocess_request(&request, None)
.await
.unwrap_err();
assert!(
format!("{error:#}").contains("Cannot continue the final message"),
"{error:#}"
);
assert_eq!(
error.downcast_ref::<DynamoError>().unwrap().error_type(),
ErrorType::InvalidArgument
);
request.inner.messages.swap(1, 2);
let ordered = preprocessor.apply_template(&request).unwrap().unwrap();
assert_eq!(ordered.as_str().matches("<|deepseek_image|>").count(), 2);
assert!(ordered.as_str().ends_with("Result B"));
request.inner.messages.swap(1, 2);
request.inner.messages.push(
serde_json::from_value(serde_json::json!({
"role":"assistant","content":"Answer:"
}))
.unwrap(),
);
let assistant = preprocessor.apply_template(&request).unwrap().unwrap();
assert_eq!(
assistant.as_str().matches("<|deepseek_image|>").count(),
2
);
assert!(assistant.as_str().ends_with("Answer:"));
request.inner.messages.pop();
request.common.continue_final_message = Some(false);
let ordinary = preprocessor.apply_template(&request).unwrap().unwrap();
assert_eq!(ordinary.as_str().matches("<|deepseek_image|>").count(), 2);
}
#[tokio::test]
async fn extra_args_messages_omit_inline_data_when_multi_modal_data_present() {
let preprocessor = test_preprocessor();
let data_url = inline_data_url();
let second_data_url = format!("{data_url}QQ");
let https_url = "https://example.com/img.png";
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{
"role": "user",
"content": [
{"type": "text", "text": "describe"},
{"type": "image_url", "image_url": {"url": data_url}},
{"type": "image_url", "image_url": {"url": second_data_url}},
{"type": "image_url", "image_url": {"url": https_url}}
]
}],
"max_tokens": 1
}))
.unwrap();
let (preprocessed, _, _) = preprocessor
.preprocess_request(&request, None)
.await
.unwrap();
let media = preprocessed
.multi_modal_data
.as_ref()
.expect("single media copy lives in multi_modal_data");
let images = media.get("image_url").expect("image_url slot");
assert_eq!(images.len(), 3);
match &images[0] {
MultimodalData::Url(url) => assert_eq!(url.as_str(), data_url),
other => panic!("expected Url for inline image, got {other:?}"),
}
match &images[1] {
MultimodalData::Url(url) => assert_eq!(url.as_str(), second_data_url),
other => panic!("expected Url for second inline image, got {other:?}"),
}
match &images[2] {
MultimodalData::Url(url) => assert_eq!(url.as_str(), https_url),
other => panic!("expected Url for HTTP image, got {other:?}"),
}
let extra_args = preprocessed
.extra_args
.as_ref()
.expect("chat-template extras must remain");
let parts = extra_args["messages"][0]["content"]
.as_array()
.expect("message content parts");
assert_eq!(parts[0]["text"], "describe");
assert_eq!(parts[1]["image_url"]["url"], "");
assert_eq!(parts[2]["image_url"]["url"], "");
assert_eq!(parts[3]["image_url"]["url"], https_url);
assert!(
extra_args.get("formatted_prompt").is_some(),
"LLaVA / TRT-LLM template path needs formatted_prompt"
);
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::protocols::common::extensions::{AGENT_CONTEXT_CONTEXT_KEY, AgentContext};
use crate::protocols::common::preprocessor::MultimodalData;
use crate::protocols::common::{OutputOptions, SamplingOptions, StopConditions};
use dynamo_protocols::types::{
ChatChoiceStream, ChatCompletionStreamResponseDelta, CreateChatCompletionStreamResponse,
FinishReason, Role,
};
use std::path::PathBuf;
#[tokio::test]
async fn eight_family_routes_and_tool_suppression() {
use crate::protocols::openai::chat_completions::unified_parser;
let card = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
let mut preprocessor = match Arc::try_unwrap(OpenAIPreprocessor::new(card).unwrap()) {
Ok(value) => value,
Err(_) => panic!("unexpected shared preprocessor"),
};
for (tool, reason, family) in [
(Some("qwen3_coder"), Some("qwen3"), "qwen3"),
(Some("kimi_k2"), Some("kimi_k25"), "kimi_k2"),
(Some("kimi-k3"), Some("kimi-k3"), "kimi_k3"),
(Some("deepseekv4"), Some("deepseek-v4"), "deepseek_v4"),
(Some("deepseek_v41"), Some("deepseek_v41"), "deepseek_v41"),
(Some("gemma-4"), Some("gemma4"), "gemma4"),
(Some("glm47"), Some("glm45"), "glm47"),
(Some("muse"), None, "muse_glimmer"),
(None, Some("muse_glimmer"), "muse_glimmer"),
] {
preprocessor.tool_call_parser = tool.map(str::to_string);
preprocessor.runtime_config.reasoning_parser = reason.map(str::to_string);
if matches!(family, "gemma4" | "muse_glimmer") {
let prompt = if family == "gemma4" {
"<|channel>"
} else {
"assistant to=self<|message|>"
};
assert!(preprocessor.configured_prompt_prefill(Some(prompt)).unified);
}
for choice in [
serde_json::Value::Null,
serde_json::json!("auto"),
serde_json::json!("none"),
serde_json::json!("required"),
serde_json::json!({"type":"function","function":{"name":"write_file"}}),
] {
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({"model":"test", "messages":[{"role":"user","content":"test"}], "tool_choice":choice, "tools":[{"type":"function","function":{"name":"write_file","parameters":{"type":"object","properties":{}}}}]})).unwrap();
let constraint = crate::preprocessor::tool_choice::guided_tool_constraint(
&request, tool, reason, false,
)
.unwrap();
let route = preprocessor
.tool_processing_route(&request, &constraint)
.unwrap();
if let Some(expected) = unified_parser::selected_family(tool, reason) {
assert!(
matches!(route, ToolProcessingRoute::Unified(actual) if actual == expected),
"{family} {choice}"
);
} else if family == "muse_glimmer"
&& (choice.is_null() || choice == "auto" || choice == "none")
{
assert!(matches!(route, ToolProcessingRoute::MuseUnified(_)));
} else {
assert!(!matches!(route, ToolProcessingRoute::Unified(_)));
}
}
if unified_parser::selected_family(tool, reason).is_none() {
continue;
}
for no_tools in [false, true] {
let mut request = serde_json::json!({"model":"test","messages":[{"role":"user","content":"test"}],"tool_choice":"none"});
if no_tools {
request.as_object_mut().unwrap().remove("tool_choice");
} else {
request["tools"] = serde_json::json!([{"type":"function","function":{"name":"write_file","parameters":{"type":"object","properties":{}}}}]);
}
let request: NvCreateChatCompletionRequest =
serde_json::from_value(request).unwrap();
let source: NvCreateChatCompletionStreamResponse = serde_json::from_value(serde_json::json!({"id":"test","object":"chat.completion.chunk","created":0,"model":"test","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call-1","type":"function","function":{"name":"write_file","arguments":"{}"}}]},"finish_reason":"tool_calls"}]})).unwrap();
let result = preprocessor
.postprocessor_parsing_stream(
futures::stream::iter([Annotated::from_data(source)]),
&request,
false,
false,
)
.unwrap()
.collect::<Vec<_>>()
.await;
assert!(
result
.iter()
.filter_map(|r| r.data.as_ref())
.flat_map(|r| &r.inner.choices)
.all(|c| c.delta.tool_calls.is_none()
&& c.finish_reason != Some(FinishReason::ToolCalls)),
"{family} no_tools={no_tools}"
);
}
}
}
#[tokio::test]
async fn eight_family_request_initialization() {
use crate::protocols::openai::chat_completions::unified_parser;
for (family, native) in unified_parser::tests::eight_family_calls("é initialized") {
let (tool, reason) = match family {
"qwen3" => ("qwen3_coder", "qwen3"),
"kimi_k2" => ("kimi_k2", "kimi_k25"),
"glm47" => ("glm47", "glm45"),
other => (other, other),
};
if unified_parser::selected_family(Some(tool), Some(reason)).is_none() {
continue;
}
let card = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
let mut preprocessor = match Arc::try_unwrap(OpenAIPreprocessor::new(card).unwrap()) {
Ok(value) => value,
Err(_) => panic!("unexpected shared preprocessor"),
};
preprocessor.tool_call_parser = Some(tool.into());
preprocessor.runtime_config.reasoning_parser = Some(reason.into());
let (open, close) = unified_parser::tests::reasoning_markers(family);
for thinking in [false, true] {
for prefilled in [false, true] {
for choice in [
serde_json::Value::Null,
serde_json::json!("auto"),
serde_json::json!("required"),
serde_json::json!({"type":"function","function":{"name":"write_file"}}),
] {
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({"model":"test","messages":[{"role":"user","content":"write it"}], "chat_template_kwargs":{"thinking":thinking},"tool_choice":choice,"tools":[{"type":"function","function":{"name":"write_file","parameters":{"type":"object","properties":{"count":{"type":"integer"},"content":{"type":"string"}}}}}]})).unwrap();
let constraint = crate::preprocessor::tool_choice::guided_tool_constraint(
&request,
Some(tool),
Some(reason),
false,
)
.unwrap();
let args = serde_json::json!({"count":2,"content":"é initialized"});
let payload = match &constraint {
crate::protocols::openai::GuidedToolConstraint::GuidedJsonNamed {
..
} => args.to_string(),
crate::protocols::openai::GuidedToolConstraint::GuidedJsonRequired => {
serde_json::json!([{"name":"write_file","arguments":args}])
.to_string()
}
_ => native.clone(),
};
let prompt = if thinking && prefilled {
open
} else {
"assistant"
};
let injected = preprocessor.configured_prompt_prefill(Some(prompt));
assert_eq!(injected.unified, thinking && prefilled, "{family}");
let raw = if thinking {
format!(
"{}private{close}{payload}",
if prefilled { "" } else { open }
)
} else {
payload
};
let chunks = raw.chars().map(|ch| Annotated::from_data(serde_json::from_value::<NvCreateChatCompletionStreamResponse>(serde_json::json!({"id":"init","object":"chat.completion.chunk","created":0,"model":"test","choices":[{"index":0,"delta":{"content":ch.to_string()},"finish_reason":null}]})).unwrap())).collect::<Vec<_>>();
let result = preprocessor
.postprocessor_parsing_stream_with_constraint(
futures::stream::iter(chunks),
&request,
injected,
constraint.clone(),
preprocessor
.tool_processing_route(&request, &constraint)
.unwrap(),
)
.unwrap()
.collect::<Vec<_>>()
.await;
let choices = result
.iter()
.filter_map(|r| r.data.as_ref())
.flat_map(|r| &r.inner.choices)
.collect::<Vec<_>>();
let reasoning = choices
.iter()
.filter_map(|c| c.delta.reasoning_content.as_deref())
.collect::<String>();
assert_eq!(
reasoning,
if thinking { "private" } else { "" },
"{family} thinking={thinking} prefilled={prefilled} choice={choice}"
);
let content = choices
.iter()
.filter_map(|choice| match &choice.delta.content {
Some(ChatCompletionMessageContent::Text(text)) => {
Some(text.as_str())
}
_ => None,
})
.collect::<String>();
assert_eq!(
content,
String::new(),
"visible reasoning {family}/{choice}"
);
let arguments = choices
.iter()
.filter_map(|c| c.delta.tool_calls.as_ref())
.flatten()
.filter_map(|c| c.function.as_ref()?.arguments.as_deref())
.collect::<String>();
assert_eq!(
serde_json::from_str::<serde_json::Value>(&arguments).unwrap(),
args,
"{family} {choice}"
);
}
}
}
}
}
#[tokio::test]
async fn kimi2_plain_response_without_prompt_prefill_stays_content() {
use crate::protocols::openai::chat_completions::unified_parser;
let mut card = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
card.runtime_config.tool_call_parser = Some("kimi_k2".into());
card.runtime_config.reasoning_parser = Some("kimi_k25".into());
let preprocessor = OpenAIPreprocessor::new(card).unwrap();
for thinking in [None, Some(true), Some(false)] {
let mut request = serde_json::json!({"model":"test", "messages":[{"role":"user", "content":"answer"}], "stream":true});
if let Some(thinking) = thinking {
request["chat_template_args"] = serde_json::json!({"enable_thinking":thinking});
}
let request: NvCreateChatCompletionRequest = serde_json::from_value(request).unwrap();
let source: NvCreateChatCompletionStreamResponse = serde_json::from_value(serde_json::json!({"id":"plain", "object":"chat.completion.chunk", "created":0,"model":"test","choices":[{"index":0,"delta":{"content":"answer"},"finish_reason":"stop"}]})).unwrap();
let output = preprocessor
.postprocessor_parsing_stream_with_constraint(
futures::stream::iter([Annotated::from_data(source)]),
&request,
false.into(),
crate::protocols::openai::GuidedToolConstraint::None,
ToolProcessingRoute::Unified(unified_parser::KIMI_K2_UNIFIED_FAMILY),
)
.unwrap()
.collect::<Vec<_>>()
.await;
let choices = output
.iter()
.filter_map(|response| response.data.as_ref())
.flat_map(|data| &data.inner.choices)
.collect::<Vec<_>>();
let text: String = choices
.iter()
.filter_map(|choice| match &choice.delta.content {
Some(ChatCompletionMessageContent::Text(text)) => Some(text.as_str()),
_ => None,
})
.collect();
assert_eq!(text, "answer", "thinking={thinking:?}");
assert!(
choices
.iter()
.all(|choice| choice.delta.reasoning_content.is_none()),
"thinking={thinking:?}"
);
}
}
#[tokio::test]
async fn gemma_prompt_channel_preserves_optional_label_ownership() {
use crate::protocols::openai::chat_completions::unified_parser;
if unified_parser::selected_family(Some("gemma4"), Some("gemma4")).is_none() {
return;
}
let card = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
let mut preprocessor = match Arc::try_unwrap(OpenAIPreprocessor::new(card).unwrap()) {
Ok(value) => value,
Err(_) => panic!("unexpected shared preprocessor"),
};
preprocessor.tool_call_parser = Some("gemma4".into());
preprocessor.runtime_config.reasoning_parser = Some("gemma4".into());
let request: NvCreateChatCompletionRequest = serde_json::from_value(
serde_json::json!({"model":"test","messages":[{"role":"user","content":"answer"}], "chat_template_args": {"enable_thinking": true}}),
)
.unwrap();
for (prompt, raw, expected) in [
("<|channel>", "private<channel|>answer", "private"),
("<|channel>thought", "\nprivate<channel|>answer", "private"),
("<|channel>thou", "ght\nprivate<channel|>answer", "private"),
(
"<|channel> ",
"thought\nprivate<channel|>answer",
"thought\nprivate",
),
("<|channel>", "thought\nprivate<channel|>answer", "private"),
(
"<|channel>thought\n",
"thought\nprivate<channel|>answer",
"thought\nprivate",
),
] {
let prefill = preprocessor.configured_prompt_prefill(Some(prompt));
for split in raw
.char_indices()
.map(|(at, _)| at)
.chain(std::iter::once(raw.len()))
{
let chunks = [&raw[..split], &raw[split..]].into_iter().map(|content| Annotated::from_data(serde_json::from_value::<NvCreateChatCompletionStreamResponse>(serde_json::json!({"id":"gemma-init","object":"chat.completion.chunk","created":0,"model":"test","choices":[{"index":0,"delta":{"content":content},"finish_reason":null}]})).unwrap())).collect::<Vec<_>>();
let result = preprocessor
.postprocessor_parsing_stream_with_constraint(
futures::stream::iter(chunks),
&request,
prefill,
crate::protocols::openai::GuidedToolConstraint::None,
ToolProcessingRoute::Unified("gemma4"),
)
.unwrap()
.collect::<Vec<_>>()
.await;
let choices = result
.iter()
.filter_map(|r| r.data.as_ref())
.flat_map(|r| &r.inner.choices)
.collect::<Vec<_>>();
assert_eq!(
choices
.iter()
.filter_map(|c| c.delta.reasoning_content.as_deref())
.collect::<String>(),
expected,
"prompt={prompt} split={split}"
);
assert_eq!(
choices
.iter()
.filter_map(|c| match &c.delta.content {
Some(ChatCompletionMessageContent::Text(text)) => Some(text.as_str()),
_ => None,
})
.collect::<String>(),
"answer",
"prompt={prompt} split={split}"
);
}
}
}
struct UnifiedReplayEngine {
preprocessor: Arc<OpenAIPreprocessor>,
native: String,
arguments: String,
release: Option<Arc<tokio::sync::Semaphore>>,
}
#[async_trait]
impl
dynamo_runtime::pipeline::AsyncEngine<
SingleIn<NvCreateChatCompletionRequest>,
ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>,
Error,
> for UnifiedReplayEngine
{
async fn generate(
&self,
request: SingleIn<NvCreateChatCompletionRequest>,
) -> Result<ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>, Error> {
use dynamo_runtime::pipeline::{AsyncEngineContextProvider, ResponseStream};
let (mut request, context) = request.transfer(());
OpenAIPreprocessor::normalize_kimi_k3_named_tool_choice(
&mut request,
self.preprocessor.tool_call_parser.as_deref(),
);
let mut backend_request = preprocessed_budget_request(None);
let constraint = self.preprocessor.apply_tool_choice_guided_decoding(
&request,
&mut backend_request,
false,
)?;
eprintln!("replay installed tool constraint: {constraint:?}");
let raw = match &request.inner.tool_choice {
Some(ChatCompletionToolChoiceOption::Required)
if constraint.installs_guided_json() =>
{
format!(
"[{{\"name\":\"write_file\",\"arguments\":{}}}]",
self.arguments
)
}
Some(ChatCompletionToolChoiceOption::Named(_))
if constraint.installs_guided_json() =>
{
self.arguments.clone()
}
_ => self.native.clone(),
};
let mut chunks = Vec::new();
for (index, character) in raw.chars().enumerate() {
let response = serde_json::from_value(
serde_json::json!({"id":"replay", "object":"chat.completion.chunk", "created":0,"model":"replay","choices":[{"index":0,"delta":{"role":if index == 0 {Some("assistant")} else {None},"content":character.to_string()},"finish_reason":null}]}),
)?;
chunks.push(Annotated::from_data(response));
}
chunks.push(Annotated::from_data(serde_json::from_value(serde_json::json!({"id":"replay","object":"chat.completion.chunk","created":0,"model":"replay","choices":[{"index":0,"delta":{"content":""},"finish_reason":"stop"}]}))?));
let pauses: Vec<usize> = [" STAGE_TWO ", " END"]
.into_iter()
.filter_map(|marker| raw.find(marker).map(|at| raw[..at].chars().count()))
.collect();
let release = self.release.clone();
let source = async_stream::stream! {
for (index, chunk) in chunks.into_iter().enumerate() {
if pauses.contains(&index) && let Some(release) = &release {
release.acquire().await.expect("replay barrier closed").forget();
}
yield chunk;
}
};
let route = self
.preprocessor
.tool_processing_route(&request, &constraint)?;
let parsed = self
.preprocessor
.postprocessor_parsing_stream_with_constraint(
source,
&request,
false.into(),
constraint,
route,
)?;
Ok(ResponseStream::new(Box::pin(parsed), context.context()))
}
}
async fn start_unified_replay(
family: &str,
native: String,
arguments: String,
release: Option<Arc<tokio::sync::Semaphore>>,
) -> (
u16,
dynamo_runtime::CancellationToken,
tokio::task::JoinHandle<anyhow::Result<()>>,
) {
use crate::http::service::service_v2::HttpService;
let (tool, reason) = match family {
"qwen3" => ("qwen3_coder", "qwen3"),
"kimi_k2" => ("kimi_k2", "kimi_k25"),
"glm47" => ("glm47", "glm45"),
other => (other, other),
};
let mut card = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
card.runtime_config.tool_call_parser = Some(tool.into());
card.runtime_config.reasoning_parser = Some(reason.into());
card.worker_type = Some(crate::worker_type::WorkerType::Aggregated);
card.needs.clear();
let preprocessor = OpenAIPreprocessor::new(card.clone()).unwrap();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let service = HttpService::builder()
.host("127.0.0.1")
.port(port)
.enable_chat_endpoints(true)
.enable_cmpl_endpoints(false)
.build()
.unwrap();
let mut workers = crate::discovery::WorkerSet::new(
"replay-workers".into(),
"replay-checksum".into(),
card,
);
workers.chat_engine = Some(Arc::new(UnifiedReplayEngine {
preprocessor,
native,
arguments,
release,
}));
assert!(
service
.model_manager()
.add_worker_set("replay", "replay-workers", workers)
);
let cancel = dynamo_runtime::CancellationToken::new();
let join = service.spawn_with_listener(cancel.clone(), listener).await;
(port, cancel, join)
}
async fn paused_http_tool_progress(guided: bool) {
use crate::protocols::openai::chat_completions::unified_parser;
assert!(
unified_parser::parsers_v2_selected(),
"run with DYN_PARSER_VERSION=2"
);
let client = reqwest::Client::builder().no_proxy().build().unwrap();
let mut failures = Vec::new();
for (shape, fields) in unified_parser::tests::write_file_progress_cases() {
let expected_arguments: serde_json::Value =
serde_json::from_str(&unified_parser::tests::ordered_arguments(&fields)).unwrap();
for (family, native) in unified_parser::tests::native_write_file_calls(&fields) {
for named in if guided {
vec![false, true]
} else {
vec![false]
} {
let release = Arc::new(tokio::sync::Semaphore::new(0));
let (port, cancel, join) = start_unified_replay(
family,
native.clone(),
unified_parser::tests::ordered_arguments(&fields),
Some(release.clone()),
)
.await;
let choice = if !guided {
serde_json::json!("auto")
} else if named {
serde_json::json!({"type":"function","function":{"name":"write_file"}})
} else {
serde_json::json!("required")
};
let request = serde_json::json!({"model":"replay","messages":[{"role":"user","content":"write it"}],"stream":true,"tool_choice":choice,"tools":[{"type":"function","function":{"name":"write_file","parameters":{"type":"object","properties":{"count":{"type":"integer"},"path":{"type":"string"},"content":{"type":"string"}}}}}]});
let response = client
.post(format!("http://127.0.0.1:{port}/v1/chat/completions"))
.json(&request)
.send()
.await
.unwrap();
assert!(response.status().is_success());
let mut source = response.bytes_stream();
let mut body = String::new();
let mut pending = Vec::new();
let mut arguments = String::new();
let mut names = String::new();
let mut ids = Vec::new();
let mut finishes = Vec::new();
let mut progress = Vec::new();
for expected in ["first content", "second content"] {
let reached =
tokio::time::timeout(std::time::Duration::from_secs(2), async {
while let Some(bytes) = source.next().await {
pending.extend_from_slice(&bytes.unwrap());
while let Some(end) =
pending.iter().position(|byte| *byte == b'\n')
{
let line =
String::from_utf8(pending.drain(..=end).collect())
.unwrap();
body.push_str(&line);
read_replay_delta(
&line,
&mut names,
&mut arguments,
&mut ids,
&mut finishes,
);
}
if names == "write_file" && arguments.contains(expected) {
return true;
}
}
false
})
.await
.unwrap_or(false);
progress.push(reached);
release.add_permits(1);
}
tokio::time::timeout(std::time::Duration::from_secs(10), async {
while let Some(bytes) = source.next().await {
pending.extend_from_slice(&bytes.unwrap());
while let Some(end) = pending.iter().position(|byte| *byte == b'\n') {
let line =
String::from_utf8(pending.drain(..=end).collect()).unwrap();
body.push_str(&line);
read_replay_delta(
&line,
&mut names,
&mut arguments,
&mut ids,
&mut finishes,
);
}
}
})
.await
.unwrap();
cancel.cancel();
tokio::time::timeout(std::time::Duration::from_secs(5), join)
.await
.unwrap()
.unwrap()
.unwrap();
assert!(body.contains("[DONE]"));
assert_eq!(names, "write_file", "{family}/{choice}");
assert_eq!(ids.len(), 1, "{family}/{choice}");
assert_eq!(finishes, ["tool_calls"], "{family}/{choice}");
assert_eq!(
serde_json::from_str::<serde_json::Value>(&arguments).unwrap(),
expected_arguments,
"{family}/{choice}"
);
eprintln!(
"paused HTTP {family}/{shape}/{choice}: first={} continued={} final=true",
progress[0], progress[1]
);
if progress != [true, true] {
failures.push(format!("{family}/{shape}/{choice}: {progress:?}"));
}
}
}
}
assert!(
failures.is_empty(),
"paused HTTP progress failed: {failures:?}"
);
}
fn read_replay_delta(
line: &str,
names: &mut String,
arguments: &mut String,
ids: &mut Vec<String>,
finishes: &mut Vec<String>,
) {
let Some(data) = line.trim_end().strip_prefix("data: ") else {
return;
};
if data == "[DONE]" {
return;
}
let event: serde_json::Value = serde_json::from_str(data).unwrap();
for choice in event["choices"].as_array().unwrap() {
if let Some(finish) = choice["finish_reason"].as_str() {
finishes.push(finish.into());
}
if let Some(calls) = choice["delta"]["tool_calls"].as_array() {
for call in calls {
assert_eq!(call["index"], 0);
if let Some(id) = call["id"].as_str() {
ids.push(id.into());
}
if let Some(name) = call["function"]["name"].as_str() {
names.push_str(name);
}
if let Some(value) = call["function"]["arguments"].as_str() {
arguments.push_str(value);
}
}
}
}
}
#[tokio::test]
#[ignore = "Run explicitly with DYN_PARSER_VERSION=2; measures upstream native buffering"]
async fn eight_family_paused_native_http_progress() {
paused_http_tool_progress(false).await;
}
#[tokio::test]
#[ignore = "Run explicitly with DYN_PARSER_VERSION=2; measures installed forced-tool modes"]
async fn eight_family_paused_forced_http_progress() {
paused_http_tool_progress(true).await;
}
#[tokio::test]
async fn eight_family_http_sse_replay() {
use crate::protocols::openai::chat_completions::unified_parser;
if unified_parser::selected_family(Some("qwen3_coder"), Some("qwen3")).is_none() {
return;
}
let client = reqwest::Client::builder().no_proxy().build().unwrap();
for (family, native) in unified_parser::tests::eight_family_calls("é replay") {
let (port, cancel, join) = start_unified_replay(
family,
native,
serde_json::json!({"count":2,"content":"é replay"}).to_string(),
None,
)
.await;
for streaming in [false, true] {
for choice in [
serde_json::Value::Null,
serde_json::json!("auto"),
serde_json::json!("required"),
serde_json::json!({"type":"function","function":{"name":"write_file"}}),
serde_json::json!("none"),
] {
let response = client.post(format!("http://127.0.0.1:{port}/v1/chat/completions")).json(&serde_json::json!({"model":"replay","messages":[{"role":"user","content":"write it"}],"stream":streaming,"tool_choice":choice,"tools":[{"type":"function","function":{"name":"write_file","parameters":{"type":"object","properties":{"count":{"type":"integer"},"content":{"type":"string"}}}}}]})).send().await.unwrap();
let status = response.status();
let body = response.text().await.unwrap();
assert!(status.is_success(), "{family} {choice} {status}: {body}");
let mut arguments = String::new();
if streaming {
assert!(body.contains("[DONE]"), "{family}");
for line in body
.lines()
.filter_map(|line| line.strip_prefix("data: "))
.filter(|line| *line != "[DONE]")
{
let event: serde_json::Value = serde_json::from_str(line).unwrap();
if let Some(calls) =
event["choices"][0]["delta"]["tool_calls"].as_array()
{
for call in calls {
arguments.push_str(
call["function"]["arguments"].as_str().unwrap_or(""),
);
}
}
}
} else {
let result: serde_json::Value = serde_json::from_str(&body).unwrap();
if let Some(calls) =
result["choices"][0]["message"]["tool_calls"].as_array()
{
assert_eq!(calls.len(), 1, "{family}");
arguments.push_str(calls[0]["function"]["arguments"].as_str().unwrap());
}
}
if choice == "none" {
assert!(arguments.is_empty(), "{family}");
} else {
assert_eq!(
serde_json::from_str::<serde_json::Value>(&arguments).unwrap_or_else(
|error| panic!(
"{family} {choice} stream={streaming}: {error}; body={body}"
)
),
serde_json::json!({"count":2,"content":"é replay"}),
"{family} {choice}"
);
}
}
}
cancel.cancel();
tokio::time::timeout(std::time::Duration::from_secs(5), join)
.await
.unwrap()
.unwrap()
.unwrap();
}
}
#[test]
fn deepseek_v41_preserves_markers_and_initializes_backend_reasoning() {
for (tool, reasoning) in [(Some("deepseek_v41"), None), (None, Some("deepseek_v41"))] {
assert!(OpenAIPreprocessor::parser_requires_special_tokens(tool, reasoning).unwrap());
}
assert_eq!(
OpenAIPreprocessor::prompt_injected_reasoning_ended_arg(
Some("deepseek_v41"),
Some("<|Assistant|><think>"),
),
Some(false)
);
assert_eq!(
OpenAIPreprocessor::prompt_injected_reasoning_ended_arg(
Some("deepseek_v41"),
Some("<|Assistant|></think>"),
),
None
);
let disabled = HashMap::from([("thinking".to_string(), serde_json::json!(false))]);
assert!(OpenAIPreprocessor::is_reasoning_disabled_by_request(
Some("deepseek_v41"),
Some(&disabled)
));
}
#[test]
fn lifecycle_preprocessing_requires_frontend_capture() {
const CHILD: &str = "DYNAMO_PREPROCESSOR_LIFECYCLE_TEST_CHILD";
if std::env::var_os(CHILD).is_none() {
let output = crate::test_utils::isolated_command(
"preprocessor::tests::lifecycle_preprocessing_requires_frontend_capture",
)
.env(CHILD, "1")
.env("DYN_LIFECYCLE_TRACE_ENABLED", "true")
.output()
.unwrap();
crate::test_utils::assert_isolated_success(&output);
return;
}
assert!(
LifecycleTrace::from_request_id_with_role(
"knob-probe",
dynamo_runtime::telemetry::LifecycleOperationRole::Worker,
)
.is_enabled()
);
let mut context = PipelineContext::new(());
assert!(!preprocessing_lifecycle(&context).is_enabled());
context.insert(LIFECYCLE_TRACE_CONTEXT_KEY, LifecycleTrace::new(false));
assert!(!preprocessing_lifecycle(&context).is_enabled());
context.insert(LIFECYCLE_TRACE_CONTEXT_KEY, LifecycleTrace::new(true));
assert!(preprocessing_lifecycle(&context).is_enabled());
}
#[test]
fn guided_tool_streaming_release_only_when_guided_json_and_not_rolled_back() {
assert!(
OpenAIPreprocessor::guided_tool_streaming_release(true, false),
"a forced tool_choice with an installed grammar releases incrementally by default"
);
assert!(
!OpenAIPreprocessor::guided_tool_streaming_release(true, true),
"DYN_ENABLE_GUIDED_TOOL_STREAMING=false must fall back to buffer-to-completion \
even when guided JSON is installed"
);
assert!(
!OpenAIPreprocessor::guided_tool_streaming_release(false, false),
"no installed grammar means nothing to release incrementally, rollback or not"
);
assert!(
!OpenAIPreprocessor::guided_tool_streaming_release(false, true),
"no installed grammar means nothing to release incrementally, rollback or not"
);
}
fn parser_route_test_request(
tool_choice: ChatCompletionToolChoiceOption,
) -> NvCreateChatCompletionRequest {
NvCreateChatCompletionRequest {
inner: dynamo_protocols::types::CreateChatCompletionRequestArgs::default()
.model("test")
.messages(vec![dynamo_protocols::types::ChatCompletionRequestMessage::User(
dynamo_protocols::types::ChatCompletionRequestUserMessage {
content: dynamo_protocols::types::ChatCompletionRequestUserMessageContent::Text(
"test".to_string(),
),
name: None,
},
)])
.tools(vec![dynamo_protocols::types::ChatCompletionTool {
r#type: dynamo_protocols::types::ChatCompletionToolType::Function,
function: dynamo_protocols::types::FunctionObject {
name: "get_weather".to_string(),
description: None,
parameters: None,
strict: None,
},
}])
.tool_choice(tool_choice)
.build()
.unwrap(),
common: Default::default(),
nvext: None,
chat_template_args: None,
thinking: None,
thinking_token_budget: None,
media_io_kwargs: None,
return_tokens_as_token_ids: None,
unsupported_fields: Default::default(),
}
}
#[tokio::test]
async fn muse_forced_choice_preserves_calls_after_native_prompt_prefill() {
if crate::test_utils::run_isolated(
concat!(
module_path!(),
"::muse_forced_choice_preserves_calls_after_native_prompt_prefill"
),
&[],
) {
return;
}
let model_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("tests/data/sample-models/mock-llama-3.1-8b-instruct");
let mut card = ModelDeploymentCard::load_from_disk(model_path, None).unwrap();
card.runtime_config.reasoning_parser = Some("muse_glimmer".to_string());
let preprocessor = OpenAIPreprocessor::new(card).unwrap();
let request = parser_route_test_request(ChatCompletionToolChoiceOption::Required);
let prefill = OpenAIPreprocessor::prompt_reasoning_prefill_for_parsers(
None,
Some("muse_glimmer"),
Some("assistant to=self<|message|>"),
);
let raw = r#"[{"name":"get_weather","parameters":{"city":"Tokyo"}}]"#;
let chunks = [
serde_json::json!({"id":"test", "object":"chat.completion.chunk", "created":0,
"model":"test", "choices":[{"index":0,"delta":{"content":raw},"finish_reason":null}]}),
serde_json::json!({"id":"test", "object":"chat.completion.chunk", "created":0,
"model":"test", "choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}),
].into_iter().map(|chunk| Annotated::from_data(
serde_json::from_value::<NvCreateChatCompletionStreamResponse>(chunk).unwrap()
));
let output = preprocessor
.postprocessor_parsing_stream(stream::iter(chunks), &request, prefill, false)
.unwrap();
let result =
crate::protocols::openai::chat_completions::aggregator::DeltaAggregator::apply(
output,
crate::protocols::openai::ParsingOptions::default(),
)
.await
.unwrap();
assert_eq!(
result.inner.choices[0].finish_reason,
Some(dynamo_protocols::types::FinishReason::ToolCalls)
);
let calls = result.inner.choices[0].message.tool_calls.as_ref().unwrap();
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].function.name, "get_weather");
assert_eq!(
serde_json::from_str::<serde_json::Value>(&calls[0].function.arguments).unwrap(),
serde_json::json!({"city":"Tokyo"})
);
}
#[tokio::test]
async fn muse_prompt_prefill_follows_request_route() {
const CHILD: &str = "DYNAMO_MUSE_PREFILL_ROUTE_CHILD";
if std::env::var_os(CHILD).is_none() {
for version in ["auto", "2"] {
let output = crate::test_utils::isolated_command(
"preprocessor::tests::muse_prompt_prefill_follows_request_route",
)
.env(CHILD, version)
.env(env_llm::DYN_PARSER_VERSION, version)
.output()
.unwrap();
crate::test_utils::assert_isolated_success(&output);
}
return;
}
let explicit_v2 = std::env::var(CHILD).unwrap() == "2";
let model_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("tests/data/sample-models/mock-llama-3.1-8b-instruct");
let named = serde_json::from_value::<ChatCompletionToolChoiceOption>(
serde_json::json!({"type":"function","function":{"name":"get_weather"}}),
)
.unwrap();
for alias in ["muse_glimmer", "muse"] {
for (tool, reasoning) in [
(None, Some(alias)),
(Some("muse_glimmer"), Some(alias)),
(Some("muse_glimmer"), None),
] {
for choice in [
ChatCompletionToolChoiceOption::Required,
named.clone(),
ChatCompletionToolChoiceOption::Auto,
ChatCompletionToolChoiceOption::None,
] {
let forced = matches!(
choice,
ChatCompletionToolChoiceOption::Required
| ChatCompletionToolChoiceOption::Named(_)
);
let request = parser_route_test_request(choice.clone());
for structural_tag in [false, true] {
let mut card =
ModelDeploymentCard::load_from_disk(model_path.clone(), None).unwrap();
card.runtime_config.tool_call_parser = tool.map(str::to_string);
card.runtime_config.reasoning_parser = reasoning.map(str::to_string);
let mut preprocessor = OpenAIPreprocessor::new(card).unwrap();
let constraint = crate::preprocessor::tool_choice::guided_tool_constraint(
&request,
tool,
reasoning,
structural_tag,
)
.unwrap();
let route = preprocessor
.tool_processing_route(&request, &constraint)
.unwrap();
let unified = explicit_v2 || (!structural_tag && !forced);
assert_eq!(
matches!(
route,
ToolProcessingRoute::MuseUnified(_)
| ToolProcessingRoute::Unified(_)
),
unified
);
for (prompt, native_open, generic_open) in [
("assistant to=self<|message|>", true, false),
("assistant<think>", false, true),
("assistant", false, false),
] {
Arc::get_mut(&mut preprocessor).unwrap().formatter =
test_prompt_formatter(prompt);
let (_, _, prefill) = preprocessor
.preprocess_request(&request, None)
.await
.unwrap();
let expected_prefill = if unified {
native_open
} else if reasoning.is_some() {
generic_open
} else {
native_open
};
assert_eq!(prefill.for_route(&route), expected_prefill);
let body = match (&choice, structural_tag) {
(ChatCompletionToolChoiceOption::Required, _)
if !unified || !structural_tag =>
{
r#"[{"name":"get_weather","parameters":{"city":"Tokyo"}}]"#
}
(ChatCompletionToolChoiceOption::Named(_), _)
if !unified || !structural_tag =>
{
r#"{"city":"Tokyo"}"#
}
_ => "answer",
};
let raw = if !unified && reasoning.is_some() && generic_open {
format!("thought</think>{body}")
} else if unified && native_open && !forced {
format!(
"thought<|eom|><|start|>assistant to=user<|message|>{body}<|eot|>"
)
} else {
body.to_string()
};
for split in 0..=raw.len() {
let input = || {
[Some(&raw[..split]), Some(&raw[split..]), None].into_iter().map(|text| {
Annotated::from_data(serde_json::from_value::<NvCreateChatCompletionStreamResponse>(serde_json::json!({
"id":"test", "object":"chat.completion.chunk", "created":0, "model":"test",
"choices":[{"index":0,"delta":{"content":text},"finish_reason": if text.is_none() { Some("stop") } else { None }}]
})).unwrap())
}).collect::<Vec<_>>()
};
let actual: Vec<_> = preprocessor
.postprocessor_parsing_stream(
stream::iter(input()),
&request,
prefill,
structural_tag,
)
.unwrap()
.collect()
.await;
let expected: Vec<_> = preprocessor
.postprocessor_parsing_stream(
stream::iter(input()),
&request,
expected_prefill,
structural_tag,
)
.unwrap()
.collect()
.await;
let snapshot = |chunks: &[Annotated<
NvCreateChatCompletionStreamResponse,
>]| {
chunks
.iter()
.map(|chunk| {
let mut value = serde_json::to_value(chunk).unwrap();
if let Some(data) = value.get_mut("data") {
for choice in
data["choices"].as_array_mut().unwrap()
{
if let Some(calls) =
choice["delta"]["tool_calls"].as_array_mut()
{
for call in calls {
if call.get("id").is_some() {
call["id"] = serde_json::json!(
"generated-id"
);
}
}
}
}
}
value
})
.collect::<Vec<_>>()
};
assert_eq!(
snapshot(&actual),
snapshot(&expected),
"{alias} {choice:?} tag={structural_tag} prompt={prompt} split={split}"
);
assert!(actual.iter().all(|chunk| chunk.error.is_none()));
if !unified
&& reasoning.is_some()
&& generic_open
&& split != 0
&& split != raw.len()
{
continue;
}
assert!(
actual
.iter()
.filter_map(|chunk| chunk.data.as_ref())
.flat_map(|chunk| &chunk.inner.choices)
.any(|choice| choice.finish_reason.is_some()),
"missing finish: {alias} {choice:?} tag={structural_tag} prompt={prompt} split={split} chunks={:?}",
snapshot(&actual)
);
let batch = crate::protocols::openai::chat_completions::aggregator::DeltaAggregator::apply(
stream::iter(actual), crate::protocols::openai::ParsingOptions::default(),
).await.unwrap();
if forced && !structural_tag {
let calls =
batch.inner.choices[0].message.tool_calls.as_ref().unwrap();
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].function.name, "get_weather");
assert_eq!(
serde_json::from_str::<serde_json::Value>(
&calls[0].function.arguments
)
.unwrap(),
serde_json::json!({"city":"Tokyo"})
);
assert_eq!(
batch.inner.choices[0].finish_reason,
Some(dynamo_protocols::types::FinishReason::ToolCalls)
);
}
}
}
}
}
}
}
}
#[test]
fn all_requested_families_use_unified_serving_route_for_every_tool_policy() {
const CHILD: &str = "DYNAMO_ALL_UNIFIED_SERVING_ROUTE_CHILD";
if std::env::var_os(CHILD).is_none() {
for (version, experimental) in [(Some("2"), None), (None, Some("true"))] {
let mut command = crate::test_utils::isolated_command(
"preprocessor::tests::all_requested_families_use_unified_serving_route_for_every_tool_policy",
);
command.arg("--test-threads=1").env(CHILD, "1");
if let Some(version) = version {
command.env(env_llm::DYN_PARSER_VERSION, version);
} else {
command.env_remove(env_llm::DYN_PARSER_VERSION);
}
if let Some(experimental) = experimental {
command.env(env_llm::DYN_ENABLE_EXPERIMENTAL_PARSERS_V2, experimental);
} else {
command.env_remove(env_llm::DYN_ENABLE_EXPERIMENTAL_PARSERS_V2);
}
let output = command.output().unwrap();
crate::test_utils::assert_isolated_success(&output);
}
return;
}
let model_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("tests/data/sample-models/mock-llama-3.1-8b-instruct");
for (tool, reasoning, family) in [
("deepseek_v4", "deepseek_v4", "deepseek_v4"),
("deepseek_v41", "deepseek_v41", "deepseek_v41"),
("gemma4", "gemma4", "gemma4"),
("glm47", "glm45", "glm47"),
("kimi_k2", "kimi_k25", "kimi_k2"),
("kimi_k3", "kimi_k3", "kimi_k3"),
("muse", "muse", "muse_glimmer"),
("qwen3_coder", "qwen3", "qwen3"),
] {
let mut card = ModelDeploymentCard::load_from_disk(model_path.clone(), None).unwrap();
card.runtime_config.tool_call_parser = Some(tool.to_string());
card.runtime_config.reasoning_parser = Some(reasoning.to_string());
let preprocessor = OpenAIPreprocessor::new(card).unwrap();
for choice in [
ChatCompletionToolChoiceOption::Auto,
ChatCompletionToolChoiceOption::None,
ChatCompletionToolChoiceOption::Required,
serde_json::from_value(
serde_json::json!({"type":"function","function":{"name":"get_weather"}}),
)
.unwrap(),
] {
for structural_tag in [false, true] {
let mut request = parser_route_test_request(choice.clone());
for has_tools in [true, false] {
if !has_tools {
request.inner.tools = None;
}
let constraint = crate::preprocessor::tool_choice::guided_tool_constraint(
&request,
Some(tool),
Some(reasoning),
structural_tag,
);
if !has_tools
&& !matches!(
choice,
ChatCompletionToolChoiceOption::Auto
| ChatCompletionToolChoiceOption::None
)
{
continue;
}
let constraint = constraint.unwrap();
assert_eq!(
preprocessor
.tool_processing_route(&request, &constraint)
.unwrap(),
ToolProcessingRoute::Unified(family),
"{tool}/{reasoning} {choice:?} structural_tag={structural_tag} has_tools={has_tools}"
);
}
}
}
}
}
#[test]
fn qwen3_coder_uses_unified_route_only_after_explicit_opt_in() {
const CHILD: &str = "DYNAMO_PARSER_VERSION_TEST_QWEN3_OPT_IN";
const MODE: &str = "DYNAMO_PARSER_VERSION_TEST_QWEN3_MODE";
if std::env::var_os(CHILD).is_none() {
for (mode, experimental) in [("default", None), ("opt-in", Some("true"))] {
let mut command = crate::test_utils::isolated_command(
"preprocessor::tests::qwen3_coder_uses_unified_route_only_after_explicit_opt_in",
);
command
.arg("--test-threads=1")
.env(CHILD, "1")
.env(MODE, mode)
.env_remove(env_llm::DYN_PARSER_VERSION);
if let Some(value) = experimental {
command.env(env_llm::DYN_ENABLE_EXPERIMENTAL_PARSERS_V2, value);
} else {
command.env_remove(env_llm::DYN_ENABLE_EXPERIMENTAL_PARSERS_V2);
}
let output = command.output().unwrap();
crate::test_utils::assert_isolated_success(&output);
}
return;
}
let model_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("tests/data/sample-models/mock-llama-3.1-8b-instruct");
let mut card = ModelDeploymentCard::load_from_disk(model_path, None).unwrap();
card.runtime_config.tool_call_parser = Some("qwen3_coder".to_string());
let preprocessor = OpenAIPreprocessor::new(card).unwrap();
let request = parser_route_test_request(ChatCompletionToolChoiceOption::Auto);
let constraint = crate::preprocessor::tool_choice::guided_tool_constraint(
&request,
Some("qwen3_coder"),
None,
false,
)
.unwrap();
let route = preprocessor
.tool_processing_route(&request, &constraint)
.unwrap();
match std::env::var(MODE).as_deref() {
Ok("default") => assert_eq!(
route,
ToolProcessingRoute::LegacyJail(Some("qwen3_coder".to_string()))
),
Ok("opt-in") => assert_eq!(route, ToolProcessingRoute::Unified("qwen3")),
value => panic!("unexpected test mode: {value:?}"),
}
}
#[test]
fn experimental_flag_uses_explicit_v2_compatibility_validation() {
if crate::test_utils::run_isolated(
concat!(
module_path!(),
"::experimental_flag_uses_explicit_v2_compatibility_validation"
),
&[(env_llm::DYN_ENABLE_EXPERIMENTAL_PARSERS_V2, "true")],
) {
return;
}
for (tool, reasoning) in [("hermes", None), ("qwen3_coder", Some("hermes"))] {
let model_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("tests/data/sample-models/mock-llama-3.1-8b-instruct");
let mut card = ModelDeploymentCard::load_from_disk(model_path, None).unwrap();
card.model_type = crate::model_type::ModelType::Chat;
card.runtime_config.tool_call_parser = Some(tool.to_string());
card.runtime_config.reasoning_parser = reasoning.map(str::to_string);
let error = OpenAIPreprocessor::new(card)
.err()
.expect("unsupported V2 configuration must fail");
assert!(error.to_string().contains("DYN_PARSER_VERSION=2"));
}
}
#[test]
fn family_defaults_keep_original_routes_for_every_request_mode() {
const CHILD: &str = "DYNAMO_PARSER_FAMILY_DEFAULT_MODE_CHILD";
if std::env::var_os(CHILD).is_none() {
let output = crate::test_utils::isolated_command(
"preprocessor::tests::family_defaults_keep_original_routes_for_every_request_mode",
)
.env(CHILD, "1")
.env_remove(env_llm::DYN_PARSER_VERSION)
.env_remove(env_llm::DYN_ENABLE_EXPERIMENTAL_PARSERS_V2)
.output()
.unwrap();
crate::test_utils::assert_isolated_success(&output);
return;
}
let model_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("tests/data/sample-models/mock-llama-3.1-8b-instruct");
let families = [
("deepseek_v4", "deepseek_v4", "deepseek_v4"),
("deepseek_v41", "deepseek_v41", "deepseek_v41"),
("gemma4", "gemma4", "gemma4"),
("glm47", "glm45", "glm47"),
("kimi_k2", "kimi_k25", "kimi_k2"),
("kimi_k3", "kimi_k3", "kimi_k3"),
("muse", "muse", "muse_glimmer"),
("qwen3_coder", "qwen3", "qwen3"),
];
let choices = [
ChatCompletionToolChoiceOption::Auto,
ChatCompletionToolChoiceOption::None,
ChatCompletionToolChoiceOption::Required,
serde_json::from_value(
serde_json::json!({"type":"function","function":{"name":"get_weather"}}),
)
.unwrap(),
];
for (tool, reasoning, family) in families {
let mut card = ModelDeploymentCard::load_from_disk(model_path.clone(), None).unwrap();
card.runtime_config.tool_call_parser = Some(tool.to_string());
card.runtime_config.reasoning_parser = Some(reasoning.to_string());
let preprocessor = OpenAIPreprocessor::new(card).unwrap();
let is_default_v2 = family == "deepseek_v41";
for choice in &choices {
for structural_tag in [false, true] {
for has_tools in [true, false] {
let mut request = parser_route_test_request(choice.clone());
if !has_tools {
request.inner.tools = None;
}
let constraint = crate::preprocessor::tool_choice::guided_tool_constraint(
&request,
Some(tool),
Some(reasoning),
structural_tag,
);
let Ok(constraint) = constraint else {
continue;
};
let route = preprocessor
.tool_processing_route(&request, &constraint)
.unwrap();
let is_original_muse_unified = family == "muse_glimmer"
&& !structural_tag
&& matches!(
choice,
ChatCompletionToolChoiceOption::Auto
| ChatCompletionToolChoiceOption::None
);
if is_original_muse_unified {
assert_eq!(
route,
ToolProcessingRoute::MuseUnified("muse_glimmer"),
"{tool}/{reasoning} {choice:?} structural_tag={structural_tag} has_tools={has_tools}"
);
continue;
}
if is_default_v2 {
assert_eq!(
route,
ToolProcessingRoute::Unified(family),
"{tool}/{reasoning} {choice:?} structural_tag={structural_tag} has_tools={has_tools}"
);
} else {
assert!(
!matches!(
route,
ToolProcessingRoute::Unified(_)
| ToolProcessingRoute::MuseUnified(_)
),
"{tool}/{reasoning} {choice:?} structural_tag={structural_tag} has_tools={has_tools}: {route:?}"
);
}
}
}
}
}
}
#[test]
fn explicit_v2_rejects_modes_that_require_the_v1_jail() {
if crate::test_utils::run_isolated(
concat!(
module_path!(),
"::explicit_v2_rejects_modes_that_require_the_v1_jail"
),
&[(env_llm::DYN_PARSER_VERSION, "2")],
) {
return;
}
let model_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("tests/data/sample-models/mock-llama-3.1-8b-instruct");
for (parser, structural_tag) in [
("muse_glimmer", false),
("muse_glimmer", true),
("qwen3_coder", false),
("qwen3_coder", true),
] {
let mut card = ModelDeploymentCard::load_from_disk(model_path.clone(), None).unwrap();
card.runtime_config.tool_call_parser = Some(parser.to_string());
let preprocessor = OpenAIPreprocessor::new(card).unwrap();
let request = parser_route_test_request(ChatCompletionToolChoiceOption::Required);
let constraint = crate::preprocessor::tool_choice::guided_tool_constraint(
&request,
Some(parser),
None,
structural_tag,
)
.unwrap();
assert_eq!(
preprocessor
.tool_processing_route(&request, &constraint)
.unwrap(),
ToolProcessingRoute::Unified(if parser == "muse_glimmer" {
"muse_glimmer"
} else {
"qwen3"
}),
);
}
let card = ModelDeploymentCard::load_from_disk(model_path, None).unwrap();
let preprocessor = OpenAIPreprocessor::new(card).unwrap();
let request = parser_route_test_request(ChatCompletionToolChoiceOption::Required);
let constraint =
crate::preprocessor::tool_choice::guided_tool_constraint(&request, None, None, false)
.unwrap();
let error = preprocessor
.tool_processing_route(&request, &constraint)
.expect_err("explicit v2 must not construct an unconfigured immediate jail");
let dynamo_error = error
.downcast_ref::<DynamoError>()
.expect("request-route errors must retain their InvalidArgument classification");
assert_eq!(dynamo_error.error_type(), ErrorType::InvalidArgument);
assert!(
dynamo_error
.message()
.contains("requires the v1 tool-call jail")
);
match crate::http::service::error::http_action_for_error(dynamo_error) {
crate::http::service::error::ClientErrorAction::Respond { status, .. } => {
assert_eq!(status.as_u16(), 400);
}
crate::http::service::error::ClientErrorAction::NoDelivery => {
panic!("invalid parser-route requests must return an HTTP response");
}
}
}
#[test]
fn legacy_jail_rejects_multiple_choices_for_choice_specific_nvext() {
let engine_data = vec!["engine_data".to_string()];
let request_level = vec!["timing".to_string(), "worker_id".to_string()];
assert!(validate_legacy_jail_nvext_choice_count(2, Some(&engine_data), true).is_err());
assert!(validate_legacy_jail_nvext_choice_count(1, Some(&engine_data), true).is_ok());
assert!(validate_legacy_jail_nvext_choice_count(2, Some(&request_level), true).is_ok());
for route in [
ToolProcessingRoute::Unified("muse_glimmer"),
ToolProcessingRoute::Unified("qwen3"),
ToolProcessingRoute::PassThrough,
] {
assert!(
validate_legacy_jail_nvext_choice_count(
2,
Some(&engine_data),
route.uses_legacy_jail(),
)
.is_ok()
);
}
}
fn chat_stream_chunk(
index: u32,
role: Option<Role>,
) -> Annotated<NvCreateChatCompletionStreamResponse> {
#[allow(deprecated)]
let choice = ChatChoiceStream {
index,
delta: ChatCompletionStreamResponseDelta {
role,
content: Some(ChatCompletionMessageContent::Text("content".to_string())),
tool_calls: None,
function_call: None,
refusal: None,
reasoning_content: None,
},
finish_reason: None,
logprobs: None,
};
Annotated::from_data(NvCreateChatCompletionStreamResponse {
inner: CreateChatCompletionStreamResponse {
id: "test".to_string(),
choices: vec![choice],
created: 0,
model: "test".to_string(),
system_fingerprint: None,
object: "chat.completion.chunk".to_string(),
usage: None,
service_tier: None,
},
nvext: None,
prompt_logprobs: None,
llm_metrics: None,
tool_call_completion: Vec::new(),
})
}
fn reasoning_usage_chunk(
reasoning: Option<&str>,
content: Option<&str>,
chunk_tokens: usize,
) -> Annotated<NvCreateChatCompletionStreamResponse> {
let mut chunk = chat_stream_chunk(0, None);
let data = chunk.data.as_mut().unwrap();
data.inner.choices[0].delta.reasoning_content = reasoning.map(str::to_string);
data.inner.choices[0].delta.content = content
.map(str::to_string)
.map(ChatCompletionMessageContent::Text);
data.llm_metrics = Some(LLMMetricAnnotation {
chunk_tokens,
..Default::default()
});
chunk
}
fn reasoning_usage_trailer(
completion_tokens: u32,
backend_reasoning_tokens: Option<u32>,
) -> Annotated<NvCreateChatCompletionStreamResponse> {
let mut chunk = chat_stream_chunk(0, None);
let data = chunk.data.as_mut().unwrap();
data.inner.choices.clear();
let mut usage = dynamo_protocols::types::CompletionUsage {
prompt_tokens: 5,
completion_tokens,
total_tokens: 5 + completion_tokens,
prompt_tokens_details: None,
completion_tokens_details: None,
};
if let Some(reasoning_tokens) = backend_reasoning_tokens {
usage
.completion_tokens_details
.get_or_insert_default()
.reasoning_tokens = Some(reasoning_tokens);
}
data.inner.usage = Some(usage);
data.llm_metrics = Some(LLMMetricAnnotation {
output_tokens: completion_tokens as usize,
..Default::default()
});
chunk
}
#[tokio::test]
async fn reasoning_usage_estimator_stamps_shared_chat_usage() {
let output = annotate_reasoning_usage(stream::iter(vec![
reasoning_usage_chunk(Some("think"), None, 1),
reasoning_usage_chunk(Some(" more"), None, 1),
reasoning_usage_chunk(None, None, 1),
reasoning_usage_chunk(None, Some("answer"), 1),
reasoning_usage_trailer(4, None),
]))
.collect::<Vec<_>>()
.await;
let usage = output
.last()
.and_then(|response| response.data.as_ref())
.and_then(|chunk| chunk.inner.usage.as_ref())
.unwrap();
assert_eq!(
usage
.completion_tokens_details
.as_ref()
.and_then(|details| details.reasoning_tokens),
Some(3)
);
}
#[tokio::test]
async fn reasoning_usage_estimator_documents_mixed_chunk_overcount() {
let output = annotate_reasoning_usage(stream::iter(vec![
reasoning_usage_chunk(Some("think"), Some("answer"), 4),
reasoning_usage_trailer(4, None),
]))
.collect::<Vec<_>>()
.await;
let usage = output
.last()
.and_then(|response| response.data.as_ref())
.and_then(|chunk| chunk.inner.usage.as_ref())
.unwrap();
assert_eq!(
usage
.completion_tokens_details
.as_ref()
.and_then(|details| details.reasoning_tokens),
Some(4),
"the chunk-granular estimate intentionally attributes the mixed chunk to reasoning"
);
}
#[tokio::test]
async fn reasoning_usage_estimator_preserves_positive_backend_count() {
let output = annotate_reasoning_usage(stream::iter(vec![
reasoning_usage_chunk(Some("two estimated tokens"), None, 2),
reasoning_usage_trailer(2, Some(1)),
]))
.collect::<Vec<_>>()
.await;
let usage = output
.last()
.and_then(|response| response.data.as_ref())
.and_then(|chunk| chunk.inner.usage.as_ref())
.unwrap();
assert_eq!(
usage
.completion_tokens_details
.as_ref()
.and_then(|details| details.reasoning_tokens),
Some(1)
);
}
#[tokio::test]
async fn disabled_unified_reasoning_keeps_interleaved_text_order_in_stream_and_batch() {
use crate::protocols::openai::GuidedToolConstraint;
use crate::protocols::openai::chat_completions::unified_parser;
for (family, input) in [
(
"kimi_k3",
"<|open|>response<|sep|>before<|close|>response<|sep|><|open|>think<|sep|>hidden<|close|>think<|sep|><|open|>response<|sep|>after<|close|>response<|sep|>",
),
(
"muse_glimmer",
"<|start|>assistant to=user<|message|>before<|eom|><|start|>assistant to=self<|message|>hidden<|eom|><|start|>assistant to=user<|message|>after<|eot|>",
),
] {
let batch = unified_parser::parse_complete_with_policy(
family,
input,
&GuidedToolConstraint::None,
&[],
unified_parser::UnifiedRequestPolicy {
reasoning_disabled: true,
..Default::default()
},
)
.unwrap();
assert_eq!(batch.text, "beforehiddenafter", "{family}");
assert!(batch.reasoning.is_empty());
for split in input.char_indices().map(|(offset, _)| offset) {
let (first, second) = input.split_at(split);
let mut terminal = chat_stream_chunk(0, None);
terminal.data.as_mut().unwrap().inner.choices[0]
.delta
.content = Some(ChatCompletionMessageContent::Text(second.to_string()));
let mut first_chunk = chat_stream_chunk(0, None);
first_chunk.data.as_mut().unwrap().inner.choices[0]
.delta
.content = Some(ChatCompletionMessageContent::Text(first.to_string()));
terminal.data.as_mut().unwrap().inner.choices[0].finish_reason =
Some(FinishReason::Stop);
let raw = unified_parser::apply_stream_with_constraint(
stream::iter([first_chunk, terminal]),
None,
GuidedToolConstraint::None,
dynamo_parsers_v2::UnifiedParserStartingState::None,
family,
true,
);
let output =
OpenAIPreprocessor::apply_unified_response_policies(raw, true, false, true)
.collect::<Vec<_>>()
.await;
let mut text = String::new();
for choice in output
.iter()
.filter_map(|r| r.data.as_ref())
.flat_map(|r| &r.inner.choices)
{
assert!(
choice.delta.reasoning_content.is_none(),
"{family} split={split}"
);
if let Some(ChatCompletionMessageContent::Text(t)) = &choice.delta.content {
text.push_str(t);
}
}
assert_eq!(text, batch.text, "{family} split={split}");
}
}
}
#[tokio::test]
async fn force_nonempty_deferral_preserves_reasoning_usage() {
let reasoning = reasoning_usage_chunk(Some("deferred thought"), None, 3);
let mut terminal = chat_stream_chunk(0, None);
let terminal_choice = &mut terminal.data.as_mut().unwrap().inner.choices[0];
terminal_choice.delta.content = None;
terminal_choice.finish_reason = Some(FinishReason::Stop);
let output = OpenAIPreprocessor::apply_unified_response_policies(
stream::iter(vec![reasoning, terminal, reasoning_usage_trailer(3, None)]),
true,
true,
false,
)
.collect::<Vec<_>>()
.await;
let usage = output
.last()
.and_then(|response| response.data.as_ref())
.and_then(|chunk| chunk.inner.usage.as_ref())
.expect("held usage trailer");
assert_eq!(
usage
.completion_tokens_details
.as_ref()
.and_then(|details| details.reasoning_tokens),
Some(3),
"deferral must not erase the parser's reasoning classification"
);
assert_eq!(
output
.iter()
.filter_map(|response| response.data.as_ref())
.filter(|chunk| chunk.llm_metrics.is_some())
.count(),
2,
"one source metric and one usage-trailer metric must survive"
);
}
#[tokio::test]
async fn test_normalize_chat_stream_roles_recovers_missing_first_role_per_choice() {
let input = stream::iter(vec![
chat_stream_chunk(0, None),
chat_stream_chunk(1, None),
chat_stream_chunk(0, Some(Role::Assistant)),
chat_stream_chunk(1, Some(Role::Assistant)),
]);
let output: Vec<_> = OpenAIPreprocessor::normalize_chat_stream_roles(input)
.collect()
.await;
let roles: Vec<_> = output
.iter()
.map(|response| response.data.as_ref().unwrap().inner.choices[0].delta.role)
.collect();
assert_eq!(
roles,
vec![Some(Role::Assistant), Some(Role::Assistant), None, None,]
);
}
#[tokio::test]
async fn legacy_jail_preserves_prompt_logprobs_metadata() {
let expected: crate::protocols::common::llm_backend::PromptLogprobs =
serde_json::from_value(serde_json::json!([
null,
{"17": {"logprob": -0.25, "rank": 1, "decoded_token": " hello"}}
]))
.expect("valid prompt logprobs");
let mut chunk = terminal_chat_stream_chunk();
chunk.data.as_mut().unwrap().prompt_logprobs = Some(Arc::new(expected.clone()));
let output = OpenAIPreprocessor::apply_tool_calling_jail(
None,
None,
None,
false,
false,
stream::iter(vec![chunk]),
)
.collect::<Vec<_>>()
.await;
assert_eq!(
output
.iter()
.filter_map(|response| response.data.as_ref())
.find_map(|data| data.prompt_logprobs.as_deref()),
Some(&expected)
);
}
fn kimi_k3_reasoning_chunk(reasoning: &str) -> Annotated<NvCreateChatCompletionStreamResponse> {
let mut chunk = chat_stream_chunk(0, Some(Role::Assistant));
let choice = &mut chunk.data.as_mut().unwrap().inner.choices[0];
choice.delta.content = None;
choice.delta.reasoning_content = Some(reasoning.to_string());
chunk
}
fn terminal_chat_stream_chunk() -> Annotated<NvCreateChatCompletionStreamResponse> {
let mut chunk = chat_stream_chunk(0, None);
let choice = &mut chunk.data.as_mut().unwrap().inner.choices[0];
choice.delta.content = None;
choice.finish_reason = Some(FinishReason::Stop);
chunk
}
fn glm47_stream_chunk(
content: &str,
finish_reason: Option<FinishReason>,
) -> Annotated<NvCreateChatCompletionStreamResponse> {
let mut chunk = chat_stream_chunk(0, Some(Role::Assistant));
let choice = &mut chunk.data.as_mut().unwrap().inner.choices[0];
choice.delta.content = Some(ChatCompletionMessageContent::Text(content.to_string()));
choice.finish_reason = finish_reason;
chunk
}
async fn apply_glm47_streaming_with_terminal(
chunks: &[&str],
terminal: FinishReason,
) -> Vec<Annotated<NvCreateChatCompletionStreamResponse>> {
let chunks: Vec<String> = chunks.iter().map(|chunk| (*chunk).to_string()).collect();
let chunk_count = chunks.len();
OpenAIPreprocessor::apply_tool_calling_jail(
Some("glm47".to_string()),
None,
None,
false,
false,
stream::iter(chunks.into_iter().enumerate().map(move |(index, content)| {
glm47_stream_chunk(&content, (index + 1 == chunk_count).then_some(terminal))
})),
)
.collect()
.await
}
async fn apply_glm47_streaming_length(
chunks: &[&str],
) -> Vec<Annotated<NvCreateChatCompletionStreamResponse>> {
apply_glm47_streaming_with_terminal(chunks, FinishReason::Length).await
}
fn has_finish_reason(
output: &[Annotated<NvCreateChatCompletionStreamResponse>],
reason: FinishReason,
) -> bool {
output
.iter()
.flat_map(|response| response.data.iter())
.flat_map(|data| data.inner.choices.iter())
.any(|choice| choice.finish_reason == Some(reason))
}
fn emitted_tool_call_count(
output: &[Annotated<NvCreateChatCompletionStreamResponse>],
) -> usize {
output
.iter()
.flat_map(|response| response.data.iter())
.flat_map(|data| data.inner.choices.iter())
.filter_map(|choice| choice.delta.tool_calls.as_ref())
.map(|tool_calls| tool_calls.len())
.sum()
}
fn stream_content(output: &[Annotated<NvCreateChatCompletionStreamResponse>]) -> String {
output
.iter()
.flat_map(|response| response.data.iter())
.flat_map(|data| data.inner.choices.iter())
.filter_map(|choice| match &choice.delta.content {
Some(ChatCompletionMessageContent::Text(content)) => Some(content.as_str()),
_ => None,
})
.collect()
}
fn assert_glm47_streaming_length_output(
output: &[Annotated<NvCreateChatCompletionStreamResponse>],
expected_content: &str,
split: usize,
) {
assert_eq!(
stream_content(output),
expected_content,
"split at byte {split} must reconstruct the complete visible response"
);
assert!(
output
.iter()
.flat_map(|response| response.data.iter())
.flat_map(|data| data.inner.choices.iter())
.any(|choice| choice.finish_reason == Some(FinishReason::Length)),
"split at byte {split} must preserve the terminal length delta"
);
}
#[tokio::test]
async fn glm47_streaming_length_preserves_prose_and_suppresses_incomplete_marker() {
let output = apply_glm47_streaming_length(&[
"I can help. <tool_call>get_weather<arg_key>city</arg_key><arg_value>Par",
])
.await;
assert_eq!(stream_content(&output), "I can help. ");
assert!(
output
.iter()
.flat_map(|response| response.data.iter())
.flat_map(|data| data.inner.choices.iter())
.any(|choice| choice.finish_reason == Some(FinishReason::Length))
);
}
#[tokio::test]
async fn glm47_streaming_length_preserves_quoted_marker_prose() {
let content = r#"The literal "<tool_call>" marker is part of the explanation."#;
for split in content.char_indices().map(|(index, _)| index).skip(1) {
let output =
apply_glm47_streaming_length(&[&content[..split], &content[split..]]).await;
assert_glm47_streaming_length_output(&output, content, split);
}
}
#[tokio::test]
async fn glm47_streaming_length_suppresses_incomplete_marker_at_every_split() {
let input = "I can help. <tool_call>get_weather<arg_key>city</arg_key><arg_value>Par";
for split in input.char_indices().map(|(index, _)| index).skip(1) {
let output = apply_glm47_streaming_length(&[&input[..split], &input[split..]]).await;
assert_glm47_streaming_length_output(&output, "I can help. ", split);
}
}
#[tokio::test]
async fn glm47_streaming_stop_after_bare_tool_call_marker_reports_length() {
let output = apply_glm47_streaming_with_terminal(
&["I'll check. ", "<tool_call>"],
FinishReason::Stop,
)
.await;
assert_eq!(stream_content(&output), "I'll check. ");
assert!(has_finish_reason(&output, FinishReason::Length));
assert!(!has_finish_reason(&output, FinishReason::Stop));
assert_eq!(emitted_tool_call_count(&output), 0);
}
#[tokio::test]
async fn glm47_streaming_stop_after_partial_function_name_reports_length() {
let output = apply_glm47_streaming_with_terminal(
&["I'll check. <tool_call>ipy"],
FinishReason::Stop,
)
.await;
assert_eq!(stream_content(&output), "I'll check. ");
assert!(has_finish_reason(&output, FinishReason::Length));
assert!(!has_finish_reason(&output, FinishReason::Stop));
}
#[tokio::test]
async fn glm47_streaming_stop_after_partial_arguments_reports_length_without_markup() {
let output = apply_glm47_streaming_with_terminal(
&["<tool_call>get_weather<arg_key>city</arg_key><arg_value>Par"],
FinishReason::Stop,
)
.await;
assert!(stream_content(&output).is_empty());
assert!(has_finish_reason(&output, FinishReason::Length));
assert!(!has_finish_reason(&output, FinishReason::Stop));
assert_eq!(emitted_tool_call_count(&output), 0);
}
#[tokio::test]
async fn glm47_streaming_stop_without_marker_keeps_stop() {
let output = apply_glm47_streaming_with_terminal(
&["Done. ", "The worker exited normally."],
FinishReason::Stop,
)
.await;
assert_eq!(stream_content(&output), "Done. The worker exited normally.");
assert!(has_finish_reason(&output, FinishReason::Stop));
assert!(!has_finish_reason(&output, FinishReason::Length));
}
#[tokio::test]
async fn glm47_streaming_stop_with_quoted_marker_keeps_stop() {
let content = r#"The literal "<tool_call>" marker is part of the explanation."#;
let output = apply_glm47_streaming_with_terminal(&[content], FinishReason::Stop).await;
assert!(has_finish_reason(&output, FinishReason::Stop));
assert!(!has_finish_reason(&output, FinishReason::Length));
assert_eq!(emitted_tool_call_count(&output), 0);
}
#[tokio::test]
async fn glm47_streaming_complete_call_on_stop_keeps_tool_calls() {
let output = apply_glm47_streaming_with_terminal(
&["<tool_call>get_weather<arg_key>city</arg_key><arg_value>Paris</arg_value></tool_call>"],
FinishReason::Stop,
)
.await;
assert_eq!(emitted_tool_call_count(&output), 1);
assert!(has_finish_reason(&output, FinishReason::ToolCalls));
assert!(!has_finish_reason(&output, FinishReason::Length));
}
#[tokio::test]
async fn glm47_streaming_prose_only_round_trips_without_retaining_the_response() {
let chunk_text = "the quick brown fox ";
let chunks: Vec<&str> = std::iter::repeat_n(chunk_text, 400).collect();
let output = apply_glm47_streaming_length(&chunks).await;
assert_eq!(stream_content(&output), chunk_text.repeat(400));
}
#[tokio::test]
async fn glm47_streaming_long_prose_survives_buffer_compaction() {
let chunks: Vec<String> = (0..400).map(|i| format!("chunk {i} of prose. ")).collect();
let expected: String = chunks.concat();
let borrowed: Vec<&str> = chunks.iter().map(String::as_str).collect();
let output = apply_glm47_streaming_length(&borrowed).await;
assert_eq!(stream_content(&output), expected);
}
#[tokio::test]
async fn glm47_streaming_length_keeps_an_empty_safe_terminal_recovery_delta() {
let output = apply_glm47_streaming_length(&[
"<tool_call>get_weather<arg_key>city</arg_key><arg_value>Par",
])
.await;
assert!(stream_content(&output).is_empty());
let terminal_choices: Vec<_> = output
.iter()
.flat_map(|response| response.data.iter())
.flat_map(|data| data.inner.choices.iter())
.filter(|choice| choice.finish_reason == Some(FinishReason::Length))
.collect();
assert_eq!(
terminal_choices.len(),
1,
"the recovery event must be observable"
);
assert!(
terminal_choices[0].delta.content.is_none()
&& terminal_choices[0].delta.tool_calls.is_none(),
"the observable recovery delta must not expose raw markup or partial arguments"
);
}
#[test]
fn effective_tool_predicates_match_protocol_definition() {
let cases = [
(
"no tools",
serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "test"}]
}),
),
(
"empty top-level tools",
serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "test"}],
"tools": []
}),
),
(
"top-level tools",
serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "test"}],
"tools": [{
"type": "function",
"function": {
"name": "lookup",
"parameters": {"type": "object"}
}
}]
}),
),
(
"empty system tools",
serde_json::json!({
"model": "test-model",
"messages": [
{"role": "system", "content": "", "tools": []},
{"role": "user", "content": "test"}
]
}),
),
(
"system tools",
serde_json::json!({
"model": "test-model",
"messages": [
{"role": "system", "content": "", "tools": [{"name": "lookup"}]},
{"role": "user", "content": "test"}
]
}),
),
(
"later system tools",
serde_json::json!({
"model": "test-model",
"messages": [
{"role": "system", "content": "", "tools": []},
{"role": "user", "content": "test"},
{"role": "system", "content": "", "tools": [{"name": "lookup"}]}
]
}),
),
];
for (case, value) in cases {
let request: NvCreateChatCompletionRequest =
serde_json::from_value(value).expect("request must deserialize");
assert_eq!(
OpenAIPreprocessor::request_has_effective_tools(&request),
request.inner.has_effective_tools(),
"effective-tool predicates diverged for {case}"
);
}
let developer_tools =
serde_json::from_value::<NvCreateChatCompletionRequest>(serde_json::json!({
"model": "test-model",
"messages": [
{"role": "developer", "content": "policy", "tools": [{"name": "lookup"}]},
{"role": "user", "content": "test"}
]
}));
match developer_tools {
Ok(request) => assert_eq!(
OpenAIPreprocessor::request_has_effective_tools(&request),
request.inner.has_effective_tools(),
"generic predicate must track newly supported developer tools"
),
Err(error) => assert!(
error
.to_string()
.contains("`tools` is only accepted on system messages, not on role developer"),
"unexpected developer-tools rejection: {error}"
),
}
}
const HARMONY_CALL_TOKEN_ID: TokenIdType = 42;
struct HiddenStopTokenizer;
impl crate::tokenizers::traits::Encoder for HiddenStopTokenizer {
fn encode(&self, input: &str) -> anyhow::Result<Encoding> {
if input != "<|call|>" {
anyhow::bail!("unexpected Harmony tool-call end token {input:?}");
}
Ok(Encoding::Sp(vec![HARMONY_CALL_TOKEN_ID]))
}
fn encode_batch(&self, inputs: &[&str]) -> anyhow::Result<Vec<Encoding>> {
inputs.iter().map(|input| self.encode(input)).collect()
}
}
impl crate::tokenizers::traits::Decoder for HiddenStopTokenizer {
fn decode(
&self,
_token_ids: &[TokenIdType],
_skip_special_tokens: bool,
) -> anyhow::Result<crate::tokenizers::traits::DecodeResult> {
Ok(crate::tokenizers::traits::DecodeResult::Complete(
String::new(),
))
}
}
impl Tokenizer for HiddenStopTokenizer {}
fn hidden_stop_preprocessor() -> OpenAIPreprocessor {
let mdc = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
let mut preprocessor = match Arc::try_unwrap(OpenAIPreprocessor::new(mdc).unwrap()) {
Ok(preprocessor) => preprocessor,
Err(_) => panic!("test preprocessor unexpectedly shared"),
};
preprocessor.tool_call_parser = Some("harmony".to_string());
preprocessor.tokenizer = Arc::new(HiddenStopTokenizer);
preprocessor
}
#[test]
fn hidden_stop_path_recognizes_dynamic_system_tools_and_tool_choice_none() {
let dynamic: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "moonshotai/Kimi-K3",
"messages": [
{"role": "system", "content": "", "tools": [{"name": "lookup"}]},
{"role": "user", "content": "test"}
]
}))
.unwrap();
let none: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "moonshotai/Kimi-K3",
"messages": [
{"role": "system", "content": "", "tools": [{"name": "lookup"}]},
{"role": "user", "content": "test"}
],
"tool_choice": "none"
}))
.unwrap();
assert!(dynamic.inner.tools.is_none());
assert!(none.inner.tools.is_none());
let preprocessor = hidden_stop_preprocessor();
let mut dynamic_hidden = vec![HARMONY_CALL_TOKEN_ID, 7];
let dynamic_visible = preprocessor
.remove_tool_parser_end_tokens_from_hidden_stops(&dynamic, &mut dynamic_hidden)
.unwrap();
assert_eq!(dynamic_visible, [HARMONY_CALL_TOKEN_ID]);
assert_eq!(dynamic_hidden, [7]);
let mut none_hidden = vec![HARMONY_CALL_TOKEN_ID, 7];
let none_visible = preprocessor
.remove_tool_parser_end_tokens_from_hidden_stops(&none, &mut none_hidden)
.unwrap();
assert!(none_visible.is_empty());
assert_eq!(none_hidden, [HARMONY_CALL_TOKEN_ID, 7]);
}
async fn apply_kimi_k3_no_tools(
leaked_reasoning: &str,
) -> Vec<Annotated<NvCreateChatCompletionStreamResponse>> {
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "moonshotai/Kimi-K3",
"messages": [{"role": "user", "content": "test"}]
}))
.unwrap();
let tool_call_parsing_enabled = OpenAIPreprocessor::tool_call_parsing_enabled(&request);
assert!(!tool_call_parsing_enabled, "request has no tools");
let jailed = OpenAIPreprocessor::apply_tool_calling_jail(
Some("kimi_k3".to_string()),
request.inner.tool_choice.clone(),
None,
false,
false,
stream::iter(vec![
kimi_k3_reasoning_chunk(leaked_reasoning),
terminal_chat_stream_chunk(),
]),
);
OpenAIPreprocessor::apply_tool_call_response_policy(jailed, tool_call_parsing_enabled)
.collect()
.await
}
async fn apply_kimi_k3_dynamic_tools(
leaked_reasoning: &str,
) -> Vec<Annotated<NvCreateChatCompletionStreamResponse>> {
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "moonshotai/Kimi-K3",
"messages": [
{
"role": "system",
"content": "",
"tools": [{
"name": "lookup",
"parameters": {
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"]
}
}]
},
{"role": "user", "content": "test"}
]
}))
.unwrap();
let tool_call_parsing_enabled = OpenAIPreprocessor::tool_call_parsing_enabled(&request);
assert!(
tool_call_parsing_enabled,
"dynamic system tools grant tool-call response permission"
);
let tool_definitions =
crate::preprocessor::tool_choice::effective_tool_definitions(&request.inner).unwrap();
let jailed = OpenAIPreprocessor::apply_tool_calling_jail(
Some("kimi_k3".to_string()),
request.inner.tool_choice.clone(),
Some(tool_definitions),
false,
false,
stream::iter(vec![
kimi_k3_reasoning_chunk(leaked_reasoning),
terminal_chat_stream_chunk(),
]),
);
OpenAIPreprocessor::apply_tool_call_response_policy(jailed, tool_call_parsing_enabled)
.collect()
.await
}
#[tokio::test]
async fn test_kimi_k3_no_tools_preserves_decoded_content_and_reasoning() {
let responses = apply_kimi_k3_no_tools(concat!(
"The user requested an exact integer.",
"<|open|>response<|sep|>",
"323",
"<|close|>response<|sep|>",
"<|close|>message<|sep|>",
"<|end_of_msg|>"
))
.await;
let choices: Vec<_> = responses
.iter()
.flat_map(|response| response.data.iter())
.flat_map(|data| data.inner.choices.iter())
.collect();
let content: String = choices
.iter()
.filter_map(|choice| match &choice.delta.content {
Some(ChatCompletionMessageContent::Text(text)) => Some(text.as_str()),
_ => None,
})
.collect();
let reasoning: String = choices
.iter()
.filter_map(|choice| choice.delta.reasoning_content.as_deref())
.collect();
assert_eq!(content, "323");
assert_eq!(reasoning, "The user requested an exact integer.");
assert!(!content.contains("<|"));
assert!(!reasoning.contains("<|"));
}
#[tokio::test]
async fn test_kimi_k3_no_tools_suppresses_structured_calls_for_stream_and_batch() {
let responses = apply_kimi_k3_no_tools(concat!(
"Use the calculator.",
"<|open|>tools<|sep|>",
"<|open|>call tool=\"calc\" index=\"1\"<|sep|>",
"<|open|>argument key=\"x\" type=\"number\"<|sep|>323",
"<|close|>argument<|sep|>",
"<|close|>call<|sep|>",
"<|close|>tools<|sep|>",
"<|close|>message<|sep|>",
"<|end_of_msg|>"
))
.await;
let choices: Vec<_> = responses
.iter()
.flat_map(|response| response.data.iter())
.flat_map(|data| data.inner.choices.iter())
.collect();
assert!(
choices
.iter()
.all(|choice| choice.delta.tool_calls.is_none()),
"streaming output must not expose parser-produced tool calls"
);
assert!(
choices
.iter()
.all(|choice| choice.finish_reason != Some(FinishReason::ToolCalls))
);
assert!(
choices
.iter()
.any(|choice| choice.finish_reason == Some(FinishReason::Stop))
);
let content = choices
.iter()
.filter_map(|choice| match &choice.delta.content {
Some(ChatCompletionMessageContent::Text(text)) => Some(text.as_str()),
_ => None,
})
.collect::<String>();
assert!(
content.is_empty(),
"a suppressed tool-only turn deliberately fails closed instead of exposing raw XTML"
);
assert_eq!(
choices
.iter()
.filter_map(|choice| choice.delta.reasoning_content.as_deref())
.collect::<String>(),
"Use the calculator."
);
let response =
crate::protocols::openai::chat_completions::aggregator::DeltaAggregator::apply(
stream::iter(responses),
crate::protocols::openai::ParsingOptions::new(None, None),
)
.await
.unwrap();
let choice = &response.inner.choices[0];
assert!(
choice.message.content.as_ref().is_none_or(|content| {
matches!(content, ChatCompletionMessageContent::Text(text) if text.is_empty())
}),
"batch output must not reconstruct the suppressed call as assistant content"
);
assert!(choice.message.tool_calls.is_none());
assert_eq!(choice.finish_reason, Some(FinishReason::Stop));
assert_eq!(
choice.message.reasoning_content.as_deref(),
Some("Use the calculator.")
);
}
#[tokio::test]
async fn test_kimi_k3_dynamic_only_tools_survive_stream_and_batch_response_policy() {
let responses = apply_kimi_k3_dynamic_tools(concat!(
"Use the dynamic lookup tool.",
"<|open|>tools<|sep|>",
"<|open|>call tool=\"lookup\" index=\"1\"<|sep|>",
"<|open|>argument key=\"query\" type=\"string\"<|sep|>weather",
"<|close|>argument<|sep|>",
"<|close|>call<|sep|>",
"<|close|>tools<|sep|>",
"<|close|>message<|sep|>",
"<|end_of_msg|>"
))
.await;
let choices: Vec<_> = responses
.iter()
.flat_map(|response| response.data.iter())
.flat_map(|data| data.inner.choices.iter())
.collect();
assert!(
choices
.iter()
.any(|choice| choice.delta.tool_calls.is_some()),
"streaming output must retain a dynamically declared tool call"
);
assert!(
choices
.iter()
.any(|choice| choice.finish_reason == Some(FinishReason::ToolCalls)),
"streaming finish reason must remain tool_calls"
);
let response =
crate::protocols::openai::chat_completions::aggregator::DeltaAggregator::apply(
stream::iter(responses),
crate::protocols::openai::ParsingOptions::new(None, None),
)
.await
.unwrap();
let choice = &response.inner.choices[0];
let calls = choice
.message
.tool_calls
.as_ref()
.expect("batch output must retain the dynamic call");
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].function.name, "lookup");
assert_eq!(choice.finish_reason, Some(FinishReason::ToolCalls));
}
#[test]
fn prompt_invalid_request_maps_to_invalid_argument() {
let error = PromptRenderError::invalid_request("unsupported model parameter").into();
let mapped = OpenAIPreprocessor::map_prompt_render_error(error);
let mapped = mapped
.downcast_ref::<DynamoError>()
.expect("prompt validation should map to a DynamoError");
assert!(matches!(mapped.error_type(), ErrorType::InvalidArgument));
assert_eq!(mapped.message(), "unsupported model parameter");
}
#[test]
fn ordinary_prompt_error_maps_to_invalid_argument() {
let mapped = OpenAIPreprocessor::map_prompt_render_error(anyhow::anyhow!(
"template configuration failed"
));
let mapped = mapped
.downcast_ref::<DynamoError>()
.expect("any prompt render failure should map to a DynamoError");
assert!(matches!(mapped.error_type(), ErrorType::InvalidArgument));
assert_eq!(mapped.message(), "template configuration failed");
}
fn url_entry(u: &str) -> MultimodalData {
MultimodalData::Url(url::Url::parse(u).unwrap())
}
fn preprocessed_with_media(media: Option<MultimodalDataMap>) -> PreprocessedRequest {
let mut b = PreprocessedRequest::builder();
b.model("m".to_string())
.token_ids(vec![1, 2, 3])
.stop_conditions(StopConditions::default())
.sampling_options(SamplingOptions::default())
.output_options(OutputOptions::default());
if let Some(m) = media {
b.multi_modal_data(Some(m));
}
b.build().unwrap()
}
#[test]
fn test_multimodal_counts_from_preprocessed_mixed() {
let mut map: MultimodalDataMap = HashMap::new();
map.insert(
"image_url".to_string(),
vec![url_entry("http://x/a.png"), url_entry("http://x/b.png")],
);
map.insert("video_url".to_string(), vec![url_entry("http://x/c.mp4")]);
let req = preprocessed_with_media(Some(map));
let counts = MultimodalCounts::from_preprocessed(&req);
assert_eq!(counts.image, 2);
assert_eq!(counts.video, 1);
assert_eq!(counts.audio, 0);
}
#[test]
fn test_multimodal_counts_from_preprocessed_text_only() {
let req = preprocessed_with_media(None);
let counts = MultimodalCounts::from_preprocessed(&req);
assert_eq!(counts.image, 0);
assert_eq!(counts.video, 0);
assert_eq!(counts.audio, 0);
}
#[test]
fn tool_message_exposes_multimodal_content_parts() {
let message: ChatCompletionRequestMessage = serde_json::from_value(serde_json::json!({
"role": "tool",
"tool_call_id": "call_media",
"content": [
{
"type": "image_url",
"image_url": {
"url": "data:image/png;base64,aGVsbG8="
}
},
{
"type": "video_url",
"video_url": {
"url": "https://example.com/clip.mp4"
}
},
{
"type": "audio_url",
"audio_url": {
"url": "https://example.com/audio.wav"
}
}
]
}))
.unwrap();
let parts: Vec<_> = multimodal_content_parts(&message).unwrap().collect();
assert!(matches!(
parts[0].as_user().as_ref(),
ChatCompletionRequestUserMessageContentPart::ImageUrl(_)
));
assert!(matches!(
parts[1].as_user().as_ref(),
ChatCompletionRequestUserMessageContentPart::VideoUrl(_)
));
assert!(matches!(
parts[2].as_user().as_ref(),
ChatCompletionRequestUserMessageContentPart::AudioUrl(_)
));
}
#[cfg(feature = "mm-routing")]
#[test]
fn image_token_aggregate_requires_complete_trustworthy_counts() {
assert_eq!(aggregate_image_tokens(Some(300), 2, 2, false), Some(300));
assert_eq!(aggregate_image_tokens(Some(300), 2, 3, false), None);
assert_eq!(aggregate_image_tokens(Some(300), 2, 2, true), None);
assert_eq!(aggregate_image_tokens(None, 2, 2, false), None);
assert_eq!(aggregate_image_tokens(Some(0), 0, 0, false), None);
assert_eq!(
checked_add_image_tokens(Some(usize::MAX), 1),
None,
"overflow must omit the aggregate rather than wrap"
);
}
#[cfg(feature = "mm-routing")]
struct RoutingTestTokenizer {
atomic_controls: bool,
fail_plain_text: bool,
}
#[cfg(feature = "mm-routing")]
impl crate::tokenizers::traits::Encoder for RoutingTestTokenizer {
fn encode(&self, input: &str) -> anyhow::Result<Encoding> {
let ids = match input {
"<image>" => vec![18],
"<img>" => vec![19],
"</img>" => vec![20],
_ => input.bytes().map(|byte| 1000 + u32::from(byte)).collect(),
};
Ok(Encoding::Sp(ids))
}
fn encode_batch(&self, inputs: &[&str]) -> anyhow::Result<Vec<Encoding>> {
inputs.iter().map(|input| self.encode(input)).collect()
}
fn encode_segments(
&self,
segments: &[crate::tokenizers::EncodeSegment<'_>],
) -> anyhow::Result<Encoding> {
let mut ids = Vec::new();
for segment in segments {
if segment.allow_special {
if !self.atomic_controls {
ids.extend([9000, 9001]);
continue;
}
ids.push(match segment.text {
"<|media_begin|>" => 163602,
"<|media_content|>" => 163603,
"<|media_end|>" => 163604,
"<img>" => 19,
"<image>" => 18,
"</img>" => 20,
other => anyhow::bail!("unexpected control segment {other:?}"),
});
} else {
if self.fail_plain_text {
anyhow::bail!("injected plain-text routing encode failure");
}
ids.extend(segment.text.bytes().map(|byte| 1000 + u32::from(byte)));
}
}
Ok(Encoding::Sp(ids))
}
}
#[cfg(feature = "mm-routing")]
impl crate::tokenizers::traits::Decoder for RoutingTestTokenizer {
fn decode(
&self,
_token_ids: &[TokenIdType],
_skip_special_tokens: bool,
) -> anyhow::Result<crate::tokenizers::traits::DecodeResult> {
Ok(crate::tokenizers::traits::DecodeResult::Complete(
String::new(),
))
}
}
#[cfg(feature = "mm-routing")]
impl Tokenizer for RoutingTestTokenizer {}
#[cfg(feature = "mm-routing")]
struct ReferenceK3Tokenizer;
#[cfg(feature = "mm-routing")]
impl llm_tokenizer::traits::Encoder for ReferenceK3Tokenizer {
fn encode(
&self,
input: &str,
_add_special_tokens: bool,
) -> anyhow::Result<llm_tokenizer::Encoding> {
Ok(llm_tokenizer::Encoding::Plain(
input.bytes().map(|byte| 1000 + u32::from(byte)).collect(),
))
}
fn encode_batch(
&self,
inputs: &[&str],
add_special_tokens: bool,
) -> anyhow::Result<Vec<llm_tokenizer::Encoding>> {
inputs
.iter()
.map(|input| self.encode(input, add_special_tokens))
.collect()
}
}
#[cfg(feature = "mm-routing")]
impl llm_tokenizer::traits::Decoder for ReferenceK3Tokenizer {
fn decode(&self, _token_ids: &[u32], _skip_special_tokens: bool) -> anyhow::Result<String> {
Ok(String::new())
}
}
#[cfg(feature = "mm-routing")]
impl llm_tokenizer::traits::Tokenizer for ReferenceK3Tokenizer {
fn vocab_size(&self) -> usize {
0
}
fn get_special_tokens(&self) -> &llm_tokenizer::SpecialTokens {
static EMPTY: std::sync::LazyLock<llm_tokenizer::SpecialTokens> =
std::sync::LazyLock::new(llm_tokenizer::SpecialTokens::default);
&EMPTY
}
fn token_to_id(&self, token: &str) -> Option<u32> {
match token {
"<|media_begin|>" => Some(163602),
"<|media_content|>" => Some(163603),
"<|media_end|>" => Some(163604),
"<|media_pad|>" => Some(163605),
_ => None,
}
}
fn id_to_token(&self, _id: u32) -> Option<String> {
None
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
}
#[cfg(feature = "mm-routing")]
#[test]
fn kimi_k3_routing_replacement_includes_dimensions_and_structural_tokens() {
let tokenizer = RoutingTestTokenizer {
atomic_controls: true,
fail_plain_text: false,
};
let layout = resolve_routing_image_prompt_layout(
&tokenizer,
mm_routing::image::ImagePromptKind::KimiK3,
None,
)
.unwrap();
let image = MmImageEntry {
mm_hash: 0x1234,
width: 320,
height: 240,
};
let fill = dynamo_kv_router::protocols::pad_value_for_mm_hash(image.mm_hash);
let dimension_ids: Vec<TokenIdType> = "image 320x240"
.bytes()
.map(|byte| 1000 + u32::from(byte))
.collect();
let mut expanded = vec![7];
append_mm_routing_replacement(&mut expanded, &tokenizer, layout, image, 3).unwrap();
let mut expected = vec![7, 163602];
expected.extend(dimension_ids);
expected.push(163603);
expected.extend([fill; 3]);
expected.push(163604);
assert_eq!(expanded, expected);
}
#[cfg(feature = "mm-routing")]
#[test]
fn kimi_k2_routing_replacement_remains_repeated_pad_only() {
let tokenizer = RoutingTestTokenizer {
atomic_controls: true,
fail_plain_text: false,
};
let image = MmImageEntry {
mm_hash: 0x5678,
width: 320,
height: 240,
};
let fill = dynamo_kv_router::protocols::pad_value_for_mm_hash(image.mm_hash);
let mut expanded = vec![7];
append_mm_routing_replacement(
&mut expanded,
&tokenizer,
RoutingImagePromptLayout::RepeatedPad,
image,
3,
)
.unwrap();
assert_eq!(expanded, vec![7, fill, fill, fill]);
}
#[cfg(feature = "mm-routing")]
#[test]
fn nemotron_routing_replacement_preserves_image_wrappers() {
let tokenizer = RoutingTestTokenizer {
atomic_controls: false,
fail_plain_text: false,
};
let layout = resolve_routing_image_prompt_layout(
&tokenizer,
mm_routing::image::ImagePromptKind::Nemotron,
Some(18),
)
.unwrap();
let image = MmImageEntry {
mm_hash: 0x1234,
width: 224,
height: 224,
};
let fill = dynamo_kv_router::protocols::pad_value_for_mm_hash(image.mm_hash);
let mut expanded = vec![7];
append_mm_routing_replacement(&mut expanded, &tokenizer, layout, image, 3).unwrap();
assert_eq!(expanded, vec![7, 19, fill, fill, fill, 20]);
}
#[cfg(feature = "mm-routing")]
#[test]
fn nemotron_layout_rejects_tokenizer_and_config_id_mismatch() {
let tokenizer = RoutingTestTokenizer {
atomic_controls: true,
fail_plain_text: false,
};
let error = resolve_routing_image_prompt_layout(
&tokenizer,
mm_routing::image::ImagePromptKind::Nemotron,
Some(99),
)
.unwrap_err();
assert!(error.to_string().contains("does not match config id"));
}
#[cfg(feature = "mm-routing")]
#[test]
fn kimi_k3_layout_resolution_rejects_non_atomic_control_tokens() {
let tokenizer = RoutingTestTokenizer {
atomic_controls: false,
fail_plain_text: false,
};
let error = resolve_routing_image_prompt_layout(
&tokenizer,
mm_routing::image::ImagePromptKind::KimiK3,
None,
)
.unwrap_err();
assert!(error.to_string().contains("expected exactly one"));
}
#[cfg(feature = "mm-routing")]
#[test]
fn kimi_k3_expansion_reports_reference_exact_prompt_length() {
let tokenizer = RoutingTestTokenizer {
atomic_controls: true,
fail_plain_text: false,
};
let layout = resolve_routing_image_prompt_layout(
&tokenizer,
mm_routing::image::ImagePromptKind::KimiK3,
None,
)
.unwrap();
let image = MmImageEntry {
mm_hash: 0x1234,
width: 320,
height: 240,
};
let dimension_token_count = "image 320x240".len();
let image_token_count = 3;
let (expanded, expanded_prompt_len) = expand_mm_routing_tokens(
&tokenizer,
layout,
None,
163605,
&[image],
&[image_token_count],
&[7, 163605, 8],
)
.unwrap();
let reference_len = 2 + 3 + dimension_token_count + image_token_count;
assert_eq!(expanded_prompt_len, reference_len);
assert_eq!(expanded.len(), reference_len);
}
#[cfg(feature = "mm-routing")]
#[test]
fn kimi_k3_exact_prompt_length_matches_upstream_processor_replacement() {
use llm_multimodal::vision::{PreProcessorConfig, VisionProcessorRegistry};
let config = serde_json::json!({
"model_type": "kimi_k3",
"media_placeholder_token_id": 163605
});
let reference_tokenizer = ReferenceK3Tokenizer;
let metadata = llm_multimodal::ModelMetadata {
model_id: "moonshotai/Kimi-K3",
tokenizer: &reference_tokenizer,
config: &config,
};
let registry = llm_multimodal::ModelRegistry::new();
let spec = registry.lookup(&metadata).unwrap();
let vision_registry = VisionProcessorRegistry::with_defaults();
let processor = vision_registry
.find("moonshotai/Kimi-K3", Some("kimi_k3"))
.unwrap();
let image = image::DynamicImage::new_rgb8(320, 240);
let processed = processor
.preprocess(&[image], &PreProcessorConfig::default())
.unwrap();
let reference_replacements = spec.prompt_replacements(&metadata, &processed).unwrap();
let reference_replacement = &reference_replacements[0];
let tokenizer = RoutingTestTokenizer {
atomic_controls: true,
fail_plain_text: false,
};
let layout = resolve_routing_image_prompt_layout(
&tokenizer,
mm_routing::image::ImagePromptKind::KimiK3,
None,
)
.unwrap();
let image_entry = MmImageEntry {
mm_hash: 0x1234,
width: 320,
height: 240,
};
let (expanded, expanded_prompt_len) = expand_mm_routing_tokens(
&tokenizer,
layout,
None,
163605,
&[image_entry],
&processed.feature_token_counts,
&[7, 163605, 8],
)
.unwrap();
let reference_prompt_len = 2 + reference_replacement.tokens.len();
assert_eq!(expanded_prompt_len, reference_prompt_len);
assert_eq!(expanded.len(), reference_prompt_len);
}
#[cfg(feature = "mm-routing")]
#[test]
fn kimi_k3_plain_text_encode_failure_is_routing_only() {
let tokenizer = RoutingTestTokenizer {
atomic_controls: true,
fail_plain_text: true,
};
let layout = resolve_routing_image_prompt_layout(
&tokenizer,
mm_routing::image::ImagePromptKind::KimiK3,
None,
)
.unwrap();
let image = MmImageEntry {
mm_hash: 0x1234,
width: 320,
height: 240,
};
let result = try_expand_mm_routing_tokens(
&tokenizer,
layout,
None,
163605,
&[image],
&[3],
&[7, 163605, 8],
"moonshotai/Kimi-K3",
);
assert!(result.is_none());
}
#[cfg(feature = "mm-routing")]
#[test]
fn image_expansion_applies_placeholders_in_request_order() {
let tokenizer = RoutingTestTokenizer {
atomic_controls: true,
fail_plain_text: false,
};
let images = [
MmImageEntry {
mm_hash: 0x1234,
width: 320,
height: 240,
},
MmImageEntry {
mm_hash: 0x5678,
width: 640,
height: 480,
},
];
let first_fill = dynamo_kv_router::protocols::pad_value_for_mm_hash(images[0].mm_hash);
let second_fill = dynamo_kv_router::protocols::pad_value_for_mm_hash(images[1].mm_hash);
let (expanded, prompt_len) = expand_mm_routing_tokens(
&tokenizer,
RoutingImagePromptLayout::RepeatedPad,
Some(1),
10,
&images,
&[2, 3],
&[7, 10, 8, 10, 9],
)
.unwrap();
assert_eq!(
expanded,
[
1,
7,
first_fill,
first_fill,
8,
second_fill,
second_fill,
second_fill,
9,
]
);
assert_eq!(prompt_len, expanded.len());
}
#[cfg(feature = "mm-routing")]
#[test]
fn exact_routing_rejects_image_placeholder_count_mismatch() {
let model_dir = tempfile::tempdir().unwrap();
std::fs::write(model_dir.path().join("preprocessor_config.json"), "{}").unwrap();
let counter = mm_routing::image::ImageRoutingProcessor::try_new(
"Qwen/Qwen3-VL-2B-Instruct",
Some("qwen3_vl"),
model_dir.path(),
)
.unwrap();
let mdc = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
let mut preprocessor = match Arc::try_unwrap(OpenAIPreprocessor::new(mdc).unwrap()) {
Ok(preprocessor) => preprocessor,
Err(_) => panic!("test preprocessor unexpectedly shared"),
};
preprocessor.image_token_counter = Some(counter);
preprocessor.routing_image_token_id = Some(10);
preprocessor.routing_image_prompt_layout = Some(RoutingImagePromptLayout::RepeatedPad);
preprocessor.kv_cache_block_size = 16;
let image = [MmRoutingEntry::Image {
mm_hash: 0x1234,
width: 320,
height: 240,
}];
assert!(
preprocessor
.build_mm_exact_routing_info(&image, &[7, 8], None)
.is_none(),
"missing placeholder must fail closed at the production boundary"
);
assert!(
preprocessor
.build_mm_exact_routing_info(&image, &[7, 10, 10, 8], None)
.is_none(),
"extra placeholder must fail closed at the production boundary"
);
}
#[cfg(feature = "mm-routing")]
#[tokio::test]
async fn nemotron_exact_routing_uses_vllm_dummy_text_budget() {
let model_dir = tempfile::tempdir().unwrap();
std::fs::write(
model_dir.path().join("config.json"),
serde_json::json!({
"architectures": [mm_routing::nemotron::MODEL_TYPE],
"model_type": mm_routing::nemotron::MODEL_TYPE,
"force_image_size": 512,
"patch_size": 16,
"downsample_ratio": 0.5,
"img_context_token": mm_routing::nemotron::IMAGE_CONTEXT,
"img_context_token_id": 18,
"img_start_token": mm_routing::nemotron::IMAGE_START,
"img_end_token": mm_routing::nemotron::IMAGE_END,
"vision_config": {
"args": {
"min_num_patches": 1024,
"max_num_patches": 13312
}
}
})
.to_string(),
)
.unwrap();
std::fs::write(
model_dir.path().join("preprocessor_config.json"),
serde_json::json!({
"image_processor_type": "NemotronH_Nano_Omni_Reasoning_V3ImageProcessor",
"image_size": 512,
"patch_size": 16,
"downsample_ratio": 0.5
})
.to_string(),
)
.unwrap();
let counter = mm_routing::image::ImageRoutingProcessor::try_new(
"nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-FP8",
Some(mm_routing::nemotron::MODEL_TYPE),
model_dir.path(),
)
.unwrap();
let image_text_prompt_len = counter
.context_budget_text_len(
Arc::new(RoutingTestTokenizer {
atomic_controls: true,
fail_plain_text: false,
}),
Some("before<image>middle<image>after<image>"),
3,
)
.await;
assert_eq!(image_text_prompt_len, Some(0));
let mdc = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
let mut preprocessor = match Arc::try_unwrap(OpenAIPreprocessor::new(mdc).unwrap()) {
Ok(preprocessor) => preprocessor,
Err(_) => panic!("test preprocessor unexpectedly shared"),
};
preprocessor.image_token_counter = Some(counter);
preprocessor.routing_image_token_id = Some(18);
preprocessor.routing_image_prompt_layout = Some(
resolve_routing_image_prompt_layout(
&RoutingTestTokenizer {
atomic_controls: true,
fail_plain_text: false,
},
mm_routing::image::ImagePromptKind::Nemotron,
Some(18),
)
.unwrap(),
);
preprocessor.context_length = 4096;
preprocessor.kv_cache_block_size = 16;
let entries = [0x1234, 0x5678, 0x9abc].map(|mm_hash| MmRoutingEntry::Image {
mm_hash,
width: 1920,
height: 1080,
});
let routing = preprocessor
.build_mm_exact_routing_info(
&entries,
&[7, 18, 8, 18, 9, 18, 10],
image_text_prompt_len,
)
.expect("Nemotron routing info should be exact");
assert_eq!(routing.expanded_prompt_len, 4 + 3 * (1323 + 2));
assert_eq!(routing.routing_token_ids.len(), 3984);
assert_eq!(routing.routing_token_ids[0], 7);
assert_eq!(routing.routing_token_ids[1], 19);
assert_eq!(routing.routing_token_ids[1325], 20);
assert_eq!(routing.routing_token_ids[1326], 8);
assert!(routing.block_mm_infos.is_empty());
assert!(
preprocessor
.build_mm_exact_routing_info(&entries, &[7, 18, 8, 18, 9, 18, 10], None)
.is_none(),
"missing rendered-text length must fail closed"
);
}
#[cfg(feature = "mm-routing")]
#[test]
fn mm_processor_override_is_conservative() {
assert!(!has_mm_processor_override(None));
assert!(!has_mm_processor_override(Some(&serde_json::Value::Null)));
assert!(!has_mm_processor_override(Some(&serde_json::json!({}))));
assert!(has_mm_processor_override(Some(&serde_json::json!([]))));
assert!(has_mm_processor_override(Some(
&serde_json::json!({"min_pixels": 64})
)));
}
#[cfg(feature = "mm-routing")]
#[test]
fn processor_kwargs_disable_exact_mm_routing() {
assert!(exact_mm_routing_preconditions_met(false, 1, 1, false));
assert!(!exact_mm_routing_preconditions_met(false, 1, 1, true));
assert!(!exact_mm_routing_preconditions_met(true, 1, 1, false));
assert!(!exact_mm_routing_preconditions_met(false, 0, 1, false));
}
#[cfg(all(feature = "mm-routing", feature = "media-ffmpeg"))]
#[test]
fn decoded_video_hash_requires_exact_routing_eligibility() {
assert!(should_hash_decoded_video(true, false, false, true));
assert!(!should_hash_decoded_video(false, false, false, true));
assert!(!should_hash_decoded_video(true, true, false, true));
assert!(!should_hash_decoded_video(true, false, true, true));
assert!(!should_hash_decoded_video(true, false, false, false));
}
#[cfg(all(
feature = "mm-routing",
feature = "media-ffmpeg",
feature = "testing-nixl"
))]
#[tokio::test]
async fn adjacent_videos_reach_media_loader_with_hashing_disabled() {
let video_bytes = include_bytes!("../tests/data/media/240p_10.mp4");
let mut server = mockito::Server::new_async().await;
let video_mock = server
.mock("GET", "/video.mp4")
.with_status(200)
.with_header("content-type", "video/mp4")
.with_body(&video_bytes[..])
.expect(2)
.create_async()
.await;
let mdc = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
let mut preprocessor = Arc::try_unwrap(OpenAIPreprocessor::new(mdc).unwrap())
.unwrap_or_else(|_| panic!("test preprocessor unexpectedly shared"));
let media_decoder: MediaDecoder = serde_json::from_value(serde_json::json!({
"video": {"num_frames": 2}
}))
.unwrap();
let media_fetcher = MediaFetcher {
allow_direct_ip: true,
allow_direct_port: true,
allow_private_ips: true,
..Default::default()
};
preprocessor.media_loader = Some(
MediaLoader::new(media_decoder, Some(media_fetcher))
.expect("test media loader must initialize"),
);
preprocessor.video_routing_processor = Some(mm_routing::VideoRoutingProcessor::test_stub());
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{
"role": "user",
"content": [
{
"type": "video_url",
"video_url": {"url": format!("{}/video.mp4", server.url())}
},
{
"type": "video_url",
"video_url": {"url": format!("{}/video.mp4", server.url())}
}
]
}],
"max_tokens": 1
}))
.unwrap();
let mut builder = PreprocessedRequest::builder();
builder
.model("test-model".to_string())
.token_ids(Vec::new())
.stop_conditions(Default::default())
.sampling_options(Default::default())
.output_options(Default::default());
let (routing_entries, _) = preprocessor
.gather_multi_modal_data_with_image_tokens(&request, &mut builder, None, &[])
.await
.unwrap();
let preprocessed = builder.build().unwrap();
let videos = preprocessed
.multi_modal_data
.as_ref()
.and_then(|media| media.get("video_url"))
.expect("both decoded videos must be forwarded");
assert_eq!(videos.len(), 2);
for video in videos {
let MultimodalData::Decoded(descriptor) = video else {
panic!("video must reach the media loader and be decoded");
};
assert!(
descriptor.content_hash().is_none(),
"adjacent videos must reach the media loader with hashing disabled"
);
}
assert!(routing_entries.is_empty());
assert!(preprocessed.mm_routing_info.is_none());
video_mock.assert_async().await;
}
#[test]
fn replace_reserved_media_slot_preserves_alignment_and_returns_errors() {
let mut map = HashMap::from([(
"image_url".to_string(),
vec![
url_entry("http://x/a.png"),
MultimodalData::UuidOnly("cached-b".to_string()),
url_entry("http://x/c.png"),
],
)]);
OpenAIPreprocessor::replace_reserved_media_slot(
&mut map,
"image_url",
2,
MultimodalData::RawUrl("decoded-c".to_string()),
)
.unwrap();
let images = &map["image_url"];
assert!(matches!(images[0], MultimodalData::Url(_)));
assert!(matches!(
&images[1],
MultimodalData::UuidOnly(uuid) if uuid == "cached-b"
));
assert!(matches!(
&images[2],
MultimodalData::RawUrl(value) if value == "decoded-c"
));
let error = OpenAIPreprocessor::replace_reserved_media_slot(
&mut map,
"image_url",
3,
MultimodalData::RawUrl("out-of-range".to_string()),
)
.expect_err("an out-of-range reserved slot must return an error");
assert!(error.to_string().contains("image_url[3]"));
}
#[test]
fn routing_priorities_keep_strict_tier_independent() {
let hints = crate::protocols::common::extensions::AgentHints {
priority: Some(-3),
strict_priority: Some(7),
..Default::default()
};
assert_eq!(
routing_priorities(Some(&hints)),
(Some(0.0), Some(7), Some(-3))
);
assert_eq!(routing_priorities(None), (None, None, None));
}
fn test_llm_metrics_annotation() -> LLMMetricAnnotation {
LLMMetricAnnotation {
input_tokens: 10,
output_tokens: 20,
chunk_tokens: 3,
cached_tokens: Some(4),
image_tokens: Some(512),
prefill_worker_id: Some(1),
prefill_dp_rank: Some(2),
prefill_worker_type: Some("prefill".to_string()),
decode_worker_id: Some(3),
decode_dp_rank: Some(4),
decode_worker_type: Some("decode".to_string()),
tokenize_latency: Some(std::time::Duration::from_millis(5)),
detokenize_total_latency: Some(std::time::Duration::from_micros(50)),
detokenize_count: Some(6),
..Default::default()
}
}
#[test]
fn llm_metrics_from_annotation_recognizes_both_metric_event_tags() {
let base = test_llm_metrics_annotation()
.to_annotation::<()>()
.expect("metrics annotation serializes");
for tag in [ANNOTATION_LLM_METRICS, ANNOTATION_PAYLOAD_USAGE] {
let tagged = Annotated::<()> {
id: None,
data: None,
event: Some(tag.to_string()),
comment: base.comment.clone(),
error: None,
};
let metrics = LLMMetricAnnotation::from_annotation(&tagged)
.expect("metrics annotation parses")
.unwrap_or_else(|| panic!("metrics recognized for tag {tag}"));
assert_eq!(metrics.input_tokens, 10);
assert_eq!(metrics.output_tokens, 20);
assert_eq!(metrics.image_tokens, Some(512));
assert_eq!(metrics.detokenize_count, Some(6));
}
}
#[test]
fn llm_metrics_from_annotation_ignores_untagged_and_other_events() {
let untagged = Annotated::<()> {
id: None,
data: None,
event: None,
comment: Some(vec!["{\"input_tokens\":1}".to_string()]),
error: None,
};
assert!(
LLMMetricAnnotation::from_annotation(&untagged)
.expect("untagged chunk is not an error")
.is_none()
);
let other = Annotated::<()> {
id: None,
data: None,
event: Some(ANNOTATION_TOKEN_IDS.to_string()),
comment: None,
error: None,
};
assert!(
LLMMetricAnnotation::from_annotation(&other)
.expect("other event is not an error")
.is_none()
);
}
#[test]
fn special_token_requirements_are_resolved_once_per_model() {
if crate::test_utils::run_isolated(
concat!(
module_path!(),
"::special_token_requirements_are_resolved_once_per_model"
),
&[(env_llm::DYN_PARSER_VERSION, "2")],
) {
return;
}
static CREATED: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
dynamo_parsers_v2::register_unified_parser("kimi_k3", |tools| {
CREATED.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
dynamo_parsers_v2::create_unified_parser_for_family("gemma4", tools)
});
CREATED.store(0, std::sync::atomic::Ordering::SeqCst);
let mut card = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
card.model_type = crate::model_type::ModelType::Chat;
card.runtime_config.tool_call_parser = Some("kimi_k3".into());
let preprocessor = OpenAIPreprocessor::new(card).unwrap();
assert_eq!(CREATED.load(std::sync::atomic::Ordering::SeqCst), 1);
for skip in [None, Some(false), Some(true)] {
let mut request = parser_route_test_request(ChatCompletionToolChoiceOption::Auto);
request.common.skip_special_tokens = skip;
for lora in [None, Some("adapter".to_string())] {
let output = preprocessor
.builder_with_lora(&request, lora)
.unwrap()
.token_ids(vec![])
.build()
.unwrap();
assert_eq!(
output.output_options.skip_special_tokens,
skip.or(Some(false))
);
assert_eq!(CREATED.load(std::sync::atomic::Ordering::SeqCst), 1);
}
}
dynamo_parsers_v2::unregister_unified_parser("kimi_k3");
}
#[test]
fn special_token_factory_failure_rejects_model_setup() {
if crate::test_utils::run_isolated(
concat!(
module_path!(),
"::special_token_factory_failure_rejects_model_setup"
),
&[(env_llm::DYN_PARSER_VERSION, "2")],
) {
return;
}
dynamo_parsers_v2::register_unified_parser("kimi_k3", |_| {
anyhow::bail!("factory unavailable")
});
let mut card = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
card.model_type = crate::model_type::ModelType::Chat;
card.runtime_config.tool_call_parser = Some("kimi_k3".into());
let error = OpenAIPreprocessor::new(card)
.err()
.expect("factory failure must reject model setup");
assert!(
error
.to_string()
.contains("special-token requirements for kimi_k3")
);
assert!(format!("{error:#}").contains("factory unavailable"));
dynamo_parsers_v2::unregister_unified_parser("kimi_k3");
}
#[test]
fn test_parser_requires_special_tokens() {
let cases: &[(Option<&str>, Option<&str>, bool, &str)] = &[
(
Some("gemma4"),
None,
true,
"gemma4 tool-call only → required",
),
(
None,
Some("gemma4"),
true,
"gemma4 reasoning only → required",
),
(
Some("gemma-4"),
None,
true,
"gemma-4 hyphen alias (tool) → required",
),
(
None,
Some("gemma-4"),
true,
"gemma-4 hyphen alias (reasoning) → required",
),
(
Some("gemma4"),
Some("gemma4"),
true,
"gemma4 paired → required",
),
(Some("hermes"), None, false, "hermes → not required"),
(
Some("harmony"),
None,
true,
"harmony tool-call only → required",
),
(
None,
Some("gpt_oss"),
true,
"gpt_oss reasoning only → required",
),
(
Some("harmony"),
Some("gpt_oss"),
true,
"harmony + gpt_oss paired → required",
),
(
Some("kimi_k2"),
Some("kimi_k25"),
true,
"kimi_k2 + kimi_k25 paired → required \
(tool-call markers `<|tool_calls_section_*|>` and reasoning \
marker `</think>` are special tokens that get stripped under \
the default skip_special_tokens=true)",
),
(
None,
Some("kimi_k25"),
true,
"kimi_k25 reasoning only → required (`</think>` is special token id 163607)",
),
(
Some("kimi_k3"),
Some("kimi_k3"),
true,
"kimi_k3 paired XTML parsers → required",
),
(
Some("kimi-k3"),
Some("kimi-k3"),
true,
"kimi-k3 aliases → required",
),
(
None,
Some("mistral"),
true,
"mistral reasoning only → required (`[THINK]` / `[/THINK]` are special tokens)",
),
(
Some("kimi_k2"),
None,
true,
"kimi_k2 tool-call only → required \
(`<|tool_calls_section_begin|>` / `<|tool_calls_section_end|>` are special)",
),
(
Some("minimax_m3"),
Some("minimax_m3"),
true,
"minimax_m3 paired → required",
),
(
Some("minimax-m3-nom"),
Some("minimax-m3"),
true,
"MiniMax M3 SGLang aliases → required",
),
(
Some("muse_glimmer"),
None,
true,
"muse_glimmer tool-call only → required \
(`<|start|>` / `<|message|>` / `<|eom|>` / `<|eot|>` are special)",
),
(
None,
Some("muse_glimmer"),
true,
"muse_glimmer reasoning-name only → required",
),
(
Some("muse"),
None,
true,
"muse (SGLang's registered name, tool) → required",
),
(None, None, false, "no parsers → not required"),
];
for (tool, reasoning, expected, desc) in cases {
assert_eq!(
OpenAIPreprocessor::parser_requires_special_tokens(*tool, *reasoning).unwrap(),
*expected,
"FAILED: {desc}",
);
}
}
#[test]
fn test_special_tokens_will_be_stripped() {
for skip in [None, Some(false), Some(true)] {
for required in [false, true] {
assert_eq!(
OpenAIPreprocessor::special_tokens_will_be_stripped(skip, required),
skip == Some(true) && required,
);
}
}
}
#[test]
fn test_kimi_k3_jail_always_unwraps_xtml_response_channels() {
let parser = "kimi_k3".to_string();
assert!(OpenAIPreprocessor::should_apply_tool_jail(Some(&parser), None, false).unwrap());
let alias = "kimi-k3".to_string();
assert!(
OpenAIPreprocessor::should_apply_tool_jail(
Some(&alias),
Some(&ChatCompletionToolChoiceOption::None),
true,
)
.unwrap()
);
}
#[test]
fn test_reasoning_before_guided_json_parser_allowlist() {
for parser in [
"deepseek_r1",
"deepseek_v3",
"deepseek_v3_1",
"deepseek_v3_2",
"step3",
"kimi_k25",
"mistral",
"minimax_m2",
"nemotron_nano",
"nemotron3",
"nemotron_v3",
] {
assert!(
OpenAIPreprocessor::supports_reasoning_before_guided_json(Some(parser)),
"{parser} should inspect guided output shape"
);
}
assert!(
!OpenAIPreprocessor::supports_reasoning_before_guided_json(Some(
"minimax_append_think"
)),
"minimax_append_think must retain the guided-JSON bypass"
);
}
#[test]
fn test_guided_json_reasoning_start_tokens() {
assert_eq!(
OpenAIPreprocessor::guided_json_reasoning_start_token(Some("mistral")),
Some("[THINK]")
);
assert_eq!(
OpenAIPreprocessor::guided_json_reasoning_start_token(Some("nemotron_v3")),
None
);
}
#[test]
fn test_prompt_injected_reasoning_start_by_parser() {
let cases = [
(
Some("minimax_m3"),
Some("...<mm:think>\n"),
true,
"MiniMax M3 starts from <mm:think>",
),
(
Some("minimax-m3"),
Some("...</mm:think>\n"),
false,
"MiniMax M3 prefilled end marker means not in reasoning",
),
(
Some("minimax_m3"),
Some("...<think>\n"),
false,
"MiniMax M3 must not use generic <think>",
),
(
Some("qwen3"),
Some("...<think>\n"),
true,
"Qwen-style templated <think> behavior remains",
),
(
Some("kimi_k3"),
Some("...<|open|>think<|sep|>\n"),
true,
"Kimi K3 canonical XTML prompt retains legacy behavior",
),
(
Some("kimi-k3"),
Some("...<|open|> think <|sep|>\n"),
true,
"Kimi K3 ignores whitespace inside the suffix",
),
(
Some("kimi-k3"),
Some("...<|open|>\u{2003}think\t<|sep|>\n"),
true,
"Kimi K3 ignores Unicode whitespace inside the suffix",
),
(
Some("kimi_k3"),
Some("think<|sep|>"),
false,
"Kimi K3 requires the complete suffix",
),
(
Some("kimi-k3"),
Some("...<think>\n"),
false,
"Kimi K3 must not use the generic think marker",
),
(
Some("deepseek_v4"),
Some("...<think>\n"),
true,
"existing <think> parser behavior remains",
),
(
None,
Some("...<think>\n"),
true,
"legacy no-parser detection",
),
(Some("minimax_m3"), None, false, "no prompt"),
];
for (parser, prompt, expected, desc) in cases {
assert_eq!(
OpenAIPreprocessor::prompt_injected_reasoning_start(parser, prompt),
expected,
"FAILED: {desc}",
);
}
}
#[test]
fn tool_only_unified_family_detects_its_prompt_injected_reasoning_marker() {
let cases = [
(
"gemma4",
"...<|channel>thought",
"Gemma 4 tool-only routing uses its thought channel",
),
(
"kimi_k3",
"...<|open|>think<|sep|>",
"Kimi K3 tool-only routing uses its think channel",
),
(
"muse_glimmer",
"...assistant to=self<|message|>",
"Muse tool-only routing uses its thought message",
),
];
for (family, prompt, description) in cases {
assert!(
OpenAIPreprocessor::prompt_injected_reasoning_start_for_parsers(
Some(family),
None,
Some(prompt),
),
"{description}"
);
}
assert!(
!OpenAIPreprocessor::prompt_injected_reasoning_start_for_parsers(
Some("gemma4"),
Some("qwen3"),
Some("...<|channel>thought"),
),
"mismatched parser families must keep the configured reasoning parser"
);
}
#[test]
fn test_prompt_injected_reasoning_ended_backend_arg_by_parser() {
let cases = [
(
Some("minimax_m2"),
Some("...<think>\n"),
Some(false),
"MiniMax M2 needs native backend reasoning state aligned with the prompt",
),
(
Some("minimax_m3"),
Some("...<mm:think>\n"),
Some(false),
"MiniMax M3 needs native backend reasoning state aligned with the prompt",
),
(
Some("kimi_k3"),
Some("...<|open|>think<|sep|>\n"),
Some(false),
"Kimi K3 canonical prompt keeps its backend reasoning state",
),
(
Some("kimi-k3"),
Some("...<|open|> think <|sep|>\n"),
Some(false),
"Kimi K3 spaced header keeps its backend reasoning state",
),
(
Some("deepseek_v4"),
Some("...<think>\n"),
None,
"DeepSeek V4 native guided JSON must not be forced into reasoning mode",
),
(
Some("deepseek-v4"),
Some("...<think>\n"),
None,
"DeepSeek V4 alias must not receive reasoning_ended=false",
),
(
Some("qwen3"),
Some("...<think>\n"),
None,
"Qwen-style prompt injection is handled by Dynamo postprocessing only",
),
(
Some("minimax_m2"),
Some("plain prompt"),
None,
"no injected reasoning opener means no backend state override",
),
];
for (parser, prompt, expected, desc) in cases {
assert_eq!(
OpenAIPreprocessor::prompt_injected_reasoning_ended_arg(parser, prompt),
expected,
"FAILED: {desc}",
);
}
}
#[test]
fn test_backend_extra_args_preserves_nvext_and_sampling_extensions() {
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hi"}],
"detokenize": false,
"allowed_token_ids": [10, 11],
"bad_words_token_ids": [[12, 13]],
"logprob_token_ids": [14, 15],
"nvext": {
"cache_salt": "step_7",
"extra_fields": ["completion_token_ids"],
"metadata_upload": {
"url": "s3://bucket/root/rollouts"
}
}
}))
.unwrap();
let extra_args = OpenAIPreprocessor::backend_extra_args(&request, false, None).unwrap();
assert_eq!(extra_args["nvext"]["cache_salt"], "step_7");
assert_eq!(
extra_args["nvext"]["extra_fields"],
serde_json::json!(["completion_token_ids"])
);
assert_eq!(
extra_args["nvext"]["metadata_upload"],
serde_json::json!({
"url": "s3://bucket/root/rollouts"
})
);
assert_eq!(extra_args["sampling_options"]["detokenize"], false);
assert_eq!(
extra_args["sampling_options"]["allowed_token_ids"],
serde_json::json!([10, 11])
);
assert_eq!(
extra_args["sampling_options"]["bad_words_token_ids"],
serde_json::json!([[12, 13]])
);
assert_eq!(
extra_args["sampling_options"]["logprob_token_ids"],
serde_json::json!([14, 15])
);
}
#[test]
fn capture_prompt_token_ids_is_exact_and_opt_in() {
let requested: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hi"}],
"nvext": {"extra_fields": ["prompt_token_ids"]}
}))
.unwrap();
let tracker = RequestTracker::new();
OpenAIPreprocessor::capture_prompt_token_ids(&requested, Some(&tracker), &[101, 102, 103]);
assert_eq!(tracker.prompt_token_ids(), Some(&[101u32, 102, 103][..]));
let ordinary: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hi"}]
}))
.unwrap();
let ordinary_tracker = RequestTracker::new();
OpenAIPreprocessor::capture_prompt_token_ids(
&ordinary,
Some(&ordinary_tracker),
&[201, 202],
);
assert!(ordinary_tracker.prompt_token_ids().is_none());
}
fn chat_request_with_args(
chat_template_args: Option<HashMap<String, serde_json::Value>>,
) -> NvCreateChatCompletionRequest {
let mut request: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hi"}]
}))
.unwrap();
request.chat_template_args = chat_template_args;
request
}
fn runtime_config_with_default_thinking_mode(
mode: &str,
) -> crate::local_model::runtime_config::ModelRuntimeConfig {
let mut runtime_config = crate::local_model::runtime_config::ModelRuntimeConfig::new();
runtime_config
.set_engine_specific(DEFAULT_THINKING_MODE_RUNTIME_KEY, mode)
.unwrap();
runtime_config
}
#[test]
fn test_default_thinking_mode_disabled_adds_template_args() {
let runtime_config = runtime_config_with_default_thinking_mode("disabled");
let mut request = chat_request_with_args(None);
OpenAIPreprocessor::apply_default_thinking_mode_from_runtime_config(
&runtime_config,
&mut request,
);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(args.get("thinking"), Some(&serde_json::json!(false)));
assert_eq!(args.get("enable_thinking"), Some(&serde_json::json!(false)));
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("disabled"))
);
}
#[test]
fn test_default_thinking_mode_enabled_adds_template_args() {
let runtime_config = runtime_config_with_default_thinking_mode("enabled");
let mut request = chat_request_with_args(None);
OpenAIPreprocessor::apply_default_thinking_mode_from_runtime_config(
&runtime_config,
&mut request,
);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(args.get("thinking"), Some(&serde_json::json!(true)));
assert_eq!(args.get("enable_thinking"), Some(&serde_json::json!(true)));
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("enabled"))
);
}
#[test]
fn test_gemma4_default_thinking_mode_controls_reasoning_parser() {
for (mode, expected_disabled) in [("enabled", false), ("disabled", true)] {
let runtime_config = runtime_config_with_default_thinking_mode(mode);
let mut request = chat_request_with_args(None);
OpenAIPreprocessor::apply_default_thinking_mode_from_runtime_config(
&runtime_config,
&mut request,
);
for parser in ["gemma4", "gemma-4"] {
assert_eq!(
OpenAIPreprocessor::is_reasoning_disabled_by_request(
Some(parser),
request.chat_template_args.as_ref(),
),
expected_disabled,
"parser={parser}, default_thinking_mode={mode}",
);
}
}
}
#[test]
fn test_default_thinking_mode_does_not_override_request() {
let runtime_config = runtime_config_with_default_thinking_mode("disabled");
let mut request = chat_request_with_args(Some(HashMap::from([(
"thinking_mode".to_string(),
serde_json::json!("enabled"),
)])));
OpenAIPreprocessor::apply_default_thinking_mode_from_runtime_config(
&runtime_config,
&mut request,
);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("enabled"))
);
assert!(!args.contains_key("thinking"));
assert!(!args.contains_key("enable_thinking"));
}
#[test]
fn test_default_thinking_mode_precedes_parser_implicit_default() {
let runtime_config = runtime_config_with_default_thinking_mode("disabled");
let mut request = chat_request_with_args(None);
OpenAIPreprocessor::apply_default_thinking_mode_from_runtime_config(
&runtime_config,
&mut request,
);
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some("kimi_k25"), None);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(args.get("thinking"), Some(&serde_json::json!(false)));
assert_eq!(args.get("enable_thinking"), Some(&serde_json::json!(false)));
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("disabled"))
);
}
#[test]
fn test_default_thinking_mode_does_not_override_reasoning_effort() {
let runtime_config = runtime_config_with_default_thinking_mode("disabled");
let mut request = chat_request_with_args(Some(HashMap::from([(
"reasoning_effort".to_string(),
serde_json::json!("high"),
)])));
OpenAIPreprocessor::apply_default_thinking_mode_from_runtime_config(
&runtime_config,
&mut request,
);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(
args.get("reasoning_effort"),
Some(&serde_json::json!("high"))
);
assert!(!args.contains_key("thinking"));
assert!(!args.contains_key("enable_thinking"));
assert!(!args.contains_key("thinking_mode"));
}
#[test]
fn test_backend_extra_args_forwards_reasoning_template_args() {
for enable_thinking in [true, false] {
let request: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hi"}],
"chat_template_kwargs": {
"enable_thinking": enable_thinking,
"reasoning_effort": "high"
}
}))
.unwrap();
let extra_args =
OpenAIPreprocessor::backend_extra_args(&request, true, Some(false)).unwrap();
assert_eq!(
extra_args["reasoning_parser_kwargs"]["chat_template_kwargs"],
serde_json::json!({
"enable_thinking": enable_thinking,
"reasoning_effort": "high"
})
);
assert_eq!(extra_args["reasoning_ended"], false);
}
}
#[test]
fn test_backend_extra_args_omits_reasoning_metadata_without_configured_parser() {
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hi"}],
"detokenize": false,
"chat_template_kwargs": {
"enable_thinking": true
}
}))
.unwrap();
let extra_args =
OpenAIPreprocessor::backend_extra_args(&request, false, Some(false)).unwrap();
assert_eq!(extra_args["sampling_options"]["detokenize"], false);
assert!(extra_args.get("reasoning_parser_kwargs").is_none());
assert!(extra_args.get("reasoning_ended").is_none());
}
#[test]
fn test_backend_extra_args_omits_reasoning_metadata_without_template_args() {
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hi"}]
}))
.unwrap();
assert!(OpenAIPreprocessor::backend_extra_args(&request, true, None).is_none());
}
#[test]
fn test_sglang_gpt_oss_structured_output_rejects_thinking_budget() {
use crate::local_model::runtime_config::SGLANG_GENERATE_CAPABILITY;
let mut mdc = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
mdc.runtime_config.reasoning_parser = Some("gpt_oss".to_string());
mdc.runtime_config
.set_engine_specific(SGLANG_GENERATE_CAPABILITY, true)
.unwrap();
let preprocessor = OpenAIPreprocessor::new(mdc).unwrap();
for response_format in [
serde_json::json!({"type": "json_object"}),
serde_json::json!({"type": "json_schema", "json_schema": {
"name": "result", "schema": {"type": "object", "properties": {}}
}}),
] {
for budget_fields in [
serde_json::json!({"thinking_token_budget": 32}),
serde_json::json!({"thinking_token_budget": 0}),
serde_json::json!({"nvext": {"max_thinking_tokens": 16}}),
serde_json::json!({"thinking_token_budget": 0, "nvext": {"max_thinking_tokens": 16}}),
] {
let mut value = serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hi"}],
"response_format": response_format,
});
value
.as_object_mut()
.unwrap()
.extend(budget_fields.as_object().unwrap().clone());
let request: NvCreateChatCompletionRequest = serde_json::from_value(value).unwrap();
let error = preprocessor
.builder(&request)
.err()
.expect("SGLang cannot honor GPT-OSS structured output with a thinking budget");
assert_eq!(
error.downcast_ref::<DynamoError>().unwrap().error_type(),
ErrorType::InvalidArgument,
);
assert!(error.to_string().contains("thinking_token_budget"));
}
}
}
#[test]
fn test_gpt_oss_thinking_budget_preserves_supported_output_paths() {
use crate::local_model::runtime_config::{
SGLANG_GENERATE_CAPABILITY, VLLM_INFERENCE_V1_GENERATE_CAPABILITY,
};
for (capability, budget, structured, expected_reasoning) in [
(SGLANG_GENERATE_CAPABILITY, Some(32), false, true),
(SGLANG_GENERATE_CAPABILITY, None, true, false),
(VLLM_INFERENCE_V1_GENERATE_CAPABILITY, Some(0), true, true),
] {
let mut mdc = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
mdc.runtime_config.reasoning_parser = Some("gpt_oss".to_string());
mdc.runtime_config
.set_engine_specific(capability, true)
.unwrap();
let preprocessor = OpenAIPreprocessor::new(mdc).unwrap();
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hi"}],
"thinking_token_budget": budget,
"response_format": if structured { serde_json::json!({"type": "json_object"}) } else { serde_json::Value::Null },
})).unwrap();
let preprocessed = preprocessor
.builder(&request)
.unwrap()
.token_ids(vec![1])
.build()
.unwrap();
assert_eq!(preprocessed.require_reasoning, expected_reasoning);
assert_eq!(preprocessed.stop_conditions.max_thinking_tokens, budget);
}
}
#[test]
fn test_request_requires_reasoning_with_thinking_budget() {
let cases = [
(
serde_json::json!({"thinking_token_budget": 32}),
Some("qwen3"),
true,
),
(
serde_json::json!({"thinking_token_budget": 0}),
Some("qwen3"),
true,
),
(
serde_json::json!({"nvext": {"max_thinking_tokens": 16}}),
Some("qwen3"),
true,
),
(serde_json::json!({}), Some("qwen3"), false),
(
serde_json::json!({"thinking_token_budget": 32}),
None,
false,
),
(
serde_json::json!({"thinking_token_budget": 32}),
Some("gemma4"),
false,
),
(
serde_json::json!({"thinking_token_budget": 32, "chat_template_kwargs": {"enable_thinking": true}}),
Some("gemma4"),
true,
),
(
serde_json::json!({"thinking_token_budget": 32, "chat_template_kwargs": {"enable_thinking": false}}),
Some("qwen3"),
false,
),
(
serde_json::json!({"thinking_token_budget": 32, "response_format": {"type": "json_object"}}),
Some("qwen3"),
true,
),
(
serde_json::json!({"response_format": {"type": "json_object"}}),
Some("qwen3"),
true,
),
(
serde_json::json!({"response_format": {"type": "json_object"}}),
Some("gpt_oss"),
false,
),
];
for (fields, parser, expected) in cases {
let mut value = serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hi"}]
});
value
.as_object_mut()
.unwrap()
.extend(fields.as_object().unwrap().clone());
let request: NvCreateChatCompletionRequest = serde_json::from_value(value).unwrap();
let has_budget = request
.extract_stop_conditions()
.unwrap()
.max_thinking_tokens
.is_some();
assert_eq!(
OpenAIPreprocessor::request_requires_reasoning(&request, parser, has_budget),
expected,
"parser={parser:?}, fields={fields}",
);
}
}
#[test]
fn test_guided_output_requires_reasoning() {
let request = |tool_choice: serde_json::Value, enable_thinking: Option<bool>| {
let mut value = serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "use the tool"}],
"tools": [{
"type": "function",
"function": {
"name": "lookup",
"parameters": {"type": "object", "properties": {}}
}
}],
"tool_choice": tool_choice
});
if let Some(enabled) = enable_thinking {
value["chat_template_kwargs"] = serde_json::json!({
"enable_thinking": enabled
});
}
serde_json::from_value::<NvCreateChatCompletionRequest>(value).unwrap()
};
let required = request(serde_json::json!("required"), Some(true));
assert!(OpenAIPreprocessor::guided_output_requires_reasoning(
&required,
Some("nemotron_v3")
));
let named = request(
serde_json::json!({
"type": "function",
"function": {"name": "lookup"}
}),
None,
);
assert!(OpenAIPreprocessor::guided_output_requires_reasoning(
&named,
Some("nemotron_v3")
));
let disabled = request(serde_json::json!("required"), Some(false));
assert!(!OpenAIPreprocessor::guided_output_requires_reasoning(
&disabled,
Some("nemotron_v3")
));
assert!(!OpenAIPreprocessor::guided_output_requires_reasoning(
&required, None
));
let automatic = request(serde_json::json!("auto"), Some(true));
assert!(!OpenAIPreprocessor::guided_output_requires_reasoning(
&automatic,
Some("nemotron_v3")
));
let gemma_without_opt_in = request(serde_json::json!("required"), None);
assert!(!OpenAIPreprocessor::guided_output_requires_reasoning(
&gemma_without_opt_in,
Some("gemma4")
));
let gemma_with_opt_in = request(serde_json::json!("required"), Some(true));
assert!(OpenAIPreprocessor::guided_output_requires_reasoning(
&gemma_with_opt_in,
Some("gemma4")
));
let deepseek_default = request(serde_json::json!("required"), None);
assert!(OpenAIPreprocessor::guided_output_requires_reasoning(
&deepseek_default,
Some("deepseek_v4")
));
let deepseek_disabled = request(serde_json::json!("required"), Some(false));
assert!(!OpenAIPreprocessor::guided_output_requires_reasoning(
&deepseek_disabled,
Some("deepseek_v4")
));
let structured_request = |enable_thinking: bool| {
serde_json::from_value::<NvCreateChatCompletionRequest>(serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "return json"}],
"chat_template_kwargs": {"enable_thinking": enable_thinking},
"response_format": {
"type": "json_schema",
"json_schema": {
"name": "result",
"schema": {"type": "object"}
}
}
}))
.unwrap()
};
let json_object_request = |enable_thinking: bool| {
serde_json::from_value::<NvCreateChatCompletionRequest>(serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "return json"}],
"chat_template_kwargs": {"enable_thinking": enable_thinking},
"response_format": {"type": "json_object"}
}))
.unwrap()
};
assert!(OpenAIPreprocessor::guided_output_requires_reasoning(
&structured_request(true),
Some("qwen3")
));
assert!(OpenAIPreprocessor::guided_output_requires_reasoning(
&json_object_request(true),
Some("qwen3")
));
assert!(!OpenAIPreprocessor::guided_output_requires_reasoning(
&structured_request(false),
Some("qwen3")
));
assert!(!OpenAIPreprocessor::guided_output_requires_reasoning(
&structured_request(true),
Some("gpt_oss")
));
assert!(!OpenAIPreprocessor::guided_output_requires_reasoning(
&json_object_request(true),
Some("gpt_oss")
));
assert!(!OpenAIPreprocessor::guided_output_requires_reasoning(
&structured_request(false),
Some("gpt_oss")
));
assert!(OpenAIPreprocessor::guided_output_requires_reasoning(
&request(serde_json::json!("required"), None),
Some("gpt_oss")
));
}
#[cfg(feature = "mm-routing")]
#[test]
fn exact_mm_routing_rejects_backend_owned_modalities() {
assert!(exact_mm_routing_supports_modality("image_url", false));
assert!(exact_mm_routing_supports_modality("video_url", true));
assert!(!exact_mm_routing_supports_modality("video_url", false));
assert!(!exact_mm_routing_supports_modality("audio_url", true));
}
#[test]
fn test_sglang_effective_reasoning_enabled() {
let cases = [
("qwen3", serde_json::json!({}), true),
(
"qwen3",
serde_json::json!({"enable_thinking": false}),
false,
),
("nemotron_v3", serde_json::json!({}), true),
("gemma4", serde_json::json!({}), false),
("gemma4", serde_json::json!({"enable_thinking": true}), true),
("deepseek_v4", serde_json::json!({}), true),
("deepseek_v4", serde_json::json!({"thinking": true}), true),
(
"deepseek_v4",
serde_json::json!({"enable_thinking": true}),
true,
),
(
"deepseek_v4",
serde_json::json!({"thinking_mode": "thinking"}),
true,
),
(
"deepseek_v4",
serde_json::json!({"enable_thinking": false}),
false,
),
(
"deepseek_v4",
serde_json::json!({"thinking_mode": "chat"}),
false,
),
(
"deepseek_v4",
serde_json::json!({"thinking": false, "thinking_mode": "thinking"}),
false,
),
("deepseek_v3_2", serde_json::json!({}), true),
("deepseek_v3_1", serde_json::json!({}), false),
(
"deepseek_v3_1",
serde_json::json!({"enable_thinking": true}),
true,
),
("kimi_k25", serde_json::json!({"thinking": false}), false),
("kimi_k3", serde_json::json!({}), true),
("kimi-k3", serde_json::json!({"thinking": false}), false),
("minimax_m2", serde_json::json!({}), true),
("minimax_m2", serde_json::json!({"thinking": false}), false),
("mistral", serde_json::json!({}), false),
(
"mistral",
serde_json::json!({"reasoning_effort": "none"}),
false,
),
(
"mistral",
serde_json::json!({"reasoning_effort": "high"}),
true,
),
("minimax_m3", serde_json::json!({}), true),
(
"minimax_m3",
serde_json::json!({"thinking_mode": "disabled"}),
false,
),
("gpt_oss", serde_json::json!({}), true),
(
"gpt_oss",
serde_json::json!({"enable_thinking": true}),
true,
),
(
"gpt_oss",
serde_json::json!({"enable_thinking": false}),
true,
),
("deepseek_r1", serde_json::json!({}), true),
("minimax_append_think", serde_json::json!({}), false),
("basic", serde_json::json!({}), false),
];
for (parser, request_args, expected) in cases {
let request_args = serde_json::from_value(request_args).unwrap();
assert_eq!(
OpenAIPreprocessor::sglang_effective_reasoning_enabled(
Some(parser),
Some(&request_args),
),
expected,
"parser={parser}, args={request_args:?}",
);
}
assert!(!OpenAIPreprocessor::sglang_effective_reasoning_enabled(
None, None
));
}
#[test]
fn test_internal_preserve_omitted_max_tokens_option() {
assert_eq!(
OpenAIPreprocessor::omitted_max_tokens_default(
10,
Some(100),
PreprocessRequestOptions::default()
),
Some(90)
);
assert_eq!(
OpenAIPreprocessor::omitted_max_tokens_default(
10,
Some(100),
PreprocessRequestOptions {
preserve_omitted_max_tokens: true,
},
),
None
);
assert_eq!(
OpenAIPreprocessor::omitted_max_tokens_default(
10,
None,
PreprocessRequestOptions::default()
),
None
);
assert_eq!(
OpenAIPreprocessor::omitted_max_tokens_default(
10,
Some(0),
PreprocessRequestOptions::default()
),
Some(0)
);
}
#[test]
fn test_exact_prompt_len() {
let images = MultimodalDataMap::from([(
"image_url".to_string(),
vec![MultimodalData::Url(
url::Url::parse("https://example.com/image.png").unwrap(),
)],
)]);
let videos = MultimodalDataMap::from([(
"video_url".to_string(),
vec![MultimodalData::Url(
url::Url::parse("https://example.com/video.mp4").unwrap(),
)],
)]);
let mixed = MultimodalDataMap::from([
(
"image_url".to_string(),
vec![MultimodalData::Url(
url::Url::parse("https://example.com/image.png").unwrap(),
)],
),
(
"audio_url".to_string(),
vec![MultimodalData::Url(
url::Url::parse("https://example.com/audio.wav").unwrap(),
)],
),
]);
assert_eq!(
OpenAIPreprocessor::exact_prompt_len(Some(500), Some(&images), 12),
Some(500)
);
assert_eq!(
OpenAIPreprocessor::exact_prompt_len(Some(0), Some(&images), 12),
None
);
assert_eq!(
OpenAIPreprocessor::exact_prompt_len(None, Some(&images), 12),
None
);
assert_eq!(
OpenAIPreprocessor::exact_prompt_len(None, Some(&videos), 12),
None
);
assert_eq!(
OpenAIPreprocessor::exact_prompt_len(Some(500), Some(&mixed), 12),
None
);
assert_eq!(
OpenAIPreprocessor::exact_prompt_len(None, None, 12),
Some(12)
);
assert_eq!(
OpenAIPreprocessor::exact_prompt_len(Some(0), None, 12),
Some(12)
);
}
fn token_budget(
combined_limit: u32,
reject_prompt_overflow: bool,
reject_total_overflow: bool,
) -> TokenBudget {
TokenBudget {
combined_limit,
reject_prompt_overflow,
reject_total_overflow,
}
}
#[test]
fn test_requested_token_budget_policy_matrix() {
let cases = [
(3, Some(7), true, true, false),
(3, Some(8), true, true, true),
(10, None, true, true, true),
(3, Some(8), true, false, false),
(10, Some(8), false, true, false),
];
for (
prompt_len,
max_tokens,
reject_prompt_overflow,
reject_total_overflow,
should_reject,
) in cases
{
let budget = token_budget(10, reject_prompt_overflow, reject_total_overflow);
assert_eq!(
OpenAIPreprocessor::validate_requested_token_budget(
prompt_len,
max_tokens,
Some(&budget),
)
.is_err(),
should_reject,
);
}
}
fn preprocessed_budget_request(max_tokens: Option<u32>) -> PreprocessedRequest {
let stop_conditions = crate::protocols::common::StopConditions {
max_tokens,
..Default::default()
};
PreprocessedRequest::builder()
.model("test-model".to_string())
.token_ids(vec![1, 2, 3])
.stop_conditions(stop_conditions)
.sampling_options(crate::protocols::common::SamplingOptions::default())
.output_options(crate::protocols::common::OutputOptions::default())
.build()
.unwrap()
}
#[test]
fn attach_agent_context_forwards_opaque_headers() {
let agent_context = AgentContext {
session_id: "codex-thread".to_string(),
parent_session_id: None,
session_final: None,
agent_headers: std::collections::BTreeMap::from([(
"x-claude-code-future".into(),
vec!["unknown".into(), "second".into()],
)])
.into(),
input_trigger: None,
};
let mut context = PipelineContext::new(());
context.insert(AGENT_CONTEXT_CONTEXT_KEY, agent_context.clone());
let mut request = preprocessed_budget_request(None);
attach_agent_context_from_context(&mut request, &context);
assert_eq!(request.agent_context.as_ref(), Some(&agent_context));
let wire = serde_json::to_value(&request).unwrap();
assert_eq!(
wire["agent_context"]["agent_headers"]["x-claude-code-future"],
serde_json::json!(["unknown", "second"])
);
let restored: AgentContext = serde_json::from_value(wire["agent_context"].clone()).unwrap();
assert_eq!(restored, agent_context);
assert!(std::sync::Arc::ptr_eq(
&request.agent_context.as_ref().unwrap().agent_headers,
&agent_context.agent_headers
));
}
#[test]
fn attach_request_context_metadata_keeps_affinity_separate_from_agent_context() {
use crate::protocols::common::extensions::{
SESSION_AFFINITY_CONTEXT_KEY, SessionAffinityId,
};
let agent_context = AgentContext {
session_id: "agent-session".to_string(),
parent_session_id: Some("agent-parent".to_string()),
session_final: None,
agent_headers: Default::default(),
input_trigger: None,
};
let mut context = PipelineContext::new(());
context.insert(AGENT_CONTEXT_CONTEXT_KEY, agent_context.clone());
context.insert(
SESSION_AFFINITY_CONTEXT_KEY,
SessionAffinityId::new("routing-session"),
);
let mut request = preprocessed_budget_request(None);
attach_request_context_metadata(&mut request, &context);
assert_eq!(request.agent_context.as_ref(), Some(&agent_context));
assert_eq!(
request.image_cache_scope.as_deref(),
Some("routing-session")
);
let wire = serde_json::to_value(&request).unwrap();
assert_eq!(wire["agent_context"]["session_id"], "agent-session");
assert_eq!(wire["image_cache_scope"], "routing-session");
let mut affinity_only_context = PipelineContext::new(());
affinity_only_context.insert(
SESSION_AFFINITY_CONTEXT_KEY,
SessionAffinityId::new("routing-only"),
);
let mut affinity_only_request = preprocessed_budget_request(None);
attach_request_context_metadata(&mut affinity_only_request, &affinity_only_context);
assert!(affinity_only_request.agent_context.is_none());
assert_eq!(
affinity_only_request.image_cache_scope.as_deref(),
Some("routing-only")
);
}
#[test]
fn test_preprocessed_completion_budget_validation_and_deferral() {
let text_request = preprocessed_budget_request(Some(8));
let reject = token_budget(10, true, true);
assert!(
OpenAIPreprocessor::validate_preprocessed_token_budget(&text_request, Some(&reject))
.is_err()
);
let defer_total = token_budget(10, true, false);
assert_eq!(
OpenAIPreprocessor::validate_preprocessed_token_budget(
&text_request,
Some(&defer_total),
)
.unwrap(),
Some(3)
);
let mut embeddings_request = text_request.clone();
embeddings_request.prompt_embeds = Some("opaque-tensor".to_string());
assert_eq!(
OpenAIPreprocessor::validate_preprocessed_token_budget(
&embeddings_request,
Some(&reject),
)
.unwrap(),
None
);
}
struct UnreachableBackend;
#[async_trait]
impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<BackendOutput>>, Error>
for UnreachableBackend
{
async fn generate(
&self,
_request: SingleIn<PreprocessedRequest>,
) -> Result<ManyOut<Annotated<BackendOutput>>, Error> {
panic!("request must be rejected before backend dispatch")
}
}
#[tokio::test]
async fn test_completion_operator_rejects_token_budget_overflow() {
let mut mdc = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
mdc.runtime_config.context_length = Some(100);
mdc.runtime_config
.set_engine_specific(TOKEN_BUDGET_RUNTIME_KEY, token_budget(10, true, true))
.unwrap();
let preprocessor = OpenAIPreprocessor::new(mdc).unwrap();
let request = NvCreateCompletionRequest {
inner: dynamo_protocols::types::CreateCompletionRequest {
model: "test-model".to_string(),
prompt: dynamo_protocols::types::Prompt::IntegerArray(vec![1, 2, 3]),
max_tokens: Some(8),
..Default::default()
},
common: Default::default(),
nvext: None,
metadata: None,
return_tokens_as_token_ids: None,
no_stop_trim: None,
unsupported_fields: Default::default(),
};
let next: Arc<
dyn AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<BackendOutput>>, Error>,
> = Arc::new(UnreachableBackend);
let result =
Operator::generate(preprocessor.as_ref(), PipelineContext::new(request), next).await;
let Err(err) = result else {
panic!("over-budget completion should fail admission");
};
let dynamo_err = err
.downcast_ref::<DynamoError>()
.expect("error should preserve the DynamoError type");
assert_eq!(dynamo_err.error_type(), ErrorType::InvalidArgument);
}
#[tokio::test]
async fn test_completion_operator_rejects_out_of_range_token_prompt() {
let mdc = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
let preprocessor = OpenAIPreprocessor::new(mdc).unwrap();
let request: NvCreateCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"prompt": [1, 128256]
}))
.unwrap();
let next: Arc<
dyn AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<BackendOutput>>, Error>,
> = Arc::new(UnreachableBackend);
let result =
Operator::generate(preprocessor.as_ref(), PipelineContext::new(request), next).await;
let Err(err) = result else {
panic!("an out-of-range token prompt should fail before dispatch");
};
let dynamo_err = err
.downcast_ref::<DynamoError>()
.expect("error should preserve the DynamoError type");
assert_eq!(dynamo_err.error_type(), ErrorType::InvalidArgument);
}
#[tokio::test]
async fn chat_operator_rejects_invalid_legacy_jail_request_before_dispatch() {
let mut mdc = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
mdc.runtime_config.tool_call_parser = Some("hermes".to_string());
let preprocessor = OpenAIPreprocessor::new(mdc).unwrap();
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "What is the weather?"}],
"n": 2,
"max_tokens": 4,
"tools": [{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the weather",
"parameters": {"type": "object", "properties": {}}
}
}],
"tool_choice": "auto",
"nvext": {"extra_fields": ["engine_data"]}
}))
.unwrap();
let next: Arc<
dyn AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<BackendOutput>>, Error>,
> = Arc::new(UnreachableBackend);
let result =
Operator::generate(preprocessor.as_ref(), PipelineContext::new(request), next).await;
let Err(err) = result else {
panic!("invalid legacy-jail request should fail admission");
};
let dynamo_err = err
.downcast_ref::<DynamoError>()
.expect("error should preserve the DynamoError type");
assert_eq!(dynamo_err.error_type(), ErrorType::InvalidArgument);
assert!(dynamo_err.to_string().contains("legacy tool-call parsing"));
}
fn test_prompt_formatter(template: &str) -> Arc<dyn OAIPromptFormatter> {
let template: dynamo_renderer::ChatTemplate = serde_json::from_value(serde_json::json!({
"chat_template": template
}))
.unwrap();
match dynamo_renderer::PromptFormatter::from_parts(
template,
dynamo_renderer::ContextMixins::default(),
false,
)
.unwrap()
{
dynamo_renderer::PromptFormatter::OAI(formatter) => formatter,
}
}
fn assistant_only_request() -> NvCreateChatCompletionRequest {
serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{"role": "assistant", "content": "prefill"}]
}))
.unwrap()
}
const REQUIRES_USER_TEMPLATE: &str = "\
{% set ns = namespace(has_user=false) %}\
{% for message in messages %}\
{% if message['role'] == 'user' %}{% set ns.has_user = true %}{% endif %}\
{% endfor %}\
{% if not ns.has_user %}{{ raise_exception('No user query found in messages.') }}{% endif %}\
{{ messages[0]['content'] }}";
const REQUIRES_LEADING_SYSTEM_TEMPLATE: &str = "\
{%- for message in messages -%}\
{%- if message['role'] == 'system' and not loop.first -%}\
{{- raise_exception('System message must be at the beginning.') -}}\
{%- endif -%}\
{{- message['role'] }}:{{ message['content'] }}\n\
{%- endfor -%}";
fn render_through_preprocessor(
formatter: &dyn OAIPromptFormatter,
request: &dyn OAIChatLikeRequest,
) -> Result<RenderedPrompt> {
formatter
.render_prompt(request)
.map_err(OpenAIPreprocessor::map_prompt_render_error)
}
#[test]
fn test_assistant_only_request_accepted_when_template_accepts_it() {
let formatter = test_prompt_formatter(
"{% for message in messages %}{{ message['role'] }}:{{ message['content'] }}{% endfor %}",
);
let rendered =
render_through_preprocessor(formatter.as_ref(), &assistant_only_request()).unwrap();
assert_eq!(rendered.as_str(), "assistant:prefill");
}
#[test]
fn test_assistant_only_template_error_is_invalid_argument() {
let formatter = test_prompt_formatter(REQUIRES_USER_TEMPLATE);
let error = render_through_preprocessor(formatter.as_ref(), &assistant_only_request())
.context("Failed to apply prompt template")
.unwrap_err();
let dynamo_error = error
.chain()
.find_map(|cause| cause.downcast_ref::<DynamoError>())
.expect("template render error should be classified as a DynamoError");
assert_eq!(dynamo_error.error_type(), ErrorType::InvalidArgument);
assert!(
dynamo_error
.message()
.contains("No user query found in messages.")
);
}
#[test]
fn test_restrictive_template_accepts_request_with_user_message() {
let formatter = test_prompt_formatter(REQUIRES_USER_TEMPLATE);
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hello"}]
}))
.unwrap();
let rendered = render_through_preprocessor(formatter.as_ref(), &request).unwrap();
assert_eq!(rendered.as_str(), "hello");
}
#[test]
fn test_nonleading_system_message_normalized_for_strict_template() {
let formatter = test_prompt_formatter(REQUIRES_LEADING_SYSTEM_TEMPLATE);
let anthropic_request: crate::protocols::anthropic::AnthropicCreateMessageRequest =
serde_json::from_value(serde_json::json!({
"model": "test-model",
"max_tokens": 100,
"system": "You are Claude Code.",
"messages": [
{"role": "user", "content": "Run make test."},
{"role": "system", "content": "Available agent types and skills."}
],
"tools": [{
"name": "Bash",
"description": "Run a shell command",
"input_schema": {
"type": "object",
"properties": {
"command": {"type": "string"}
},
"required": ["command"]
}
}]
}))
.unwrap();
let request: NvCreateChatCompletionRequest = anthropic_request.try_into().unwrap();
let rendered = render_through_preprocessor(formatter.as_ref(), &request).unwrap();
assert_eq!(
rendered.as_str(),
"system:You are Claude Code.\
user:Run make test.\
user:Available agent types and skills."
);
}
#[test]
fn continue_final_message_leaves_last_assistant_open_on_llama_template() {
let mut mdc = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
mdc.set_name("test-model");
let preprocessor = OpenAIPreprocessor::new(mdc).unwrap();
let default_request: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [
{"role": "user", "content": "Continue this sentence"},
{"role": "assistant", "content": "LLM-Native Interaction"}
]
}))
.unwrap();
let continue_request: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [
{"role": "user", "content": "Continue this sentence"},
{"role": "assistant", "content": "LLM-Native Interaction"}
],
"add_generation_prompt": false,
"continue_final_message": true
}))
.unwrap();
let default_prompt = preprocessor
.apply_template(&default_request)
.unwrap()
.unwrap();
let continue_prompt = preprocessor
.apply_template(&continue_request)
.unwrap()
.unwrap();
assert!(
default_prompt
.as_str()
.ends_with("<|start_header_id|>assistant<|end_header_id|>\n\n"),
"default render should start a new assistant turn, got {:?}",
default_prompt.as_str()
);
assert!(
continue_prompt.as_str().ends_with("LLM-Native Interaction"),
"continue_final_message should leave the last assistant open, got {:?}",
continue_prompt.as_str()
);
assert!(
!continue_prompt.as_str().contains(
"LLM-Native Interaction<|eot_id|><|start_header_id|>assistant<|end_header_id|>"
),
"continue_final_message must not close the last assistant and start a new turn, got {:?}",
continue_prompt.as_str()
);
}
#[test]
fn should_add_generation_prompt_defaults_true_and_continue_forces_false() {
use dynamo_renderer::OAIChatLikeRequest;
let unset: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hi"}]
}))
.unwrap();
assert!(unset.should_add_generation_prompt());
let explicit_false: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{"role": "user", "content": "hi"}],
"add_generation_prompt": false
}))
.unwrap();
assert!(!explicit_false.should_add_generation_prompt());
let continue_with_false: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": "partial"}
],
"add_generation_prompt": false,
"continue_final_message": true
}))
.unwrap();
assert!(!continue_with_false.should_add_generation_prompt());
}
const QWEN_STYLE_TEMPLATE: &str = "\
{%- for message in messages -%}\
{%- if message.role == 'user' -%}{{ '<|im_start|>user\n' + message.content + '<|im_end|>\n' }}\
{%- elif message.role == 'assistant' -%}{{ '<|im_start|>assistant\n' + message.content + '<|im_end|>\n' }}\
{%- endif -%}\
{%- endfor -%}\
{%- if add_generation_prompt -%}{{ '<|im_start|>assistant\n' }}{%- endif -%}";
fn continue_request() -> NvCreateChatCompletionRequest {
serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [
{"role": "user", "content": "Continue this sentence"},
{"role": "assistant", "content": "LLM-Native Interaction"}
],
"add_generation_prompt": false,
"continue_final_message": true
}))
.unwrap()
}
#[test]
fn normalized_args_request_preserves_render_transformations() {
use crate::preprocessor::prompt::CONTINUE_FINAL_MESSAGE_TAG;
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{"role": "assistant", "content": "prefix", "tool_calls": [{
"id": "c1", "type": "function",
"function": {"name": "lookup", "arguments": "{\"query\":\"hello\"}"}
}]}]
}))
.unwrap();
let normalized = NormalizedArgsRequest {
inner: &request,
normalize_tool_call_args: true,
continue_final_message: true,
};
assert!(request.typed_messages().is_some());
assert!(normalized.typed_messages().is_none());
let messages = serde_json::to_value(normalized.messages()).unwrap();
assert_eq!(
messages[0]["tool_calls"][0]["function"]["arguments"],
serde_json::json!({"query": "hello"})
);
assert_eq!(
messages[0]["content"],
format!("prefix{CONTINUE_FINAL_MESSAGE_TAG}")
);
let formatter = test_prompt_formatter(
"{% for m in messages %}{% for call in m.tool_calls %}{% if call.function.arguments is string %}RAW{% else %}{{ call.function.arguments.query }}{% endif %}{% endfor %}{{ m.content }}{% endfor %}",
);
let rendered = formatter.render_prompt(&normalized).unwrap();
assert_eq!(
rendered.as_str(),
format!("helloprefix{CONTINUE_FINAL_MESSAGE_TAG}")
);
assert_eq!(
apply_continue_final_message(rendered).unwrap().as_str(),
"helloprefix"
);
}
fn render_with_continue_final_message(
formatter: &dyn OAIPromptFormatter,
request: &NvCreateChatCompletionRequest,
) -> RenderedPrompt {
use crate::protocols::openai::common_ext::CommonExtProvider;
let continue_final = request.get_continue_final_message() == Some(true);
let rendered = if continue_final {
formatter
.render_prompt(&NormalizedArgsRequest {
inner: request,
normalize_tool_call_args: false,
continue_final_message: true,
})
.unwrap()
} else {
formatter.render_prompt(request).unwrap()
};
if continue_final {
apply_continue_final_message(rendered).unwrap()
} else {
rendered
}
}
#[test]
fn continue_final_message_strips_qwen_style_closing_tokens() {
let formatter = test_prompt_formatter(QWEN_STYLE_TEMPLATE);
let default_request: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [
{"role": "user", "content": "Continue this sentence"},
{"role": "assistant", "content": "LLM-Native Interaction"}
]
}))
.unwrap();
let closed_only: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [
{"role": "user", "content": "Continue this sentence"},
{"role": "assistant", "content": "LLM-Native Interaction"}
],
"add_generation_prompt": false
}))
.unwrap();
let default_prompt =
render_with_continue_final_message(formatter.as_ref(), &default_request);
let closed_prompt = render_with_continue_final_message(formatter.as_ref(), &closed_only);
let continue_prompt =
render_with_continue_final_message(formatter.as_ref(), &continue_request());
assert!(
default_prompt.as_str().ends_with("<|im_start|>assistant\n"),
"default Qwen render should start a new assistant turn, got {:?}",
default_prompt.as_str()
);
assert!(
closed_prompt
.as_str()
.ends_with("LLM-Native Interaction<|im_end|>\n"),
"add_generation_prompt=false alone must still close the last assistant, got {:?}",
closed_prompt.as_str()
);
assert_eq!(
continue_prompt.as_str(),
"<|im_start|>user\nContinue this sentence<|im_end|>\n<|im_start|>assistant\nLLM-Native Interaction"
);
}
#[test]
fn continue_final_message_marker_survives_rewritten_final_turn() {
const TEMPLATE: &str = "\
{%- for message in messages -%}\
{%- if loop.last -%}{{ message.content | upper }}|{%- else -%}{{ message.content }}|{%- endif -%}\
{%- endfor -%}closed";
let formatter = test_prompt_formatter(TEMPLATE);
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [
{"role": "user", "content": "same"},
{"role": "user", "content": "previous"},
{"role": "assistant", "content": "same"}
],
"add_generation_prompt": false,
"continue_final_message": true
}))
.unwrap();
let prompt = render_with_continue_final_message(formatter.as_ref(), &request);
assert_eq!(prompt.as_str(), "same|previous|SAME");
assert!(
!prompt.as_str().contains("CONTINUE_FINAL_MESSAGE_TAG"),
"marker must be stripped from the prompt sent to the model, got {:?}",
prompt.as_str()
);
}
#[test]
fn test_kimi_thinking_normalization_keeps_template_and_gates_in_sync() {
let template: dynamo_renderer::ChatTemplate = serde_json::from_value(serde_json::json!({
"chat_template": "{% if thinking == true and enable_thinking == true %}<think>{% elif thinking == false and enable_thinking == false %}<chat>{% else %}<mismatch>{% endif %}"
}))
.unwrap();
let formatter = dynamo_renderer::PromptFormatter::from_parts(
template,
dynamo_renderer::ContextMixins::default(),
false,
)
.unwrap();
let cases = [
(serde_json::json!(true), true, "true"),
(serde_json::json!("true"), true, "string true"),
(serde_json::json!("TRUE"), true, "uppercase true"),
(serde_json::json!("1"), true, "string one"),
(serde_json::json!("yes"), true, "yes"),
(serde_json::json!("on"), true, "on"),
(serde_json::json!(1), true, "one"),
(serde_json::json!(false), false, "false"),
(serde_json::json!("false"), false, "string false"),
(serde_json::json!("no"), false, "no"),
(serde_json::json!("off"), false, "off"),
(serde_json::json!(0), false, "zero"),
];
let assert_case =
|key: Option<&str>, value: Option<&serde_json::Value>, expected, description| {
let mut request: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"messages": [{"role": "user", "content": "hello"}],
"model": "moonshotai/Kimi-K2.5-Instruct"
}))
.unwrap();
if let (Some(key), Some(value)) = (key, value) {
request.chat_template_args = Some(std::collections::HashMap::from([(
key.to_string(),
value.clone(),
)]));
}
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some("kimi_k25"), None);
let args = request.chat_template_args.as_ref();
assert_eq!(
dynamo_renderer::thinking_bool_from_args(args),
Some(expected),
"renderer helper mismatch for {description}"
);
assert_eq!(
OpenAIPreprocessor::sglang_effective_reasoning_enabled(Some("kimi_k25"), args),
expected,
"SGLang gate mismatch for {description}"
);
assert_eq!(
!OpenAIPreprocessor::is_reasoning_disabled_by_request(Some("kimi_k25"), args,),
expected,
"postprocessor gate mismatch for {description}"
);
let rendered = match &formatter {
dynamo_renderer::PromptFormatter::OAI(formatter) => {
formatter.render(&request).unwrap()
}
};
assert_eq!(
rendered,
if expected { "<think>" } else { "<chat>" },
"rendered prompt mismatch for {description}"
);
};
for key in ["thinking", "enable_thinking"] {
for (value, expected, description) in &cases {
assert_case(Some(key), Some(value), *expected, *description);
}
}
assert_case(None, None, true, "omitted Kimi default");
for parser in ["kimi_k3", "kimi-k3"] {
let mut request: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"messages": [{"role": "user", "content": "hello"}],
"model": "moonshotai/Kimi-K3"
}))
.unwrap();
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some(parser), None);
assert_eq!(
dynamo_renderer::thinking_bool_from_args(request.chat_template_args.as_ref()),
Some(true),
"{parser} should default to thinking mode",
);
}
let mut conflicting_request: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"messages": [{"role": "user", "content": "hello"}],
"model": "moonshotai/Kimi-K2.5-Instruct",
"chat_template_kwargs": {
"thinking": "true",
"enable_thinking": false
}
}))
.unwrap();
OpenAIPreprocessor::normalize_thinking_arg(
&mut conflicting_request,
Some("kimi_k25"),
None,
);
let args = conflicting_request.chat_template_args.as_ref().unwrap();
assert_eq!(args.get("thinking"), Some(&serde_json::Value::Bool(true)));
assert_eq!(
args.get("enable_thinking"),
Some(&serde_json::Value::Bool(true))
);
let rendered = match &formatter {
dynamo_renderer::PromptFormatter::OAI(formatter) => {
formatter.render(&conflicting_request).unwrap()
}
};
assert_eq!(rendered, "<think>");
}
#[test]
fn test_named_kimi_k3_normalizes_every_reasoning_consumer_to_disabled() {
for parser in ["kimi_k3", "kimi-k3"] {
let mut request: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"messages": [{"role": "user", "content": "Weather in Berlin?"}],
"model": "moonshotai/Kimi-K3",
"tools": [{
"type": "function",
"function": {
"name": "get_weather",
"parameters": {
"type": "object",
"properties": {"city": {"type": "string"}}
}
}
}],
"tool_choice": {
"type": "function",
"function": {"name": "get_weather"}
},
"reasoning_effort": "high",
"chat_template_args": {
"thinking": true,
"enable_thinking": true
}
}))
.unwrap();
request.normalize_reasoning_template_args().unwrap();
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some(parser), None);
OpenAIPreprocessor::normalize_kimi_k3_named_tool_choice(&mut request, Some(parser));
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(args.get("thinking"), Some(&serde_json::Value::Bool(false)));
assert_eq!(
args.get("enable_thinking"),
Some(&serde_json::Value::Bool(false))
);
assert_eq!(
args.get("reasoning_effort"),
Some(&serde_json::Value::String("high".to_string())),
"the public effort value is preserved while the named-call exception disables thinking"
);
assert!(OpenAIPreprocessor::is_reasoning_disabled_by_request(
Some(parser),
Some(args)
));
}
}
#[test]
fn test_named_kimi_k3_override_is_scoped_to_k3_tool_parser() {
let request_json = serde_json::json!({
"messages": [{"role": "user", "content": "Weather?"}],
"model": "test",
"tools": [{
"type": "function",
"function": {
"name": "get_weather",
"parameters": {"type": "object"}
}
}],
"tool_choice": {
"type": "function",
"function": {"name": "get_weather"}
},
"chat_template_args": {"thinking": true}
});
for parser in [None, Some("hermes"), Some("kimi_k2")] {
let mut request: NvCreateChatCompletionRequest =
serde_json::from_value(request_json.clone()).unwrap();
OpenAIPreprocessor::normalize_kimi_k3_named_tool_choice(&mut request, parser);
assert_eq!(
dynamo_renderer::thinking_bool_from_args(request.chat_template_args.as_ref()),
Some(true),
"parser {parser:?} must retain its existing policy"
);
}
}
fn minimax_m3_request(body: serde_json::Value) -> NvCreateChatCompletionRequest {
let mut base = serde_json::json!({
"model": "MiniMaxAI/MiniMax-M3",
"messages": [{"role": "user", "content": "hi"}],
});
let base_obj = base.as_object_mut().unwrap();
for (k, v) in body.as_object().unwrap() {
base_obj.insert(k.clone(), v.clone());
}
serde_json::from_value(base).unwrap()
}
#[test]
fn test_normalize_thinking_arg_m3_defaults_disabled_for_json_schema() {
for parser in ["minimax_m3", "minimax-m3"] {
let mut request = minimax_m3_request(serde_json::json!({
"response_format": {
"type": "json_schema",
"json_schema": {"name": "s", "schema": {"type": "object"}}
}
}));
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some(parser), None);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("disabled"))
);
assert_eq!(args.get("thinking"), Some(&serde_json::json!(false)));
assert_eq!(args.get("enable_thinking"), Some(&serde_json::json!(false)));
}
}
#[test]
fn test_normalize_thinking_arg_m3_defaults_disabled_for_tool_choice_required() {
let mut request = minimax_m3_request(serde_json::json!({
"tools": [{"type": "function", "function": {"name": "f", "parameters": {"type": "object"}}}],
"tool_choice": "required"
}));
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some("minimax_m3"), None);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("disabled"))
);
}
#[test]
fn test_normalize_thinking_arg_m3_defaults_disabled_for_named_tool_choice() {
let mut request = minimax_m3_request(serde_json::json!({
"tools": [{"type": "function", "function": {"name": "f", "parameters": {"type": "object"}}}],
"tool_choice": {"type": "function", "function": {"name": "f"}}
}));
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some("minimax_m3"), None);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("disabled"))
);
}
#[test]
fn test_normalize_thinking_arg_m3_tool_call_only_deployment_gets_default() {
for tool_call_parser in [
"minimax_m3",
"minimax-m3",
"minimax_m3_nom",
"minimax-m3-nom",
] {
let mut request = minimax_m3_request(serde_json::json!({
"tools": [{"type": "function", "function": {"name": "f", "parameters": {"type": "object"}}}],
"tool_choice": "required"
}));
OpenAIPreprocessor::normalize_thinking_arg(&mut request, None, Some(tool_call_parser));
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("disabled")),
"tool_call_parser={tool_call_parser:?} + tool_choice=required with no reasoning parser must still default thinking_mode=disabled"
);
}
}
#[test]
fn test_normalize_thinking_arg_m3_preserves_client_adaptive() {
let mut request = minimax_m3_request(serde_json::json!({
"response_format": {
"type": "json_schema",
"json_schema": {"name": "s", "schema": {"type": "object"}}
},
"chat_template_kwargs": {"thinking_mode": "adaptive"}
}));
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some("minimax_m3"), None);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("adaptive"))
);
assert!(!args.contains_key("thinking"));
assert!(!args.contains_key("enable_thinking"));
}
#[test]
fn test_normalize_thinking_arg_m3_bridges_thinking_bool_to_mode() {
let mut request = minimax_m3_request(serde_json::json!({
"chat_template_kwargs": {"thinking": true}
}));
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some("minimax_m3"), None);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(args.get("thinking"), Some(&serde_json::json!(true)));
assert_eq!(args.get("enable_thinking"), Some(&serde_json::json!(true)));
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("enabled"))
);
}
#[test]
fn test_normalize_thinking_arg_m3_bridges_enable_thinking_to_mode() {
let mut request = minimax_m3_request(serde_json::json!({
"chat_template_kwargs": {"enable_thinking": false}
}));
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some("minimax_m3"), None);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(args.get("thinking"), Some(&serde_json::json!(false)));
assert_eq!(args.get("enable_thinking"), Some(&serde_json::json!(false)));
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("disabled"))
);
}
#[test]
fn test_normalize_thinking_arg_thinking_mode_enabled_is_thinking_on() {
for value in ["enabled", "ENABLED", "Enabled"] {
let mut request = minimax_m3_request(serde_json::json!({
"chat_template_kwargs": {"thinking_mode": value}
}));
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some("minimax_m3"), None);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(
args.get("thinking"),
Some(&serde_json::json!(true)),
"thinking_mode={value:?} must map to thinking=true"
);
assert_eq!(args.get("enable_thinking"), Some(&serde_json::json!(true)));
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("enabled"))
);
}
}
#[test]
fn test_normalize_thinking_arg_thinking_mode_disabled_is_thinking_off() {
for value in ["disabled", "DISABLED"] {
let mut request = minimax_m3_request(serde_json::json!({
"chat_template_kwargs": {"thinking_mode": value}
}));
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some("minimax_m3"), None);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(args.get("thinking"), Some(&serde_json::json!(false)));
assert_eq!(args.get("enable_thinking"), Some(&serde_json::json!(false)));
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("disabled"))
);
}
}
#[test]
fn test_normalize_thinking_arg_m3_canonicalizes_non_canonical_thinking_mode() {
for value in [
serde_json::json!(false),
serde_json::json!(0),
serde_json::json!("false"),
serde_json::json!("no"),
serde_json::json!("off"),
] {
let mut request = minimax_m3_request(serde_json::json!({
"chat_template_kwargs": {"thinking_mode": value.clone()}
}));
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some("minimax_m3"), None);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("disabled")),
"input={value} must canonicalize to \"disabled\""
);
assert_eq!(args.get("thinking"), Some(&serde_json::json!(false)));
assert_eq!(args.get("enable_thinking"), Some(&serde_json::json!(false)));
}
for value in [
serde_json::json!(true),
serde_json::json!(1),
serde_json::json!("true"),
serde_json::json!("yes"),
serde_json::json!("on"),
] {
let mut request = minimax_m3_request(serde_json::json!({
"chat_template_kwargs": {"thinking_mode": value.clone()}
}));
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some("minimax_m3"), None);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("enabled")),
"input={value} must canonicalize to \"enabled\""
);
assert_eq!(args.get("thinking"), Some(&serde_json::json!(true)));
assert_eq!(args.get("enable_thinking"), Some(&serde_json::json!(true)));
}
}
#[test]
fn test_normalize_thinking_arg_deepseek_thinking_mode_words_preserved() {
for value in ["thinking", "chat"] {
let mut request: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"model": "deepseek/DeepSeek-V3.1",
"messages": [{"role": "user", "content": "hi"}],
"chat_template_kwargs": {"thinking_mode": value}
}))
.unwrap();
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some("deepseek_v3_1"), None);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(
args.get("thinking_mode").and_then(|v| v.as_str()),
Some(value),
"DeepSeek {value:?} must survive verbatim"
);
assert!(
!args.contains_key("thinking"),
"no boolean coercion for DeepSeek {value:?}"
);
assert!(!args.contains_key("enable_thinking"));
}
}
#[test]
fn test_normalize_thinking_arg_thinking_mode_not_written_for_non_m3_parsers() {
for parser in [
"deepseek_v3_1",
"deepseek_v3_2",
"minimax_m2",
"qwen3",
"kimi_k25",
"kimi_k3",
] {
let mut request: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"model": "test",
"messages": [{"role": "user", "content": "hi"}],
"chat_template_kwargs": {"thinking": true}
}))
.unwrap();
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some(parser), None);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(args.get("thinking"), Some(&serde_json::json!(true)));
assert_eq!(args.get("enable_thinking"), Some(&serde_json::json!(true)));
assert!(
!args.contains_key("thinking_mode"),
"parser {parser:?} must not receive an injected thinking_mode value"
);
}
}
#[test]
fn test_normalize_thinking_arg_kimi_or_else_does_not_write_thinking_mode() {
for parser in ["kimi_k25", "kimi_k3", "kimi-k3"] {
let mut request: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"model": "test",
"messages": [{"role": "user", "content": "hi"}]
}))
.unwrap();
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some(parser), None);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(args.get("thinking"), Some(&serde_json::json!(true)));
assert!(!args.contains_key("thinking_mode"));
}
}
#[test]
fn test_normalize_thinking_arg_m3_leaves_plain_chat_untouched() {
let mut request = minimax_m3_request(serde_json::json!({}));
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some("minimax_m3"), None);
assert!(request.chat_template_args.is_none());
}
#[test]
fn test_normalize_thinking_arg_non_m3_adaptive_is_no_op() {
let mut request: NvCreateChatCompletionRequest =
serde_json::from_value(serde_json::json!({
"model": "test",
"messages": [{"role": "user", "content": "hi"}],
"chat_template_kwargs": {"thinking_mode": "adaptive"}
}))
.unwrap();
OpenAIPreprocessor::normalize_thinking_arg(&mut request, Some("qwen3"), None);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("adaptive"))
);
assert!(!args.contains_key("thinking"));
assert!(!args.contains_key("enable_thinking"));
}
#[test]
fn test_normalize_thinking_arg_operator_default_forces_m3_constrained_disabled() {
let runtime_config = runtime_config_with_default_thinking_mode("enabled");
let mut request = minimax_m3_request(serde_json::json!({
"response_format": {
"type": "json_schema",
"json_schema": {"name": "s", "schema": {"type": "object"}}
}
}));
let original_args = request.chat_template_args.clone();
assert!(OpenAIPreprocessor::request_disables_reasoning_with_family(
&request,
None,
Some("minimax_m3"),
Some("enabled"),
Some("minimax_m3"),
));
assert_eq!(request.chat_template_args, original_args);
let thinking_control_from_client =
OpenAIPreprocessor::request_has_client_thinking_control(&request);
assert!(
!thinking_control_from_client,
"no chat_template_kwargs → not client intent"
);
OpenAIPreprocessor::apply_default_thinking_mode_from_runtime_config(
&runtime_config,
&mut request,
);
OpenAIPreprocessor::normalize_thinking_arg_with_source(
&mut request,
Some("minimax_m3"),
None,
thinking_control_from_client,
);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!("disabled"))
);
assert_eq!(args.get("thinking"), Some(&serde_json::json!(false)));
assert_eq!(args.get("enable_thinking"), Some(&serde_json::json!(false)));
}
#[test]
fn test_normalize_thinking_arg_m3_preserves_explicit_thinking_for_constrained() {
for (body, expected_mode, desc) in [
(
serde_json::json!({
"response_format": {
"type": "json_schema",
"json_schema": {"name": "s", "schema": {"type": "object"}}
},
"chat_template_kwargs": {"thinking_mode": "enabled"}
}),
"enabled",
"json_schema + explicit thinking_mode=enabled",
),
(
serde_json::json!({
"tools": [{"type": "function", "function": {"name": "f", "parameters": {"type": "object"}}}],
"tool_choice": "required",
"chat_template_kwargs": {"thinking_mode": "adaptive"}
}),
"adaptive",
"required tool + explicit thinking_mode=adaptive",
),
(
serde_json::json!({
"tools": [{"type": "function", "function": {"name": "f", "parameters": {"type": "object"}}}],
"tool_choice": {"type": "function", "function": {"name": "f"}},
"chat_template_kwargs": {"thinking": true}
}),
"enabled",
"named tool + explicit thinking=true",
),
] {
let mut request = minimax_m3_request(body);
assert!(!OpenAIPreprocessor::request_disables_reasoning_with_family(
&request,
None,
Some("minimax_m3"),
Some("disabled"),
Some("minimax_m3"),
));
let thinking_control_from_client =
OpenAIPreprocessor::request_has_client_thinking_control(&request);
assert!(
thinking_control_from_client,
"{desc} must be classified as client thinking intent"
);
OpenAIPreprocessor::normalize_thinking_arg_with_source(
&mut request,
Some("minimax_m3"),
None,
thinking_control_from_client,
);
let args = request.chat_template_args.as_ref().unwrap();
assert_eq!(
args.get("thinking_mode"),
Some(&serde_json::json!(expected_mode)),
"{desc} must preserve explicit client thinking"
);
}
}
#[test]
fn test_is_reasoning_disabled_by_request() {
let thinking_true = {
let mut m = std::collections::HashMap::new();
m.insert("thinking".to_string(), serde_json::Value::Bool(true));
m
};
let thinking_false = {
let mut m = std::collections::HashMap::new();
m.insert("thinking".to_string(), serde_json::Value::Bool(false));
m
};
let enable_thinking_true = {
let mut m = std::collections::HashMap::new();
m.insert("enable_thinking".to_string(), serde_json::Value::Bool(true));
m
};
let enable_thinking_false = {
let mut m = std::collections::HashMap::new();
m.insert(
"enable_thinking".to_string(),
serde_json::Value::Bool(false),
);
m
};
let force_nonempty_content_true = {
let mut m = std::collections::HashMap::new();
m.insert(
"force_nonempty_content".to_string(),
serde_json::Value::Bool(true),
);
m
};
let thinking_mode_chat = {
let mut m = std::collections::HashMap::new();
m.insert(
"thinking_mode".to_string(),
serde_json::Value::String("chat".to_string()),
);
m
};
let thinking_mode_thinking = {
let mut m = std::collections::HashMap::new();
m.insert(
"thinking_mode".to_string(),
serde_json::Value::String("thinking".to_string()),
);
m
};
let thinking_mode_disabled = {
let mut m = std::collections::HashMap::new();
m.insert(
"thinking_mode".to_string(),
serde_json::Value::String("disabled".to_string()),
);
m
};
let reasoning_effort_none = {
let mut m = std::collections::HashMap::new();
m.insert(
"reasoning_effort".to_string(),
serde_json::Value::String("none".to_string()),
);
m
};
let reasoning_effort_high = {
let mut m = std::collections::HashMap::new();
m.insert(
"reasoning_effort".to_string(),
serde_json::Value::String("high".to_string()),
);
m
};
let empty_args = std::collections::HashMap::new();
let cases = [
(
Some("kimi_k25"),
Some(&thinking_false),
true,
"kimi_k25 + thinking=false → disabled",
),
(
Some("kimi_k25"),
Some(&thinking_true),
false,
"kimi_k25 + thinking=true → enabled",
),
(
Some("kimi_k25"),
None,
false,
"kimi_k25 + no args → enabled",
),
(
Some("kimi_k25"),
Some(&empty_args),
false,
"kimi_k25 + empty args → enabled",
),
(
Some("kimi_k3"),
Some(&thinking_false),
true,
"kimi_k3 + thinking=false → disabled",
),
(
Some("kimi-k3"),
Some(&thinking_true),
false,
"kimi-k3 + thinking=true → enabled",
),
(
Some("deepseek_r1"),
Some(&thinking_false),
true,
"deepseek_r1 + thinking=false → disabled",
),
(
Some("deepseek_r1"),
Some(&thinking_true),
false,
"deepseek_r1 + thinking=true → enabled",
),
(
Some("deepseek_r1"),
Some(&thinking_mode_chat),
true,
"deepseek_r1 + thinking_mode=chat → disabled",
),
(
Some("deepseek_r1"),
Some(&thinking_mode_thinking),
false,
"deepseek_r1 + thinking_mode=thinking → enabled",
),
(
Some("deepseek_r1"),
None,
false,
"deepseek_r1 + no args → enabled",
),
(
Some("deepseek_r1"),
Some(&empty_args),
false,
"deepseek_r1 + empty args → enabled",
),
(
Some("deepseek_v3"),
None,
true,
"deepseek_v3 + no args → disabled",
),
(
Some("deepseek_v3"),
Some(&thinking_true),
false,
"deepseek_v3 + thinking=true → enabled",
),
(
Some("deepseek_v3_1"),
Some(&enable_thinking_true),
false,
"deepseek_v3_1 + enable_thinking=true → enabled",
),
(
Some("deepseek_v3_2"),
None,
false,
"deepseek_v3_2 + no args → enabled",
),
(
Some("deepseek_v3_2"),
Some(&thinking_false),
true,
"deepseek_v3_2 + thinking=false → disabled",
),
(
Some("minimax_m2"),
Some(&thinking_false),
true,
"minimax_m2 + thinking=false → disabled",
),
(
Some("minimax_m2"),
Some(&thinking_true),
false,
"minimax_m2 + thinking=true → enabled",
),
(
Some("minimax_m2"),
None,
false,
"minimax_m2 + no args → enabled",
),
(
Some("basic"),
Some(&thinking_false),
false,
"basic → never disabled",
),
(
None,
Some(&thinking_false),
false,
"no parser → never disabled",
),
(
Some("nemotron_nano"),
Some(&enable_thinking_false),
true,
"nemotron_nano + enable_thinking=false → disabled",
),
(
Some("nemotron_nano"),
Some(&enable_thinking_true),
false,
"nemotron_nano + enable_thinking=true → enabled",
),
(
Some("nemotron_nano"),
None,
false,
"nemotron_nano + no args → enabled",
),
(
Some("nemotron_nano"),
Some(&empty_args),
false,
"nemotron_nano + empty args → enabled",
),
(
Some("nemotron3"),
Some(&force_nonempty_content_true),
false,
"nemotron3 + force_nonempty_content=true → NOT disabled (parser stays on)",
),
(
Some("nemotron_v3"),
Some(&enable_thinking_false),
true,
"nemotron_v3 + enable_thinking=false → disabled",
),
(
Some("nemotron_v3"),
Some(&force_nonempty_content_true),
false,
"nemotron_v3 + force_nonempty_content=true → NOT disabled (parser stays on)",
),
(
Some("deepseek_v4"),
Some(&thinking_false),
true,
"deepseek_v4 + thinking=false → disabled",
),
(
Some("deepseek_v4"),
Some(&thinking_true),
false,
"deepseek_v4 + thinking=true → enabled",
),
(
Some("deepseek_v4"),
Some(&thinking_mode_chat),
true,
"deepseek_v4 + thinking_mode=chat → disabled",
),
(
Some("deepseek_v4"),
Some(&thinking_mode_thinking),
false,
"deepseek_v4 + thinking_mode=thinking → enabled",
),
(
Some("deepseek_v4"),
None,
false,
"deepseek_v4 + no args → enabled",
),
(
Some("deepseek-v4"),
Some(&thinking_false),
true,
"deepseek-v4 (hyphen alias) + thinking=false → disabled",
),
(
Some("deepseekv4"),
Some(&thinking_mode_chat),
true,
"deepseekv4 (joined alias) + thinking_mode=chat → disabled",
),
(
Some("deepseek_v4"),
Some(&enable_thinking_false),
true,
"deepseek_v4 + enable_thinking=false → disabled (vLLM alias)",
),
(
Some("deepseek_v4"),
Some(&enable_thinking_true),
false,
"deepseek_v4 + enable_thinking=true → enabled (vLLM alias)",
),
(
Some("gemma4"),
Some(&enable_thinking_false),
true,
"gemma4 + enable_thinking=false → disabled",
),
(
Some("gemma4"),
Some(&enable_thinking_true),
false,
"gemma4 + enable_thinking=true → enabled",
),
(
Some("gemma4"),
None,
true,
"gemma4 + no args → disabled (reasoning is opt-in)",
),
(
Some("gemma-4"),
Some(&enable_thinking_false),
true,
"gemma-4 (hyphen alias) + enable_thinking=false → disabled",
),
(
Some("gemma-4"),
None,
true,
"gemma-4 (hyphen alias) + no args → disabled (reasoning is opt-in)",
),
(Some("mistral"), None, true, "mistral + no args → disabled"),
(
Some("mistral"),
Some(&reasoning_effort_none),
true,
"mistral + reasoning_effort=none → disabled",
),
(
Some("mistral"),
Some(&reasoning_effort_high),
false,
"mistral + reasoning_effort=high → enabled",
),
(
Some("minimax_m3"),
Some(&thinking_mode_disabled),
true,
"minimax_m3 + thinking_mode=disabled → disabled",
),
(
Some("minimax-m3"),
Some(&thinking_mode_disabled),
true,
"minimax-m3 + thinking_mode=disabled → disabled",
),
(
Some("minimax_m3"),
Some(&thinking_mode_thinking),
false,
"minimax_m3 + thinking_mode=thinking → enabled",
),
(
Some("minimax_m3"),
None,
false,
"minimax_m3 + no args → enabled",
),
];
for (parser, args, expected, desc) in cases {
assert_eq!(
OpenAIPreprocessor::is_reasoning_disabled_by_request(parser, args),
expected,
"FAILED: {desc}",
);
}
assert!(
!OpenAIPreprocessor::is_reasoning_disabled_by_request(
Some("nemotron3"),
Some(&force_nonempty_content_true),
),
"nemotron3 + force_nonempty_content=true → NOT disabled",
);
assert!(
!OpenAIPreprocessor::is_reasoning_disabled_by_request(
Some("nemotron_v3"),
Some(&force_nonempty_content_true),
),
"nemotron_v3 + force_nonempty_content=true → NOT disabled",
);
assert!(
OpenAIPreprocessor::is_reasoning_disabled_by_request(
Some("nemotron3"),
Some(&enable_thinking_false),
),
"nemotron3 + enable_thinking=false → disabled",
);
}
#[test]
fn muse_force_nonempty_stream_reports_that_it_can_defer_all_output() {
let mut args = std::collections::HashMap::new();
args.insert(
"force_nonempty_content".to_string(),
serde_json::Value::Bool(true),
);
assert!(OpenAIPreprocessor::stream_can_defer_all_output(
None,
Some("muse_glimmer"),
Some(&args),
));
}
#[cfg(feature = "mm-routing")]
#[test]
fn hash_image_url_distinguishes_query_strings() {
let base = "https://cdn.example.com/img.jpg";
let v1 = OpenAIPreprocessor::hash_image_url(&format!("{base}?v=1"));
let v2 = OpenAIPreprocessor::hash_image_url(&format!("{base}?v=2"));
let no_q = OpenAIPreprocessor::hash_image_url(base);
assert_ne!(v1, v2, "different query values must hash differently");
assert_ne!(v1, no_q, "presence of a query string must change the hash");
}
#[cfg(feature = "mm-routing")]
#[test]
fn image_dim_cache_preserves_terminal_policy_classification() {
let terminal = ImageDimFetchFailure::from_error(invalid_argument_error(
"media destination rejected by SSRF policy",
));
let terminal = terminal.to_error();
assert!(MediaFetcher::is_policy_rejection(&terminal));
let recoverable =
ImageDimFetchFailure::from_error(anyhow::anyhow!("image header was truncated"));
let recoverable = recoverable.to_error();
assert!(!MediaFetcher::is_policy_rejection(&recoverable));
}
#[tokio::test]
async fn gather_rejects_data_url_over_the_size_cap() {
let mdc = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
let preprocessor = OpenAIPreprocessor::new(mdc).unwrap();
assert!(preprocessor.media_loader.is_none());
let request_with_data_url = |size: usize| -> NvCreateChatCompletionRequest {
let prefix = "data:audio/wav;base64,";
let url = format!("{prefix}{}", "A".repeat(size - prefix.len()));
serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{
"role": "user",
"content": [{"type": "audio_url", "audio_url": {"url": url}}]
}]
}))
.unwrap()
};
temp_env::async_with_vars([("DYN_MM_MAX_DATA_URL_MB", Some("1"))], async {
let limit = 1024 * 1024;
let mut builder = PreprocessedRequest::builder();
preprocessor
.gather_multi_modal_data(&request_with_data_url(limit), &mut builder, None, &[])
.await
.expect("a data: URL at the cap is accepted");
let mut builder = PreprocessedRequest::builder();
let error = preprocessor
.gather_multi_modal_data(&request_with_data_url(limit + 1), &mut builder, None, &[])
.await
.expect_err("a data: URL over the cap is rejected");
let dynamo_error = error
.downcast_ref::<DynamoError>()
.expect("error should preserve the DynamoError type");
assert_eq!(dynamo_error.error_type(), ErrorType::InvalidArgument);
assert!(
dynamo_error.message().contains(
"audio_url data: URL is 1048577 bytes, exceeds the 1048576-byte limit"
),
"{}",
dynamo_error.message()
);
assert!(
dynamo_error.message().contains(
"To raise the limit, set DYN_MM_MAX_DATA_URL_MB (in megabytes) on both \
the frontend and the workers."
),
"{}",
dynamo_error.message()
);
assert_eq!(dynamo_error.public_message(), Some(dynamo_error.message()));
})
.await;
}
#[cfg(feature = "mm-routing")]
#[tokio::test]
async fn url_passthrough_policy_rejection_is_terminal() {
if std::env::var("DYN_MM_ALLOW_INTERNAL").as_deref() == Ok("1") {
return;
}
let mdc = ModelDeploymentCard::load_from_disk(
"tests/data/sample-models/mock-llama-3.1-8b-instruct",
None,
)
.unwrap();
let preprocessor = OpenAIPreprocessor::new(mdc).unwrap();
assert!(preprocessor.media_loader.is_none());
let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "test-model",
"messages": [{
"role": "user",
"content": [{
"type": "image_url",
"image_url": {
"url": "https://169.254.169.254/latest/meta-data/private-image.png"
}
}]
}]
}))
.unwrap();
let mut builder = PreprocessedRequest::builder();
let error = preprocessor
.gather_multi_modal_data(&request, &mut builder, None, &[])
.await
.expect_err("URL-passthrough must stop at the policy rejection");
assert!(MediaFetcher::is_policy_rejection(&error));
}
#[cfg(feature = "mm-routing")]
#[tokio::test]
async fn dim_fetch_declines_passthrough_schemes_without_rejecting_them() {
for url in [
"s3://bucket/private-image.png",
"gs://bucket/private-image.png",
"file:///tmp/private-image.png",
] {
let error = OpenAIPreprocessor::fetch_image_dims_uncached(
url,
RoutingImageDimensionPolicy::Encoded,
)
.await
.expect_err("a non-fetchable scheme cannot yield dimensions");
assert!(
!MediaFetcher::is_policy_rejection(&error),
"{url} must be a skippable dim-fetch failure, not a terminal 4xx: {error:#}"
);
}
}
#[cfg(feature = "mm-routing")]
#[test]
fn exif_dimension_policy_is_limited_to_kimi_k3_vllm_url_passthrough() {
use crate::local_model::runtime_config::{
ModelRuntimeConfig, SGLANG_GENERATE_CAPABILITY, VLLM_INFERENCE_V1_GENERATE_CAPABILITY,
};
let kimi_k3 = Some(RoutingImagePromptLayout::KimiK3 {
media_begin: 1,
media_content: 2,
media_end: 3,
});
let mut runtime_config = ModelRuntimeConfig::default();
assert_eq!(
routing_image_dimension_policy(&runtime_config, false, kimi_k3),
RoutingImageDimensionPolicy::Encoded
);
let mut sglang_config = ModelRuntimeConfig::default();
sglang_config
.set_engine_specific(SGLANG_GENERATE_CAPABILITY, true)
.unwrap();
assert_eq!(
routing_image_dimension_policy(&sglang_config, false, kimi_k3),
RoutingImageDimensionPolicy::Encoded
);
runtime_config
.set_engine_specific(VLLM_INFERENCE_V1_GENERATE_CAPABILITY, true)
.unwrap();
assert_eq!(
routing_image_dimension_policy(&runtime_config, false, kimi_k3),
RoutingImageDimensionPolicy::ExifTransposed
);
assert_eq!(
routing_image_dimension_policy(&runtime_config, true, kimi_k3),
RoutingImageDimensionPolicy::Encoded
);
assert_eq!(
routing_image_dimension_policy(
&runtime_config,
false,
Some(RoutingImagePromptLayout::RepeatedPad),
),
RoutingImageDimensionPolicy::Encoded
);
}
#[cfg(all(feature = "mm-routing", feature = "media-ffmpeg"))]
#[test]
fn dual_qwen_video_contracts_disable_exact_routing() {
let mut runtime_config = ModelRuntimeConfig::default();
runtime_config
.set_engine_specific(
VLLM_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY,
serde_json::json!({
"placeholder_target": "bare_video_token",
"resize_mode": "legacy_ceil"
}),
)
.unwrap();
runtime_config
.set_engine_specific(
SGLANG_QWEN_VIDEO_PROCESSOR_CONTRACT_RUNTIME_KEY,
serde_json::json!({
"placeholder_target": "bare_video_token",
"resize_mode": "legacy_ceil",
"runless_boundary_hash": "tokens_only",
"sglang_preprocess": {
"image_factor": 28,
"video_min_pixels": 100352,
"video_max_pixels": 602112,
"video_total_pixels": 90316800,
"frame_factor": 2,
"fps": 2.0,
"min_frames": 4,
"max_frames": 768
}
}),
)
.unwrap();
let error = resolve_qwen_video_processor_contract(&runtime_config).unwrap_err();
assert!(
error
.to_string()
.contains("multiple Qwen video processor contracts")
);
}
#[cfg(feature = "mm-routing")]
#[test]
fn exif_transposed_dimensions_match_vllm_image_loading() {
use image::{ExtendedColorType, ImageEncoder, codecs::jpeg::JpegEncoder};
let exif = vec![
b'I', b'I', 42, 0, 8, 0, 0, 0, 1, 0, 0x12, 0x01, 3, 0, 1, 0, 0, 0, 6, 0, 0, 0, 0, 0, 0,
0,
];
let mut jpeg = Vec::new();
let mut encoder = JpegEncoder::new(&mut jpeg);
encoder.set_exif_metadata(exif).unwrap();
encoder
.encode(&[255, 0, 0, 0, 255, 0], 2, 1, ExtendedColorType::Rgb8)
.unwrap();
assert_eq!(
OpenAIPreprocessor::dimensions_from_image_bytes(
&jpeg,
RoutingImageDimensionPolicy::Encoded,
)
.unwrap(),
(2, 1)
);
assert_eq!(
OpenAIPreprocessor::dimensions_from_image_bytes(
&jpeg,
RoutingImageDimensionPolicy::ExifTransposed,
)
.unwrap(),
(1, 2)
);
}
#[cfg(feature = "mm-routing")]
#[test]
fn hash_image_url_distinguishes_rotating_signatures() {
let base = "https://bucket.s3.amazonaws.com/img.jpg";
let a = OpenAIPreprocessor::hash_image_url(&format!(
"{base}?X-Amz-Signature=AAA&X-Amz-Date=20260101T000000Z&X-Amz-Expires=600"
));
let b = OpenAIPreprocessor::hash_image_url(&format!(
"{base}?X-Amz-Signature=BBB&X-Amz-Date=20260101T010000Z&X-Amz-Expires=900"
));
assert_ne!(
a, b,
"rotating presign params produce a different URL and must hash differently"
);
}
#[cfg(feature = "mm-routing")]
#[test]
fn hash_image_url_is_deterministic_for_identical_urls() {
let url = "https://cdn.example.com/img.jpg?width=256";
assert_eq!(
OpenAIPreprocessor::hash_image_url(url),
OpenAIPreprocessor::hash_image_url(url),
);
}
#[cfg(feature = "mm-routing")]
#[test]
fn hash_image_url_data_uri_content_addressed() {
let same = "data:image/png;base64,AAAA";
let other = "data:image/png;base64,BBBB";
assert_eq!(
OpenAIPreprocessor::hash_image_url(same),
OpenAIPreprocessor::hash_image_url(same)
);
assert_ne!(
OpenAIPreprocessor::hash_image_url(same),
OpenAIPreprocessor::hash_image_url(other),
"different data URI payloads must hash differently"
);
}
#[cfg(feature = "mm-routing")]
#[test]
fn hash_image_url_other_schemes_passthrough() {
let s3a = OpenAIPreprocessor::hash_image_url("s3://bucket/key?v=1");
let s3b = OpenAIPreprocessor::hash_image_url("s3://bucket/key?v=2");
assert_ne!(
s3a, s3b,
"s3:// query params identify objects and must not collide"
);
}
#[cfg(feature = "mm-routing")]
#[test]
fn tracked_video_boundary_uses_native_metadata_only_when_needed() {
use dynamo_kv_router::protocols::pad_value_for_mm_hash;
let video_token_id = 100;
let mm_hash = 41;
let video_pad = pad_value_for_mm_hash(mm_hash);
let replacement = TrackedMmRoutingReplacement {
mm_hash,
target_tokens: vec![9],
worker_tokens: vec![3, 4, 5, 6, video_token_id, video_token_id],
routing_tokens: vec![3, 4, 5, 6, video_pad, video_pad],
};
let (tokens, prompt_len, infos) = apply_tracked_mm_replacements(
None,
&[replacement],
&[1, 9, 2],
4,
Some(99),
Some(video_token_id),
mm_routing::KvEventMmIdentity::MmMetadata,
)
.unwrap();
assert_eq!(prompt_len, 8);
assert_eq!(&tokens[..4], &[1, 3, 4, 5]);
assert_eq!(&tokens[4..], &[6, video_pad, video_pad, 2]);
assert_eq!(infos[0].as_ref().unwrap().mm_objects[0].mm_hash, mm_hash);
assert!(infos[1].is_none());
}
#[cfg(feature = "mm-routing")]
#[test]
fn tracked_video_boundary_matches_token_only_worker_contract() {
use dynamo_kv_router::protocols::pad_value_for_mm_hash;
let video_token_id = 100;
let mm_hash = 41;
let video_pad = pad_value_for_mm_hash(mm_hash);
let replacement = TrackedMmRoutingReplacement {
mm_hash,
target_tokens: vec![9],
worker_tokens: vec![3, 4, 5, 6, video_token_id, video_token_id],
routing_tokens: vec![3, 4, 5, 6, video_pad, video_pad],
};
let (tokens, prompt_len, infos) = apply_tracked_mm_replacements(
None,
&[replacement],
&[1, 9, 2],
4,
Some(99),
Some(video_token_id),
mm_routing::KvEventMmIdentity::PadValueTokens,
)
.unwrap();
assert_eq!(prompt_len, 8);
assert_eq!(&tokens[..4], &[1, 3, 4, 5]);
assert_eq!(&tokens[4..], &[6, video_pad, video_pad, 2]);
assert!(infos.iter().all(Option::is_none));
}
#[cfg(feature = "mm-routing")]
#[test]
fn tracked_token_only_worker_still_validates_normalizable_blocks() {
let video_token_id = 100;
let replacement = TrackedMmRoutingReplacement {
mm_hash: 41,
target_tokens: vec![9],
worker_tokens: vec![video_token_id, video_token_id],
routing_tokens: vec![1, 2],
};
let error = apply_tracked_mm_replacements(
None,
&[replacement],
&[9],
4,
Some(99),
Some(video_token_id),
mm_routing::KvEventMmIdentity::PadValueTokens,
)
.unwrap_err();
assert!(
error
.to_string()
.contains("frontend MM replacement differs from KV-event normalization")
);
}
#[cfg(feature = "mm-routing")]
#[test]
fn tracked_mixed_boundary_preserves_worker_hash_fallback() {
use dynamo_kv_router::protocols::pad_value_for_mm_hash;
let image_token_id = 99;
let video_token_id = 100;
let image_hash = 41;
let video_hash = 42;
let replacements = [
TrackedMmRoutingReplacement {
mm_hash: image_hash,
target_tokens: vec![image_token_id],
worker_tokens: vec![image_token_id, image_token_id, 7],
routing_tokens: vec![
pad_value_for_mm_hash(image_hash),
pad_value_for_mm_hash(image_hash),
7,
],
},
TrackedMmRoutingReplacement {
mm_hash: video_hash,
target_tokens: vec![video_token_id],
worker_tokens: vec![8, video_token_id, video_token_id, 9],
routing_tokens: vec![
8,
pad_value_for_mm_hash(video_hash),
pad_value_for_mm_hash(video_hash),
9,
],
},
];
let (tokens, prompt_len, infos) = apply_tracked_mm_replacements(
None,
&replacements,
&[image_token_id, video_token_id],
4,
Some(image_token_id),
Some(video_token_id),
mm_routing::KvEventMmIdentity::MmMetadata,
)
.unwrap();
assert_eq!(prompt_len, 7);
assert_eq!(&tokens[..4], &[image_token_id, image_token_id, 7, 8]);
assert_eq!(
&tokens[4..],
&[
pad_value_for_mm_hash(video_hash),
pad_value_for_mm_hash(video_hash),
9,
0
]
);
assert_eq!(
infos[0]
.as_ref()
.unwrap()
.mm_objects
.iter()
.map(|object| object.mm_hash)
.collect::<Vec<_>>(),
[image_hash, video_hash]
);
assert!(infos[1].is_none());
}
#[cfg(feature = "mm-routing")]
#[test]
fn tracked_mixed_token_only_boundary_preserves_canonical_pads() {
use dynamo_kv_router::protocols::pad_value_for_mm_hash;
let image_token_id = 99;
let video_token_id = 100;
let image_hash = 41;
let video_hash = 42;
let image_pad = pad_value_for_mm_hash(image_hash);
let video_pad = pad_value_for_mm_hash(video_hash);
let replacements = [
TrackedMmRoutingReplacement {
mm_hash: image_hash,
target_tokens: vec![image_token_id],
worker_tokens: vec![image_token_id; 14],
routing_tokens: vec![image_pad; 14],
},
TrackedMmRoutingReplacement {
mm_hash: video_hash,
target_tokens: vec![video_token_id],
worker_tokens: vec![7, 8, video_token_id, video_token_id],
routing_tokens: vec![7, 8, video_pad, video_pad],
},
];
let (tokens, prompt_len, infos) = apply_tracked_mm_replacements(
None,
&replacements,
&[image_token_id, video_token_id],
16,
Some(image_token_id),
Some(video_token_id),
mm_routing::KvEventMmIdentity::PadValueTokens,
)
.unwrap();
assert_eq!(prompt_len, 18);
assert_eq!(&tokens[..14], &[image_pad; 14]);
assert_eq!(&tokens[14..18], &[7, 8, video_pad, video_pad]);
assert!(tokens[18..].iter().all(|token| *token == 0));
assert!(infos.iter().all(Option::is_none));
}
#[cfg(feature = "mm-routing")]
#[test]
fn tracked_replacements_preserve_image_video_image_order() {
use dynamo_kv_router::protocols::pad_value_for_mm_hash;
let image_token_id = 99;
let video_token_id = 100;
let replacement = |mm_hash, target_token| TrackedMmRoutingReplacement {
mm_hash,
target_tokens: vec![target_token],
worker_tokens: vec![target_token, target_token],
routing_tokens: vec![
pad_value_for_mm_hash(mm_hash),
pad_value_for_mm_hash(mm_hash),
],
};
let replacements = [
replacement(41, image_token_id),
replacement(42, video_token_id),
replacement(43, image_token_id),
];
let (tokens, prompt_len, infos) = apply_tracked_mm_replacements(
None,
&replacements,
&[1, image_token_id, 2, video_token_id, 3, image_token_id, 4],
16,
Some(image_token_id),
Some(video_token_id),
mm_routing::KvEventMmIdentity::MmMetadata,
)
.unwrap();
assert_eq!(prompt_len, 10);
assert_eq!(
&tokens[..prompt_len],
&[
1,
pad_value_for_mm_hash(41),
pad_value_for_mm_hash(41),
2,
pad_value_for_mm_hash(42),
pad_value_for_mm_hash(42),
3,
pad_value_for_mm_hash(43),
pad_value_for_mm_hash(43),
4,
]
);
assert!(infos[0].is_none());
}
#[cfg(feature = "mm-routing")]
#[test]
fn tracked_replacements_reject_misordered_missing_and_extra_targets() {
let replacement = |mm_hash, target| TrackedMmRoutingReplacement {
mm_hash,
target_tokens: vec![target],
worker_tokens: vec![target],
routing_tokens: vec![target],
};
let replacements = [replacement(41, 10), replacement(42, 20)];
for token_ids in [&[20, 10][..], &[10][..], &[10, 20, 20][..]] {
assert!(
apply_tracked_mm_replacements(
None,
&replacements,
token_ids,
4,
Some(10),
Some(20),
mm_routing::KvEventMmIdentity::MmMetadata,
)
.is_err(),
"invalid target sequence {token_ids:?} must fail closed"
);
}
}
#[cfg(feature = "mm-routing")]
#[test]
fn shared_placeholder_video_matches_legacy_worker_normalization() {
use dynamo_kv_router::protocols::pad_value_for_mm_hash;
let mm_hash = 41;
let pad = pad_value_for_mm_hash(mm_hash);
let replacement = TrackedMmRoutingReplacement {
mm_hash,
target_tokens: vec![7],
worker_tokens: vec![100, 19, 18, 18, 20, 101, 19, 18, 20],
routing_tokens: vec![100, 19, pad, pad, 20, 101, 19, pad, 20],
};
let (tokens, prompt_len, block_infos) = apply_tracked_mm_replacements(
None,
&[replacement],
&[7],
4,
Some(18),
None,
mm_routing::KvEventMmIdentity::MmMetadata,
)
.unwrap();
assert_eq!(prompt_len, 9);
assert_eq!(tokens, [100, 19, pad, pad, 20, 101, 19, pad, 20, 0, 0, 0]);
assert_eq!(block_infos.len(), 3);
assert!(block_infos[0].is_none());
assert!(block_infos[1].is_none());
assert_eq!(
block_infos[2]
.as_ref()
.unwrap()
.mm_objects
.iter()
.map(|object| object.mm_hash)
.collect::<Vec<_>>(),
[mm_hash]
);
}
#[cfg(feature = "mm-routing")]
#[test]
fn routing_bos_preserves_image_behavior_and_skips_video_only_requests() {
assert_eq!(routing_bos_to_prepend(Some(1), true), Some(1));
assert_eq!(routing_bos_to_prepend(None, true), None);
assert_eq!(routing_bos_to_prepend(Some(1), false), None);
}
#[cfg(feature = "mm-routing")]
#[test]
fn consecutive_video_entries_are_not_exactly_routable() {
let video = |mm_hash| MmRoutingEntry::Video {
mm_hash,
placeholder_token_id: 3,
event_video_token_id: Some(3),
target_tokens: vec![3],
replacement_tokens: vec![3],
};
let image = MmRoutingEntry::Image {
mm_hash: 2,
width: 1,
height: 1,
};
assert!(exact_mm_routing_entries_are_unambiguous(&[video(1)]));
assert!(!exact_mm_routing_entries_are_unambiguous(&[
video(1),
video(2)
]));
assert!(exact_mm_routing_entries_are_unambiguous(&[
video(1),
image,
video(2)
]));
}
}