Skip to main content

dynamo_renderer/deepseek/
v4.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! DeepSeek V4 native prompt formatting
5//!
6//! Native Rust port of DeepSeek V4's chat encoding (encoding_dsv4.py).
7//!
8//! Reference: DeepSeek-V4-Pro/encoding/encoding_dsv4.py
9
10use anyhow::{Context, Result};
11use serde_json::Value as JsonValue;
12
13use super::common::{
14    NormalizeNonText, REASONING_EFFORT_HIGH, REASONING_EFFORT_MAX, RESPONSE_FORMAT_TEMPLATE,
15    TOOL_CALLS_BLOCK_NAME, TOOLS_TEMPLATE, drop_thinking_messages, encode_arguments_to_dsml,
16    find_last_user_index, merge_tool_messages, normalize_message_contents, render_tools,
17    sort_tool_results_by_call_order, task_token, to_json,
18};
19pub use super::common::{ReasoningEffort, ThinkingMode, tokens};
20
21/// Render a single message at the given index.
22fn 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    // Reasoning effort prefix (only at index 0 in thinking mode). Low is the
40    // reference encoder's no-prefix baseline.
41    if index == 0 && thinking_mode == ThinkingMode::Thinking {
42        match reasoning_effort {
43            Some(ReasoningEffort::High) => prompt.push_str(REASONING_EFFORT_HIGH),
44            Some(ReasoningEffort::Max) => prompt.push_str(REASONING_EFFORT_MAX),
45            None => {}
46        }
47    }
48
49    match role {
50        "system" => {
51            let content = msg.get("content").and_then(|c| c.as_str()).unwrap_or("");
52            prompt.push_str(content);
53            if let Some(tools) = msg.get("tools").and_then(|t| t.as_array()) {
54                prompt.push_str("\n\n");
55                prompt.push_str(&render_tools(TOOLS_TEMPLATE, tools));
56            }
57            if let Some(response_format) = msg.get("response_format") {
58                prompt.push_str("\n\n");
59                prompt.push_str(
60                    &RESPONSE_FORMAT_TEMPLATE.replace("{schema}", &to_json(response_format)),
61                );
62            }
63        }
64
65        "developer" => {
66            let content = msg
67                .get("content")
68                .and_then(|c| c.as_str())
69                .filter(|s| !s.is_empty())
70                .context("Developer role requires content")?;
71
72            let mut content_developer = String::from(tokens::USER_START);
73            content_developer.push_str(content);
74
75            if let Some(tools) = msg.get("tools").and_then(|t| t.as_array()) {
76                content_developer.push_str("\n\n");
77                content_developer.push_str(&render_tools(TOOLS_TEMPLATE, tools));
78            }
79            if let Some(response_format) = msg.get("response_format") {
80                content_developer.push_str("\n\n");
81                content_developer.push_str(
82                    &RESPONSE_FORMAT_TEMPLATE.replace("{schema}", &to_json(response_format)),
83                );
84            }
85            prompt.push_str(&content_developer);
86        }
87
88        "user" => {
89            prompt.push_str(tokens::USER_START);
90            if let Some(blocks) = msg.get("content_blocks").and_then(|b| b.as_array()) {
91                let mut parts: Vec<String> = Vec::with_capacity(blocks.len());
92                for block in blocks {
93                    let block_type = block.get("type").and_then(|v| v.as_str()).unwrap_or("");
94                    match block_type {
95                        "text" => {
96                            let text = block.get("text").and_then(|v| v.as_str()).unwrap_or("");
97                            parts.push(text.to_string());
98                        }
99                        "tool_result" => {
100                            let rendered = render_tool_result_content(
101                                block.get("content").unwrap_or(&JsonValue::Null),
102                            );
103                            parts.push(format!("<tool_result>{}</tool_result>", rendered));
104                        }
105                        other => {
106                            parts.push(format!("[Unsupported {}]", other));
107                        }
108                    }
109                }
110                prompt.push_str(&parts.join("\n\n"));
111            } else {
112                let content = msg.get("content").and_then(|c| c.as_str()).unwrap_or("");
113                prompt.push_str(content);
114            }
115        }
116
117        "latest_reminder" => {
118            let content = msg.get("content").and_then(|c| c.as_str()).unwrap_or("");
119            prompt.push_str(tokens::LATEST_REMINDER);
120            prompt.push_str(content);
121        }
122
123        "tool" => {
124            anyhow::bail!(
125                "deepseek_v4 merges tool messages into user; preprocess with merge_tool_messages()"
126            );
127        }
128
129        "assistant" => {
130            let content = msg.get("content").and_then(|c| c.as_str()).unwrap_or("");
131            let reasoning = msg
132                .get("reasoning_content")
133                .and_then(|c| c.as_str())
134                .unwrap_or("");
135            let wo_eos = msg.get("wo_eos").and_then(|v| v.as_bool()).unwrap_or(false);
136
137            let prev_has_task = index > 0
138                && messages[index - 1]
139                    .get("task")
140                    .map(|v| !v.is_null())
141                    .unwrap_or(false);
142
143            let mut thinking_part = String::new();
144            if thinking_mode == ThinkingMode::Thinking && !prev_has_task {
145                let render_thinking = !drop_thinking || last_user_idx.is_none_or(|u| index > u);
146                if render_thinking {
147                    thinking_part.push_str(reasoning);
148                    thinking_part.push_str(tokens::THINKING_END);
149                }
150            }
151
152            prompt.push_str(&thinking_part);
153            prompt.push_str(content);
154
155            if let Some(tool_calls) = msg.get("tool_calls").and_then(|t| t.as_array())
156                && !tool_calls.is_empty()
157            {
158                prompt.push_str("\n\n");
159                prompt.push_str(&format!(
160                    "<{}{}>\n",
161                    tokens::DSML_TOKEN,
162                    TOOL_CALLS_BLOCK_NAME
163                ));
164
165                let mut invocations = Vec::with_capacity(tool_calls.len());
166                for tc in tool_calls {
167                    // Accept both OpenAI-format (nested `function`) and internal
168                    // `{name, arguments}` shape, matching Python's `tool_calls_from_openai_format`.
169                    let fn_obj = tc.get("function").unwrap_or(tc);
170                    let name = fn_obj
171                        .get("name")
172                        .and_then(|n| n.as_str())
173                        .context("Missing tool call name")?;
174                    let arguments = encode_arguments_to_dsml(fn_obj)?;
175                    invocations.push(format!(
176                        "<{}invoke name=\"{}\">\n{}\n</{}invoke>",
177                        tokens::DSML_TOKEN,
178                        name,
179                        arguments,
180                        tokens::DSML_TOKEN
181                    ));
182                }
183                prompt.push_str(&invocations.join("\n"));
184                prompt.push_str(&format!(
185                    "\n</{}{}>",
186                    tokens::DSML_TOKEN,
187                    TOOL_CALLS_BLOCK_NAME
188                ));
189            }
190
191            if !wo_eos {
192                prompt.push_str(tokens::EOS);
193            }
194        }
195
196        other => anyhow::bail!("Unknown role: {}", other),
197    }
198
199    // Early return if the next message is not assistant/latest_reminder — no transition appended.
200    if index + 1 < messages.len() {
201        let next_role = messages[index + 1].get("role").and_then(|r| r.as_str());
202        if !matches!(next_role, Some("assistant") | Some("latest_reminder")) {
203            return Ok(prompt);
204        }
205    }
206
207    // Transition tokens based on task field and role.
208    let task = msg.get("task").and_then(|v| v.as_str());
209    if let Some(task) = task {
210        let sp = task_token(task).with_context(|| format!("Invalid task: '{}'", task))?;
211        if task != "action" {
212            prompt.push_str(sp);
213        } else {
214            prompt.push_str(tokens::ASSISTANT_START);
215            prompt.push_str(if thinking_mode != ThinkingMode::Thinking {
216                tokens::THINKING_END
217            } else {
218                tokens::THINKING_START
219            });
220            prompt.push_str(sp);
221        }
222    } else if matches!(role, "user" | "developer") {
223        prompt.push_str(tokens::ASSISTANT_START);
224        let seed_thinking = thinking_mode == ThinkingMode::Thinking
225            && (!drop_thinking || last_user_idx.is_none_or(|u| index >= u));
226        prompt.push_str(if seed_thinking {
227            tokens::THINKING_START
228        } else {
229            tokens::THINKING_END
230        });
231    }
232
233    Ok(prompt)
234}
235
236/// Render a tool_result `content` payload (string or content-block list).
237fn render_tool_result_content(content: &JsonValue) -> String {
238    match content {
239        JsonValue::String(s) => s.clone(),
240        JsonValue::Array(items) => {
241            let mut parts: Vec<String> = Vec::with_capacity(items.len());
242            for item in items {
243                let item_type = item.get("type").and_then(|v| v.as_str()).unwrap_or("");
244                if item_type == "text" {
245                    parts.push(
246                        item.get("text")
247                            .and_then(|v| v.as_str())
248                            .unwrap_or("")
249                            .to_string(),
250                    );
251                } else {
252                    parts.push(format!("[Unsupported {}]", item_type));
253                }
254            }
255            parts.join("\n\n")
256        }
257        JsonValue::Null => String::new(),
258        _ => to_json(content),
259    }
260}
261
262/// Encode messages to prompt string with default options.
263///
264/// Equivalent to `encode_messages_with_options(.., drop_thinking=true, reasoning_effort=None)`.
265pub fn encode_messages(
266    messages: &[JsonValue],
267    thinking_mode: ThinkingMode,
268    add_bos_token: bool,
269) -> Result<String> {
270    encode_messages_with_options(messages, thinking_mode, add_bos_token, true, None)
271}
272
273/// Encode messages to prompt string.
274///
275/// # Arguments
276/// * `messages` - Array of messages in OpenAI format
277/// * `thinking_mode` - Chat or Thinking
278/// * `add_bos_token` - Whether to prepend BOS token
279/// * `drop_thinking` - Drop reasoning_content from earlier turns (auto-disabled if tools present)
280/// * `reasoning_effort` - Optional reasoning effort level (High and Max prepend distinct verbatim blocks)
281pub fn encode_messages_with_options(
282    messages: &[JsonValue],
283    thinking_mode: ThinkingMode,
284    add_bos_token: bool,
285    drop_thinking: bool,
286    reasoning_effort: Option<ReasoningEffort>,
287) -> Result<String> {
288    let merged = merge_tool_messages(messages);
289    let mut full = sort_tool_results_by_call_order(merged);
290
291    let mut prompt = String::new();
292    if add_bos_token {
293        prompt.push_str(tokens::BOS);
294    }
295
296    // Auto-disable drop_thinking when any message carries a `tools` field.
297    let has_tools = full.iter().any(|m| {
298        m.get("tools")
299            .map(|v| match v {
300                JsonValue::Array(a) => !a.is_empty(),
301                JsonValue::Null => false,
302                _ => true,
303            })
304            .unwrap_or(false)
305    });
306    let effective_drop_thinking = drop_thinking && !has_tools;
307
308    if thinking_mode == ThinkingMode::Thinking && effective_drop_thinking {
309        full = drop_thinking_messages(full);
310    }
311
312    let last_user_idx = find_last_user_index(&full);
313    for idx in 0..full.len() {
314        let part = render_message(
315            idx,
316            &full,
317            thinking_mode,
318            effective_drop_thinking,
319            reasoning_effort,
320            last_user_idx,
321        )?;
322        prompt.push_str(&part);
323    }
324
325    Ok(prompt)
326}
327
328/// DeepSeek V4 Prompt Formatter
329#[derive(Debug)]
330pub struct DeepSeekV4Formatter {
331    thinking_mode: ThinkingMode,
332}
333
334impl DeepSeekV4Formatter {
335    pub fn new(thinking_mode: ThinkingMode) -> Self {
336        Self { thinking_mode }
337    }
338
339    /// Create formatter with thinking mode enabled (default for DSV4)
340    pub fn new_thinking() -> Self {
341        Self::new(ThinkingMode::Thinking)
342    }
343
344    /// Create formatter with chat mode
345    pub fn new_chat() -> Self {
346        Self::new(ThinkingMode::Chat)
347    }
348
349    fn resolve_reasoning_effort(v: Option<&JsonValue>) -> (bool, Option<ReasoningEffort>) {
350        match v.and_then(JsonValue::as_str) {
351            Some("none") => (true, None),
352            Some("max") => (false, Some(ReasoningEffort::Max)),
353            Some("high") | Some("medium") | Some("xhigh") => (false, Some(ReasoningEffort::High)),
354            Some("low") | Some("minimal") => (false, None),
355            None if v.is_none() => (false, Some(ReasoningEffort::High)),
356            _ => {
357                tracing::warn!(
358                    value = ?v,
359                    "reasoning_effort must be one of \"none\", \"minimal\", \"low\", \"medium\", \"high\", \"xhigh\", \"max\"; ignoring and using API default (high)"
360                );
361                (false, Some(ReasoningEffort::High))
362            }
363        }
364    }
365
366    fn resolve_drop_thinking(
367        args: Option<&std::collections::HashMap<String, serde_json::Value>>,
368    ) -> bool {
369        let Some(args) = args else { return true };
370        let Some(v) = args.get("drop_thinking") else {
371            return true;
372        };
373        if let Some(b) = v.as_bool() {
374            return b;
375        }
376        tracing::warn!(
377            value = ?v,
378            "chat_template_args.drop_thinking must be a bool; ignoring and using default (true)"
379        );
380        true
381    }
382}
383
384impl crate::OAIPromptFormatter for DeepSeekV4Formatter {
385    fn supports_add_generation_prompt(&self) -> bool {
386        true
387    }
388
389    fn render(&self, req: &dyn crate::OAIChatLikeRequest) -> Result<String> {
390        let args = req.chat_template_args();
391        let effort_value = req
392            .reasoning_effort()
393            .map(|value| serde_json::to_value(value).context("serialize reasoning_effort"))
394            .transpose()?
395            .or_else(|| args.and_then(|args| args.get("reasoning_effort").cloned()));
396        let (disable_thinking, reasoning_effort) =
397            Self::resolve_reasoning_effort(effort_value.as_ref());
398        let mut thinking_mode = super::common::resolve_thinking_mode(args, self.thinking_mode);
399        if disable_thinking {
400            thinking_mode = ThinkingMode::Chat;
401        }
402        let drop_thinking = Self::resolve_drop_thinking(args);
403
404        let messages_value = req.messages();
405        let messages_json =
406            serde_json::to_value(&messages_value).context("Failed to convert messages to JSON")?;
407
408        let mut messages_array = messages_json
409            .as_array()
410            .context("Messages is not an array")?
411            .clone();
412
413        normalize_message_contents(&mut messages_array, NormalizeNonText::LeaveUntouched);
414
415        super::common::inject_tools_and_response_format(&mut messages_array, req)?;
416
417        encode_messages_with_options(
418            &messages_array,
419            thinking_mode,
420            true,
421            drop_thinking,
422            reasoning_effort,
423        )
424    }
425}
426
427#[cfg(test)]
428mod tests {
429    use super::*;
430    use serde_json::json;
431
432    #[test]
433    fn test_simple_conversation() {
434        let messages = json!([
435            {"role": "system", "content": "You are a helpful assistant."},
436            {"role": "user", "content": "Hello"},
437            {"role": "assistant", "reasoning_content": "greet", "content": "Hi!"},
438            {"role": "user", "content": "What is 2+2?"}
439        ]);
440        let out =
441            encode_messages(messages.as_array().unwrap(), ThinkingMode::Thinking, true).unwrap();
442        assert!(out.starts_with(tokens::BOS));
443        assert!(out.ends_with(&format!(
444            "{}{}",
445            tokens::ASSISTANT_START,
446            tokens::THINKING_START
447        )));
448        // drop_thinking default true → earlier reasoning stripped
449        assert!(!out.contains("greet"));
450    }
451
452    #[test]
453    fn test_reasoning_effort_prefixes() {
454        let messages = json!([
455            {"role": "system", "content": "hi"},
456            {"role": "user", "content": "hello"}
457        ]);
458
459        let high = encode_messages_with_options(
460            messages.as_array().unwrap(),
461            ThinkingMode::Thinking,
462            true,
463            true,
464            Some(ReasoningEffort::High),
465        )
466        .unwrap();
467        let max = encode_messages_with_options(
468            messages.as_array().unwrap(),
469            ThinkingMode::Thinking,
470            true,
471            true,
472            Some(ReasoningEffort::Max),
473        )
474        .unwrap();
475        let low = encode_messages_with_options(
476            messages.as_array().unwrap(),
477            ThinkingMode::Thinking,
478            true,
479            true,
480            None,
481        )
482        .unwrap();
483
484        assert_eq!(
485            high,
486            concat!(
487                "<|begin▁of▁sentence|>Reasoning Effort: Absolute maximum with no shortcuts permitted.\n",
488                "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n",
489                "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n",
490                "hi<|User|>hello<|Assistant|><think>"
491            )
492        );
493        assert_eq!(
494            max,
495            concat!(
496                "<|begin▁of▁sentence|>Reasoning Effort: Beyond maximum — exhaustive, relentless, and uncompromising.\n",
497                "You MUST reason with the utmost depth and rigor, leaving absolutely nothing to chance: exhaustively decompose the problem into its most fundamental components, trace every causal chain to its root, and resolve the underlying cause rather than any surface symptom.\n",
498                "Do not stop reasoning until you have independently verified the solution from multiple angles and are certain that no assumption remains unchecked and no error remains undiscovered.\n\n",
499                "hi<|User|>hello<|Assistant|><think>"
500            )
501        );
502        assert_eq!(
503            low,
504            "<|begin▁of▁sentence|>hi<|User|>hello<|Assistant|><think>"
505        );
506    }
507
508    #[test]
509    fn test_content_blocks_with_tool_result() {
510        // `merge_tool_messages` turns a `tool` role followed by a plain user text
511        // into a single user turn whose `content_blocks` interleave the tool result
512        // with the text, joined by "\n\n" at render time. Users don't construct
513        // `content_blocks` directly — both the Python reference and this port
514        // overwrite any user-supplied `content_blocks` with a single text block.
515        let messages = json!([
516            {"role": "user", "content": "call tool"},
517            {"role": "assistant", "content": "", "tool_calls": [{
518                "id": "c1", "type": "function",
519                "function": {"name": "f", "arguments": "{}"}
520            }]},
521            {"role": "tool", "tool_call_id": "c1", "content": "RESULT"},
522            {"role": "user", "content": "thanks"}
523        ]);
524        let out = encode_messages(messages.as_array().unwrap(), ThinkingMode::Chat, true).unwrap();
525        assert!(
526            out.contains("<tool_result>RESULT</tool_result>\n\nthanks"),
527            "expected tool_result block followed by 'thanks' in the merged user turn, got:\n{}",
528            out
529        );
530    }
531
532    #[test]
533    fn test_user_task_preserved_when_merged_after_tool_result() {
534        let messages = json!([
535            {"role": "assistant", "content": "", "tool_calls": [{
536                "id": "c1", "type": "function",
537                "function": {"name": "search", "arguments": "{}"}
538            }]},
539            {"role": "tool", "tool_call_id": "c1", "content": "RESULT"},
540            {"role": "user", "content": "Search", "task": "action"},
541            {"role": "assistant", "content": "OK"}
542        ]);
543
544        let out = encode_messages(messages.as_array().unwrap(), ThinkingMode::Chat, true).unwrap();
545        assert!(
546            out.contains(&format!(
547                "{}Search{}{}{}OK",
548                "<tool_result>RESULT</tool_result>\n\n",
549                tokens::ASSISTANT_START,
550                tokens::THINKING_END,
551                tokens::TASK_ACTION
552            )),
553            "expected merged user text to keep the action task transition, got:\n{}",
554            out
555        );
556    }
557
558    #[test]
559    fn test_drop_thinking_auto_disable_when_tools_present() {
560        let messages = json!([
561            {"role": "system", "content": "s", "tools": [{
562                "type": "function",
563                "function": {"name": "f", "description": "", "parameters": {"type": "object", "properties": {}}}
564            }]},
565            {"role": "user", "content": "hi"},
566            {"role": "assistant", "reasoning_content": "PRIOR_REASONING", "content": "reply"},
567            {"role": "user", "content": "again"}
568        ]);
569        let out =
570            encode_messages(messages.as_array().unwrap(), ThinkingMode::Thinking, true).unwrap();
571        // Tools present → drop_thinking auto-disabled → earlier reasoning preserved.
572        assert!(out.contains("PRIOR_REASONING"));
573    }
574
575    // ---- Regression tests for known divergences from the Python reference ----
576
577    /// Bug: `last_user_idx = None` (no user/developer in history) should behave
578    /// like Python's `-1` sentinel — `index >= -1` / `idx >= -1` always true, so
579    /// earlier reasoning is preserved and the assistant's reasoning block is
580    /// rendered. Rust defaulting `None` to `usize::MAX` / `is_some_and` silently
581    /// stripped reasoning instead.
582    ///
583    /// Byte-equivalent to Python reference with the same input:
584    /// `<BOS>sysREASONING_BLOCK</think>hello<EOS>`
585    #[test]
586    fn test_assistant_reasoning_preserved_when_no_user_in_history() {
587        let messages = json!([
588            {"role": "system", "content": "sys"},
589            {"role": "assistant", "content": "hello", "reasoning_content": "REASONING_BLOCK"}
590        ]);
591        let out =
592            encode_messages(messages.as_array().unwrap(), ThinkingMode::Thinking, true).unwrap();
593        assert_eq!(
594            out, "<|begin▁of▁sentence|>sysREASONING_BLOCK</think>hello<|end▁of▁sentence|>",
595            "Output must match Python reference byte-for-byte when no user/developer in history"
596        );
597    }
598
599    /// Bug: `to_json` tracks in-string state via `prev_char != '\\'` which
600    /// mis-handles consecutive backslashes. A value containing `\\` (one literal
601    /// backslash in JSON) makes the helper think the closing `"` is escaped,
602    /// so it stops inserting Python-compatible spaces after subsequent `:`/`,`.
603    ///
604    /// Python `json.dumps({"path": "\\", "count": 5}, ensure_ascii=False)`
605    /// emits `{"path": "\\", "count": 5}` — space after every `:` and `,`.
606    #[test]
607    fn test_to_json_preserves_spacing_past_escaped_backslash() {
608        let v = json!({"path": "\\", "count": 5});
609        let got = to_json(&v);
610        assert_eq!(
611            got, r#"{"path": "\\", "count": 5}"#,
612            "to_json must match Python's json.dumps formatting past an escaped backslash"
613        );
614    }
615
616    #[test]
617    fn test_resolve_drop_thinking_warns_on_malformed_value() {
618        use std::collections::HashMap;
619        // String "false" where a bool is expected → fall back to default (true) and warn.
620        let mut args = HashMap::new();
621        args.insert(
622            "drop_thinking".to_string(),
623            serde_json::Value::String("false".to_string()),
624        );
625        assert!(DeepSeekV4Formatter::resolve_drop_thinking(Some(&args)));
626        // Malformed reasoning_effort falls back to the API default (high).
627        let malformed = serde_json::Value::String("HIGH".to_string());
628        assert_eq!(
629            DeepSeekV4Formatter::resolve_reasoning_effort(Some(&malformed)),
630            (false, Some(ReasoningEffort::High))
631        );
632    }
633
634    #[test]
635    fn test_resolve_thinking_mode_honors_enable_thinking() {
636        use std::collections::HashMap;
637        let mut args = HashMap::new();
638        args.insert(
639            "enable_thinking".to_string(),
640            serde_json::Value::Bool(false),
641        );
642        assert_eq!(
643            super::super::common::resolve_thinking_mode(Some(&args), ThinkingMode::Thinking),
644            ThinkingMode::Chat
645        );
646        args.insert("enable_thinking".to_string(), serde_json::Value::Bool(true));
647        assert_eq!(
648            super::super::common::resolve_thinking_mode(Some(&args), ThinkingMode::Thinking),
649            ThinkingMode::Thinking
650        );
651    }
652
653    struct MockRequest {
654        messages: JsonValue,
655        chat_template_args: Option<std::collections::HashMap<String, JsonValue>>,
656        reasoning_effort: Option<JsonValue>,
657        tools: Option<JsonValue>,
658        tool_choice: Option<JsonValue>,
659        response_format: Option<JsonValue>,
660    }
661
662    impl MockRequest {
663        fn new(messages: JsonValue) -> Self {
664            Self {
665                messages,
666                chat_template_args: None,
667                reasoning_effort: None,
668                tools: None,
669                tool_choice: None,
670                response_format: None,
671            }
672        }
673
674        fn with_chat_template_args(
675            mut self,
676            args: std::collections::HashMap<String, JsonValue>,
677        ) -> Self {
678            self.chat_template_args = Some(args);
679            self
680        }
681
682        fn with_reasoning_effort(mut self, reasoning_effort: JsonValue) -> Self {
683            self.reasoning_effort = Some(reasoning_effort);
684            self
685        }
686
687        fn with_tools(mut self, tools: JsonValue) -> Self {
688            self.tools = Some(tools);
689            self
690        }
691
692        fn with_tool_choice(mut self, tool_choice: JsonValue) -> Self {
693            self.tool_choice = Some(tool_choice);
694            self
695        }
696
697        fn with_response_format(mut self, response_format: JsonValue) -> Self {
698            self.response_format = Some(response_format);
699            self
700        }
701    }
702
703    impl crate::OAIChatLikeRequest for MockRequest {
704        fn model(&self) -> String {
705            "deepseek-v4".to_string()
706        }
707
708        fn messages(&self) -> minijinja::value::Value {
709            minijinja::value::Value::from_serialize(&self.messages)
710        }
711
712        fn should_add_generation_prompt(&self) -> bool {
713            true
714        }
715
716        fn chat_template_args(
717            &self,
718        ) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
719            self.chat_template_args.as_ref()
720        }
721
722        fn reasoning_effort(&self) -> Option<minijinja::value::Value> {
723            self.reasoning_effort
724                .as_ref()
725                .map(minijinja::value::Value::from_serialize)
726        }
727
728        fn tools(&self) -> Option<minijinja::value::Value> {
729            self.tools
730                .as_ref()
731                .map(minijinja::value::Value::from_serialize)
732        }
733
734        fn tool_choice(&self) -> Option<minijinja::value::Value> {
735            self.tool_choice
736                .as_ref()
737                .map(minijinja::value::Value::from_serialize)
738        }
739
740        fn response_format(&self) -> Option<minijinja::value::Value> {
741            self.response_format
742                .as_ref()
743                .map(minijinja::value::Value::from_serialize)
744        }
745    }
746
747    fn weather_tool() -> JsonValue {
748        json!([{
749            "type": "function",
750            "function": {
751                "name": "get_current_weather",
752                "description": "Get the current weather in a given location",
753                "parameters": {
754                    "type": "object",
755                    "properties": {"location": {"type": "string"}},
756                    "required": ["location"]
757                }
758            }
759        }])
760    }
761
762    #[test]
763    fn test_render_tool_choice_none_strips_tools_keeps_response_format() {
764        use crate::OAIPromptFormatter;
765
766        let req = MockRequest::new(json!([
767            {"role": "system", "content": "sys"},
768            {"role": "user", "content": "weather in Boston?"}
769        ]))
770        .with_tools(weather_tool())
771        .with_tool_choice(json!("none"))
772        .with_response_format(json!({"type": "json_object"}));
773
774        let formatter = DeepSeekV4Formatter::new_chat();
775        let out = formatter.render(&req).unwrap();
776
777        assert!(
778            !out.contains("## Tools"),
779            "tool_choice=none must strip the tools block, got: {out}"
780        );
781        assert!(
782            !out.contains("get_current_weather"),
783            "tool schema leaked into prompt despite tool_choice=none: {out}"
784        );
785        assert!(
786            out.contains("## Response Format"),
787            "response_format must survive tool_choice=none: {out}"
788        );
789    }
790
791    #[test]
792    fn test_render_tool_choice_auto_keeps_tools() {
793        use crate::OAIPromptFormatter;
794
795        let req = MockRequest::new(json!([
796            {"role": "system", "content": "sys"},
797            {"role": "user", "content": "weather in Boston?"}
798        ]))
799        .with_tools(weather_tool())
800        .with_tool_choice(json!("auto"));
801
802        let formatter = DeepSeekV4Formatter::new_chat();
803        let out = formatter.render(&req).unwrap();
804
805        assert!(out.contains("## Tools"));
806        assert!(out.contains("get_current_weather"));
807    }
808
809    #[test]
810    fn test_render_absent_tool_choice_keeps_tools() {
811        use crate::OAIPromptFormatter;
812
813        let req = MockRequest::new(json!([
814            {"role": "system", "content": "sys"},
815            {"role": "user", "content": "weather in Boston?"}
816        ]))
817        .with_tools(weather_tool());
818
819        let formatter = DeepSeekV4Formatter::new_chat();
820        let out = formatter.render(&req).unwrap();
821
822        assert!(out.contains("## Tools"));
823        assert!(out.contains("get_current_weather"));
824    }
825
826    #[test]
827    fn test_resolve_reasoning_effort_accepts_full_range() {
828        let effort = |v: &str| {
829            let value = json!(v);
830            DeepSeekV4Formatter::resolve_reasoning_effort(Some(&value))
831        };
832
833        assert_eq!(effort("max"), (false, Some(ReasoningEffort::Max)));
834        assert_eq!(effort("xhigh"), (false, Some(ReasoningEffort::High)));
835        assert_eq!(effort("high"), (false, Some(ReasoningEffort::High)));
836        assert_eq!(effort("minimal"), (false, None));
837        assert_eq!(effort("low"), (false, None));
838        assert_eq!(effort("medium"), (false, Some(ReasoningEffort::High)));
839        assert_eq!(effort("none"), (true, None));
840        assert_eq!(effort("bogus"), (false, Some(ReasoningEffort::High)));
841        assert_eq!(
842            DeepSeekV4Formatter::resolve_reasoning_effort(None),
843            (false, Some(ReasoningEffort::High))
844        );
845    }
846
847    #[test]
848    fn test_render_leaves_null_assistant_tool_content_empty() {
849        use crate::OAIPromptFormatter;
850
851        let req = MockRequest::new(json!([
852            {"role": "user", "content": "call tool"},
853            {"role": "assistant", "content": null, "tool_calls": [{
854                "id": "c1", "type": "function",
855                "function": {"name": "f", "arguments": "{}"}
856            }]}
857        ]));
858
859        let formatter = DeepSeekV4Formatter::new_chat();
860        let out = formatter.render(&req).unwrap();
861
862        assert!(out.contains(&format!(
863            "<{}{}>",
864            tokens::DSML_TOKEN,
865            TOOL_CALLS_BLOCK_NAME
866        )));
867        assert!(!out.contains("null"));
868    }
869
870    #[test]
871    fn test_render_wires_reasoning_effort_from_chat_template_args() {
872        use crate::OAIPromptFormatter;
873        use std::collections::HashMap;
874
875        for (effort, expected) in [
876            ("high", REASONING_EFFORT_HIGH),
877            ("max", REASONING_EFFORT_MAX),
878        ] {
879            let mut args = HashMap::new();
880            args.insert("reasoning_effort".to_string(), json!(effort));
881
882            let req = MockRequest::new(json!([
883                {"role": "system", "content": "sys"},
884                {"role": "user", "content": "hi"}
885            ]))
886            .with_chat_template_args(args);
887
888            let formatter = DeepSeekV4Formatter::new_thinking();
889            let out = formatter.render(&req).unwrap();
890
891            assert!(out.starts_with(tokens::BOS));
892            assert!(
893                out[tokens::BOS.len()..].starts_with(expected),
894                "{effort} preamble should appear after BOS, got:\n{out}"
895            );
896        }
897    }
898
899    #[test]
900    fn test_render_wires_top_level_reasoning_effort_and_none_disables_thinking() {
901        use crate::OAIPromptFormatter;
902
903        let formatter = DeepSeekV4Formatter::new_thinking();
904        for (effort, expected_prefix) in [
905            ("high", "Reasoning Effort: Absolute maximum"),
906            ("max", "Reasoning Effort: Beyond maximum"),
907        ] {
908            let req: dynamo_protocols::types::CreateChatCompletionRequest =
909                serde_json::from_value(json!({
910                    "model": "deepseek-v4",
911                    "messages": [{"role": "user", "content": "hi"}],
912                    "reasoning_effort": effort
913                }))
914                .unwrap();
915            let out = formatter.render(&req).unwrap();
916
917            assert!(
918                out[tokens::BOS.len()..].starts_with(expected_prefix),
919                "top-level {effort} did not select its prefix: {out}"
920            );
921            assert!(out.ends_with(tokens::THINKING_START));
922        }
923
924        let req: dynamo_protocols::types::CreateChatCompletionRequest =
925            serde_json::from_value(json!({
926                "model": "deepseek-v4",
927                "messages": [{"role": "user", "content": "hi"}],
928                "reasoning_effort": "none"
929            }))
930            .unwrap();
931        let out = formatter.render(&req).unwrap();
932
933        assert_eq!(
934            out,
935            "<|begin▁of▁sentence|><|User|>hi<|Assistant|></think>"
936        );
937    }
938
939    #[test]
940    fn test_top_level_reasoning_effort_precedes_template_argument() {
941        use crate::OAIPromptFormatter;
942        use std::collections::HashMap;
943
944        let mut args = HashMap::new();
945        args.insert("reasoning_effort".to_string(), json!("max"));
946        let req = MockRequest::new(json!([{"role": "user", "content": "hi"}]))
947            .with_chat_template_args(args)
948            .with_reasoning_effort(json!("low"));
949
950        let out = DeepSeekV4Formatter::new_thinking().render(&req).unwrap();
951
952        assert_eq!(
953            out,
954            "<|begin▁of▁sentence|><|User|>hi<|Assistant|><think>"
955        );
956    }
957
958    #[test]
959    fn test_render_drop_thinking_override_from_chat_template_args() {
960        use crate::OAIPromptFormatter;
961        use std::collections::HashMap;
962
963        let messages = json!([
964            {"role": "user", "content": "first"},
965            {"role": "assistant", "reasoning_content": "PRIOR", "content": "reply"},
966            {"role": "user", "content": "again"}
967        ]);
968
969        // Default (drop_thinking=true): prior reasoning stripped.
970        let req_default = MockRequest::new(messages.clone());
971        let formatter = DeepSeekV4Formatter::new_thinking();
972        let out_default = formatter.render(&req_default).unwrap();
973        assert!(
974            !out_default.contains("PRIOR"),
975            "default drop_thinking=true should strip prior reasoning, got:\n{}",
976            out_default
977        );
978
979        // drop_thinking=false override: prior reasoning survives.
980        let mut args = HashMap::new();
981        args.insert("drop_thinking".to_string(), json!(false));
982        let req_keep = MockRequest::new(messages).with_chat_template_args(args);
983        let out_keep = formatter.render(&req_keep).unwrap();
984        assert!(
985            out_keep.contains("PRIOR"),
986            "drop_thinking=false override should preserve prior reasoning, got:\n{}",
987            out_keep
988        );
989    }
990
991    // N4: developer-role interactions with drop_thinking.
992    // find_last_user_index returns the index of user OR developer messages; the
993    // drop_thinking reasoning cutoff and the thinking-seed insertion treat
994    // user and developer identically.
995
996    #[test]
997    fn test_developer_only_conversation_renders_developer_content() {
998        let messages = json!([
999            {"role": "system", "content": "sys"},
1000            {"role": "developer", "content": "x"},
1001            {"role": "assistant", "reasoning_content": "R", "content": "ok"}
1002        ]);
1003        let out =
1004            encode_messages(messages.as_array().unwrap(), ThinkingMode::Thinking, true).unwrap();
1005        assert!(
1006            out.contains("x"),
1007            "developer content should appear in output, got:\n{}",
1008            out
1009        );
1010    }
1011
1012    #[test]
1013    fn test_developer_as_last_user_index_controls_reasoning_cutoff() {
1014        // Indices: 0=user, 1=assistant(FIRST), 2=developer(y), 3=assistant(SECOND).
1015        // find_last_user_index = 2 (developer). With drop_thinking=true:
1016        //   - assistant idx=1 < 2  → reasoning_content stripped.
1017        //   - assistant idx=3 >= 2 → reasoning_content preserved.
1018        let messages = json!([
1019            {"role": "user", "content": "a"},
1020            {"role": "assistant", "reasoning_content": "FIRST", "content": "r1"},
1021            {"role": "developer", "content": "y"},
1022            {"role": "assistant", "reasoning_content": "SECOND", "content": "r2"}
1023        ]);
1024        let out =
1025            encode_messages(messages.as_array().unwrap(), ThinkingMode::Thinking, true).unwrap();
1026        assert!(
1027            !out.contains("FIRST"),
1028            "reasoning before last user/developer (idx 1 < 2) should be stripped, got:\n{}",
1029            out
1030        );
1031        assert!(
1032            out.contains("SECOND"),
1033            "reasoning at/after last user/developer (idx 3 > 2) should survive, got:\n{}",
1034            out
1035        );
1036    }
1037}