Skip to main content

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 template;
41
42pub use template::{
43    ChatTemplate, ChatTemplateValue, ContextMixins, deepseek_formatter_for, kimi_k3_formatter_for,
44    may_be_fix_tool_schema, native_formatter_for,
45};
46
47/// Selects which context-mixin behaviors a template renders with.
48///
49/// Carried on the model deployment card (`prompt_context`) and consumed by the
50/// chat-template renderer via [`ContextMixins`].
51#[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq, Hash)]
52#[serde(rename_all = "snake_case")]
53pub enum PromptContextMixin {
54    /// Support OAI Chat Messages and Tools
55    OaiChat,
56
57    /// Enables templates with `{{datetime}}` to be rendered with the current date and time.
58    Llama3DateTime,
59}
60
61/// Shared helper: extract a boolean thinking toggle from `chat_template_args`.
62///
63/// Reads the two equivalent keys (`thinking`, `enable_thinking` — vLLM's
64/// canonical kwarg) in order and returns the first bool value found, or `None`
65/// if neither key is present (or neither carries a bool). Used by the V4
66/// formatter's `resolve_thinking_mode` and by reasoning-parser gating in
67/// consumers so both paths agree on the signal interpretation.
68pub fn thinking_bool_from_args(args: Option<&HashMap<String, serde_json::Value>>) -> Option<bool> {
69    let args = args?;
70    for key in ["thinking", "enable_thinking"] {
71        if let Some(v) = args.get(key).and_then(|x| x.as_bool()) {
72            return Some(v);
73        }
74    }
75    None
76}
77
78#[derive(Debug)]
79pub enum TokenInput {
80    Single(Vec<u32>),
81    Batch(Vec<Vec<u32>>),
82}
83
84#[derive(Debug)]
85pub enum TextInput {
86    Single(String),
87    Batch(Vec<String>),
88}
89
90#[derive(Debug)]
91pub enum PromptInput {
92    Tokens(TokenInput),
93    Text(TextInput),
94}
95
96/// One owned prompt segment with an explicit special-token trust boundary.
97#[derive(Debug, Clone, PartialEq, Eq)]
98pub struct RenderedSegment {
99    pub text: String,
100    pub allow_special: bool,
101}
102
103impl RenderedSegment {
104    pub fn new(text: impl Into<String>, allow_special: bool) -> Self {
105        Self {
106            text: text.into(),
107            allow_special,
108        }
109    }
110
111    pub fn as_encode_segment(&self) -> dynamo_tokenizers::EncodeSegment<'_> {
112        dynamo_tokenizers::EncodeSegment::new(&self.text, self.allow_special)
113    }
114}
115
116/// A rendered prompt plus its optional tokenization boundaries.
117///
118/// The prompt owns its segment text while `dynamo-tokenizers` borrows that text
119/// during encoding. Keeping the types separate preserves the tokenizer crate's
120/// published zero-copy `EncodeSegment<'_>` API.
121#[derive(Debug, Clone, PartialEq, Eq)]
122pub struct RenderedPrompt {
123    text: String,
124    segments: Option<Vec<RenderedSegment>>,
125}
126
127impl RenderedPrompt {
128    pub fn text(text: String) -> Self {
129        Self {
130            text,
131            segments: None,
132        }
133    }
134
135    pub fn segmented(segments: Vec<RenderedSegment>) -> Self {
136        let text = segments
137            .iter()
138            .map(|segment| segment.text.as_str())
139            .collect();
140        Self {
141            text,
142            segments: Some(segments),
143        }
144    }
145
146    pub fn as_str(&self) -> &str {
147        &self.text
148    }
149
150    pub fn segments(&self) -> Option<&[RenderedSegment]> {
151        self.segments.as_deref()
152    }
153
154    pub fn encode_segments(&self) -> Option<Vec<dynamo_tokenizers::EncodeSegment<'_>>> {
155        Some(
156            self.segments()?
157                .iter()
158                .map(RenderedSegment::as_encode_segment)
159                .collect(),
160        )
161    }
162
163    pub fn into_text(self) -> String {
164        self.text
165    }
166}
167
168/// A prompt-rendering failure caused by the request rather than server state.
169///
170/// Callers can downcast an [`anyhow::Error`] to this type and map it to their
171/// protocol's invalid-request status without treating every template failure as
172/// a client error.
173#[derive(Debug, Clone, PartialEq, Eq)]
174pub enum PromptRenderError {
175    InvalidRequest(String),
176}
177
178impl PromptRenderError {
179    pub fn invalid_request(message: impl Into<String>) -> Self {
180        Self::InvalidRequest(message.into())
181    }
182}
183
184impl std::fmt::Display for PromptRenderError {
185    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
186        match self {
187            Self::InvalidRequest(message) => f.write_str(message),
188        }
189    }
190}
191
192impl std::error::Error for PromptRenderError {}
193
194/// Trait that defines a request that can map to an OpenAI-like request.
195///
196/// Implement this for your request type to render it through a
197/// [`PromptFormatter`]. Media/multimodal IO config is intentionally *not* part
198/// of this trait — it is a preprocessing concern owned by the consumer, kept
199/// off the rendering surface so this crate stays runtime-free.
200pub trait OAIChatLikeRequest {
201    fn model(&self) -> String;
202    fn messages(&self) -> Value;
203    /// Optional fast path for rendering without calling [`Self::messages`].
204    ///
205    /// Returning `Some` requires the same serialized messages, including all
206    /// transformations, as `messages()`. Adapters that rewrite messages must
207    /// return `None` unless this slice includes those same changes.
208    fn typed_messages(&self) -> Option<&[dynamo_protocols::types::ChatCompletionRequestMessage]> {
209        None
210    }
211    fn tools(&self) -> Option<Value> {
212        None
213    }
214    fn tool_choice(&self) -> Option<Value> {
215        None
216    }
217    fn response_format(&self) -> Option<Value> {
218        None
219    }
220
221    /// OpenAI-compatible reasoning-effort control, when the request type
222    /// exposes it as a top-level field.
223    fn reasoning_effort(&self) -> Option<Value> {
224        None
225    }
226
227    fn should_add_generation_prompt(&self) -> bool;
228
229    /// Optional additional args to merge into the chat template context
230    fn chat_template_args(&self) -> Option<&HashMap<String, serde_json::Value>> {
231        None
232    }
233
234    /// Returns the type of input for the prompt. Default is Text.
235    fn prompt_input_type(&self) -> PromptInput {
236        PromptInput::Text(TextInput::Single(String::new()))
237    }
238
239    /// Extract tokens if the input is pre-tokenized
240    fn extract_tokens(&self) -> Option<TokenInput> {
241        None
242    }
243
244    fn extract_text(&self) -> Option<TextInput> {
245        None
246    }
247
248    fn mm_processor_kwargs(&self) -> Option<&serde_json::Value> {
249        None
250    }
251}
252
253/// Native renderers and normalization work on JSON, so avoid a MiniJinja
254/// round trip when the request exposes its typed messages.
255pub(crate) fn messages_to_json(req: &dyn OAIChatLikeRequest) -> Result<serde_json::Value> {
256    if let Some(messages) = req.typed_messages() {
257        serde_json::to_value(messages)
258    } else {
259        serde_json::to_value(req.messages())
260    }
261    .context("Failed to convert messages to JSON")
262}
263
264pub trait OAIPromptFormatter: Send + Sync + 'static {
265    fn supports_add_generation_prompt(&self) -> bool;
266    fn render(&self, req: &dyn OAIChatLikeRequest) -> Result<String>;
267
268    /// Source message indices in rendered media order; `None` preserves arrival order.
269    /// An explicit order must be a permutation of all source-message indices.
270    /// Parts within each message retain their original order.
271    fn media_message_order(&self, _request: &dyn OAIChatLikeRequest) -> Option<Vec<usize>> {
272        None
273    }
274
275    fn render_prompt(&self, req: &dyn OAIChatLikeRequest) -> Result<RenderedPrompt> {
276        self.render(req).map(RenderedPrompt::text)
277    }
278}
279
280/// Reject Kimi-style Partial Mode in a formatter that cannot leave the final
281/// assistant turn open for continuation.
282///
283/// `partial: false` and `partial: null` are ordinary message metadata and are
284/// intentionally ignored. Supporting formatters (currently Kimi K3) do not
285/// call this helper and implement the open-turn rendering themselves.
286pub(crate) fn reject_unsupported_partial_assistant(messages: &serde_json::Value) -> Result<()> {
287    let has_partial =
288        messages.as_array().into_iter().flatten().any(|message| {
289            message.get("partial").and_then(serde_json::Value::as_bool) == Some(true)
290        });
291    if has_partial {
292        return Err(PromptRenderError::invalid_request(
293            "assistant `partial: true` is not supported by this model's prompt formatter",
294        )
295        .into());
296    }
297    Ok(())
298}
299
300/// Reject non-empty message-level tool declarations on roles where a formatter
301/// does not support that field. `tools: null` and `tools: []` declare nothing.
302pub(crate) fn reject_unsupported_message_tools(
303    messages: &serde_json::Value,
304    supported_tool_roles: &[&str],
305) -> Result<()> {
306    let offending = messages.as_array().into_iter().flatten().find(|message| {
307        let declares_tools = message
308            .get("tools")
309            .is_some_and(|tools| !tools.is_null() && !tools.as_array().is_some_and(Vec::is_empty));
310        let role_is_supported = message
311            .get("role")
312            .and_then(serde_json::Value::as_str)
313            .is_some_and(|role| supported_tool_roles.contains(&role));
314        declares_tools && !role_is_supported
315    });
316
317    if let Some(message) = offending {
318        let role = message
319            .get("role")
320            .and_then(serde_json::Value::as_str)
321            .unwrap_or("<missing>");
322        return Err(PromptRenderError::invalid_request(format!(
323            "message-level `tools` on role {role:?} are not supported by this model's prompt \
324             formatter"
325        ))
326        .into());
327    }
328    Ok(())
329}
330
331#[derive(Clone)]
332pub enum PromptFormatter {
333    OAI(Arc<dyn OAIPromptFormatter>),
334}
335
336// No-op formatter: used for models without chat_template
337#[derive(Debug, Default)]
338pub struct NoOpFormatter;
339
340impl OAIPromptFormatter for NoOpFormatter {
341    fn supports_add_generation_prompt(&self) -> bool {
342        false
343    }
344
345    fn render(&self, req: &dyn OAIChatLikeRequest) -> Result<String> {
346        let messages = req.messages();
347        let messages_json = serde_json::to_value(&messages)?;
348        reject_unsupported_partial_assistant(&messages_json)?;
349        reject_unsupported_message_tools(&messages_json, &[])?;
350
351        let first_message = messages
352            .get_item_by_index(0)
353            .map_err(|_| anyhow::Error::msg("No message at index 0 or messages array is empty"))?;
354
355        let content = first_message
356            .get_attr("content")
357            .map_err(|_| anyhow::Error::msg("First message has no 'content' field"))?;
358
359        let content_str = content
360            .as_str()
361            .ok_or_else(|| anyhow::Error::msg("Message content is not a string"))?
362            .to_string();
363        Ok(content_str)
364    }
365}
366
367impl PromptFormatter {
368    pub fn no_op() -> Self {
369        Self::OAI(Arc::new(NoOpFormatter))
370    }
371}
372
373#[cfg(test)]
374mod rendered_prompt_tests {
375    use super::{
376        NoOpFormatter, OAIPromptFormatter, PromptRenderError, RenderedPrompt, RenderedSegment,
377    };
378
379    #[test]
380    fn owned_segments_borrow_into_tokenizer_segments() {
381        let prompt = RenderedPrompt::segmented(vec![
382            RenderedSegment::new("<|open|>", true),
383            RenderedSegment::new("user text", false),
384        ]);
385
386        let segments = prompt.encode_segments().expect("segmented prompt");
387        assert_eq!(segments[0].text, "<|open|>");
388        assert!(segments[0].allow_special);
389        assert_eq!(segments[1].text, "user text");
390        assert!(!segments[1].allow_special);
391        assert_eq!(prompt.as_str(), "<|open|>user text");
392    }
393
394    #[test]
395    fn no_op_formatter_rejects_unsupported_partial_assistant() {
396        let request: dynamo_protocols::types::CreateChatCompletionRequest =
397            serde_json::from_value(serde_json::json!({
398                "model": "test",
399                "messages": [
400                    {"role": "user", "content": "Continue"},
401                    {"role": "assistant", "content": "prefix", "partial": true}
402                ]
403            }))
404            .unwrap();
405
406        let error = NoOpFormatter.render(&request).unwrap_err();
407        assert!(matches!(
408            error.downcast_ref::<PromptRenderError>(),
409            Some(PromptRenderError::InvalidRequest(message))
410                if message.contains("`partial: true` is not supported")
411        ));
412    }
413
414    #[test]
415    fn no_op_formatter_rejects_message_level_tools() {
416        let request: dynamo_protocols::types::CreateChatCompletionRequest =
417            serde_json::from_value(serde_json::json!({
418                "model": "test",
419                "messages": [
420                    {"role": "system", "tools": [{"name": "lookup"}]},
421                    {"role": "user", "content": "Continue"}
422                ]
423            }))
424            .unwrap();
425
426        let error = NoOpFormatter.render(&request).unwrap_err();
427        assert!(matches!(
428            error.downcast_ref::<PromptRenderError>(),
429            Some(PromptRenderError::InvalidRequest(message))
430                if message.contains("message-level `tools`")
431        ));
432    }
433}