Skip to main content

deepseek_recipe_encoding/v4/
mod.rs

1use deepseek_recipe_core::conversation::{Conversation, ReasoningEffort, ResponseFormat};
2use deepseek_recipe_core::messages::{InputMessage, ToolCall};
3use deepseek_recipe_core::tools::{ToolChoice, ToolDefinition};
4use deepseek_recipe_core::util::json_formatter::stringify_python_style;
5
6use crate::EncodingError;
7use crate::PromptEncoding;
8use crate::RenderedPrompt;
9use crate::TokenizerEncoder;
10
11pub mod dsv4;
12pub mod dsv41;
13
14/// Marks the start of the prompt.
15pub const BOS_TOKEN: &str = "<|begin▁of▁sentence|>";
16/// Starts the reasoning content of an assistant turn.
17pub const THINKING_START_TOKEN: &str = "<think>";
18/// Ends the reasoning content of an assistant turn.
19pub const THINKING_END_TOKEN: &str = "</think>";
20
21/// Starts a system message in a V4.1 prompt.
22pub const SYSTEM_SP_TOKEN: &str = "<|System|>";
23/// Starts a user message.
24pub const USER_SP_TOKEN: &str = "<|User|>";
25/// Starts an assistant message.
26pub const ASSISTANT_SP_TOKEN: &str = "<|Assistant|>";
27/// Starts the latest reminder message.
28pub const LATEST_REMINDER_SP_TOKEN: &str = "<|latest_reminder|>";
29/// Terminates a message.
30pub const EOS_TOKEN: &str = "<|end▁of▁sentence|>";
31/// Tag name prefix of the markup that structures tool calls. The angle brackets
32/// come from the surrounding template, as in `<|DSML|tool_calls>` for V4 and
33/// `<|DSML| calls>` for V4.1.
34pub const DSML_SP_TOKEN: &str = "|DSML|";
35
36fn parameter_template(
37    dsml_token: &str,
38    tool_parameter_tag_name: &str,
39    key: &str,
40    is_str: &str,
41    value: &str,
42) -> String {
43    format!(
44        "<{dsml_token}{tool_parameter_tag_name} name=\"{key}\" string=\"{is_str}\">{value}</{dsml_token}{tool_parameter_tag_name}>"
45    )
46}
47
48fn render_tool_arguments(tool_call: &ToolCall, tool_parameter_tag_name: &str) -> String {
49    let arguments =
50        serde_json::from_str::<serde_json::Map<String, serde_json::Value>>(&tool_call.arguments)
51            .unwrap_or_else(|err| {
52                tracing::warn!(?err, "invalid tool call arguments");
53                serde_json::Map::from_iter([(
54                    "arguments".to_owned(),
55                    tool_call.arguments.clone().into(),
56                )])
57            });
58    arguments
59        .iter()
60        .map(|(key, value)| {
61            let (is_str, kv_str) = match value.as_str() {
62                Some(s) => ("true", s.to_owned()),
63                None => ("false", stringify_python_style(value)),
64            };
65            parameter_template(DSML_SP_TOKEN, tool_parameter_tag_name, key, is_str, &kv_str)
66        })
67        .collect::<Vec<_>>()
68        .join("\n")
69}
70
71pub(crate) trait EncodingV4 {
72    fn tokenizer(&self) -> Option<&dyn TokenizerEncoder>;
73
74    fn supports_mid_conversation_system(&self) -> bool;
75
76    fn system_token(&self) -> &'static str;
77
78    fn tool_calls_block_name(&self) -> &'static str;
79
80    fn tool_call_tag_name(&self) -> &'static str;
81
82    fn tool_parameter_tag_name(&self) -> &'static str;
83
84    fn render_reasoning_effort(
85        &self,
86        index: usize,
87        thinking_mode: bool,
88        effort: Option<ReasoningEffort>,
89    ) -> String;
90}
91
92fn tool_call_template(encoding: &impl EncodingV4, name: &str, arguments: &str) -> String {
93    let tool_call_tag_name = encoding.tool_call_tag_name();
94    format!(
95        "<{DSML_SP_TOKEN}{tool_call_tag_name} name=\"{name}\">\n{arguments}\n</{DSML_SP_TOKEN}{tool_call_tag_name}>"
96    )
97}
98
99fn tool_calls_template(encoding: &impl EncodingV4, tool_calls: &str) -> String {
100    let tc_block_name = encoding.tool_calls_block_name();
101    format!("<{DSML_SP_TOKEN}{tc_block_name}>\n{tool_calls}\n</{DSML_SP_TOKEN}{tc_block_name}>")
102}
103
104fn render_tool_calls(encoding: &impl EncodingV4, tool_calls: &[ToolCall]) -> String {
105    tool_calls
106        .iter()
107        .map(|tool_call| {
108            tool_call_template(
109                encoding,
110                &tool_call.name,
111                &render_tool_arguments(tool_call, encoding.tool_parameter_tag_name()),
112            )
113        })
114        .collect::<Vec<_>>()
115        .join("\n")
116}
117
118fn render_message(
119    encoding: &impl EncodingV4,
120    messages: &[InputMessage],
121    index: usize,
122    thinking_mode: bool,
123    reasoning_effort: Option<ReasoningEffort>,
124) -> String {
125    let msg = &messages[index];
126    let prev = messages[..index].last();
127    let reasoning_effort_prompt =
128        encoding.render_reasoning_effort(index, thinking_mode, reasoning_effort);
129    let mut prompt = if index == 0
130        && (!reasoning_effort_prompt.is_empty() || matches!(msg, InputMessage::System { .. }))
131    {
132        encoding.system_token().to_string()
133    } else {
134        String::new()
135    };
136    prompt += &reasoning_effort_prompt;
137    match msg {
138        InputMessage::System { content } => {
139            if index > 0 && encoding.supports_mid_conversation_system() {
140                prompt += encoding.system_token();
141            }
142            prompt += content;
143        }
144        InputMessage::User { content, .. } => {
145            if matches!(
146                prev,
147                Some(InputMessage::User { .. } | InputMessage::Tool { .. })
148            ) {
149                prompt += "\n\n";
150            } else {
151                prompt += USER_SP_TOKEN;
152            }
153            prompt += content;
154        }
155        InputMessage::LatestReminder { content } => {
156            prompt += LATEST_REMINDER_SP_TOKEN;
157            prompt += content;
158        }
159        InputMessage::Tool { content, .. } => {
160            if matches!(
161                prev,
162                Some(InputMessage::User { .. } | InputMessage::Tool { .. })
163            ) {
164                prompt += "\n\n";
165            } else {
166                prompt += USER_SP_TOKEN;
167            }
168            prompt += &format!("<tool_result>{content}</tool_result>");
169        }
170        InputMessage::Assistant {
171            content,
172            reasoning_content,
173            tool_calls,
174        } => {
175            let tool_calls_content = match tool_calls {
176                Some(tool_calls) if !tool_calls.is_empty() => {
177                    format!(
178                        "\n\n{}",
179                        tool_calls_template(encoding, &render_tool_calls(encoding, tool_calls))
180                    )
181                }
182                _ => String::new(),
183            };
184            let mut thinking_part = String::new();
185            if thinking_mode && index > 0 {
186                if let Some(reasoning_content) = reasoning_content {
187                    thinking_part += reasoning_content;
188                }
189                thinking_part += THINKING_END_TOKEN;
190            }
191            prompt += ASSISTANT_SP_TOKEN;
192            prompt += if !thinking_part.is_empty() {
193                THINKING_START_TOKEN
194            } else {
195                THINKING_END_TOKEN
196            };
197            prompt += &thinking_part;
198            prompt += content;
199            prompt += &tool_calls_content;
200            prompt += EOS_TOKEN;
201        }
202    }
203    prompt
204}
205
206impl<T: EncodingV4> PromptEncoding for T {
207    fn encode(&self, conversation: &Conversation) -> Result<Vec<u32>, EncodingError> {
208        let tokenizer = self.tokenizer().ok_or(EncodingError::MissingTokenizer)?;
209        let rendered = self.render_conversation(conversation);
210        tokenizer
211            .encode_ids(&rendered.prompt)
212            .map_err(EncodingError::Encode)
213    }
214
215    fn render_conversation(&self, conversation: &Conversation) -> RenderedPrompt {
216        let mut messages = normalize_messages(self, &conversation.messages);
217        let has_tools =
218            conversation.tool_choice != ToolChoice::None && !conversation.tools.is_empty();
219        let format_schema = match &conversation.response_format {
220            ResponseFormat::Text => None,
221            ResponseFormat::JsonObject => Some(stringify_python_style(&serde_json::json!({
222                "type": "json_object"
223            }))),
224        };
225        if has_tools || format_schema.is_some() {
226            if !matches!(messages.first(), Some(InputMessage::System { .. })) {
227                messages.insert(
228                    0,
229                    InputMessage::System {
230                        content: String::new(),
231                    },
232                );
233            }
234            if let Some(InputMessage::System { content }) = messages.first_mut() {
235                if has_tools {
236                    content.push_str("\n\n");
237                    content.push_str(&render_tool_prompt(self, &conversation.tools));
238                }
239                if let Some(schema) = format_schema {
240                    content.push_str("\n\n## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n");
241                    content.push_str(&schema);
242                }
243            }
244        }
245        let mut prompt = BOS_TOKEN.to_string();
246        for index in 0..messages.len() {
247            prompt += &render_message(
248                self,
249                &messages,
250                index,
251                conversation.thinking_mode,
252                conversation.reasoning_effort,
253            );
254        }
255        prompt.push_str(ASSISTANT_SP_TOKEN);
256        prompt.push_str(if conversation.thinking_mode {
257            THINKING_START_TOKEN
258        } else {
259            THINKING_END_TOKEN
260        });
261        if conversation.tool_choice == ToolChoice::Required && !conversation.tools.is_empty() {
262            prompt.push_str(&format!(
263                "\n\n<{DSML_SP_TOKEN}{}>\n",
264                self.tool_calls_block_name()
265            ));
266        }
267        let image_sources = messages
268            .iter()
269            .filter_map(|message| match message {
270                InputMessage::User { image_sources, .. }
271                | InputMessage::Tool { image_sources, .. } => Some(image_sources.as_slice()),
272                _ => None,
273            })
274            .flatten()
275            .cloned()
276            .collect();
277        RenderedPrompt {
278            prompt,
279            image_sources,
280        }
281    }
282}
283
284fn render_tool_prompt(encoding: &impl EncodingV4, tools: &[ToolDefinition]) -> String {
285    let tool_schemas = tools
286        .iter()
287        .map(|tool| {
288            stringify_python_style(&serde_json::json!({
289              "name": tool.name,
290              "description": tool.description.as_deref().unwrap_or_default(),
291              "parameters": tool.parameters,
292            }))
293        })
294        .collect::<Vec<_>>()
295        .join("\n");
296    let dsml_token = DSML_SP_TOKEN;
297    let tc_block_name = encoding.tool_calls_block_name();
298    let tool_call_tag_name = encoding.tool_call_tag_name();
299    let tool_parameter_tag_name = encoding.tool_parameter_tag_name();
300    let thinking_start_token = THINKING_START_TOKEN;
301    let thinking_end_token = THINKING_END_TOKEN;
302    format!(
303        r#"## Tools
304
305You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}{tc_block_name}>" block like the following:
306
307<{dsml_token}{tc_block_name}>
308<{dsml_token}{tool_call_tag_name} name="$TOOL_NAME">
309<{dsml_token}{tool_parameter_tag_name} name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}{tool_parameter_tag_name}>
310...
311</{dsml_token}{tool_call_tag_name}>
312<{dsml_token}{tool_call_tag_name} name="$TOOL_NAME2">
313...
314</{dsml_token}{tool_call_tag_name}>
315</{dsml_token}{tc_block_name}>
316
317String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
318
319If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response.
320
321Otherwise, output directly after {thinking_end_token} with tool calls or final response.
322
323### Available Tool Schemas
324
325{tool_schemas}
326
327You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
328"#
329    )
330}
331
332fn normalize_messages(encoding: &impl EncodingV4, messages: &[InputMessage]) -> Vec<InputMessage> {
333    let mut normalized = Vec::new();
334    let mut has_non_system = false;
335    for message in messages.iter().cloned() {
336        match message {
337            InputMessage::System { content } if !encoding.supports_mid_conversation_system() => {
338                if !has_non_system {
339                    if let Some(InputMessage::System { content: head }) = normalized.last_mut() {
340                        if !head.is_empty() && !content.is_empty() {
341                            head.push_str("\n\n");
342                        }
343                        head.push_str(&content);
344                    } else {
345                        normalized.push(InputMessage::System { content });
346                    }
347                } else if !content.is_empty() {
348                    normalized.push(InputMessage::User {
349                        content,
350                        image_sources: Vec::new(),
351                    });
352                }
353            }
354            InputMessage::User {
355                content,
356                image_sources,
357            } => {
358                has_non_system = true;
359                if let Some(InputMessage::User {
360                    content: previous,
361                    image_sources: previous_image_sources,
362                }) = normalized.last_mut()
363                {
364                    previous.push_str("\n\n");
365                    previous.push_str(&content);
366                    previous_image_sources.extend(image_sources);
367                } else {
368                    normalized.push(InputMessage::User {
369                        content,
370                        image_sources,
371                    });
372                }
373            }
374            message => {
375                has_non_system |= !matches!(message, InputMessage::System { .. });
376                normalized.push(message);
377            }
378        }
379    }
380    sort_tool_results_by_call_order(&mut normalized);
381    normalized
382}
383
384fn sort_tool_results_by_call_order(messages: &mut [InputMessage]) {
385    let mut order: Vec<String> = Vec::new();
386    let mut idx = 0;
387    while idx < messages.len() {
388        match &messages[idx] {
389            InputMessage::Assistant {
390                tool_calls: Some(tool_calls),
391                ..
392            } if !tool_calls.is_empty() => {
393                order = tool_calls.iter().map(|tc| tc.id.clone()).collect();
394                idx += 1;
395            }
396            InputMessage::User { .. } | InputMessage::Tool { .. } => {
397                let start = idx;
398                while idx < messages.len()
399                    && matches!(
400                        messages[idx],
401                        InputMessage::User { .. } | InputMessage::Tool { .. }
402                    )
403                {
404                    idx += 1;
405                }
406                let tool_idxs: Vec<usize> = (start..idx)
407                    .filter(|&i| matches!(messages[i], InputMessage::Tool { .. }))
408                    .collect();
409                if tool_idxs.len() > 1 && !order.is_empty() {
410                    let mut tools: Vec<InputMessage> = tool_idxs
411                        .iter()
412                        .map(|&i| {
413                            std::mem::replace(
414                                &mut messages[i],
415                                InputMessage::LatestReminder {
416                                    content: String::new(),
417                                },
418                            )
419                        })
420                        .collect();
421                    tools.sort_by_key(|m| match m {
422                        InputMessage::Tool { tool_call_id, .. } => {
423                            order.iter().position(|id| id == tool_call_id).unwrap_or(0)
424                        }
425                        _ => 0,
426                    });
427                    for (&i, tool) in tool_idxs.iter().zip(tools) {
428                        messages[i] = tool;
429                    }
430                }
431            }
432            _ => idx += 1,
433        }
434    }
435}