1use anyhow::Result;
28use minijinja::value::Value;
29use std::collections::HashMap;
30use std::sync::Arc;
31
32pub 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#[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq, Hash)]
52#[serde(rename_all = "snake_case")]
53pub enum PromptContextMixin {
54 OaiChat,
56
57 Llama3DateTime,
59}
60
61pub 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#[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#[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#[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
194pub trait OAIChatLikeRequest {
201 fn model(&self) -> String;
202 fn messages(&self) -> Value;
203 fn typed_messages(&self) -> Option<&[dynamo_protocols::types::ChatCompletionRequestMessage]> {
204 None
205 }
206 fn tools(&self) -> Option<Value> {
207 None
208 }
209 fn tool_choice(&self) -> Option<Value> {
210 None
211 }
212 fn response_format(&self) -> Option<Value> {
213 None
214 }
215
216 fn reasoning_effort(&self) -> Option<Value> {
219 None
220 }
221
222 fn should_add_generation_prompt(&self) -> bool;
223
224 fn chat_template_args(&self) -> Option<&HashMap<String, serde_json::Value>> {
226 None
227 }
228
229 fn prompt_input_type(&self) -> PromptInput {
231 PromptInput::Text(TextInput::Single(String::new()))
232 }
233
234 fn extract_tokens(&self) -> Option<TokenInput> {
236 None
237 }
238
239 fn extract_text(&self) -> Option<TextInput> {
240 None
241 }
242
243 fn mm_processor_kwargs(&self) -> Option<&serde_json::Value> {
244 None
245 }
246}
247
248pub trait OAIPromptFormatter: Send + Sync + 'static {
249 fn supports_add_generation_prompt(&self) -> bool;
250 fn render(&self, req: &dyn OAIChatLikeRequest) -> Result<String>;
251
252 fn render_prompt(&self, req: &dyn OAIChatLikeRequest) -> Result<RenderedPrompt> {
253 self.render(req).map(RenderedPrompt::text)
254 }
255}
256
257pub(crate) fn reject_unsupported_partial_assistant(messages: &serde_json::Value) -> Result<()> {
264 let has_partial =
265 messages.as_array().into_iter().flatten().any(|message| {
266 message.get("partial").and_then(serde_json::Value::as_bool) == Some(true)
267 });
268 if has_partial {
269 return Err(PromptRenderError::invalid_request(
270 "assistant `partial: true` is not supported by this model's prompt formatter",
271 )
272 .into());
273 }
274 Ok(())
275}
276
277pub(crate) fn reject_unsupported_message_tools(
280 messages: &serde_json::Value,
281 supported_tool_roles: &[&str],
282) -> Result<()> {
283 let offending = messages.as_array().into_iter().flatten().find(|message| {
284 let declares_tools = message
285 .get("tools")
286 .is_some_and(|tools| !tools.is_null() && !tools.as_array().is_some_and(Vec::is_empty));
287 let role_is_supported = message
288 .get("role")
289 .and_then(serde_json::Value::as_str)
290 .is_some_and(|role| supported_tool_roles.contains(&role));
291 declares_tools && !role_is_supported
292 });
293
294 if let Some(message) = offending {
295 let role = message
296 .get("role")
297 .and_then(serde_json::Value::as_str)
298 .unwrap_or("<missing>");
299 return Err(PromptRenderError::invalid_request(format!(
300 "message-level `tools` on role {role:?} are not supported by this model's prompt \
301 formatter"
302 ))
303 .into());
304 }
305 Ok(())
306}
307
308#[derive(Clone)]
309pub enum PromptFormatter {
310 OAI(Arc<dyn OAIPromptFormatter>),
311}
312
313#[derive(Debug, Default)]
315pub struct NoOpFormatter;
316
317impl OAIPromptFormatter for NoOpFormatter {
318 fn supports_add_generation_prompt(&self) -> bool {
319 false
320 }
321
322 fn render(&self, req: &dyn OAIChatLikeRequest) -> Result<String> {
323 let messages = req.messages();
324 let messages_json = serde_json::to_value(&messages)?;
325 reject_unsupported_partial_assistant(&messages_json)?;
326 reject_unsupported_message_tools(&messages_json, &[])?;
327
328 let first_message = messages
329 .get_item_by_index(0)
330 .map_err(|_| anyhow::Error::msg("No message at index 0 or messages array is empty"))?;
331
332 let content = first_message
333 .get_attr("content")
334 .map_err(|_| anyhow::Error::msg("First message has no 'content' field"))?;
335
336 let content_str = content
337 .as_str()
338 .ok_or_else(|| anyhow::Error::msg("Message content is not a string"))?
339 .to_string();
340 Ok(content_str)
341 }
342}
343
344impl PromptFormatter {
345 pub fn no_op() -> Self {
346 Self::OAI(Arc::new(NoOpFormatter))
347 }
348}
349
350#[cfg(test)]
351mod rendered_prompt_tests {
352 use super::{
353 NoOpFormatter, OAIPromptFormatter, PromptRenderError, RenderedPrompt, RenderedSegment,
354 };
355
356 #[test]
357 fn owned_segments_borrow_into_tokenizer_segments() {
358 let prompt = RenderedPrompt::segmented(vec![
359 RenderedSegment::new("<|open|>", true),
360 RenderedSegment::new("user text", false),
361 ]);
362
363 let segments = prompt.encode_segments().expect("segmented prompt");
364 assert_eq!(segments[0].text, "<|open|>");
365 assert!(segments[0].allow_special);
366 assert_eq!(segments[1].text, "user text");
367 assert!(!segments[1].allow_special);
368 assert_eq!(prompt.as_str(), "<|open|>user text");
369 }
370
371 #[test]
372 fn no_op_formatter_rejects_unsupported_partial_assistant() {
373 let request: dynamo_protocols::types::CreateChatCompletionRequest =
374 serde_json::from_value(serde_json::json!({
375 "model": "test",
376 "messages": [
377 {"role": "user", "content": "Continue"},
378 {"role": "assistant", "content": "prefix", "partial": true}
379 ]
380 }))
381 .unwrap();
382
383 let error = NoOpFormatter.render(&request).unwrap_err();
384 assert!(matches!(
385 error.downcast_ref::<PromptRenderError>(),
386 Some(PromptRenderError::InvalidRequest(message))
387 if message.contains("`partial: true` is not supported")
388 ));
389 }
390
391 #[test]
392 fn no_op_formatter_rejects_message_level_tools() {
393 let request: dynamo_protocols::types::CreateChatCompletionRequest =
394 serde_json::from_value(serde_json::json!({
395 "model": "test",
396 "messages": [
397 {"role": "system", "tools": [{"name": "lookup"}]},
398 {"role": "user", "content": "Continue"}
399 ]
400 }))
401 .unwrap();
402
403 let error = NoOpFormatter.render(&request).unwrap_err();
404 assert!(matches!(
405 error.downcast_ref::<PromptRenderError>(),
406 Some(PromptRenderError::InvalidRequest(message))
407 if message.contains("message-level `tools`")
408 ));
409 }
410}