1use bitflags::bitflags;
76use serde::{Deserialize, Serialize};
77
78bitflags! {
79 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
84 #[repr(transparent)]
85 pub struct ModelCapabilities: u32 {
86 const SUPPORTS_SYSTEM_ROLE = 0b0000_0001;
91
92 const REQUIRES_STRICT_TURNS = 0b0000_0010;
97
98 const SUPPORTS_TOOL_CALLS = 0b0000_0100;
103
104 const SUPPORTS_REASONING = 0b0000_1000;
109 }
110}
111
112impl Default for ModelCapabilities {
113 fn default() -> Self {
118 Self::empty()
119 }
120}
121
122impl Serialize for ModelCapabilities {
123 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
124 where
125 S: serde::Serializer,
126 {
127 self.bits().serialize(serializer)
128 }
129}
130
131impl<'de> Deserialize<'de> for ModelCapabilities {
132 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
133 where
134 D: serde::Deserializer<'de>,
135 {
136 let bits = u32::deserialize(deserializer)?;
137 Ok(Self::from_bits_truncate(bits))
138 }
139}
140
141impl ModelCapabilities {
142 pub const fn supports_system_role(self) -> bool {
144 self.contains(Self::SUPPORTS_SYSTEM_ROLE)
145 }
146
147 pub const fn requires_strict_turns(self) -> bool {
149 self.contains(Self::REQUIRES_STRICT_TURNS)
150 }
151
152 pub const fn supports_tool_calls(self) -> bool {
154 self.contains(Self::SUPPORTS_TOOL_CALLS)
155 }
156
157 pub const fn supports_reasoning(self) -> bool {
159 self.contains(Self::SUPPORTS_REASONING)
160 }
161}
162
163pub fn infer_from_chat_template(
190 template: Option<&str>,
191 model_name: Option<&str>,
192) -> ModelCapabilities {
193 let mut caps = ModelCapabilities::empty();
194
195 let mut tool_detected_from_metadata = false;
200 let mut reasoning_detected_from_metadata = false;
201
202 if let Some(template) = template {
203 let supports_system_positive =
215 template.contains("[SYSTEM_PROMPT]") || template.contains("[AVAILABLE_TOOLS]");
216
217 let forbids_system = !supports_system_positive
218 && (template.contains("Only user, assistant and tool roles are supported")
219 || template.contains("got system")
220 || template.contains("Raise exception for unsupported roles"));
221
222 if !forbids_system {
223 caps |= ModelCapabilities::SUPPORTS_SYSTEM_ROLE;
224 }
225
226 let requires_alternation = template.contains("must alternate user and assistant")
229 || template.contains("conversation roles must alternate")
230 || template.contains("ns.index % 2");
231
232 if requires_alternation {
233 caps |= ModelCapabilities::REQUIRES_STRICT_TURNS;
234 }
235
236 let has_tool_patterns = template.contains("<tool_call>")
238 || template.contains("<|python_tag|>")
239 || template.contains("if tools")
240 || template.contains("tools is defined")
241 || template.contains("tool_calls")
242 || template.contains("function_call");
243
244 if has_tool_patterns {
245 caps |= ModelCapabilities::SUPPORTS_TOOL_CALLS;
246 tool_detected_from_metadata = true;
247 }
248
249 let has_reasoning_patterns = template.contains("<think>")
251 || template.contains("</think>")
252 || template.contains("<reasoning>")
253 || template.contains("</reasoning>")
254 || template.contains("enable_thinking")
255 || template.contains("thinking_forced_open")
256 || template.contains("reasoning_content");
257
258 if has_reasoning_patterns {
259 caps |= ModelCapabilities::SUPPORTS_REASONING;
260 reasoning_detected_from_metadata = true;
261 }
262 }
263
264 if let Some(name) = model_name {
272 let name_lower = name.to_lowercase();
273
274 if !tool_detected_from_metadata {
276 let has_tool_name = name_lower.contains("hermes")
277 || name_lower.contains("functionary")
278 || name_lower.contains("firefunction")
279 || name_lower.contains("gorilla");
280
281 if has_tool_name {
282 caps |= ModelCapabilities::SUPPORTS_TOOL_CALLS;
283 }
284 }
285
286 if !reasoning_detected_from_metadata {
288 let has_reasoning_name = name_lower.contains("deepseek-r1")
289 || name_lower.contains("qwq")
290 || name_lower.contains("-r1-")
291 || name_lower.contains("o1");
292
293 if has_reasoning_name {
294 caps |= ModelCapabilities::SUPPORTS_REASONING;
295 }
296 }
297 }
298
299 caps
300}
301
302#[must_use]
350pub fn capabilities_from_architecture(arch: Option<&str>) -> ModelCapabilities {
351 let Some(arch) = arch else {
352 return ModelCapabilities::empty();
353 };
354
355 match arch {
356 "mistral" => ModelCapabilities::REQUIRES_STRICT_TURNS,
360
361 "mistral3" => {
367 ModelCapabilities::REQUIRES_STRICT_TURNS | ModelCapabilities::SUPPORTS_SYSTEM_ROLE
368 }
369
370 _ => ModelCapabilities::empty(),
375 }
376}
377
378#[cfg(test)]
379mod tests {
380 use super::*;
381
382 #[test]
383 fn test_default_capabilities() {
384 let caps = ModelCapabilities::default();
385 assert!(caps.is_empty());
387 assert!(!caps.supports_system_role());
388 assert!(!caps.requires_strict_turns());
389 assert!(!caps.supports_tool_calls());
390 assert!(!caps.supports_reasoning());
391 }
392
393 #[test]
394 fn test_infer_openai_style() {
395 let template = r"
396 {% for message in messages %}
397 {{ message.role }}: {{ message.content }}
398 {% endfor %}
399 ";
400 let caps = infer_from_chat_template(Some(template), None);
401 assert!(caps.supports_system_role());
402 assert!(!caps.requires_strict_turns());
403 }
404
405 #[test]
406 fn test_infer_mistral_style() {
407 let template = r"
408 {% if message.role == 'system' %}
409 {{ raise_exception('Only user, assistant and tool roles are supported, got system.') }}
410 {% endif %}
411 {% if (message['role'] == 'user') != (ns.index % 2 == 0) %}
412 {{ raise_exception('conversation roles must alternate user and assistant') }}
413 {% endif %}
414 ";
415 let caps = infer_from_chat_template(Some(template), None);
416 assert!(!caps.supports_system_role());
417 assert!(caps.requires_strict_turns());
418 }
419
420 #[test]
421 fn test_infer_missing_template() {
422 let caps = infer_from_chat_template(None, None);
423 assert!(caps.is_empty());
425 assert!(!caps.supports_system_role());
426 }
427
428 #[test]
429 fn test_tool_calling_from_template() {
430 let template = r"
431 {% if tools %}
432 <tool_call>{{ message.tool_calls }}</tool_call>
433 {% endif %}
434 ";
435 let caps = infer_from_chat_template(Some(template), None);
436 assert!(caps.supports_tool_calls());
437 }
438
439 #[test]
440 fn test_reasoning_from_template() {
441 let template = r"
442 {% if enable_thinking %}
443 <think>{{ message.thinking }}</think>
444 {% endif %}
445 ";
446 let caps = infer_from_chat_template(Some(template), None);
447 assert!(caps.supports_reasoning());
448 }
449
450 #[test]
451 fn test_tool_calling_name_fallback() {
452 let caps = infer_from_chat_template(None, Some("hermes-2-pro-7b"));
454 assert!(caps.supports_tool_calls());
455 }
456
457 #[test]
458 fn test_reasoning_name_fallback() {
459 let caps = infer_from_chat_template(None, Some("deepseek-r1-lite"));
461 assert!(caps.supports_reasoning());
462 }
463
464 #[test]
465 fn test_metadata_plus_name_fallback() {
466 let template = "simple template with no tool markers";
468 let caps = infer_from_chat_template(Some(template), Some("hermes-model"));
469 assert!(caps.supports_tool_calls());
471 }
472
473 #[test]
474 fn test_metadata_detected_skips_name_fallback() {
475 let template = "<tool_call>detected</tool_call>";
477 let caps = infer_from_chat_template(Some(template), Some("not-a-tool-model"));
478 assert!(caps.supports_tool_calls());
480 }
481
482 #[test]
483 fn test_combined_detections() {
484 let template = r"
485 {% if tools %}<tool_call>{{ tool }}</tool_call>{% endif %}
486 <think>{{ reasoning }}</think>
487 ";
488 let caps = infer_from_chat_template(Some(template), None);
489 assert!(caps.supports_tool_calls());
490 assert!(caps.supports_reasoning());
491 }
492
493 #[test]
496 fn test_arch_none_returns_empty() {
497 assert!(capabilities_from_architecture(None).is_empty());
498 }
499
500 #[test]
501 fn test_arch_mistral_requires_strict_turns() {
502 let caps = capabilities_from_architecture(Some("mistral"));
503 assert!(caps.requires_strict_turns());
504 }
505
506 #[test]
507 fn test_arch_llama_returns_empty() {
508 assert!(capabilities_from_architecture(Some("llama")).is_empty());
509 }
510
511 #[test]
512 fn test_arch_unknown_returns_empty() {
513 assert!(capabilities_from_architecture(Some("future-arch-xyz")).is_empty());
514 }
515
516 #[test]
517 fn test_arch_mistral3_strict_turns_and_system_role() {
518 let caps = capabilities_from_architecture(Some("mistral3"));
519 assert!(
520 caps.requires_strict_turns(),
521 "mistral3 must enforce strict turns"
522 );
523 assert!(
524 caps.supports_system_role(),
525 "mistral3 supports system via [SYSTEM_PROMPT]"
526 );
527 }
528
529 #[test]
530 fn test_infer_mistral_v7_supports_system() {
531 let template = r"
534 {% if messages[0].role == 'system' %}
535 [SYSTEM_PROMPT]{{ messages[0].content }}[/SYSTEM_PROMPT]
536 {% endif %}
537 {% for message in messages %}
538 {% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
539 {{ raise_exception('conversation roles must alternate') }}
540 {% endif %}
541 {% endfor %}
542 ";
543 let caps = infer_from_chat_template(Some(template), None);
544 assert!(
545 caps.supports_system_role(),
546 "[SYSTEM_PROMPT] is positive evidence"
547 );
548 assert!(caps.requires_strict_turns(), "still enforces alternation");
549 }
550
551 #[test]
552 fn test_infer_mistral_v3_supports_system() {
553 let template = r"
556 {% if tools is defined %}[AVAILABLE_TOOLS]{{ tools | tojson }}[/AVAILABLE_TOOLS]{% endif %}
557 {% for message in messages %}
558 {% if message.role == 'user' %}[INST]{{ message.content }}[/INST]
559 {% elif message.role == 'assistant' %}{{ message.content }}</s>
560 {% endif %}
561 {% endfor %}
562 ";
563 let caps = infer_from_chat_template(Some(template), None);
564 assert!(
565 caps.supports_system_role(),
566 "[AVAILABLE_TOOLS] is positive evidence"
567 );
568 }
569
570 #[test]
571 fn test_infer_mistral_v1_forbids_system() {
572 let template = r"
575 {% if message.role == 'system' %}
576 {{ raise_exception('Only user, assistant and tool roles are supported, got system.') }}
577 {% endif %}
578 ";
579 let caps = infer_from_chat_template(Some(template), None);
580 assert!(
581 !caps.supports_system_role(),
582 "v1/v2 genuinely rejects system role"
583 );
584 }
585
586 #[test]
587 fn test_arch_or_template_additive() {
588 let template = "<tool_call>{{ tool }}</tool_call>";
591 let from_template = infer_from_chat_template(Some(template), None);
592 let from_arch = capabilities_from_architecture(Some("mistral"));
593 let combined = from_template | from_arch;
594 assert!(combined.supports_tool_calls(), "tool calls from template");
595 assert!(combined.requires_strict_turns(), "strict turns from arch");
596 }
597}
598
599#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
619#[serde(untagged)]
620pub enum MessageContent {
621 Text(String),
623 Parts(Vec<serde_json::Value>),
628}
629
630impl MessageContent {
631 pub fn as_str(&self) -> Option<&str> {
633 match self {
634 Self::Text(s) => Some(s),
635 Self::Parts(_) => None,
636 }
637 }
638
639 pub fn into_string(self) -> String {
650 match self {
651 Self::Text(s) => s,
652 Self::Parts(parts) => parts
653 .iter()
654 .filter_map(|p| p.get("text").and_then(|t| t.as_str()))
655 .collect::<Vec<_>>()
656 .join(""),
657 }
658 }
659
660 fn merge_with(self, other: Self) -> Self {
672 match (self, other) {
673 (Self::Text(mut a), Self::Text(b)) => {
674 if a.is_empty() {
675 return Self::Text(b);
676 }
677 if b.is_empty() {
678 return Self::Text(a);
679 }
680 a.push_str("\n\n");
681 a.push_str(&b);
682 Self::Text(a)
683 }
684 (Self::Parts(mut a), Self::Parts(b)) => {
685 a.extend(b);
686 Self::Parts(a)
687 }
688 (Self::Text(a), Self::Parts(b)) => {
689 let mut parts = vec![serde_json::json!({"type": "text", "text": a})];
690 parts.extend(b);
691 Self::Parts(parts)
692 }
693 (Self::Parts(mut a), Self::Text(b)) => {
694 a.push(serde_json::json!({"type": "text", "text": b}));
695 Self::Parts(a)
696 }
697 }
698 }
699}
700
701impl From<String> for MessageContent {
702 fn from(s: String) -> Self {
703 Self::Text(s)
704 }
705}
706
707impl From<&str> for MessageContent {
708 fn from(s: &str) -> Self {
709 Self::Text(s.to_string())
710 }
711}
712
713#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
734pub struct ChatMessage {
735 pub role: String,
736 #[serde(default, skip_serializing_if = "Option::is_none")]
737 pub content: Option<MessageContent>,
738 #[serde(default, skip_serializing_if = "Option::is_none")]
739 pub tool_calls: Option<serde_json::Value>,
740 #[serde(flatten)]
745 pub extra: serde_json::Map<String, serde_json::Value>,
746}
747
748impl ChatMessage {
749 fn merge_into(&mut self, other: Self) {
762 self.content = match (self.content.take(), other.content) {
763 (None, b) => b,
764 (a, None) => a,
765 (Some(a), Some(b)) => Some(a.merge_with(b)),
766 };
767 match (self.tool_calls.as_mut(), other.tool_calls) {
768 (_, None) => {}
769 (None, tc) => self.tool_calls = tc,
770 (Some(last_tc), Some(msg_tc)) => {
771 if let (Some(la), Some(ma)) = (last_tc.as_array_mut(), msg_tc.as_array()) {
772 la.extend_from_slice(ma);
773 }
774 }
775 }
776 for (key, value) in other.extra {
777 self.extra.entry(key).or_insert(value);
778 }
779 }
780}
781
782fn merge_consecutive_system_messages(messages: Vec<ChatMessage>) -> Vec<ChatMessage> {
797 if messages.is_empty() {
798 return messages;
799 }
800
801 let mut result: Vec<ChatMessage> = Vec::with_capacity(messages.len());
802
803 for msg in messages {
804 let is_system_merge = result
805 .last()
806 .is_some_and(|last| last.role == "system" && msg.role == "system");
807 if is_system_merge {
808 let last = result.last_mut().unwrap();
809 last.content = match (last.content.take(), msg.content) {
810 (None, b) => b,
811 (a, None) => a,
812 (Some(a), Some(b)) => Some(a.merge_with(b)),
813 };
814 } else {
815 result.push(msg);
816 }
817 }
818
819 result
820}
821
822pub fn transform_messages_for_capabilities(
846 mut messages: Vec<ChatMessage>,
847 capabilities: ModelCapabilities,
848) -> Vec<ChatMessage> {
849 messages = merge_consecutive_system_messages(messages);
854
855 if capabilities.is_empty() {
857 return messages;
858 }
859
860 if !capabilities.contains(ModelCapabilities::SUPPORTS_SYSTEM_ROLE) {
862 for msg in &mut messages {
863 if msg.role == "system" {
864 msg.role = "user".to_string();
865 if let Some(content) = msg.content.take() {
866 msg.content = Some(MessageContent::Text(format!(
867 "[System]: {}",
868 content.into_string()
869 )));
870 }
871 }
872 }
873 }
874
875 if capabilities.contains(ModelCapabilities::REQUIRES_STRICT_TURNS) {
877 let mut merged: Vec<ChatMessage> = Vec::new();
878 for msg in messages {
879 let is_mergeable = msg.role == "user" || msg.role == "assistant";
880 let same_role_as_last = merged.last().is_some_and(|last| last.role == msg.role);
881 if is_mergeable && same_role_as_last {
882 merged.last_mut().unwrap().merge_into(msg);
883 } else {
884 merged.push(msg);
885 }
886 }
887 return merged;
888 }
889
890 messages
891}
892
893#[cfg(test)]
894mod transform_tests {
895 use super::*;
896
897 #[test]
898 fn test_transform_unknown_passes_through_non_system() {
899 let messages = vec![
901 ChatMessage {
902 role: "user".to_string(),
903 content: Some(MessageContent::Text("Hello".to_string())),
904 tool_calls: None,
905 ..Default::default()
906 },
907 ChatMessage {
908 role: "assistant".to_string(),
909 content: Some(MessageContent::Text("Hi there".to_string())),
910 tool_calls: None,
911 ..Default::default()
912 },
913 ];
914 let original = messages.clone();
915 let result = transform_messages_for_capabilities(messages, ModelCapabilities::empty());
916 assert_eq!(result, original);
917 }
918
919 #[test]
920 fn test_merges_consecutive_system_messages_always() {
921 let messages = vec![
923 ChatMessage {
924 role: "system".to_string(),
925 content: Some(MessageContent::Text(
926 "You are a helpful assistant.".to_string(),
927 )),
928 tool_calls: None,
929 ..Default::default()
930 },
931 ChatMessage {
932 role: "system".to_string(),
933 content: Some(MessageContent::Text(
934 "WORKING_MEMORY:\n- task1 (ok): done".to_string(),
935 )),
936 tool_calls: None,
937 ..Default::default()
938 },
939 ChatMessage {
940 role: "user".to_string(),
941 content: Some(MessageContent::Text("Hello".to_string())),
942 tool_calls: None,
943 ..Default::default()
944 },
945 ];
946 let result = transform_messages_for_capabilities(messages, ModelCapabilities::empty());
947
948 assert_eq!(result.len(), 2);
949 assert_eq!(result[0].role, "system");
950 assert_eq!(
951 result[0].content.as_ref().and_then(|c| c.as_str()),
952 Some("You are a helpful assistant.\n\nWORKING_MEMORY:\n- task1 (ok): done")
953 );
954 assert_eq!(result[1].role, "user");
955 }
956
957 #[test]
958 fn test_merges_three_consecutive_system_messages() {
959 let messages = vec![
960 ChatMessage {
961 role: "system".to_string(),
962 content: Some(MessageContent::Text("First.".to_string())),
963 tool_calls: None,
964 ..Default::default()
965 },
966 ChatMessage {
967 role: "system".to_string(),
968 content: Some(MessageContent::Text("Second.".to_string())),
969 tool_calls: None,
970 ..Default::default()
971 },
972 ChatMessage {
973 role: "system".to_string(),
974 content: Some(MessageContent::Text("Third.".to_string())),
975 tool_calls: None,
976 ..Default::default()
977 },
978 ];
979 let result = transform_messages_for_capabilities(messages, ModelCapabilities::empty());
980
981 assert_eq!(result.len(), 1);
982 assert_eq!(
983 result[0].content.as_ref().and_then(|c| c.as_str()),
984 Some("First.\n\nSecond.\n\nThird.")
985 );
986 }
987
988 #[test]
989 fn test_handles_empty_system_content() {
990 let messages = vec![
991 ChatMessage {
992 role: "system".to_string(),
993 content: Some(MessageContent::Text(String::new())),
994 tool_calls: None,
995 ..Default::default()
996 },
997 ChatMessage {
998 role: "system".to_string(),
999 content: Some(MessageContent::Text("Actual content".to_string())),
1000 tool_calls: None,
1001 ..Default::default()
1002 },
1003 ];
1004 let result = transform_messages_for_capabilities(messages, ModelCapabilities::empty());
1005
1006 assert_eq!(result.len(), 1);
1007 assert_eq!(
1008 result[0].content.as_ref().and_then(|c| c.as_str()),
1009 Some("Actual content")
1010 );
1011 }
1012
1013 #[test]
1014 fn test_transform_system_to_user() {
1015 let messages = vec![
1016 ChatMessage {
1017 role: "system".to_string(),
1018 content: Some(MessageContent::Text("You are helpful".to_string())),
1019 tool_calls: None,
1020 ..Default::default()
1021 },
1022 ChatMessage {
1023 role: "user".to_string(),
1024 content: Some(MessageContent::Text("Hello".to_string())),
1025 tool_calls: None,
1026 ..Default::default()
1027 },
1028 ];
1029 let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1031 let result = transform_messages_for_capabilities(messages, caps);
1032 assert_eq!(result.len(), 1); assert_eq!(result[0].role, "user");
1035 let content_str = result[0].content.as_ref().and_then(|c| c.as_str()).unwrap();
1036 assert!(content_str.contains("[System]: You are helpful"));
1037 assert!(content_str.contains("Hello"));
1038 }
1039
1040 #[test]
1041 fn test_transform_preserves_system_when_supported() {
1042 let messages = vec![ChatMessage {
1043 role: "system".to_string(),
1044 content: Some(MessageContent::Text("You are helpful".to_string())),
1045 tool_calls: None,
1046 ..Default::default()
1047 }];
1048 let caps = ModelCapabilities::SUPPORTS_SYSTEM_ROLE;
1049 let result = transform_messages_for_capabilities(messages, caps);
1050 assert_eq!(result[0].role, "system");
1051 assert_eq!(
1052 result[0].content,
1053 Some(MessageContent::Text("You are helpful".to_string()))
1054 );
1055 }
1056
1057 #[test]
1058 fn test_transform_merges_consecutive_user_messages() {
1059 let messages = vec![
1060 ChatMessage {
1061 role: "user".to_string(),
1062 content: Some(MessageContent::Text("First".to_string())),
1063 tool_calls: None,
1064 ..Default::default()
1065 },
1066 ChatMessage {
1067 role: "user".to_string(),
1068 content: Some(MessageContent::Text("Second".to_string())),
1069 tool_calls: None,
1070 ..Default::default()
1071 },
1072 ];
1073 let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1074 let result = transform_messages_for_capabilities(messages, caps);
1075 assert_eq!(result.len(), 1);
1076 assert_eq!(
1077 result[0].content,
1078 Some(MessageContent::Text("First\n\nSecond".to_string()))
1079 );
1080 }
1081
1082 #[test]
1083 fn test_transform_does_not_merge_tool_messages() {
1084 let messages = vec![
1085 ChatMessage {
1086 role: "tool".to_string(),
1087 content: Some(MessageContent::Text("Result 1".to_string())),
1088 tool_calls: None,
1089 ..Default::default()
1090 },
1091 ChatMessage {
1092 role: "tool".to_string(),
1093 content: Some(MessageContent::Text("Result 2".to_string())),
1094 tool_calls: None,
1095 ..Default::default()
1096 },
1097 ];
1098 let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1099 let result = transform_messages_for_capabilities(messages, caps);
1100 assert_eq!(result.len(), 2); }
1102
1103 #[test]
1104 fn test_transform_combined_system_and_merge() {
1105 let messages = vec![
1106 ChatMessage {
1107 role: "system".to_string(),
1108 content: Some(MessageContent::Text("Be helpful".to_string())),
1109 tool_calls: None,
1110 ..Default::default()
1111 },
1112 ChatMessage {
1113 role: "user".to_string(),
1114 content: Some(MessageContent::Text("First".to_string())),
1115 tool_calls: None,
1116 ..Default::default()
1117 },
1118 ChatMessage {
1119 role: "user".to_string(),
1120 content: Some(MessageContent::Text("Second".to_string())),
1121 tool_calls: None,
1122 ..Default::default()
1123 },
1124 ];
1125 let caps = ModelCapabilities::REQUIRES_STRICT_TURNS; let result = transform_messages_for_capabilities(messages, caps);
1127 assert_eq!(result.len(), 1); assert_eq!(result[0].role, "user");
1129 let content_str = result[0].content.as_ref().and_then(|c| c.as_str()).unwrap();
1130 assert!(content_str.contains("[System]: Be helpful"));
1131 assert!(content_str.contains("First"));
1132 assert!(content_str.contains("Second"));
1133 }
1134
1135 #[test]
1136 fn test_merge_consecutive_assistant_with_tool_calls() {
1137 let tool_call_1 = serde_json::json!([
1140 {
1141 "id": "call_1",
1142 "type": "function",
1143 "function": {
1144 "name": "get_weather",
1145 "arguments": "{\"location\":\"Paris\"}"
1146 }
1147 }
1148 ]);
1149 let tool_call_2 = serde_json::json!([
1150 {
1151 "id": "call_2",
1152 "type": "function",
1153 "function": {
1154 "name": "get_time",
1155 "arguments": "{\"timezone\":\"UTC\"}"
1156 }
1157 }
1158 ]);
1159
1160 let messages = vec![
1161 ChatMessage {
1162 role: "user".to_string(),
1163 content: Some(MessageContent::Text("What's the weather?".to_string())),
1164 tool_calls: None,
1165 ..Default::default()
1166 },
1167 ChatMessage {
1168 role: "assistant".to_string(),
1169 content: Some(MessageContent::Text("Let me check...".to_string())),
1170 tool_calls: Some(tool_call_1),
1171 ..Default::default()
1172 },
1173 ChatMessage {
1174 role: "assistant".to_string(),
1175 content: Some(MessageContent::Text("And the time...".to_string())),
1176 tool_calls: Some(tool_call_2),
1177 ..Default::default()
1178 },
1179 ];
1180
1181 let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1182 let result = transform_messages_for_capabilities(messages, caps);
1183
1184 assert_eq!(result.len(), 2);
1186 assert_eq!(result[0].role, "user");
1187 assert_eq!(result[1].role, "assistant");
1188
1189 assert_eq!(
1191 result[1].content,
1192 Some(MessageContent::Text(
1193 "Let me check...\n\nAnd the time...".to_string()
1194 ))
1195 );
1196
1197 let merged_tool_calls = result[1].tool_calls.as_ref().unwrap();
1199 let tool_calls_array = merged_tool_calls.as_array().unwrap();
1200 assert_eq!(tool_calls_array.len(), 2);
1201 assert_eq!(tool_calls_array[0]["id"], "call_1");
1202 assert_eq!(tool_calls_array[1]["id"], "call_2");
1203 }
1204
1205 #[test]
1206 fn test_merge_assistant_messages_only_first_has_content() {
1207 let tool_call = serde_json::json!([
1209 {
1210 "id": "call_1",
1211 "type": "function",
1212 "function": {
1213 "name": "get_weather",
1214 "arguments": "{}"
1215 }
1216 }
1217 ]);
1218
1219 let messages = vec![
1220 ChatMessage {
1221 role: "assistant".to_string(),
1222 content: Some(MessageContent::Text("Let me check...".to_string())),
1223 tool_calls: None,
1224 ..Default::default()
1225 },
1226 ChatMessage {
1227 role: "assistant".to_string(),
1228 content: None,
1229 tool_calls: Some(tool_call),
1230 ..Default::default()
1231 },
1232 ];
1233
1234 let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1235 let result = transform_messages_for_capabilities(messages, caps);
1236
1237 assert_eq!(result.len(), 1);
1238 assert_eq!(
1239 result[0].content,
1240 Some(MessageContent::Text("Let me check...".to_string()))
1241 );
1242 assert!(result[0].tool_calls.is_some());
1243 }
1244
1245 #[test]
1246 fn test_merge_assistant_messages_only_second_has_content() {
1247 let tool_call = serde_json::json!([
1249 {
1250 "id": "call_1",
1251 "type": "function",
1252 "function": {
1253 "name": "get_weather",
1254 "arguments": "{}"
1255 }
1256 }
1257 ]);
1258
1259 let messages = vec![
1260 ChatMessage {
1261 role: "assistant".to_string(),
1262 content: None,
1263 tool_calls: Some(tool_call),
1264 ..Default::default()
1265 },
1266 ChatMessage {
1267 role: "assistant".to_string(),
1268 content: Some(MessageContent::Text("Result received".to_string())),
1269 tool_calls: None,
1270 ..Default::default()
1271 },
1272 ];
1273
1274 let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1275 let result = transform_messages_for_capabilities(messages, caps);
1276
1277 assert_eq!(result.len(), 1);
1278 assert_eq!(
1279 result[0].content,
1280 Some(MessageContent::Text("Result received".to_string()))
1281 );
1282 assert!(result[0].tool_calls.is_some());
1283 }
1284
1285 #[test]
1286 fn test_merge_assistant_messages_neither_has_content() {
1287 let tool_call_1 = serde_json::json!([
1289 {
1290 "id": "call_1",
1291 "type": "function",
1292 "function": {"name": "tool1", "arguments": "{}"}
1293 }
1294 ]);
1295 let tool_call_2 = serde_json::json!([
1296 {
1297 "id": "call_2",
1298 "type": "function",
1299 "function": {"name": "tool2", "arguments": "{}"}
1300 }
1301 ]);
1302
1303 let messages = vec![
1304 ChatMessage {
1305 role: "assistant".to_string(),
1306 content: None,
1307 tool_calls: Some(tool_call_1),
1308 ..Default::default()
1309 },
1310 ChatMessage {
1311 role: "assistant".to_string(),
1312 content: None,
1313 tool_calls: Some(tool_call_2),
1314 ..Default::default()
1315 },
1316 ];
1317
1318 let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1319 let result = transform_messages_for_capabilities(messages, caps);
1320
1321 assert_eq!(result.len(), 1);
1322 assert!(result[0].content.is_none());
1323
1324 let merged_tool_calls = result[0].tool_calls.as_ref().unwrap();
1325 let tool_calls_array = merged_tool_calls.as_array().unwrap();
1326 assert_eq!(tool_calls_array.len(), 2);
1327 }
1328
1329 #[test]
1330 fn test_no_merge_without_strict_turns_capability() {
1331 let messages = vec![
1333 ChatMessage {
1334 role: "assistant".to_string(),
1335 content: Some(MessageContent::Text("First".to_string())),
1336 tool_calls: None,
1337 ..Default::default()
1338 },
1339 ChatMessage {
1340 role: "assistant".to_string(),
1341 content: Some(MessageContent::Text("Second".to_string())),
1342 tool_calls: None,
1343 ..Default::default()
1344 },
1345 ];
1346
1347 let caps = ModelCapabilities::empty();
1348 let result = transform_messages_for_capabilities(messages, caps);
1349
1350 assert_eq!(result.len(), 2);
1352 }
1353
1354 #[test]
1355 fn test_merge_preserves_different_role_boundaries() {
1356 let tool_call = serde_json::json!([
1358 {
1359 "id": "call_1",
1360 "type": "function",
1361 "function": {"name": "tool1", "arguments": "{}"}
1362 }
1363 ]);
1364
1365 let messages = vec![
1366 ChatMessage {
1367 role: "user".to_string(),
1368 content: Some(MessageContent::Text("Question".to_string())),
1369 tool_calls: None,
1370 ..Default::default()
1371 },
1372 ChatMessage {
1373 role: "assistant".to_string(),
1374 content: Some(MessageContent::Text("Answer".to_string())),
1375 tool_calls: Some(tool_call),
1376 ..Default::default()
1377 },
1378 ChatMessage {
1379 role: "user".to_string(),
1380 content: Some(MessageContent::Text("Follow-up".to_string())),
1381 tool_calls: None,
1382 ..Default::default()
1383 },
1384 ];
1385
1386 let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1387 let result = transform_messages_for_capabilities(messages, caps);
1388
1389 assert_eq!(result.len(), 3);
1391 assert_eq!(result[0].role, "user");
1392 assert_eq!(result[1].role, "assistant");
1393 assert_eq!(result[2].role, "user");
1394 }
1395
1396 #[test]
1399 fn test_array_content_deserializes_to_parts() {
1400 let json = serde_json::json!({
1401 "role": "user",
1402 "content": [
1403 {"type": "text", "text": "What is in this image?"},
1404 {"type": "image_url", "image_url": {"url": "https://example.com/img.png"}}
1405 ]
1406 });
1407 let msg: ChatMessage = serde_json::from_value(json).unwrap();
1408 assert!(matches!(msg.content, Some(MessageContent::Parts(_))));
1409 let re_serialised = serde_json::to_value(&msg).unwrap();
1411 assert!(re_serialised["content"].is_array());
1412 }
1413
1414 #[test]
1417 fn test_merge_array_content_with_text_content() {
1418 let messages = vec![
1419 ChatMessage {
1420 role: "user".to_string(),
1421 content: Some(MessageContent::Parts(vec![
1422 serde_json::json!({"type": "text", "text": "Look at this:"}),
1423 serde_json::json!({"type": "image_url", "image_url": {"url": "https://example.com/img.png"}}),
1424 ])),
1425 tool_calls: None,
1426 ..Default::default()
1427 },
1428 ChatMessage {
1429 role: "user".to_string(),
1430 content: Some(MessageContent::Text("What do you see?".to_string())),
1431 tool_calls: None,
1432 ..Default::default()
1433 },
1434 ];
1435 let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1436 let result = transform_messages_for_capabilities(messages, caps);
1437
1438 assert_eq!(result.len(), 1);
1439 assert!(matches!(&result[0].content, Some(MessageContent::Parts(p)) if p.len() == 3));
1441 }
1442
1443 #[test]
1446 fn test_tool_message_with_array_content_passes_through() {
1447 let messages = vec![
1448 ChatMessage {
1449 role: "user".to_string(),
1450 content: Some(MessageContent::Text("Run the tool".to_string())),
1451 tool_calls: None,
1452 ..Default::default()
1453 },
1454 ChatMessage {
1455 role: "tool".to_string(),
1456 content: Some(MessageContent::Parts(vec![
1457 serde_json::json!({"type": "text", "text": "tool result here"}),
1458 ])),
1459 tool_calls: None,
1460 ..Default::default()
1461 },
1462 ChatMessage {
1463 role: "user".to_string(),
1464 content: Some(MessageContent::Text("Thanks".to_string())),
1465 tool_calls: None,
1466 ..Default::default()
1467 },
1468 ];
1469 let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1470 let result = transform_messages_for_capabilities(messages, caps);
1471
1472 assert_eq!(result.len(), 3);
1475 assert_eq!(result[1].role, "tool");
1476 assert!(matches!(&result[1].content, Some(MessageContent::Parts(_))));
1477 }
1478}
1479
1480#[cfg(test)]
1488mod round_trip_tests {
1489 use super::*;
1490 use serde_json::{Value, json};
1491
1492 fn round_trip(messages: Value, capabilities: ModelCapabilities) -> Vec<Value> {
1495 let parsed: Vec<ChatMessage> = serde_json::from_value(messages).expect("valid messages");
1496 let transformed = transform_messages_for_capabilities(parsed, capabilities);
1497 serde_json::to_value(&transformed)
1498 .expect("serializable")
1499 .as_array()
1500 .expect("array")
1501 .clone()
1502 }
1503
1504 #[test]
1511 fn tool_call_id_survives_strict_turn_coalescing() {
1512 let out = round_trip(
1513 json!([
1514 {"role": "user", "content": "run it"},
1515 {"role": "assistant", "tool_calls": [
1516 {"id": "call_1", "type": "function",
1517 "function": {"name": "f", "arguments": "{}"}}
1518 ]},
1519 {"role": "tool", "tool_call_id": "call_1", "content": "result"},
1520 ]),
1521 ModelCapabilities::REQUIRES_STRICT_TURNS,
1522 );
1523
1524 let tool = out
1525 .iter()
1526 .find(|m| m["role"] == "tool")
1527 .expect("tool message");
1528 assert_eq!(
1529 tool["tool_call_id"], "call_1",
1530 "tool_call_id must survive the transform"
1531 );
1532 }
1533
1534 #[test]
1536 fn arbitrary_unknown_message_keys_are_preserved() {
1537 let out = round_trip(
1538 json!([{
1539 "role": "user",
1540 "content": "hi",
1541 "name": "alice",
1542 "x_vendor_extension": {"nested": [1, 2]},
1543 }]),
1544 ModelCapabilities::REQUIRES_STRICT_TURNS,
1545 );
1546
1547 assert_eq!(out[0]["name"], "alice");
1548 assert_eq!(out[0]["x_vendor_extension"], json!({"nested": [1, 2]}));
1549 }
1550
1551 #[test]
1556 fn flatten_does_not_disturb_untagged_content_forms() {
1557 let out = round_trip(
1558 json!([
1559 {"role": "user", "content": "plain string"},
1560 {"role": "assistant", "content": [{"type": "text", "text": "part"}]},
1561 ]),
1562 ModelCapabilities::empty(),
1563 );
1564
1565 assert_eq!(out[0]["content"], json!("plain string"));
1566 assert_eq!(out[1]["content"], json!([{"type": "text", "text": "part"}]));
1567 }
1568
1569 #[test]
1571 fn merging_prefers_the_surviving_message_extras() {
1572 let out = round_trip(
1573 json!([
1574 {"role": "user", "content": "one", "name": "first"},
1575 {"role": "user", "content": "two", "name": "second", "only_on_later": true},
1576 ]),
1577 ModelCapabilities::REQUIRES_STRICT_TURNS,
1578 );
1579
1580 assert_eq!(out.len(), 1, "consecutive user messages merge");
1581 assert_eq!(out[0]["name"], "first", "the surviving message's key wins");
1582 assert_eq!(out[0]["only_on_later"], true, "absent keys are adopted");
1583 }
1584
1585 #[test]
1587 fn messages_without_extras_serialize_unchanged() {
1588 let out = round_trip(
1589 json!([{"role": "user", "content": "hi"}]),
1590 ModelCapabilities::empty(),
1591 );
1592 assert_eq!(out[0], json!({"role": "user", "content": "hi"}));
1593 }
1594}