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
257#[derive(Clone)]
258pub enum PromptFormatter {
259 OAI(Arc<dyn OAIPromptFormatter>),
260}
261
262#[derive(Debug, Default)]
264pub struct NoOpFormatter;
265
266impl OAIPromptFormatter for NoOpFormatter {
267 fn supports_add_generation_prompt(&self) -> bool {
268 false
269 }
270
271 fn render(&self, req: &dyn OAIChatLikeRequest) -> Result<String> {
272 let messages = req.messages();
273
274 let first_message = messages
275 .get_item_by_index(0)
276 .map_err(|_| anyhow::Error::msg("No message at index 0 or messages array is empty"))?;
277
278 let content = first_message
279 .get_attr("content")
280 .map_err(|_| anyhow::Error::msg("First message has no 'content' field"))?;
281
282 let content_str = content
283 .as_str()
284 .ok_or_else(|| anyhow::Error::msg("Message content is not a string"))?
285 .to_string();
286 Ok(content_str)
287 }
288}
289
290impl PromptFormatter {
291 pub fn no_op() -> Self {
292 Self::OAI(Arc::new(NoOpFormatter))
293 }
294}
295
296#[cfg(test)]
297mod rendered_prompt_tests {
298 use super::{RenderedPrompt, RenderedSegment};
299
300 #[test]
301 fn owned_segments_borrow_into_tokenizer_segments() {
302 let prompt = RenderedPrompt::segmented(vec![
303 RenderedSegment::new("<|open|>", true),
304 RenderedSegment::new("user text", false),
305 ]);
306
307 let segments = prompt.encode_segments().expect("segmented prompt");
308 assert_eq!(segments[0].text, "<|open|>");
309 assert!(segments[0].allow_special);
310 assert_eq!(segments[1].text, "user text");
311 assert!(!segments[1].allow_special);
312 assert_eq!(prompt.as_str(), "<|open|>user text");
313 }
314}