1use anyhow::{Context, Result};
11use serde_json::Value as JsonValue;
12
13use super::common::{
14 NormalizeNonText, REASONING_EFFORT_MAX, RESPONSE_FORMAT_TEMPLATE, TOOL_CALLS_BLOCK_NAME,
15 TOOLS_TEMPLATE, drop_thinking_messages, encode_arguments_to_dsml, find_last_user_index,
16 merge_tool_messages, normalize_message_contents, render_tools, sort_tool_results_by_call_order,
17 task_token, to_json,
18};
19pub use super::common::{ReasoningEffort, ThinkingMode, tokens};
20
21fn render_message(
23 index: usize,
24 messages: &[JsonValue],
25 thinking_mode: ThinkingMode,
26 drop_thinking: bool,
27 reasoning_effort: Option<ReasoningEffort>,
28 last_user_idx: Option<usize>,
29) -> Result<String> {
30 let msg = &messages[index];
31
32 let role = msg
33 .get("role")
34 .and_then(|r| r.as_str())
35 .context("Missing 'role' field")?;
36
37 let mut prompt = String::new();
38
39 if index == 0
41 && thinking_mode == ThinkingMode::Thinking
42 && reasoning_effort == Some(ReasoningEffort::Max)
43 {
44 prompt.push_str(REASONING_EFFORT_MAX);
45 }
46
47 match role {
48 "system" => {
49 let content = msg.get("content").and_then(|c| c.as_str()).unwrap_or("");
50 prompt.push_str(content);
51 if let Some(tools) = msg.get("tools").and_then(|t| t.as_array()) {
52 prompt.push_str("\n\n");
53 prompt.push_str(&render_tools(TOOLS_TEMPLATE, tools));
54 }
55 if let Some(response_format) = msg.get("response_format") {
56 prompt.push_str("\n\n");
57 prompt.push_str(
58 &RESPONSE_FORMAT_TEMPLATE.replace("{schema}", &to_json(response_format)),
59 );
60 }
61 }
62
63 "developer" => {
64 let content = msg
65 .get("content")
66 .and_then(|c| c.as_str())
67 .filter(|s| !s.is_empty())
68 .context("Developer role requires content")?;
69
70 let mut content_developer = String::from(tokens::USER_START);
71 content_developer.push_str(content);
72
73 if let Some(tools) = msg.get("tools").and_then(|t| t.as_array()) {
74 content_developer.push_str("\n\n");
75 content_developer.push_str(&render_tools(TOOLS_TEMPLATE, tools));
76 }
77 if let Some(response_format) = msg.get("response_format") {
78 content_developer.push_str("\n\n");
79 content_developer.push_str(
80 &RESPONSE_FORMAT_TEMPLATE.replace("{schema}", &to_json(response_format)),
81 );
82 }
83 prompt.push_str(&content_developer);
84 }
85
86 "user" => {
87 prompt.push_str(tokens::USER_START);
88 if let Some(blocks) = msg.get("content_blocks").and_then(|b| b.as_array()) {
89 let mut parts: Vec<String> = Vec::with_capacity(blocks.len());
90 for block in blocks {
91 let block_type = block.get("type").and_then(|v| v.as_str()).unwrap_or("");
92 match block_type {
93 "text" => {
94 let text = block.get("text").and_then(|v| v.as_str()).unwrap_or("");
95 parts.push(text.to_string());
96 }
97 "tool_result" => {
98 let rendered = render_tool_result_content(
99 block.get("content").unwrap_or(&JsonValue::Null),
100 );
101 parts.push(format!("<tool_result>{}</tool_result>", rendered));
102 }
103 other => {
104 parts.push(format!("[Unsupported {}]", other));
105 }
106 }
107 }
108 prompt.push_str(&parts.join("\n\n"));
109 } else {
110 let content = msg.get("content").and_then(|c| c.as_str()).unwrap_or("");
111 prompt.push_str(content);
112 }
113 }
114
115 "latest_reminder" => {
116 let content = msg.get("content").and_then(|c| c.as_str()).unwrap_or("");
117 prompt.push_str(tokens::LATEST_REMINDER);
118 prompt.push_str(content);
119 }
120
121 "tool" => {
122 anyhow::bail!(
123 "deepseek_v4 merges tool messages into user; preprocess with merge_tool_messages()"
124 );
125 }
126
127 "assistant" => {
128 let content = msg.get("content").and_then(|c| c.as_str()).unwrap_or("");
129 let reasoning = msg
130 .get("reasoning_content")
131 .and_then(|c| c.as_str())
132 .unwrap_or("");
133 let wo_eos = msg.get("wo_eos").and_then(|v| v.as_bool()).unwrap_or(false);
134
135 let prev_has_task = index > 0
136 && messages[index - 1]
137 .get("task")
138 .map(|v| !v.is_null())
139 .unwrap_or(false);
140
141 let mut thinking_part = String::new();
142 if thinking_mode == ThinkingMode::Thinking && !prev_has_task {
143 let render_thinking = !drop_thinking || last_user_idx.is_none_or(|u| index > u);
144 if render_thinking {
145 thinking_part.push_str(reasoning);
146 thinking_part.push_str(tokens::THINKING_END);
147 }
148 }
149
150 prompt.push_str(&thinking_part);
151 prompt.push_str(content);
152
153 if let Some(tool_calls) = msg.get("tool_calls").and_then(|t| t.as_array())
154 && !tool_calls.is_empty()
155 {
156 prompt.push_str("\n\n");
157 prompt.push_str(&format!(
158 "<{}{}>\n",
159 tokens::DSML_TOKEN,
160 TOOL_CALLS_BLOCK_NAME
161 ));
162
163 let mut invocations = Vec::with_capacity(tool_calls.len());
164 for tc in tool_calls {
165 let fn_obj = tc.get("function").unwrap_or(tc);
168 let name = fn_obj
169 .get("name")
170 .and_then(|n| n.as_str())
171 .context("Missing tool call name")?;
172 let arguments = encode_arguments_to_dsml(fn_obj)?;
173 invocations.push(format!(
174 "<{}invoke name=\"{}\">\n{}\n</{}invoke>",
175 tokens::DSML_TOKEN,
176 name,
177 arguments,
178 tokens::DSML_TOKEN
179 ));
180 }
181 prompt.push_str(&invocations.join("\n"));
182 prompt.push_str(&format!(
183 "\n</{}{}>",
184 tokens::DSML_TOKEN,
185 TOOL_CALLS_BLOCK_NAME
186 ));
187 }
188
189 if !wo_eos {
190 prompt.push_str(tokens::EOS);
191 }
192 }
193
194 other => anyhow::bail!("Unknown role: {}", other),
195 }
196
197 if index + 1 < messages.len() {
199 let next_role = messages[index + 1].get("role").and_then(|r| r.as_str());
200 if !matches!(next_role, Some("assistant") | Some("latest_reminder")) {
201 return Ok(prompt);
202 }
203 }
204
205 let task = msg.get("task").and_then(|v| v.as_str());
207 if let Some(task) = task {
208 let sp = task_token(task).with_context(|| format!("Invalid task: '{}'", task))?;
209 if task != "action" {
210 prompt.push_str(sp);
211 } else {
212 prompt.push_str(tokens::ASSISTANT_START);
213 prompt.push_str(if thinking_mode != ThinkingMode::Thinking {
214 tokens::THINKING_END
215 } else {
216 tokens::THINKING_START
217 });
218 prompt.push_str(sp);
219 }
220 } else if matches!(role, "user" | "developer") {
221 prompt.push_str(tokens::ASSISTANT_START);
222 let seed_thinking = thinking_mode == ThinkingMode::Thinking
223 && (!drop_thinking || last_user_idx.is_none_or(|u| index >= u));
224 prompt.push_str(if seed_thinking {
225 tokens::THINKING_START
226 } else {
227 tokens::THINKING_END
228 });
229 }
230
231 Ok(prompt)
232}
233
234fn render_tool_result_content(content: &JsonValue) -> String {
236 match content {
237 JsonValue::String(s) => s.clone(),
238 JsonValue::Array(items) => {
239 let mut parts: Vec<String> = Vec::with_capacity(items.len());
240 for item in items {
241 let item_type = item.get("type").and_then(|v| v.as_str()).unwrap_or("");
242 if item_type == "text" {
243 parts.push(
244 item.get("text")
245 .and_then(|v| v.as_str())
246 .unwrap_or("")
247 .to_string(),
248 );
249 } else {
250 parts.push(format!("[Unsupported {}]", item_type));
251 }
252 }
253 parts.join("\n\n")
254 }
255 JsonValue::Null => String::new(),
256 _ => to_json(content),
257 }
258}
259
260pub fn encode_messages(
264 messages: &[JsonValue],
265 thinking_mode: ThinkingMode,
266 add_bos_token: bool,
267) -> Result<String> {
268 encode_messages_with_options(messages, thinking_mode, add_bos_token, true, None)
269}
270
271pub fn encode_messages_with_options(
280 messages: &[JsonValue],
281 thinking_mode: ThinkingMode,
282 add_bos_token: bool,
283 drop_thinking: bool,
284 reasoning_effort: Option<ReasoningEffort>,
285) -> Result<String> {
286 let merged = merge_tool_messages(messages);
287 let mut full = sort_tool_results_by_call_order(merged);
288
289 let mut prompt = String::new();
290 if add_bos_token {
291 prompt.push_str(tokens::BOS);
292 }
293
294 let has_tools = full.iter().any(|m| {
296 m.get("tools")
297 .map(|v| match v {
298 JsonValue::Array(a) => !a.is_empty(),
299 JsonValue::Null => false,
300 _ => true,
301 })
302 .unwrap_or(false)
303 });
304 let effective_drop_thinking = drop_thinking && !has_tools;
305
306 if thinking_mode == ThinkingMode::Thinking && effective_drop_thinking {
307 full = drop_thinking_messages(full);
308 }
309
310 let last_user_idx = find_last_user_index(&full);
311 for idx in 0..full.len() {
312 let part = render_message(
313 idx,
314 &full,
315 thinking_mode,
316 effective_drop_thinking,
317 reasoning_effort,
318 last_user_idx,
319 )?;
320 prompt.push_str(&part);
321 }
322
323 Ok(prompt)
324}
325
326#[derive(Debug)]
328pub struct DeepSeekV4Formatter {
329 thinking_mode: ThinkingMode,
330}
331
332impl DeepSeekV4Formatter {
333 pub fn new(thinking_mode: ThinkingMode) -> Self {
334 Self { thinking_mode }
335 }
336
337 pub fn new_thinking() -> Self {
339 Self::new(ThinkingMode::Thinking)
340 }
341
342 pub fn new_chat() -> Self {
344 Self::new(ThinkingMode::Chat)
345 }
346
347 fn resolve_reasoning_effort(
348 args: Option<&std::collections::HashMap<String, serde_json::Value>>,
349 ) -> Option<ReasoningEffort> {
350 let args = args?;
351 let v = args.get("reasoning_effort")?;
352 match v.as_str() {
353 Some("max") => Some(ReasoningEffort::Max),
354 Some("high") => Some(ReasoningEffort::High),
355 _ => {
356 tracing::warn!(
357 value = ?v,
358 "chat_template_args.reasoning_effort must be a string of \"max\" or \"high\"; ignoring and using default (none)"
359 );
360 None
361 }
362 }
363 }
364
365 fn resolve_drop_thinking(
366 args: Option<&std::collections::HashMap<String, serde_json::Value>>,
367 ) -> bool {
368 let Some(args) = args else { return true };
369 let Some(v) = args.get("drop_thinking") else {
370 return true;
371 };
372 if let Some(b) = v.as_bool() {
373 return b;
374 }
375 tracing::warn!(
376 value = ?v,
377 "chat_template_args.drop_thinking must be a bool; ignoring and using default (true)"
378 );
379 true
380 }
381}
382
383impl crate::OAIPromptFormatter for DeepSeekV4Formatter {
384 fn supports_add_generation_prompt(&self) -> bool {
385 true
386 }
387
388 fn render(&self, req: &dyn crate::OAIChatLikeRequest) -> Result<String> {
389 let args = req.chat_template_args();
390 let thinking_mode = super::common::resolve_thinking_mode(args, self.thinking_mode);
391 let reasoning_effort = Self::resolve_reasoning_effort(args);
392 let drop_thinking = Self::resolve_drop_thinking(args);
393
394 let messages_value = req.messages();
395 let messages_json =
396 serde_json::to_value(&messages_value).context("Failed to convert messages to JSON")?;
397
398 let mut messages_array = messages_json
399 .as_array()
400 .context("Messages is not an array")?
401 .clone();
402
403 normalize_message_contents(&mut messages_array, NormalizeNonText::LeaveUntouched);
404
405 super::common::inject_tools_and_response_format(&mut messages_array, req)?;
406
407 encode_messages_with_options(
408 &messages_array,
409 thinking_mode,
410 true,
411 drop_thinking,
412 reasoning_effort,
413 )
414 }
415}
416
417#[cfg(test)]
418mod tests {
419 use super::*;
420 use serde_json::json;
421
422 #[test]
423 fn test_simple_conversation() {
424 let messages = json!([
425 {"role": "system", "content": "You are a helpful assistant."},
426 {"role": "user", "content": "Hello"},
427 {"role": "assistant", "reasoning_content": "greet", "content": "Hi!"},
428 {"role": "user", "content": "What is 2+2?"}
429 ]);
430 let out =
431 encode_messages(messages.as_array().unwrap(), ThinkingMode::Thinking, true).unwrap();
432 assert!(out.starts_with(tokens::BOS));
433 assert!(out.ends_with(&format!(
434 "{}{}",
435 tokens::ASSISTANT_START,
436 tokens::THINKING_START
437 )));
438 assert!(!out.contains("greet"));
440 }
441
442 #[test]
443 fn test_reasoning_effort_max_prefix() {
444 let messages = json!([
445 {"role": "system", "content": "hi"},
446 {"role": "user", "content": "hello"}
447 ]);
448 let out = encode_messages_with_options(
449 messages.as_array().unwrap(),
450 ThinkingMode::Thinking,
451 true,
452 true,
453 Some(ReasoningEffort::Max),
454 )
455 .unwrap();
456 assert!(out.contains("Reasoning Effort: Absolute maximum"));
457 let after_bos = &out[tokens::BOS.len()..];
459 assert!(after_bos.starts_with("Reasoning Effort:"));
460
461 let out2 = encode_messages_with_options(
463 messages.as_array().unwrap(),
464 ThinkingMode::Thinking,
465 true,
466 true,
467 Some(ReasoningEffort::High),
468 )
469 .unwrap();
470 assert!(!out2.contains("Reasoning Effort: Absolute maximum"));
471 }
472
473 #[test]
474 fn test_content_blocks_with_tool_result() {
475 let messages = json!([
481 {"role": "user", "content": "call tool"},
482 {"role": "assistant", "content": "", "tool_calls": [{
483 "id": "c1", "type": "function",
484 "function": {"name": "f", "arguments": "{}"}
485 }]},
486 {"role": "tool", "tool_call_id": "c1", "content": "RESULT"},
487 {"role": "user", "content": "thanks"}
488 ]);
489 let out = encode_messages(messages.as_array().unwrap(), ThinkingMode::Chat, true).unwrap();
490 assert!(
491 out.contains("<tool_result>RESULT</tool_result>\n\nthanks"),
492 "expected tool_result block followed by 'thanks' in the merged user turn, got:\n{}",
493 out
494 );
495 }
496
497 #[test]
498 fn test_user_task_preserved_when_merged_after_tool_result() {
499 let messages = json!([
500 {"role": "assistant", "content": "", "tool_calls": [{
501 "id": "c1", "type": "function",
502 "function": {"name": "search", "arguments": "{}"}
503 }]},
504 {"role": "tool", "tool_call_id": "c1", "content": "RESULT"},
505 {"role": "user", "content": "Search", "task": "action"},
506 {"role": "assistant", "content": "OK"}
507 ]);
508
509 let out = encode_messages(messages.as_array().unwrap(), ThinkingMode::Chat, true).unwrap();
510 assert!(
511 out.contains(&format!(
512 "{}Search{}{}{}OK",
513 "<tool_result>RESULT</tool_result>\n\n",
514 tokens::ASSISTANT_START,
515 tokens::THINKING_END,
516 tokens::TASK_ACTION
517 )),
518 "expected merged user text to keep the action task transition, got:\n{}",
519 out
520 );
521 }
522
523 #[test]
524 fn test_drop_thinking_auto_disable_when_tools_present() {
525 let messages = json!([
526 {"role": "system", "content": "s", "tools": [{
527 "type": "function",
528 "function": {"name": "f", "description": "", "parameters": {"type": "object", "properties": {}}}
529 }]},
530 {"role": "user", "content": "hi"},
531 {"role": "assistant", "reasoning_content": "PRIOR_REASONING", "content": "reply"},
532 {"role": "user", "content": "again"}
533 ]);
534 let out =
535 encode_messages(messages.as_array().unwrap(), ThinkingMode::Thinking, true).unwrap();
536 assert!(out.contains("PRIOR_REASONING"));
538 }
539
540 #[test]
551 fn test_assistant_reasoning_preserved_when_no_user_in_history() {
552 let messages = json!([
553 {"role": "system", "content": "sys"},
554 {"role": "assistant", "content": "hello", "reasoning_content": "REASONING_BLOCK"}
555 ]);
556 let out =
557 encode_messages(messages.as_array().unwrap(), ThinkingMode::Thinking, true).unwrap();
558 assert_eq!(
559 out, "<|begin▁of▁sentence|>sysREASONING_BLOCK</think>hello<|end▁of▁sentence|>",
560 "Output must match Python reference byte-for-byte when no user/developer in history"
561 );
562 }
563
564 #[test]
572 fn test_to_json_preserves_spacing_past_escaped_backslash() {
573 let v = json!({"path": "\\", "count": 5});
574 let got = to_json(&v);
575 assert_eq!(
576 got, r#"{"path": "\\", "count": 5}"#,
577 "to_json must match Python's json.dumps formatting past an escaped backslash"
578 );
579 }
580
581 #[test]
582 fn test_resolve_drop_thinking_warns_on_malformed_value() {
583 use std::collections::HashMap;
584 let mut args = HashMap::new();
586 args.insert(
587 "drop_thinking".to_string(),
588 serde_json::Value::String("false".to_string()),
589 );
590 assert!(DeepSeekV4Formatter::resolve_drop_thinking(Some(&args)));
591 let mut args2 = HashMap::new();
593 args2.insert(
594 "reasoning_effort".to_string(),
595 serde_json::Value::String("HIGH".to_string()),
596 );
597 assert_eq!(
598 DeepSeekV4Formatter::resolve_reasoning_effort(Some(&args2)),
599 None
600 );
601 }
602
603 #[test]
604 fn test_resolve_thinking_mode_honors_enable_thinking() {
605 use std::collections::HashMap;
606 let mut args = HashMap::new();
607 args.insert(
608 "enable_thinking".to_string(),
609 serde_json::Value::Bool(false),
610 );
611 assert_eq!(
612 super::super::common::resolve_thinking_mode(Some(&args), ThinkingMode::Thinking),
613 ThinkingMode::Chat
614 );
615 args.insert("enable_thinking".to_string(), serde_json::Value::Bool(true));
616 assert_eq!(
617 super::super::common::resolve_thinking_mode(Some(&args), ThinkingMode::Thinking),
618 ThinkingMode::Thinking
619 );
620 }
621
622 struct MockRequest {
623 messages: JsonValue,
624 chat_template_args: Option<std::collections::HashMap<String, JsonValue>>,
625 }
626
627 impl MockRequest {
628 fn new(messages: JsonValue) -> Self {
629 Self {
630 messages,
631 chat_template_args: None,
632 }
633 }
634
635 fn with_chat_template_args(
636 mut self,
637 args: std::collections::HashMap<String, JsonValue>,
638 ) -> Self {
639 self.chat_template_args = Some(args);
640 self
641 }
642 }
643
644 impl crate::OAIChatLikeRequest for MockRequest {
645 fn model(&self) -> String {
646 "deepseek-v4".to_string()
647 }
648
649 fn messages(&self) -> minijinja::value::Value {
650 minijinja::value::Value::from_serialize(&self.messages)
651 }
652
653 fn should_add_generation_prompt(&self) -> bool {
654 true
655 }
656
657 fn chat_template_args(
658 &self,
659 ) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
660 self.chat_template_args.as_ref()
661 }
662 }
663
664 #[test]
665 fn test_render_leaves_null_assistant_tool_content_empty() {
666 use crate::OAIPromptFormatter;
667
668 let req = MockRequest::new(json!([
669 {"role": "user", "content": "call tool"},
670 {"role": "assistant", "content": null, "tool_calls": [{
671 "id": "c1", "type": "function",
672 "function": {"name": "f", "arguments": "{}"}
673 }]}
674 ]));
675
676 let formatter = DeepSeekV4Formatter::new_chat();
677 let out = formatter.render(&req).unwrap();
678
679 assert!(out.contains(&format!(
680 "<{}{}>",
681 tokens::DSML_TOKEN,
682 TOOL_CALLS_BLOCK_NAME
683 )));
684 assert!(!out.contains("null"));
685 }
686
687 #[test]
688 fn test_render_wires_reasoning_effort_max_from_chat_template_args() {
689 use crate::OAIPromptFormatter;
690 use std::collections::HashMap;
691
692 let mut args = HashMap::new();
693 args.insert("reasoning_effort".to_string(), json!("max"));
694
695 let req = MockRequest::new(json!([
696 {"role": "system", "content": "sys"},
697 {"role": "user", "content": "hi"}
698 ]))
699 .with_chat_template_args(args);
700
701 let formatter = DeepSeekV4Formatter::new_thinking();
702 let out = formatter.render(&req).unwrap();
703
704 assert!(out.starts_with(tokens::BOS));
705 let after_bos = &out[tokens::BOS.len()..];
706 assert!(
707 after_bos.starts_with("Reasoning Effort:"),
708 "REASONING_EFFORT_MAX preamble should appear at start (after BOS), got:\n{}",
709 out
710 );
711 }
712
713 #[test]
714 fn test_render_drop_thinking_override_from_chat_template_args() {
715 use crate::OAIPromptFormatter;
716 use std::collections::HashMap;
717
718 let messages = json!([
719 {"role": "user", "content": "first"},
720 {"role": "assistant", "reasoning_content": "PRIOR", "content": "reply"},
721 {"role": "user", "content": "again"}
722 ]);
723
724 let req_default = MockRequest::new(messages.clone());
726 let formatter = DeepSeekV4Formatter::new_thinking();
727 let out_default = formatter.render(&req_default).unwrap();
728 assert!(
729 !out_default.contains("PRIOR"),
730 "default drop_thinking=true should strip prior reasoning, got:\n{}",
731 out_default
732 );
733
734 let mut args = HashMap::new();
736 args.insert("drop_thinking".to_string(), json!(false));
737 let req_keep = MockRequest::new(messages).with_chat_template_args(args);
738 let out_keep = formatter.render(&req_keep).unwrap();
739 assert!(
740 out_keep.contains("PRIOR"),
741 "drop_thinking=false override should preserve prior reasoning, got:\n{}",
742 out_keep
743 );
744 }
745
746 #[test]
752 fn test_developer_only_conversation_renders_developer_content() {
753 let messages = json!([
754 {"role": "system", "content": "sys"},
755 {"role": "developer", "content": "x"},
756 {"role": "assistant", "reasoning_content": "R", "content": "ok"}
757 ]);
758 let out =
759 encode_messages(messages.as_array().unwrap(), ThinkingMode::Thinking, true).unwrap();
760 assert!(
761 out.contains("x"),
762 "developer content should appear in output, got:\n{}",
763 out
764 );
765 }
766
767 #[test]
768 fn test_developer_as_last_user_index_controls_reasoning_cutoff() {
769 let messages = json!([
774 {"role": "user", "content": "a"},
775 {"role": "assistant", "reasoning_content": "FIRST", "content": "r1"},
776 {"role": "developer", "content": "y"},
777 {"role": "assistant", "reasoning_content": "SECOND", "content": "r2"}
778 ]);
779 let out =
780 encode_messages(messages.as_array().unwrap(), ThinkingMode::Thinking, true).unwrap();
781 assert!(
782 !out.contains("FIRST"),
783 "reasoning before last user/developer (idx 1 < 2) should be stripped, got:\n{}",
784 out
785 );
786 assert!(
787 out.contains("SECOND"),
788 "reasoning at/after last user/developer (idx 3 > 2) should survive, got:\n{}",
789 out
790 );
791 }
792}