Skip to main content

cortiq_engine/
dsv41_encoding.rs

1//! DeepSeek-V4.1 prompt encoding and completion parsing.
2//!
3//! The V4.1 checkpoint intentionally ships without a Jinja chat template.
4//! This module is a Rust port of the official `encoding/encoding.py` at the
5//! pinned model revision.  It keeps the protocol in one place for the CLI,
6//! OpenAI server, and any embedding application: spaced DSML tags, numeric
7//! reasoning effort, mid conversation system turns, tool result merging, and
8//! image ordering all follow the reference implementation.
9
10use serde::Serialize;
11use serde_json::{Map, Value, json};
12use std::collections::HashMap;
13use std::io;
14use thiserror::Error;
15
16pub const BOS_TOKEN: &str = "<|begin▁of▁sentence|>";
17pub const EOS_TOKEN: &str = "<|end▁of▁sentence|>";
18pub const THINKING_START_TOKEN: &str = "<think>";
19pub const THINKING_END_TOKEN: &str = "</think>";
20pub const DSML_TOKEN: &str = "|DSML|";
21pub const USER_SP_TOKEN: &str = "<|User|>";
22pub const ASSISTANT_SP_TOKEN: &str = "<|Assistant|>";
23pub const SYSTEM_SP_TOKEN: &str = "<|System|>";
24pub const LATEST_REMINDER_SP_TOKEN: &str = "<|latest_reminder|>";
25pub const IMAGE_PLACEHOLDER: &str = "<|deepseek_image|>";
26
27const TOOL_CALLS_BLOCK_NAME: &str = " calls";
28const TOOL_CALL_TAG_NAME: &str = " invoke";
29const TOOL_PARAMETER_TAG_NAME: &str = " parameter";
30
31/// Internal quick instruction tasks supported by the official formatter.
32pub const TASK_TOKENS: &[(&str, &str)] = &[
33    ("action", "<|action|>"),
34    ("query", "<|query|>"),
35    ("authority", "<|authority|>"),
36    ("domain", "<|domain|>"),
37    ("title", "<|title|>"),
38    ("read_url", "<|read_url|>"),
39];
40
41#[derive(Debug, Error)]
42pub enum EncodingError {
43    #[error("invalid thinking mode '{0}', expected 'chat' or 'thinking'")]
44    InvalidThinkingMode(String),
45    #[error("invalid reasoning effort '{0}', expected an integer in 1..=100 or low/high/max")]
46    InvalidReasoningEffort(String),
47    #[error("message {index} has unsupported role '{role}'")]
48    UnsupportedRole { index: usize, role: String },
49    #[error("invalid message: {0}")]
50    InvalidMessage(String),
51    #[error("invalid image block: {0}")]
52    InvalidImage(String),
53    #[error("invalid DSML tool call: {0}")]
54    InvalidToolCall(String),
55    #[error("invalid completion: {0}")]
56    InvalidCompletion(String),
57}
58
59#[derive(Clone, Copy, Debug, PartialEq, Eq)]
60pub enum ThinkingMode {
61    Chat,
62    Thinking,
63}
64
65impl ThinkingMode {
66    pub fn parse(value: &str) -> Result<Self, EncodingError> {
67        match value {
68            "chat" => Ok(Self::Chat),
69            "thinking" => Ok(Self::Thinking),
70            other => Err(EncodingError::InvalidThinkingMode(other.to_string())),
71        }
72    }
73
74    pub fn as_str(self) -> &'static str {
75        match self {
76            Self::Chat => "chat",
77            Self::Thinking => "thinking",
78        }
79    }
80}
81
82#[derive(Clone, Copy, Debug, PartialEq, Eq)]
83pub enum ReasoningEffort {
84    Budget(u8),
85    Low,
86    High,
87    Max,
88}
89
90impl Default for ReasoningEffort {
91    fn default() -> Self {
92        Self::High
93    }
94}
95
96impl ReasoningEffort {
97    pub fn budget(self) -> u8 {
98        match self {
99            Self::Budget(v) => v,
100            Self::Low => 50,
101            Self::High => 75,
102            Self::Max => 100,
103        }
104    }
105
106    pub fn parse(value: &str) -> Result<Self, EncodingError> {
107        match value {
108            "low" => Ok(Self::Low),
109            "high" => Ok(Self::High),
110            "max" => Ok(Self::Max),
111            other => other
112                .parse::<u8>()
113                .ok()
114                .filter(|v| (1..=100).contains(v))
115                .map(Self::Budget)
116                .ok_or_else(|| EncodingError::InvalidReasoningEffort(other.to_string())),
117        }
118    }
119
120    pub fn from_json(value: &Value) -> Result<Self, EncodingError> {
121        if let Some(s) = value.as_str() {
122            return match s {
123                "low" => Ok(Self::Low),
124                "high" => Ok(Self::High),
125                "max" => Ok(Self::Max),
126                _ => Err(EncodingError::InvalidReasoningEffort(value.to_string())),
127            };
128        }
129        // bool is deliberately rejected even though serde_json represents it
130        // as neither an integer nor a float.
131        if let Some(v) = value.as_i64() {
132            if (1..=100).contains(&v) {
133                return Ok(Self::Budget(v as u8));
134            }
135        }
136        Err(EncodingError::InvalidReasoningEffort(value.to_string()))
137    }
138}
139
140#[derive(Clone, Debug)]
141pub struct EncodeOptions {
142    pub thinking_mode: ThinkingMode,
143    pub context: Vec<Value>,
144    pub drop_thinking: bool,
145    pub add_default_bos_token: bool,
146    pub reasoning_effort: Option<ReasoningEffort>,
147}
148
149impl Default for EncodeOptions {
150    fn default() -> Self {
151        Self {
152            thinking_mode: ThinkingMode::Chat,
153            context: Vec::new(),
154            drop_thinking: true,
155            add_default_bos_token: true,
156            reasoning_effort: None,
157        }
158    }
159}
160
161#[derive(Clone, Debug, PartialEq)]
162pub struct EncodedPrompt {
163    pub prompt: String,
164    /// Image records in the order their placeholders occur in the current
165    /// messages. Records are intentionally JSON shaped so OpenAI, Anthropic,
166    /// local paths, and data URLs survive the protocol boundary unchanged.
167    pub images: Vec<Value>,
168}
169
170#[derive(Clone, Copy, Debug, Default)]
171struct PythonJsonFormatter;
172
173impl serde_json::ser::Formatter for PythonJsonFormatter {
174    fn begin_array_value<W>(&mut self, writer: &mut W, first: bool) -> io::Result<()>
175    where
176        W: ?Sized + io::Write,
177    {
178        if first {
179            Ok(())
180        } else {
181            writer.write_all(b", ")
182        }
183    }
184
185    fn begin_object_key<W>(&mut self, writer: &mut W, first: bool) -> io::Result<()>
186    where
187        W: ?Sized + io::Write,
188    {
189        if first {
190            Ok(())
191        } else {
192            writer.write_all(b", ")
193        }
194    }
195
196    fn begin_object_value<W>(&mut self, writer: &mut W) -> io::Result<()>
197    where
198        W: ?Sized + io::Write,
199    {
200        writer.write_all(b": ")
201    }
202}
203
204pub fn to_json(value: &Value) -> String {
205    let mut bytes = Vec::new();
206    let mut serializer = serde_json::Serializer::with_formatter(&mut bytes, PythonJsonFormatter);
207    if value.serialize(&mut serializer).is_err() {
208        return "null".to_string();
209    }
210    String::from_utf8(bytes).unwrap_or_else(|_| "null".to_string())
211}
212
213fn object(value: &Value) -> Result<&Map<String, Value>, EncodingError> {
214    value
215        .as_object()
216        .ok_or_else(|| EncodingError::InvalidMessage("expected a JSON object".to_string()))
217}
218
219fn role(value: &Value, index: usize) -> Result<&str, EncodingError> {
220    object(value)?
221        .get("role")
222        .and_then(Value::as_str)
223        .ok_or_else(|| EncodingError::InvalidMessage(format!("message {index} has no string role")))
224}
225
226fn text(value: Option<&Value>) -> String {
227    value.and_then(Value::as_str).unwrap_or("").to_string()
228}
229
230/// Split a qualified tool name and validate explicit namespace metadata.
231fn split_tool_name(
232    name: &str,
233    namespace: Option<&str>,
234) -> Result<(Option<String>, String), EncodingError> {
235    let (mut ns, bare) = if let Some((prefix, suffix)) = name.split_once("::") {
236        if suffix.contains("::") {
237            return Err(EncodingError::InvalidToolCall(format!(
238                "tool name contains multiple '::': {name}"
239            )));
240        }
241        if let Some(given) = namespace {
242            if given != prefix {
243                return Err(EncodingError::InvalidToolCall(format!(
244                    "conflicting tool namespaces: {given} != {prefix}"
245                )));
246            }
247        }
248        (Some(prefix.to_string()), suffix.to_string())
249    } else {
250        (namespace.map(str::to_string), name.to_string())
251    };
252    if bare.contains("::") || ns.as_deref().is_some_and(|v| v.contains("::")) {
253        return Err(EncodingError::InvalidToolCall(format!(
254            "invalid qualified tool name: {name}"
255        )));
256    }
257    if bare.is_empty() {
258        return Err(EncodingError::InvalidToolCall(
259            "tool name is empty".to_string(),
260        ));
261    }
262    Ok((ns.take(), bare))
263}
264
265fn tool_name_for_encoding(tool: &Value) -> Result<String, EncodingError> {
266    let map = object(tool)?;
267    let name = map
268        .get("name")
269        .and_then(Value::as_str)
270        .ok_or_else(|| EncodingError::InvalidToolCall("tool definition has no name".to_string()))?;
271    let namespace = match map.get("namespace") {
272        Some(Value::String(v)) => Some(v.as_str()),
273        Some(Value::Object(v)) => v.get("name").and_then(Value::as_str),
274        _ => None,
275    };
276    let (ns, bare) = split_tool_name(name, namespace)?;
277    Ok(ns.map(|v| format!("{v}::{bare}")).unwrap_or(bare))
278}
279
280fn tools_from_openai_format(tools: &Value) -> Result<Vec<Value>, EncodingError> {
281    let list = tools
282        .as_array()
283        .ok_or_else(|| EncodingError::InvalidToolCall("tools must be an array".to_string()))?;
284    let mut out = Vec::with_capacity(list.len());
285    for tool in list {
286        let map = object(tool)?;
287        let function = map
288            .get("function")
289            .and_then(Value::as_object)
290            .ok_or_else(|| {
291                EncodingError::InvalidToolCall("OpenAI tool has no function object".to_string())
292            })?;
293        let mut f = function.clone();
294        if let Some(ns) = map.get("namespace") {
295            f.insert("namespace".to_string(), ns.clone());
296        }
297        let encoded = tool_name_for_encoding(&Value::Object(f.clone()))?;
298        f.insert("name".to_string(), Value::String(encoded));
299        if let Some(Value::Object(ns)) = f.get("namespace") {
300            if let Some(description) = ns.get("description").and_then(Value::as_str) {
301                let old = f.get("description").and_then(Value::as_str).unwrap_or("");
302                f.insert(
303                    "description".to_string(),
304                    Value::String(format!("{description}\n{old}")),
305                );
306            }
307        }
308        // The internal schema carries namespace only while constructing the
309        // qualified name.  The official formatter removes it from the
310        // emitted function object regardless of whether the value was a
311        // string or an object.
312        f.remove("namespace");
313        out.push(Value::Object(f));
314    }
315    Ok(out)
316}
317
318fn tool_calls_from_openai_format(value: &Value) -> Result<Vec<Value>, EncodingError> {
319    let calls = value
320        .as_array()
321        .ok_or_else(|| EncodingError::InvalidToolCall("tool_calls must be an array".to_string()))?;
322    let mut out = Vec::with_capacity(calls.len());
323    for call in calls {
324        let map = object(call)?;
325        let function = map
326            .get("function")
327            .and_then(Value::as_object)
328            .ok_or_else(|| {
329                EncodingError::InvalidToolCall("tool call has no function object".to_string())
330            })?;
331        let name = function
332            .get("name")
333            .and_then(Value::as_str)
334            .ok_or_else(|| {
335                EncodingError::InvalidToolCall("tool call has no function name".to_string())
336            })?;
337        let namespace = map
338            .get("namespace")
339            .and_then(Value::as_str)
340            .or_else(|| function.get("namespace").and_then(Value::as_str));
341        let (ns, bare) = split_tool_name(name, namespace)?;
342        let mut result = Map::new();
343        result.insert("name".to_string(), Value::String(bare));
344        result.insert(
345            "arguments".to_string(),
346            function
347                .get("arguments")
348                .cloned()
349                .unwrap_or(Value::String(String::new())),
350        );
351        if let Some(ns) = ns {
352            result.insert("namespace".to_string(), Value::String(ns));
353        }
354        out.push(Value::Object(result));
355    }
356    Ok(out)
357}
358
359fn encode_arguments_to_dsml(tool_call: &Value) -> Result<String, EncodingError> {
360    let map = object(tool_call)?;
361    let mut arguments = map.get("arguments").cloned().unwrap_or(Value::Null);
362    if !arguments.is_object() {
363        for _ in 0..2 {
364            if let Some(s) = arguments.as_str() {
365                if let Ok(v) = serde_json::from_str::<Value>(s) {
366                    arguments = v;
367                } else {
368                    break;
369                }
370            } else {
371                break;
372            }
373        }
374    }
375    let mut fields = Vec::new();
376    if let Some(args) = arguments.as_object() {
377        for (key, value) in args {
378            let is_string = value.is_string();
379            let encoded = value
380                .as_str()
381                .map(str::to_string)
382                .unwrap_or_else(|| to_json(value));
383            fields.push(format!(
384                "<{DSML_TOKEN}{TOOL_PARAMETER_TAG_NAME} name=\"{key}\" string=\"{}\">{encoded}</{DSML_TOKEN}{TOOL_PARAMETER_TAG_NAME}>",
385                if is_string { "true" } else { "false" }
386            ));
387        }
388    } else {
389        fields.push(format!(
390            "<{DSML_TOKEN}{TOOL_PARAMETER_TAG_NAME} name=\"arguments\" string=\"false\">{}</{DSML_TOKEN}{TOOL_PARAMETER_TAG_NAME}>",
391            to_json(&arguments)
392        ));
393    }
394    Ok(fields.join("\n"))
395}
396
397fn render_tools(tools: &Value) -> Result<String, EncodingError> {
398    let functions = tools_from_openai_format(tools)?;
399    let schemas = functions.iter().map(to_json).collect::<Vec<_>>().join("\n");
400    Ok(format!(
401        "## Tools\n\nYou have access to a set of tools to help answer the user's question. You can invoke tools by writing a \"<{DSML_TOKEN}{TOOL_CALLS_BLOCK_NAME}>\" block like the following:\n\n<{DSML_TOKEN}{TOOL_CALLS_BLOCK_NAME}>\n<{DSML_TOKEN}{TOOL_CALL_TAG_NAME} name=\"$TOOL_NAME\">\n<{DSML_TOKEN}{TOOL_PARAMETER_TAG_NAME} name=\"$PARAMETER_NAME\" string=\"true|false\">$PARAMETER_VALUE</{DSML_TOKEN}{TOOL_PARAMETER_TAG_NAME}>\n...\n</{DSML_TOKEN}{TOOL_CALL_TAG_NAME}>\n<{DSML_TOKEN}{TOOL_CALL_TAG_NAME} name=\"$TOOL_NAME2\">\n...\n</{DSML_TOKEN}{TOOL_CALL_TAG_NAME}>\n</{DSML_TOKEN}{TOOL_CALLS_BLOCK_NAME}>\n\nString 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\"`.\n\nIf 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.\n\nOtherwise, output directly after {THINKING_END_TOKEN} with tool calls or final response.\n\n### Available Tool Schemas\n\n{schemas}\n\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n"
402    ))
403}
404
405fn decode_dsml_to_arguments(
406    tool_name: &str,
407    args: &[(String, String, String)],
408) -> Result<Value, EncodingError> {
409    let mut fields = Vec::with_capacity(args.len());
410    for (key, value, string) in args {
411        let encoded = if string == "true" {
412            to_json(&Value::String(value.clone()))
413        } else {
414            value.clone()
415        };
416        fields.push(format!(
417            "{}: {}",
418            to_json(&Value::String(key.clone())),
419            encoded
420        ));
421    }
422    let arguments = format!("{{{}}}", fields.join(", "));
423    let (namespace, name) = split_tool_name(tool_name, None)?;
424    let mut out = Map::new();
425    out.insert("name".to_string(), Value::String(name));
426    out.insert("arguments".to_string(), Value::String(arguments));
427    if let Some(namespace) = namespace {
428        out.insert("namespace".to_string(), Value::String(namespace));
429    }
430    Ok(Value::Object(out))
431}
432
433/// Convert compact `<image>path</image>` input to OpenAI style content blocks.
434pub fn parse_tagged_text(input: &str) -> Result<Value, EncodingError> {
435    let mut blocks = Vec::new();
436    let mut cursor = 0usize;
437    let mut found = false;
438    while let Some(rel) = input[cursor..].find("<image>") {
439        let start = cursor + rel;
440        if input[cursor..start].contains("</image>") {
441            return Err(EncodingError::InvalidImage(
442                "malformed <image>path</image> tag".to_string(),
443            ));
444        }
445        let end_rel = input[start + 7..]
446            .find("</image>")
447            .ok_or_else(|| EncodingError::InvalidImage("malformed <image> tag".to_string()))?;
448        let end = start + 7 + end_rel;
449        if start > cursor {
450            blocks.push(json!({"type":"text", "text": &input[cursor..start]}));
451        }
452        let path = &input[start + 7..end];
453        if path.is_empty() {
454            return Err(EncodingError::InvalidImage(
455                "image path must not be empty".to_string(),
456            ));
457        }
458        blocks.push(json!({"type":"image_url", "image_url":{"url":path}}));
459        cursor = end + 8;
460        found = true;
461    }
462    if input[cursor..].contains("<image>") || input[cursor..].contains("</image>") {
463        return Err(EncodingError::InvalidImage(
464            "malformed <image>path</image> tag".to_string(),
465        ));
466    }
467    if !found {
468        return Ok(Value::String(input.to_string()));
469    }
470    if cursor < input.len() {
471        blocks.push(json!({"type":"text", "text": &input[cursor..]}));
472    }
473    Ok(Value::Array(blocks))
474}
475
476fn is_image_block(block: &Value) -> bool {
477    matches!(
478        block.get("type").and_then(Value::as_str),
479        Some("image" | "image_url")
480    )
481}
482
483fn extract_image(block: &Value) -> Result<Value, EncodingError> {
484    let map = object(block)?;
485    let mut record = Map::new();
486    record.insert("type".to_string(), Value::String("image".to_string()));
487    if map.get("type").and_then(Value::as_str) == Some("image_url") {
488        match map.get("image_url") {
489            Some(Value::String(url)) => {
490                record.insert("url".to_string(), Value::String(url.clone()));
491            }
492            Some(Value::Object(image_url)) => {
493                record.insert(
494                    "url".to_string(),
495                    image_url
496                        .get("url")
497                        .cloned()
498                        .unwrap_or(Value::String(String::new())),
499                );
500            }
501            _ => {}
502        }
503    } else {
504        for key in ["source", "url", "data"] {
505            if let Some(value) = map.get(key) {
506                record.insert(key.to_string(), value.clone());
507            }
508        }
509    }
510    if !["source", "url", "data"]
511        .iter()
512        .any(|key| record.get(*key).is_some_and(json_truthy))
513    {
514        return Err(EncodingError::InvalidImage(
515            "image block has no source".to_string(),
516        ));
517    }
518    Ok(Value::Object(record))
519}
520
521fn json_truthy(value: &Value) -> bool {
522    match value {
523        Value::Null => false,
524        Value::Bool(value) => *value,
525        Value::Number(value) => value.as_f64().is_some_and(|value| value != 0.0),
526        Value::String(value) => !value.is_empty(),
527        Value::Array(value) => !value.is_empty(),
528        Value::Object(value) => !value.is_empty(),
529    }
530}
531
532fn process_image_blocks(blocks: &[Value]) -> Result<(Vec<Value>, Vec<Value>), EncodingError> {
533    let mut out = Vec::with_capacity(blocks.len());
534    let mut images = Vec::new();
535    for block in blocks {
536        if !block.is_object() {
537            out.push(block.clone());
538            continue;
539        }
540        if is_image_block(block) {
541            out.push(json!({"type":"text", "text":IMAGE_PLACEHOLDER}));
542            images.push(extract_image(block)?);
543            continue;
544        }
545        if block.get("type").and_then(Value::as_str) == Some("tool_result")
546            && block.get("content").and_then(Value::as_array).is_some()
547        {
548            let (nested, nested_images) =
549                process_image_blocks(block.get("content").and_then(Value::as_array).unwrap())?;
550            let mut copy = object(block)?.clone();
551            copy.insert("content".to_string(), Value::Array(nested));
552            out.push(Value::Object(copy));
553            images.extend(nested_images);
554            continue;
555        }
556        if block.get("type").and_then(Value::as_str) == Some("text") {
557            let value = text(block.get("text"));
558            if value.contains(IMAGE_PLACEHOLDER) {
559                return Err(EncodingError::InvalidImage(
560                    "text blocks must use image content blocks for the image placeholder"
561                        .to_string(),
562                ));
563            }
564        }
565        out.push(block.clone());
566    }
567    Ok((out, images))
568}
569
570fn validate_no_image_tokens(message: &Value) -> Result<(), EncodingError> {
571    let map = object(message)?;
572    for key in ["content", "reasoning_content"] {
573        if map
574            .get(key)
575            .and_then(Value::as_str)
576            .is_some_and(|s| s.contains(IMAGE_PLACEHOLDER))
577        {
578            return Err(EncodingError::InvalidImage(format!(
579                "message {key} contains {IMAGE_PLACEHOLDER}; use an image content block"
580            )));
581        }
582    }
583    Ok(())
584}
585
586/// Normalize image blocks and return image records in prompt order.
587pub fn process_image_messages(
588    messages: &[Value],
589) -> Result<(Vec<Value>, Vec<Value>), EncodingError> {
590    let mut processed = Vec::with_capacity(messages.len());
591    let mut images = Vec::new();
592    for original in messages {
593        validate_no_image_tokens(original)?;
594        let mut message = object(original)?.clone();
595        if message.get("content_blocks").is_none() {
596            if let Some(Value::Array(blocks)) = message.get("content") {
597                message.insert("content_blocks".to_string(), Value::Array(blocks.clone()));
598                message.remove("content");
599            }
600        }
601        if let Some(Value::Array(blocks)) = message
602            .get("content_blocks")
603            .filter(|blocks| !blocks.as_array().map_or(true, |items| items.is_empty()))
604        {
605            let (new_blocks, new_images) = process_image_blocks(blocks)?;
606            let strings = new_blocks
607                .iter()
608                .filter_map(|block| {
609                    (block.get("type").and_then(Value::as_str) == Some("text"))
610                        .then(|| text(block.get("text")))
611                })
612                .collect::<Vec<_>>();
613            message.insert("content_blocks".to_string(), Value::Array(new_blocks));
614            if !message.get("content").is_some_and(Value::is_string) {
615                message.insert("content".to_string(), Value::String(strings.join("\n\n")));
616            }
617            images.extend(new_images);
618        }
619        processed.push(Value::Object(message));
620    }
621    Ok((processed, images))
622}
623
624/// Merge standalone OpenAI `tool` messages into a preceding user content list.
625pub fn merge_tool_messages(messages: &[Value]) -> Result<Vec<Value>, EncodingError> {
626    let mut merged = Vec::new();
627    for original in messages {
628        let message = object(original)?.clone();
629        let message_role = message.get("role").and_then(Value::as_str).unwrap_or("");
630        if message_role == "tool" {
631            // A tool result may arrive as structured blocks (text and
632            // images); keep them so the images are collected in prompt
633            // order after the results are sorted.
634            let content = match (message.get("content_blocks"), message.get("content")) {
635                (Some(blocks), content)
636                    if json_truthy(blocks) && !content.is_some_and(Value::is_string) =>
637                {
638                    blocks.clone()
639                }
640                (_, content) => content.cloned().unwrap_or(Value::String(String::new())),
641            };
642            let block = json!({
643                "type":"tool_result",
644                "tool_use_id": message.get("tool_call_id").cloned().unwrap_or(Value::String(String::new())),
645                "content": content,
646            });
647            let can_append = merged.last().is_some_and(|m: &Value| {
648                m.get("role").and_then(Value::as_str) == Some("user")
649                    && m.get("content_blocks").is_some()
650            });
651            if can_append {
652                merged.last_mut().unwrap()["content_blocks"]
653                    .as_array_mut()
654                    .unwrap()
655                    .push(block);
656            } else {
657                merged.push(json!({"role":"user", "content_blocks":[block]}));
658            }
659        } else if message_role == "user" {
660            let blocks = message
661                .get("content_blocks")
662                .filter(|value| !value.is_null())
663                .cloned()
664                .unwrap_or_else(|| match message.get("content") {
665                    // OpenAI-style structured content is already a block list.
666                    Some(Value::Array(blocks)) => Value::Array(blocks.clone()),
667                    content => json!([{
668                        "type":"text",
669                        "text": content.cloned().unwrap_or(Value::String(String::new()))
670                    }]),
671                });
672            let can_append = merged.last().is_some_and(|m: &Value| {
673                m.get("role").and_then(Value::as_str) == Some("user")
674                    && m.get("content_blocks").is_some()
675                    && m.get("task").is_none()
676            });
677            if can_append {
678                let dst = merged.last_mut().unwrap()["content_blocks"]
679                    .as_array_mut()
680                    .unwrap();
681                if let Some(src) = blocks.as_array() {
682                    dst.extend(src.iter().cloned());
683                }
684            } else {
685                let mut copy = message;
686                copy.insert("content_blocks".to_string(), blocks);
687                merged.push(Value::Object(copy));
688            }
689        } else {
690            merged.push(Value::Object(message));
691        }
692    }
693    Ok(merged)
694}
695
696/// Stable tool result order follows the preceding assistant's tool call order.
697pub fn sort_tool_results_by_call_order(messages: &mut [Value]) {
698    let mut order: HashMap<String, usize> = HashMap::new();
699    for message in messages {
700        match message.get("role").and_then(Value::as_str) {
701            Some("assistant") => {
702                if let Some(calls) = message
703                    .get("tool_calls")
704                    .filter(|value| json_truthy(value))
705                    .and_then(Value::as_array)
706                {
707                    order.clear();
708                    for (index, call) in calls.iter().enumerate() {
709                        let id = call.get("id").and_then(Value::as_str).or_else(|| {
710                            call.get("function")
711                                .and_then(|f| f.get("id"))
712                                .and_then(Value::as_str)
713                        });
714                        if let Some(id) = id {
715                            order.insert(id.to_string(), index);
716                        }
717                    }
718                }
719            }
720            Some("user") => {
721                if let Some(blocks) = message
722                    .get_mut("content_blocks")
723                    .and_then(Value::as_array_mut)
724                {
725                    let mut tools = blocks
726                        .iter()
727                        .filter(|b| b.get("type").and_then(Value::as_str) == Some("tool_result"))
728                        .cloned()
729                        .collect::<Vec<_>>();
730                    if tools.len() > 1 && !order.is_empty() {
731                        tools.sort_by_key(|b| {
732                            order
733                                .get(b.get("tool_use_id").and_then(Value::as_str).unwrap_or(""))
734                                .copied()
735                                .unwrap_or(0)
736                        });
737                        let mut next = 0;
738                        for block in blocks.iter_mut() {
739                            if block.get("type").and_then(Value::as_str) == Some("tool_result") {
740                                *block = tools[next].clone();
741                                next += 1;
742                            }
743                        }
744                    }
745                }
746            }
747            _ => {}
748        }
749    }
750}
751
752/// Index of the message that opened the current turn. Mid-conversation
753/// system messages count as user turns; a user message holding only tool
754/// results does not start a new turn, so the reasoning of the assistant
755/// message that called the tools stays in the prompt.
756fn last_user_index(messages: &[Value]) -> isize {
757    messages
758        .iter()
759        .enumerate()
760        .rev()
761        .find(|(idx, m)| {
762            let role = m.get("role").and_then(Value::as_str);
763            if !(role == Some("user") || (role == Some("system") && *idx > 0)) {
764                return false;
765            }
766            match m.get("content_blocks").and_then(Value::as_array) {
767                Some(blocks) if !blocks.is_empty() => blocks.iter().any(|block| {
768                    matches!(
769                        block.get("type").and_then(Value::as_str),
770                        Some("text" | "image" | "image_url")
771                    )
772                }),
773                _ => true,
774            }
775        })
776        .map(|(i, _)| i as isize)
777        .unwrap_or(-1)
778}
779
780fn render_reasoning_effort(
781    index: usize,
782    mode: ThinkingMode,
783    effort: Option<ReasoningEffort>,
784) -> String {
785    if index == 0 && mode == ThinkingMode::Thinking {
786        let budget = effort.unwrap_or_default().budget();
787        format!(
788            "Reasoning Effort: {budget} (range 1-100, the higher the value, the more thorough the reasoning)\n\n"
789        )
790    } else {
791        String::new()
792    }
793}
794
795fn render_content_blocks(blocks: &[Value]) -> String {
796    blocks
797        .iter()
798        .map(|block| match block.get("type").and_then(Value::as_str) {
799            Some("text") => text(block.get("text")),
800            Some("tool_result") => {
801                let content = block.get("content");
802                let rendered = if let Some(parts) = content.and_then(Value::as_array) {
803                    parts
804                        .iter()
805                        .map(|part| {
806                            if part.get("type").and_then(Value::as_str) == Some("text") {
807                                text(part.get("text"))
808                            } else {
809                                format!("[Unsupported {}]", python_type_label(part))
810                            }
811                        })
812                        .collect::<Vec<_>>()
813                        .join("\n\n")
814                } else {
815                    // Python's formatter inserts a string tool result as-is;
816                    // JSON quoting here would change the prompt protocol.
817                    content
818                        .map(|value| {
819                            value
820                                .as_str()
821                                .map(str::to_string)
822                                .unwrap_or_else(|| to_json(value))
823                        })
824                        .unwrap_or_default()
825                };
826                format!("<tool_result>{rendered}</tool_result>")
827            }
828            _ => format!("[Unsupported {}]", python_type_label(block)),
829        })
830        .collect::<Vec<_>>()
831        .join("\n\n")
832}
833
834/// `f"{block.get('type')}"` for the unsupported-block placeholder.
835fn python_type_label(block: &Value) -> String {
836    match block.get("type") {
837        None | Some(Value::Null) => "None".to_string(),
838        Some(Value::String(kind)) => kind.clone(),
839        Some(other) => to_json(other),
840    }
841}
842
843fn task_token(task: &str) -> Option<&'static str> {
844    TASK_TOKENS
845        .iter()
846        .find(|(name, _)| *name == task)
847        .map(|(_, token)| *token)
848}
849
850/// Render one message exactly as the official Python implementation does.
851pub fn render_message(
852    index: usize,
853    messages: &[Value],
854    options: &EncodeOptions,
855) -> Result<String, EncodingError> {
856    let message = object(messages.get(index).ok_or_else(|| {
857        EncodingError::InvalidMessage(format!("message index {index} out of range"))
858    })?)?;
859    let message_value = Value::Object(message.clone());
860    let message_role = role(&message_value, index)?;
861    let last_user = last_user_index(messages);
862    let effort = render_reasoning_effort(index, options.thinking_mode, options.reasoning_effort);
863    let mut prompt = if index == 0 && (!effort.is_empty() || message_role == "system") {
864        SYSTEM_SP_TOKEN.to_string()
865    } else {
866        String::new()
867    };
868    prompt.push_str(&effort);
869
870    match message_role {
871        "system" => {
872            if index > 0 {
873                prompt.push_str(SYSTEM_SP_TOKEN);
874            }
875            prompt.push_str(&text(message.get("content")));
876            if let Some(tools) = message.get("tools").filter(|value| json_truthy(value)) {
877                prompt.push_str("\n\n");
878                prompt.push_str(&render_tools(tools)?);
879            }
880            if let Some(schema) = message
881                .get("response_format")
882                .filter(|value| json_truthy(value))
883            {
884                prompt.push_str("\n\n## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n");
885                prompt.push_str(&to_json(schema));
886            }
887        }
888        "user" => {
889            prompt.push_str(USER_SP_TOKEN);
890            if let Some(blocks) = message
891                .get("content_blocks")
892                .and_then(Value::as_array)
893                .filter(|blocks| !blocks.is_empty())
894            {
895                prompt.push_str(&render_content_blocks(blocks));
896            } else {
897                prompt.push_str(&text(message.get("content")));
898            }
899        }
900        "latest_reminder" => {
901            prompt.push_str(LATEST_REMINDER_SP_TOKEN);
902            prompt.push_str(&text(message.get("content")));
903        }
904        "tool" => {
905            return Err(EncodingError::UnsupportedRole {
906                index,
907                role: "tool".to_string(),
908            });
909        }
910        "assistant" => {
911            let mut reasoning = String::new();
912            let mut tool_calls = String::new();
913            if let Some(calls) = message.get("tool_calls") {
914                let calls = tool_calls_from_openai_format(calls)?;
915                if !calls.is_empty() {
916                    let mut encoded = Vec::new();
917                    for call in calls {
918                        let cm = object(&call)?;
919                        let name = tool_name_for_encoding(&call)?;
920                        encoded.push(format!(
921                            "<{DSML_TOKEN}{TOOL_CALL_TAG_NAME} name=\"{name}\">\n{}\n</{DSML_TOKEN}{TOOL_CALL_TAG_NAME}>",
922                            encode_arguments_to_dsml(&Value::Object(cm.clone()))?
923                        ));
924                    }
925                    tool_calls = format!(
926                        "\n\n<{DSML_TOKEN}{TOOL_CALLS_BLOCK_NAME}>\n{}\n</{DSML_TOKEN}{TOOL_CALLS_BLOCK_NAME}>",
927                        encoded.join("\n")
928                    );
929                }
930            }
931            let previous_has_task = index > 0
932                && messages[index - 1]
933                    .get("task")
934                    .is_some_and(|v| !v.is_null());
935            if options.thinking_mode == ThinkingMode::Thinking && !previous_has_task {
936                let keep = !options.drop_thinking || index as isize > last_user;
937                if keep {
938                    reasoning.push_str(&text(message.get("reasoning_content")));
939                    reasoning.push_str(THINKING_END_TOKEN);
940                }
941            }
942            prompt.push_str(&reasoning);
943            prompt.push_str(&text(message.get("content")));
944            prompt.push_str(&tool_calls);
945            if !message
946                .get("wo_eos")
947                .and_then(Value::as_bool)
948                .unwrap_or(false)
949            {
950                prompt.push_str(EOS_TOKEN);
951            }
952        }
953        other => {
954            return Err(EncodingError::UnsupportedRole {
955                index,
956                role: other.to_string(),
957            });
958        }
959    }
960
961    // The assistant generation header belongs after a user/mid-system turn
962    // when the following message is an assistant (or when this is the final
963    // message).  A non-assistant next turn is a completed transition and
964    // leaves the header to that later turn, matching the Python formatter.
965    let next_is_non_assistant = index + 1 < messages.len()
966        && !matches!(
967            messages[index + 1].get("role").and_then(Value::as_str),
968            Some("assistant" | "latest_reminder")
969        );
970    if next_is_non_assistant {
971        return Ok(prompt);
972    }
973    if let Some(task) = message.get("task").and_then(Value::as_str) {
974        let token = task_token(task)
975            .ok_or_else(|| EncodingError::InvalidMessage(format!("invalid task '{task}'")))?;
976        if task != "action" {
977            prompt.push_str(token);
978        } else {
979            prompt.push_str(ASSISTANT_SP_TOKEN);
980            prompt.push_str(if options.thinking_mode == ThinkingMode::Thinking {
981                THINKING_START_TOKEN
982            } else {
983                THINKING_END_TOKEN
984            });
985            prompt.push_str(token);
986        }
987    } else if message_role == "user" || (message_role == "system" && index > 0) {
988        prompt.push_str(ASSISTANT_SP_TOKEN);
989        if options.thinking_mode == ThinkingMode::Thinking
990            && (!options.drop_thinking || index as isize >= last_user)
991        {
992            prompt.push_str(THINKING_START_TOKEN);
993        } else {
994            prompt.push_str(THINKING_END_TOKEN);
995        }
996    }
997    Ok(prompt)
998}
999
1000fn drop_thinking_messages(messages: &[Value]) -> Vec<Value> {
1001    let last = last_user_index(messages);
1002    messages
1003        .iter()
1004        .enumerate()
1005        .filter_map(|(index, message)| {
1006            let role = message.get("role").and_then(Value::as_str).unwrap_or("");
1007            let keep_role = matches!(
1008                role,
1009                "user" | "system" | "tool" | "latest_reminder" | "direct_search_results"
1010            );
1011            if keep_role || index as isize >= last {
1012                return Some(message.clone());
1013            }
1014            if role == "assistant" {
1015                let mut copy = object(message).ok()?.clone();
1016                copy.remove("reasoning_content");
1017                Some(Value::Object(copy))
1018            } else {
1019                None
1020            }
1021        })
1022        .collect()
1023}
1024
1025/// Render already preprocessed messages (tool messages merged, tool
1026/// results sorted, image blocks replaced) after `context`.
1027fn encode_messages_text(
1028    messages: &[Value],
1029    options: &EncodeOptions,
1030    context: &[Value],
1031) -> Result<String, EncodingError> {
1032    let mut full = context.to_vec();
1033    full.extend(messages.iter().cloned());
1034    let mut prompt = if options.add_default_bos_token && context.is_empty() {
1035        BOS_TOKEN.to_string()
1036    } else {
1037        String::new()
1038    };
1039    // Any message carrying tools keeps every turn's reasoning.
1040    let effective_drop = options.drop_thinking
1041        && !full
1042            .iter()
1043            .any(|m| m.get("tools").is_some_and(json_truthy));
1044    let (render_messages, render_count, context_len) =
1045        if options.thinking_mode == ThinkingMode::Thinking && effective_drop {
1046            let reduced_full = drop_thinking_messages(&full);
1047            let reduced_context = drop_thinking_messages(context);
1048            let count = reduced_full.len().saturating_sub(reduced_context.len());
1049            let context_len = reduced_full.len() - count;
1050            (reduced_full, count, context_len)
1051        } else {
1052            (full, messages.len(), context.len())
1053        };
1054    let render_options = EncodeOptions {
1055        drop_thinking: effective_drop,
1056        ..options.clone()
1057    };
1058    for index in 0..render_count {
1059        prompt.push_str(&render_message(
1060            index + context_len,
1061            &render_messages,
1062            &render_options,
1063        )?);
1064    }
1065    Ok(prompt)
1066}
1067
1068/// Encode current messages. `images` contains only records belonging to the
1069/// current turn, matching the official helper's context-cache contract.
1070pub fn encode_messages(
1071    messages: &[Value],
1072    options: &EncodeOptions,
1073) -> Result<EncodedPrompt, EncodingError> {
1074    // Validate before merging can discard message-level textual fields.
1075    for message in options.context.iter().chain(messages) {
1076        validate_no_image_tokens(message)?;
1077    }
1078    // Collect images only after tool results are in their final prompt
1079    // order. The reference sorts context and current turn as one list (its
1080    // dicts are shared, so the context is sorted in place as well) and then
1081    // slices the current turn back off by the merged context length.
1082    let context = merge_tool_messages(&options.context)?;
1083    let context_len = context.len();
1084    let mut all = context;
1085    all.extend(merge_tool_messages(messages)?);
1086    sort_tool_results_by_call_order(&mut all);
1087    let current = all.split_off(context_len);
1088    let (processed_context, _) = if all.is_empty() {
1089        (Vec::new(), Vec::new())
1090    } else {
1091        process_image_messages(&all)?
1092    };
1093    let (processed, images) = process_image_messages(&current)?;
1094    let prompt = encode_messages_text(&processed, options, &processed_context)?;
1095    Ok(EncodedPrompt { prompt, images })
1096}
1097
1098pub fn encode_messages_simple(
1099    messages: &[Value],
1100    thinking_mode: ThinkingMode,
1101) -> Result<EncodedPrompt, EncodingError> {
1102    encode_messages(
1103        messages,
1104        &EncodeOptions {
1105            thinking_mode,
1106            ..EncodeOptions::default()
1107        },
1108    )
1109}
1110
1111/// Normalize the reference's JSON case files (`encoding/tests/*.json`): a
1112/// single case object, a bare message list, or a list of cases. A
1113/// case-level `tools` list moves onto the first message, as the official
1114/// `load_cases` does.
1115pub fn load_cases(data: &Value) -> Result<Vec<Value>, EncodingError> {
1116    let cases: Vec<Value> = match data {
1117        Value::Object(_) => vec![data.clone()],
1118        Value::Array(items)
1119            if items
1120                .first()
1121                .is_some_and(|first| first.get("role").is_some()) =>
1122        {
1123            vec![json!({ "messages": items })]
1124        }
1125        Value::Array(items) => items.clone(),
1126        _ => {
1127            return Err(EncodingError::InvalidMessage(
1128                "case file must hold an object or a list".to_string(),
1129            ));
1130        }
1131    };
1132    let mut out = Vec::with_capacity(cases.len());
1133    for case in cases {
1134        let map = object(&case)?;
1135        let mut messages = map
1136            .get("messages")
1137            .and_then(Value::as_array)
1138            .cloned()
1139            .ok_or_else(|| {
1140                EncodingError::InvalidMessage("case has no messages array".to_string())
1141            })?;
1142        if let Some(tools) = map.get("tools") {
1143            let first = messages.first_mut().ok_or_else(|| {
1144                EncodingError::InvalidMessage(
1145                    "a case with tools must contain at least one message".to_string(),
1146                )
1147            })?;
1148            let first = first.as_object_mut().ok_or_else(|| {
1149                EncodingError::InvalidMessage("expected a JSON object".to_string())
1150            })?;
1151            first.insert("tools".to_string(), tools.clone());
1152        }
1153        let field = |key: &str| map.get(key).cloned().unwrap_or(Value::Null);
1154        out.push(json!({
1155            "messages": messages,
1156            "context": field("context"),
1157            "thinking_mode": field("thinking_mode"),
1158            "reasoning_effort": field("reasoning_effort"),
1159        }));
1160    }
1161    Ok(out)
1162}
1163
1164pub fn encode_case(
1165    case: &Value,
1166    default_mode: ThinkingMode,
1167) -> Result<EncodedPrompt, EncodingError> {
1168    let map = object(case)?;
1169    let messages = map
1170        .get("messages")
1171        .and_then(Value::as_array)
1172        .ok_or_else(|| EncodingError::InvalidMessage("case has no messages array".to_string()))?;
1173    // `case.get("thinking_mode") or thinking_mode`: an empty value falls back.
1174    let mode = map
1175        .get("thinking_mode")
1176        .and_then(Value::as_str)
1177        .filter(|mode| !mode.is_empty())
1178        .map(ThinkingMode::parse)
1179        .transpose()?
1180        .unwrap_or(default_mode);
1181    let effort = map
1182        .get("reasoning_effort")
1183        .filter(|value| !value.is_null())
1184        .map(ReasoningEffort::from_json)
1185        .transpose()?;
1186    let context = map
1187        .get("context")
1188        .and_then(Value::as_array)
1189        .cloned()
1190        .unwrap_or_default();
1191    encode_messages(
1192        messages,
1193        &EncodeOptions {
1194            thinking_mode: mode,
1195            reasoning_effort: effort,
1196            context,
1197            ..EncodeOptions::default()
1198        },
1199    )
1200}
1201
1202fn read_until_stop(index: usize, input: &str, stops: &[&str]) -> (usize, String, Option<String>) {
1203    let mut position = input.len();
1204    let mut matched = None;
1205    for stop in stops {
1206        if let Some(relative) = input[index..].find(stop) {
1207            let candidate = index + relative;
1208            if candidate < position {
1209                position = candidate;
1210                matched = Some((*stop).to_string());
1211            }
1212        }
1213    }
1214    match matched {
1215        Some(stop) => (
1216            position + stop.len(),
1217            input[index..position].to_string(),
1218            Some(stop),
1219        ),
1220        None => (input.len(), input[index..].to_string(), None),
1221    }
1222}
1223
1224pub fn parse_tool_calls(
1225    mut index: usize,
1226    input: &str,
1227) -> Result<(usize, Option<String>, Vec<Value>), EncodingError> {
1228    let calls_end = format!("</{DSML_TOKEN}{TOOL_CALLS_BLOCK_NAME}>");
1229    let call_start = format!("<{DSML_TOKEN}{TOOL_CALL_TAG_NAME}");
1230    let call_end = format!("</{DSML_TOKEN}{TOOL_CALL_TAG_NAME}");
1231    let parameter_start = format!("<{DSML_TOKEN}{TOOL_PARAMETER_TAG_NAME}");
1232    let parameter_end = format!("/{DSML_TOKEN}{TOOL_PARAMETER_TAG_NAME}");
1233    let mut calls = Vec::new();
1234    let mut last_stop = None;
1235    while index < input.len() {
1236        let (next, content, stop) = read_until_stop(index, input, &[&call_start, &calls_end]);
1237        index = next;
1238        last_stop = stop.clone();
1239        if content != ">\n" {
1240            return Err(EncodingError::InvalidCompletion(format!(
1241                "tool call header expected '>\\n', got {content:?}"
1242            )));
1243        }
1244        if stop.as_deref() == Some(calls_end.as_str()) {
1245            break;
1246        }
1247        if stop.is_none() {
1248            return Err(EncodingError::InvalidCompletion(
1249                "missing DSML tool-call tag".to_string(),
1250            ));
1251        }
1252        let (next, name_content, mut stop) =
1253            read_until_stop(index, input, &[&parameter_start, &call_end]);
1254        index = next;
1255        let name = name_content
1256            .strip_prefix(" name=\"")
1257            .and_then(|s| s.strip_suffix("\">\n"))
1258            .ok_or_else(|| {
1259                EncodingError::InvalidCompletion(format!(
1260                    "invalid tool name header {name_content:?}"
1261                ))
1262            })?;
1263        let mut args = Vec::new();
1264        while stop.as_deref() == Some(parameter_start.as_str()) {
1265            let (next, parameter, matched) = read_until_stop(index, input, &[&parameter_end]);
1266            index = next;
1267            let body = parameter.strip_prefix(" name=\"").ok_or_else(|| {
1268                EncodingError::InvalidCompletion("invalid parameter header".to_string())
1269            })?;
1270            let (name_end, rest) = body.split_once("\" string=\"").ok_or_else(|| {
1271                EncodingError::InvalidCompletion(format!("invalid parameter {parameter:?}"))
1272            })?;
1273            let (string, value) = rest.split_once("\">").ok_or_else(|| {
1274                EncodingError::InvalidCompletion(format!("invalid parameter {parameter:?}"))
1275            })?;
1276            let (value, close) = value.rsplit_once("<").ok_or_else(|| {
1277                EncodingError::InvalidCompletion(format!("invalid parameter {parameter:?}"))
1278            })?;
1279            if close != "" || !matches!(string, "true" | "false") {
1280                return Err(EncodingError::InvalidCompletion(format!(
1281                    "invalid parameter {parameter:?}"
1282                )));
1283            }
1284            if args
1285                .iter()
1286                .any(|(key, _, _): &(String, String, String)| key == name_end)
1287            {
1288                return Err(EncodingError::InvalidCompletion(format!(
1289                    "duplicate parameter {name_end}"
1290                )));
1291            }
1292            let (next, content, next_stop) =
1293                read_until_stop(index, input, &[&parameter_start, &call_end]);
1294            index = next;
1295            if content != ">\n" {
1296                return Err(EncodingError::InvalidCompletion(
1297                    "parameter separator expected '>\\n'".to_string(),
1298                ));
1299            }
1300            args.push((name_end.to_string(), value.to_string(), string.to_string()));
1301            stop = next_stop;
1302            if matched.is_none() {
1303                return Err(EncodingError::InvalidCompletion(
1304                    "unterminated parameter".to_string(),
1305                ));
1306            }
1307        }
1308        calls.push(decode_dsml_to_arguments(name, &args)?);
1309    }
1310    Ok((index, last_stop, calls))
1311}
1312
1313fn tool_calls_to_openai(calls: &[Value]) -> Value {
1314    Value::Array(
1315        calls
1316            .iter()
1317            .map(|call| {
1318                let name = call
1319                    .get("name")
1320                    .cloned()
1321                    .unwrap_or(Value::String(String::new()));
1322                let arguments = call
1323                    .get("arguments")
1324                    .cloned()
1325                    .unwrap_or(Value::String(String::new()));
1326                let mut function = Map::new();
1327                function.insert("name".to_string(), name);
1328                function.insert("arguments".to_string(), arguments);
1329                let mut result = Map::new();
1330                result.insert("type".to_string(), Value::String("function".to_string()));
1331                result.insert("function".to_string(), Value::Object(function));
1332                if let Some(namespace) = call.get("namespace") {
1333                    result.insert("namespace".to_string(), namespace.clone());
1334                }
1335                Value::Object(result)
1336            })
1337            .collect(),
1338    )
1339}
1340
1341/// Parse a generated V4.1 completion into an OpenAI style assistant object.
1342pub fn parse_message_from_completion_text(
1343    input: &str,
1344    mode: ThinkingMode,
1345) -> Result<Value, EncodingError> {
1346    let tool_start = format!("\n\n<{DSML_TOKEN}{TOOL_CALLS_BLOCK_NAME}");
1347    let (mut index, reasoning, mut stop) = if mode == ThinkingMode::Thinking {
1348        let (next, content, matched) =
1349            read_until_stop(0, input, &[THINKING_END_TOKEN, &tool_start]);
1350        if matched.as_deref() != Some(THINKING_END_TOKEN) {
1351            return Err(EncodingError::InvalidCompletion(
1352                "thinking completion is missing </think>".to_string(),
1353            ));
1354        }
1355        (next, content, matched)
1356    } else {
1357        (0, String::new(), None)
1358    };
1359    let (next, summary, matched) = read_until_stop(index, input, &[EOS_TOKEN, &tool_start]);
1360    index = next;
1361    stop = matched;
1362    let mut calls = Vec::new();
1363    if stop.as_deref() == Some(tool_start.as_str()) {
1364        let (next, _, parsed_calls) = parse_tool_calls(index, input)?;
1365        index = next;
1366        calls = parsed_calls;
1367        let (next, tail, matched) = read_until_stop(index, input, &[EOS_TOKEN]);
1368        if !tail.is_empty() {
1369            return Err(EncodingError::InvalidCompletion(
1370                "content follows DSML calls".to_string(),
1371            ));
1372        }
1373        index = next;
1374        stop = matched;
1375    } else if stop.as_deref() != Some(EOS_TOKEN) {
1376        return Err(EncodingError::InvalidCompletion(
1377            "completion is missing EOS token".to_string(),
1378        ));
1379    }
1380    if index != input.len() {
1381        return Err(EncodingError::InvalidCompletion(
1382            "unexpected bytes after completion".to_string(),
1383        ));
1384    }
1385    for special in [
1386        BOS_TOKEN,
1387        EOS_TOKEN,
1388        THINKING_START_TOKEN,
1389        THINKING_END_TOKEN,
1390        DSML_TOKEN,
1391    ] {
1392        if summary.contains(special) || reasoning.contains(special) {
1393            return Err(EncodingError::InvalidCompletion(format!(
1394                "special token {special} leaked into content"
1395            )));
1396        }
1397    }
1398    Ok(json!({
1399        "role":"assistant",
1400        "content":summary,
1401        "reasoning_content":reasoning,
1402        "tool_calls":tool_calls_to_openai(&calls),
1403    }))
1404}
1405
1406#[cfg(test)]
1407mod tests {
1408    use super::*;
1409
1410    fn user(content: Value) -> Value {
1411        json!({"role":"user", "content":content})
1412    }
1413
1414    #[test]
1415    fn text_chat_and_numeric_reasoning_match_reference() {
1416        let chat = encode_messages(
1417            &[user(Value::String("hello".into()))],
1418            &EncodeOptions::default(),
1419        )
1420        .unwrap();
1421        assert_eq!(
1422            chat.prompt,
1423            "<|begin▁of▁sentence|><|User|>hello<|Assistant|></think>"
1424        );
1425        let thinking = encode_messages(
1426            &[user(Value::String("question".into()))],
1427            &EncodeOptions {
1428                thinking_mode: ThinkingMode::Thinking,
1429                reasoning_effort: Some(ReasoningEffort::Budget(42)),
1430                ..EncodeOptions::default()
1431            },
1432        )
1433        .unwrap();
1434        assert_eq!(
1435            thinking.prompt,
1436            "<|begin▁of▁sentence|><|System|>Reasoning Effort: 42 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<|User|>question<|Assistant|><think>"
1437        );
1438    }
1439
1440    #[test]
1441    fn images_and_tool_tags_are_ordered_and_spaced() {
1442        let prompt = encode_messages(
1443            &[user(json!([
1444                {"type":"text","text":"inspect"},
1445                {"type":"image_url","image_url":{"url":"/tmp/a.png"}},
1446            ]))],
1447            &EncodeOptions::default(),
1448        )
1449        .unwrap();
1450        assert_eq!(
1451            prompt.images,
1452            vec![json!({"type":"image","url":"/tmp/a.png"})]
1453        );
1454        assert_eq!(
1455            prompt.prompt,
1456            "<|begin▁of▁sentence|><|User|>inspect\n\n<|deepseek_image|><|Assistant|></think>"
1457        );
1458        assert!(!prompt.prompt.contains("<|DSML|tool_calls>"));
1459    }
1460
1461    #[test]
1462    fn mid_system_fixture_keeps_prior_reasoning_and_new_header() {
1463        let encoded = encode_messages(
1464            &[
1465                user(Value::String("old".into())),
1466                json!({
1467                    "role":"assistant",
1468                    "content":"old answer",
1469                    "reasoning_content":"private",
1470                    "wo_eos":true
1471                }),
1472                json!({"role":"system", "content":"new policy"}),
1473                user(Value::String("now".into())),
1474            ],
1475            &EncodeOptions {
1476                thinking_mode: ThinkingMode::Thinking,
1477                reasoning_effort: Some(ReasoningEffort::Budget(75)),
1478                ..EncodeOptions::default()
1479            },
1480        )
1481        .unwrap();
1482        assert_eq!(
1483            encoded.prompt,
1484            "<|begin▁of▁sentence|><|System|>Reasoning Effort: 75 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<|User|>old<|Assistant|></think>old answer<|System|>new policy<|User|>now<|Assistant|><think>"
1485        );
1486    }
1487
1488    #[test]
1489    fn spaced_dsml_round_trip() {
1490        let completion = "reason </think>answer\n\n<|DSML| calls>\n<|DSML| invoke name=\"lookup\">\n<|DSML| parameter name=\"q\" string=\"true\">Paris</|DSML| parameter>\n</|DSML| invoke>\n</|DSML| calls><|end▁of▁sentence|>";
1491        let parsed =
1492            parse_message_from_completion_text(completion, ThinkingMode::Thinking).unwrap();
1493        assert_eq!(parsed["tool_calls"][0]["function"]["name"], "lookup");
1494        assert_eq!(
1495            parsed["tool_calls"][0]["function"]["arguments"],
1496            "{\"q\": \"Paris\"}"
1497        );
1498    }
1499
1500    /// The official golden fixtures (`encoding/tests/test_input_*.json` ->
1501    /// `test_output_*.txt`, DeepSeek-V4.1-Flash-1017), encoded the way
1502    /// `test_examples_encoding_golden_outputs` does: `load_cases`, first case,
1503    /// `encode_case(case, "chat")`, byte-exact.
1504    #[test]
1505    fn reference_golden_fixtures_encode_byte_exact() {
1506        let cases: [(&str, &str); 5] = [
1507            (
1508                include_str!("../tests/fixtures/dsv41/encoding/test_input_1.json"),
1509                include_str!("../tests/fixtures/dsv41/encoding/test_output_1.txt"),
1510            ),
1511            (
1512                include_str!("../tests/fixtures/dsv41/encoding/test_input_2.json"),
1513                include_str!("../tests/fixtures/dsv41/encoding/test_output_2.txt"),
1514            ),
1515            (
1516                include_str!("../tests/fixtures/dsv41/encoding/test_input_3.json"),
1517                include_str!("../tests/fixtures/dsv41/encoding/test_output_3.txt"),
1518            ),
1519            (
1520                include_str!("../tests/fixtures/dsv41/encoding/test_input_4.json"),
1521                include_str!("../tests/fixtures/dsv41/encoding/test_output_4.txt"),
1522            ),
1523            (
1524                include_str!("../tests/fixtures/dsv41/encoding/test_input_5.json"),
1525                include_str!("../tests/fixtures/dsv41/encoding/test_output_5.txt"),
1526            ),
1527        ];
1528        for (i, (input, output)) in cases.iter().enumerate() {
1529            let data: Value = serde_json::from_str(input).unwrap();
1530            let case = load_cases(&data).unwrap().remove(0);
1531            let encoded = encode_case(&case, ThinkingMode::Chat).unwrap();
1532            assert_eq!(encoded.prompt, *output, "test_output_{}.txt", i + 1);
1533        }
1534        // Case 5 carries two images in prompt order.
1535        let data: Value = serde_json::from_str(cases[4].0).unwrap();
1536        let case = load_cases(&data).unwrap().remove(0);
1537        let images = encode_case(&case, ThinkingMode::Chat).unwrap().images;
1538        assert_eq!(
1539            images,
1540            vec![
1541                json!({"type":"image","url":"examples/images/carrots.jpeg"}),
1542                json!({"type":"image","url":"examples/images/corn.jpeg"}),
1543            ]
1544        );
1545    }
1546
1547    /// `test_v41_tool_result_images_follow_prompt_order` (reference commit
1548    /// 314113f8): tool results are sorted into call order before image
1549    /// records are collected, for results given as content blocks, as tool
1550    /// messages with list content, as tool messages with `content_blocks`,
1551    /// and across the context boundary.
1552    #[test]
1553    fn tool_result_images_follow_prompt_order() {
1554        let history = vec![
1555            json!({"role":"user","content":[{"type":"image_url","image_url":{"url":"history.png"}}]}),
1556            json!({"role":"assistant","tool_calls":[
1557                {"id":"a","type":"function","function":{"name":"a","arguments":"{}"}},
1558                {"id":"b","type":"function","function":{"name":"b","arguments":"{}"}},
1559            ]}),
1560        ];
1561        let results: Vec<Value> = ["b", "a"]
1562            .iter()
1563            .map(|name| {
1564                json!({"type":"tool_result","tool_use_id":name,"content":[
1565                    {"type":"text","text":format!("result {name}")},
1566                    {"type":"image_url","image_url":{"url":format!("{name}.png")}},
1567                ]})
1568            })
1569            .collect();
1570        for format in [
1571            "content_blocks",
1572            "tool_messages",
1573            "tool_content_blocks",
1574            "context",
1575        ] {
1576            let (messages, context) = match format {
1577                "tool_messages" | "tool_content_blocks" => {
1578                    let field = if format == "tool_content_blocks" {
1579                        "content_blocks"
1580                    } else {
1581                        "content"
1582                    };
1583                    let mut messages = history.clone();
1584                    for block in &results {
1585                        let mut m = Map::new();
1586                        m.insert("role".into(), json!("tool"));
1587                        m.insert("tool_call_id".into(), block["tool_use_id"].clone());
1588                        m.insert(field.into(), block["content"].clone());
1589                        messages.push(Value::Object(m));
1590                    }
1591                    (messages, Vec::new())
1592                }
1593                "context" => (
1594                    vec![json!({"role":"user","content_blocks":results})],
1595                    history.clone(),
1596                ),
1597                _ => {
1598                    let mut messages = history.clone();
1599                    messages.push(json!({"role":"user","content_blocks":results}));
1600                    (messages, Vec::new())
1601                }
1602            };
1603            let original = (messages.clone(), context.clone());
1604            let encoded = encode_messages(
1605                &messages,
1606                &EncodeOptions {
1607                    context: context.clone(),
1608                    ..EncodeOptions::default()
1609                },
1610            )
1611            .unwrap();
1612            assert!(
1613                encoded.prompt.ends_with(&format!(
1614                    "<|User|><tool_result>result a\n\n{IMAGE_PLACEHOLDER}</tool_result>\n\n\
1615                     <tool_result>result b\n\n{IMAGE_PLACEHOLDER}</tool_result><|Assistant|></think>"
1616                )),
1617                "{format}: {}",
1618                encoded.prompt
1619            );
1620            let urls: &[&str] = if context.is_empty() {
1621                &["history.png", "a.png", "b.png"]
1622            } else {
1623                &["a.png", "b.png"]
1624            };
1625            let expected: Vec<Value> = urls
1626                .iter()
1627                .map(|url| json!({"type":"image","url":url}))
1628                .collect();
1629            assert_eq!(encoded.images, expected, "{format}");
1630            assert_eq!((messages, context), original, "{format}: inputs mutated");
1631        }
1632    }
1633
1634    /// `test_v41_drop_thinking_without_tools` (314113f8): completed turns
1635    /// lose their reasoning, the current turn keeps it across tool results,
1636    /// and a following user message closes that turn.
1637    #[test]
1638    fn current_turn_thinking_survives_tool_results() {
1639        let messages = vec![
1640            json!({"role":"user","content":"q1"}),
1641            json!({"role":"assistant","content":"a1","reasoning_content":"r1"}),
1642            json!({"role":"user","content":"q2"}),
1643            json!({"role":"assistant","reasoning_content":"r2","tool_calls":[
1644                {"id":"c1","type":"function","function":{"name":"lookup","arguments":"{}"}}
1645            ]}),
1646            json!({"role":"tool","tool_call_id":"c1","content":"value"}),
1647        ];
1648        let options = EncodeOptions {
1649            thinking_mode: ThinkingMode::Thinking,
1650            drop_thinking: true,
1651            ..EncodeOptions::default()
1652        };
1653        let prompt = encode_messages(&messages, &options).unwrap().prompt;
1654        assert!(prompt.contains("<|User|>q1<|Assistant|></think>a1<|end▁of▁sentence|>"));
1655        assert!(!prompt.contains("r1"));
1656        assert!(prompt.contains("<think>r2</think>"), "{prompt}");
1657        assert!(
1658            prompt.ends_with("<tool_result>value</tool_result><|Assistant|><think>"),
1659            "{prompt}"
1660        );
1661        let mut next = messages.clone();
1662        next.push(json!({"role":"user","content":"q3"}));
1663        let next_turn = encode_messages(&next, &options).unwrap().prompt;
1664        assert!(!next_turn.contains("r2"), "{next_turn}");
1665        assert!(next_turn.ends_with("q3<|Assistant|><think>"));
1666    }
1667
1668    /// The reference reference-tests around the same change: placeholders in
1669    /// user or tool text are refused before merging; tool messages merge
1670    /// into one user message of tool_result blocks; a task token replaces the
1671    /// assistant header.
1672    #[test]
1673    fn reference_preprocessing_rules() {
1674        for role in ["user", "tool"] {
1675            let err = encode_messages(
1676                &[json!({"role":role,"content":format!("hi {IMAGE_PLACEHOLDER}")})],
1677                &EncodeOptions::default(),
1678            );
1679            assert!(err.is_err(), "{role} placeholder accepted");
1680        }
1681        let merged = merge_tool_messages(&[
1682            json!({"role":"assistant","content":"","tool_calls":[]}),
1683            json!({"role":"tool","tool_call_id":"a","content":"r1"}),
1684            json!({"role":"tool","tool_call_id":"b","content":"r2"}),
1685        ])
1686        .unwrap();
1687        assert_eq!(merged.len(), 2);
1688        assert_eq!(merged[1]["role"], "user");
1689        let kinds: Vec<&str> = merged[1]["content_blocks"]
1690            .as_array()
1691            .unwrap()
1692            .iter()
1693            .map(|b| b["type"].as_str().unwrap())
1694            .collect();
1695        assert_eq!(kinds, ["tool_result", "tool_result"]);
1696        let prompt = encode_messages(
1697            &[json!({"role":"user","content":"classify me","task":"query"})],
1698            &EncodeOptions::default(),
1699        )
1700        .unwrap()
1701        .prompt;
1702        assert!(prompt.ends_with("classify me<|query|>"));
1703        assert!(!prompt.contains(ASSISTANT_SP_TOKEN));
1704        // A user turn holding only tool results does not open a new turn.
1705        let only_tools = [
1706            json!({"role":"user","content":"q"}),
1707            json!({"role":"user","content_blocks":[{"type":"tool_result","tool_use_id":"x","content":"v"}]}),
1708        ];
1709        assert_eq!(last_user_index(&only_tools), 0);
1710    }
1711
1712    #[test]
1713    fn effort_rejects_float_and_bool_json() {
1714        assert!(ReasoningEffort::from_json(&json!(1.5)).is_err());
1715        assert!(ReasoningEffort::from_json(&json!(true)).is_err());
1716        assert!(ReasoningEffort::from_json(&json!(0)).is_err());
1717        assert!(ReasoningEffort::from_json(&json!("42")).is_err());
1718    }
1719}