dynamo_renderer/
lib.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Prompt Formatting
5//!
6//! Standalone, runtime-free chat-template / prompt formatting for
7//! OpenAI-compatible inference frontends. Renders HuggingFace `chat_template`
8//! jinja2 (via `minijinja` + `minijinja-contrib` pycompat), handles tool
9//! usage formatting and generation-prompt handling.
10//!
11//! Consumers implement [`OAIChatLikeRequest`] for their request type (or use
12//! the ready-made impl for `dynamo-protocols`' OpenAI chat request) and render
13//! with a [`PromptFormatter`] built from a HuggingFace `tokenizer_config.json`
14//! ([`ChatTemplate`]).
15//!
16//! This crate is a *bridge* between OpenAI request types ([`dynamo_protocols`])
17//! and prompt rendering. Most formatters return text; segment-sensitive native
18//! formats can preserve tokenizer policy through [`RenderedPrompt`].
19
20// TODO:
21// 1. Query if `add_generation_prompt` is present in the prompt template
22// 2. Support for models with add_generation_prompt:
23//    - PALS (Prefix-Assisted Language Sampling)
24//    - Continuation - Detected on user turns, where we can return
25//      partial assistant responses without add_generation_prompt
26
27use anyhow::{Context, Result};
28use minijinja::value::Value;
29use std::collections::HashMap;
30use std::sync::Arc;
31
32/// Re-export of `dynamo-tokenizers` as a one-import convenience: consumers that
33/// want both tokenization and chat templating can reach the tokenizer types via
34/// `dynamo_renderer::dynamo_tokenizers::*` without adding a second dependency.
35pub use dynamo_tokenizers;
36
37pub mod deepseek;
38pub mod inkling;
39pub mod kimi_k3;
40mod python;
41mod template;
42
43pub use template::{
44    ChatTemplate, ChatTemplateValue, ContextMixins, deepseek_formatter_for, kimi_k3_formatter_for,
45    may_be_fix_tool_schema, native_formatter_for,
46};
47
48/// Selects which context-mixin behaviors a template renders with.
49///
50/// Carried on the model deployment card (`prompt_context`) and consumed by the
51/// chat-template renderer via [`ContextMixins`].
52#[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq, Hash)]
53#[serde(rename_all = "snake_case")]
54pub enum PromptContextMixin {
55    /// Support OAI Chat Messages and Tools
56    OaiChat,
57
58    /// Enables templates with `{{datetime}}` to be rendered with the current date and time.
59    Llama3DateTime,
60}
61
62/// Shared helper: extract a boolean thinking toggle from `chat_template_args`.
63///
64/// Reads the two equivalent keys (`thinking`, `enable_thinking` — vLLM's
65/// canonical kwarg) in order and returns the first bool value found, or `None`
66/// if neither key is present (or neither carries a bool). Used by the V4
67/// formatter's `resolve_thinking_mode` and by reasoning-parser gating in
68/// consumers so both paths agree on the signal interpretation.
69pub fn thinking_bool_from_args(args: Option<&HashMap<String, serde_json::Value>>) -> Option<bool> {
70    let args = args?;
71    for key in ["thinking", "enable_thinking"] {
72        if let Some(v) = args.get(key).and_then(|x| x.as_bool()) {
73            return Some(v);
74        }
75    }
76    None
77}
78
79#[derive(Debug)]
80pub enum TokenInput {
81    Single(Vec<u32>),
82    Batch(Vec<Vec<u32>>),
83}
84
85#[derive(Debug)]
86pub enum TextInput {
87    Single(String),
88    Batch(Vec<String>),
89}
90
91#[derive(Debug)]
92pub enum PromptInput {
93    Tokens(TokenInput),
94    Text(TextInput),
95}
96
97/// One owned prompt segment with an explicit special-token trust boundary.
98#[derive(Debug, Clone, PartialEq, Eq)]
99pub struct RenderedSegment {
100    pub text: String,
101    pub allow_special: bool,
102}
103
104impl RenderedSegment {
105    pub fn new(text: impl Into<String>, allow_special: bool) -> Self {
106        Self {
107            text: text.into(),
108            allow_special,
109        }
110    }
111
112    pub fn as_encode_segment(&self) -> dynamo_tokenizers::EncodeSegment<'_> {
113        dynamo_tokenizers::EncodeSegment::new(&self.text, self.allow_special)
114    }
115}
116
117/// A rendered prompt plus its optional tokenization boundaries.
118///
119/// The prompt owns its segment text while `dynamo-tokenizers` borrows that text
120/// during encoding. Keeping the types separate preserves the tokenizer crate's
121/// published zero-copy `EncodeSegment<'_>` API.
122#[derive(Debug, Clone, PartialEq, Eq)]
123pub struct RenderedPrompt {
124    text: String,
125    segments: Option<Vec<RenderedSegment>>,
126    /// Number of trailing segments that form the assistant generation stub.
127    ///
128    /// Some reference APIs (Moonshot's Kimi K3) render a channel-opening stub
129    /// such as `<|open|>response<|sep|>` after the conversation so the model
130    /// starts inside the right channel, but exclude those tokens from the
131    /// reported `usage.prompt_tokens`. The renderer records how many trailing
132    /// segments belong to that stub so callers can apply the same convention
133    /// without hardcoding a token count. Zero for every other prompt.
134    pending_segments: usize,
135}
136
137impl RenderedPrompt {
138    pub fn text(text: String) -> Self {
139        Self {
140            text,
141            segments: None,
142            pending_segments: 0,
143        }
144    }
145
146    pub fn segmented(segments: Vec<RenderedSegment>) -> Self {
147        Self::segmented_with_pending(segments, 0)
148    }
149
150    /// Build a segmented prompt whose last `pending_segments` segments are the
151    /// assistant generation stub (see [`RenderedPrompt::pending_segments`]).
152    ///
153    /// # Panics
154    ///
155    /// Panics if `pending_segments` exceeds the number of segments; a renderer
156    /// that claims more pending segments than it produced is a bug.
157    pub fn segmented_with_pending(segments: Vec<RenderedSegment>, pending_segments: usize) -> Self {
158        assert!(
159            pending_segments <= segments.len(),
160            "pending_segments ({pending_segments}) exceeds rendered segments ({})",
161            segments.len()
162        );
163        let text = segments
164            .iter()
165            .map(|segment| segment.text.as_str())
166            .collect();
167        Self {
168            text,
169            segments: Some(segments),
170            pending_segments,
171        }
172    }
173
174    pub fn as_str(&self) -> &str {
175        &self.text
176    }
177
178    pub fn segments(&self) -> Option<&[RenderedSegment]> {
179        self.segments.as_deref()
180    }
181
182    pub fn encode_segments(&self) -> Option<Vec<dynamo_tokenizers::EncodeSegment<'_>>> {
183        Some(
184            self.segments()?
185                .iter()
186                .map(RenderedSegment::as_encode_segment)
187                .collect(),
188        )
189    }
190
191    /// Number of trailing segments that form the assistant generation stub.
192    ///
193    /// Zero when the prompt is unsegmented, when no generation prompt was
194    /// added, or when the renderer has no stub convention to report.
195    pub fn pending_segments(&self) -> usize {
196        self.pending_segments
197    }
198
199    /// The trailing generation-stub segments, or `None` for unsegmented prompts.
200    ///
201    /// Callers that report usage under the reference API's convention encode
202    /// these and subtract their token count from the physical prompt length.
203    pub fn pending_encode_segments(&self) -> Option<Vec<dynamo_tokenizers::EncodeSegment<'_>>> {
204        let segments = self.segments()?;
205        let start = segments.len() - self.pending_segments;
206        Some(
207            segments[start..]
208                .iter()
209                .map(RenderedSegment::as_encode_segment)
210                .collect(),
211        )
212    }
213
214    pub fn into_text(self) -> String {
215        self.text
216    }
217}
218
219/// A prompt-rendering failure caused by the request rather than server state.
220///
221/// Callers can downcast an [`anyhow::Error`] to this type and map it to their
222/// protocol's invalid-request status without treating every template failure as
223/// a client error.
224#[derive(Debug, Clone, PartialEq, Eq)]
225pub enum PromptRenderError {
226    InvalidRequest(String),
227}
228
229impl PromptRenderError {
230    pub fn invalid_request(message: impl Into<String>) -> Self {
231        Self::InvalidRequest(message.into())
232    }
233}
234
235impl std::fmt::Display for PromptRenderError {
236    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
237        match self {
238            Self::InvalidRequest(message) => f.write_str(message),
239        }
240    }
241}
242
243impl std::error::Error for PromptRenderError {}
244
245/// Trait that defines a request that can map to an OpenAI-like request.
246///
247/// Implement this for your request type to render it through a
248/// [`PromptFormatter`]. Media/multimodal IO config is intentionally *not* part
249/// of this trait — it is a preprocessing concern owned by the consumer, kept
250/// off the rendering surface so this crate stays runtime-free.
251pub trait OAIChatLikeRequest {
252    fn model(&self) -> String;
253    fn messages(&self) -> Value;
254    /// Optional fast path for rendering without calling [`Self::messages`].
255    ///
256    /// Returning `Some` requires the same serialized messages, including all
257    /// transformations, as `messages()`. Adapters that rewrite messages must
258    /// return `None` unless this slice includes those same changes.
259    fn typed_messages(&self) -> Option<&[dynamo_protocols::types::ChatCompletionRequestMessage]> {
260        None
261    }
262    fn tools(&self) -> Option<Value> {
263        None
264    }
265    fn tool_choice(&self) -> Option<Value> {
266        None
267    }
268    fn response_format(&self) -> Option<Value> {
269        None
270    }
271
272    /// OpenAI-compatible reasoning-effort control, when the request type
273    /// exposes it as a top-level field.
274    fn reasoning_effort(&self) -> Option<Value> {
275        None
276    }
277
278    fn should_add_generation_prompt(&self) -> bool;
279
280    /// Optional additional args to merge into the chat template context
281    fn chat_template_args(&self) -> Option<&HashMap<String, serde_json::Value>> {
282        None
283    }
284
285    /// Returns the type of input for the prompt. Default is Text.
286    fn prompt_input_type(&self) -> PromptInput {
287        PromptInput::Text(TextInput::Single(String::new()))
288    }
289
290    /// Extract tokens if the input is pre-tokenized
291    fn extract_tokens(&self) -> Option<TokenInput> {
292        None
293    }
294
295    fn extract_text(&self) -> Option<TextInput> {
296        None
297    }
298
299    fn mm_processor_kwargs(&self) -> Option<&serde_json::Value> {
300        None
301    }
302}
303
304/// Native renderers and normalization work on JSON, so avoid a MiniJinja
305/// round trip when the request exposes its typed messages.
306pub(crate) fn messages_to_json(req: &dyn OAIChatLikeRequest) -> Result<serde_json::Value> {
307    if let Some(messages) = req.typed_messages() {
308        serde_json::to_value(messages)
309    } else {
310        serde_json::to_value(req.messages())
311    }
312    .context("Failed to convert messages to JSON")
313}
314
315pub trait OAIPromptFormatter: Send + Sync + 'static {
316    fn supports_add_generation_prompt(&self) -> bool;
317    fn render(&self, req: &dyn OAIChatLikeRequest) -> Result<String>;
318
319    /// Source message indices in rendered media order; `None` preserves arrival order.
320    /// An explicit order must be a permutation of all source-message indices.
321    /// Parts within each message retain their original order.
322    fn media_message_order(&self, _request: &dyn OAIChatLikeRequest) -> Option<Vec<usize>> {
323        None
324    }
325
326    fn render_prompt(&self, req: &dyn OAIChatLikeRequest) -> Result<RenderedPrompt> {
327        self.render(req).map(RenderedPrompt::text)
328    }
329}
330
331/// Reject Kimi-style Partial Mode in a formatter that cannot leave the final
332/// assistant turn open for continuation.
333///
334/// `partial: false` and `partial: null` are ordinary message metadata and are
335/// intentionally ignored. Supporting formatters (currently Kimi K3) do not
336/// call this helper and implement the open-turn rendering themselves.
337pub(crate) fn reject_unsupported_partial_assistant(messages: &serde_json::Value) -> Result<()> {
338    let has_partial =
339        messages.as_array().into_iter().flatten().any(|message| {
340            message.get("partial").and_then(serde_json::Value::as_bool) == Some(true)
341        });
342    if has_partial {
343        return Err(PromptRenderError::invalid_request(
344            "assistant `partial: true` is not supported by this model's prompt formatter",
345        )
346        .into());
347    }
348    Ok(())
349}
350
351/// Reject non-empty message-level tool declarations on roles where a formatter
352/// does not support that field. `tools: null` and `tools: []` declare nothing.
353pub(crate) fn reject_unsupported_message_tools(
354    messages: &serde_json::Value,
355    supported_tool_roles: &[&str],
356) -> Result<()> {
357    let offending = messages.as_array().into_iter().flatten().find(|message| {
358        let declares_tools = message
359            .get("tools")
360            .is_some_and(|tools| !tools.is_null() && !tools.as_array().is_some_and(Vec::is_empty));
361        let role_is_supported = message
362            .get("role")
363            .and_then(serde_json::Value::as_str)
364            .is_some_and(|role| supported_tool_roles.contains(&role));
365        declares_tools && !role_is_supported
366    });
367
368    if let Some(message) = offending {
369        let role = message
370            .get("role")
371            .and_then(serde_json::Value::as_str)
372            .unwrap_or("<missing>");
373        return Err(PromptRenderError::invalid_request(format!(
374            "message-level `tools` on role {role:?} are not supported by this model's prompt \
375             formatter"
376        ))
377        .into());
378    }
379    Ok(())
380}
381
382#[derive(Clone)]
383pub enum PromptFormatter {
384    OAI(Arc<dyn OAIPromptFormatter>),
385}
386
387// No-op formatter: used for models without chat_template
388#[derive(Debug, Default)]
389pub struct NoOpFormatter;
390
391impl OAIPromptFormatter for NoOpFormatter {
392    fn supports_add_generation_prompt(&self) -> bool {
393        false
394    }
395
396    fn render(&self, req: &dyn OAIChatLikeRequest) -> Result<String> {
397        let messages = req.messages();
398        let messages_json = serde_json::to_value(&messages)?;
399        reject_unsupported_partial_assistant(&messages_json)?;
400        reject_unsupported_message_tools(&messages_json, &[])?;
401
402        let first_message = messages
403            .get_item_by_index(0)
404            .map_err(|_| anyhow::Error::msg("No message at index 0 or messages array is empty"))?;
405
406        let content = first_message
407            .get_attr("content")
408            .map_err(|_| anyhow::Error::msg("First message has no 'content' field"))?;
409
410        let content_str = content
411            .as_str()
412            .ok_or_else(|| anyhow::Error::msg("Message content is not a string"))?
413            .to_string();
414        Ok(content_str)
415    }
416}
417
418impl PromptFormatter {
419    pub fn no_op() -> Self {
420        Self::OAI(Arc::new(NoOpFormatter))
421    }
422}
423
424#[cfg(test)]
425mod rendered_prompt_tests {
426    use super::{
427        NoOpFormatter, OAIPromptFormatter, PromptRenderError, RenderedPrompt, RenderedSegment,
428    };
429
430    #[test]
431    fn owned_segments_borrow_into_tokenizer_segments() {
432        let prompt = RenderedPrompt::segmented(vec![
433            RenderedSegment::new("<|open|>", true),
434            RenderedSegment::new("user text", false),
435        ]);
436
437        let segments = prompt.encode_segments().expect("segmented prompt");
438        assert_eq!(segments[0].text, "<|open|>");
439        assert!(segments[0].allow_special);
440        assert_eq!(segments[1].text, "user text");
441        assert!(!segments[1].allow_special);
442        assert_eq!(prompt.as_str(), "<|open|>user text");
443    }
444
445    #[test]
446    fn no_op_formatter_rejects_unsupported_partial_assistant() {
447        let request: dynamo_protocols::types::CreateChatCompletionRequest =
448            serde_json::from_value(serde_json::json!({
449                "model": "test",
450                "messages": [
451                    {"role": "user", "content": "Continue"},
452                    {"role": "assistant", "content": "prefix", "partial": true}
453                ]
454            }))
455            .unwrap();
456
457        let error = NoOpFormatter.render(&request).unwrap_err();
458        assert!(matches!(
459            error.downcast_ref::<PromptRenderError>(),
460            Some(PromptRenderError::InvalidRequest(message))
461                if message.contains("`partial: true` is not supported")
462        ));
463    }
464
465    #[test]
466    fn no_op_formatter_rejects_message_level_tools() {
467        let request: dynamo_protocols::types::CreateChatCompletionRequest =
468            serde_json::from_value(serde_json::json!({
469                "model": "test",
470                "messages": [
471                    {"role": "system", "tools": [{"name": "lookup"}]},
472                    {"role": "user", "content": "Continue"}
473                ]
474            }))
475            .unwrap();
476
477        let error = NoOpFormatter.render(&request).unwrap_err();
478        assert!(matches!(
479            error.downcast_ref::<PromptRenderError>(),
480            Some(PromptRenderError::InvalidRequest(message))
481                if message.contains("message-level `tools`")
482        ));
483    }
484
485    #[test]
486    fn rendered_prompt_pending_segments_default_to_zero() {
487        let prompt = RenderedPrompt::segmented(vec![RenderedSegment::new("a", false)]);
488        assert_eq!(prompt.pending_segments(), 0);
489        assert!(prompt.pending_encode_segments().unwrap().is_empty());
490        assert_eq!(RenderedPrompt::text("a".to_string()).pending_segments(), 0);
491        assert!(
492            RenderedPrompt::text("a".to_string())
493                .pending_encode_segments()
494                .is_none()
495        );
496    }
497
498    #[test]
499    fn rendered_prompt_exposes_trailing_pending_segments() {
500        let prompt = RenderedPrompt::segmented_with_pending(
501            vec![
502                RenderedSegment::new("prompt", false),
503                RenderedSegment::new("<|open|>", true),
504                RenderedSegment::new("response", false),
505                RenderedSegment::new("<|sep|>", true),
506            ],
507            3,
508        );
509        assert_eq!(prompt.as_str(), "prompt<|open|>response<|sep|>");
510        assert_eq!(prompt.pending_segments(), 3);
511        let pending: Vec<&str> = prompt
512            .pending_encode_segments()
513            .unwrap()
514            .iter()
515            .map(|segment| segment.text)
516            .collect();
517        assert_eq!(pending, ["<|open|>", "response", "<|sep|>"]);
518    }
519
520    #[test]
521    #[should_panic(expected = "pending_segments")]
522    fn rendered_prompt_rejects_more_pending_than_segments() {
523        RenderedPrompt::segmented_with_pending(vec![RenderedSegment::new("a", false)], 2);
524    }
525}
dynamo_renderer/lib.rs

dynamo_renderer/
lib.rs
