1use 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
31pub 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 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 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
230fn 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 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
433pub 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
586pub 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
624pub 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 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 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
696pub 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
752fn 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 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
834fn 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
850pub 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 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
1025fn 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 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
1068pub fn encode_messages(
1071 messages: &[Value],
1072 options: &EncodeOptions,
1073) -> Result<EncodedPrompt, EncodingError> {
1074 for message in options.context.iter().chain(messages) {
1076 validate_no_image_tokens(message)?;
1077 }
1078 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(¤t)?;
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
1111pub 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 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, &[¶meter_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, &[¶meter_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, &[¶meter_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
1341pub 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 #[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 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]
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]
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 #[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 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}