Skip to main content

tool_parser/parsers/
minimax_m2.rs

1use std::{collections::HashMap, fmt::Write as FmtWrite};
2
3use async_trait::async_trait;
4use openai_protocol::common::Tool;
5use regex::Regex;
6use serde_json::Value;
7
8use crate::{
9    errors::{ParserError, ParserResult},
10    parsers::helpers,
11    traits::ToolParser,
12    types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
13};
14
15/// MiniMax M2 format parser for tool calls
16///
17/// Handles the MiniMax M2 specific format:
18/// `<minimax:tool_call><invoke name="func"><parameter name="key">value</parameter></invoke></minimax:tool_call>`
19///
20/// Features:
21/// - Namespaced XML tags (`minimax:tool_call`)
22/// - Function wrapped in `<invoke name="...">` tags
23/// - Parameters as `<parameter name="key">value</parameter>`
24/// - Incremental JSON streaming for parameters
25///
26/// Reference: https://huggingface.co/MiniMaxAI/MiniMax-M2?chat_template=default
27pub struct MinimaxM2Parser {
28    // Regex patterns
29    tool_call_extractor: Regex,
30    invoke_extractor: Regex,
31    param_extractor: Regex,
32
33    // Streaming state
34    buffer: String,
35    prev_tool_call_arr: Vec<Value>,
36    current_tool_id: i32,
37    streamed_args_for_tool: Vec<String>,
38    current_function_name: String,
39    current_parameters: HashMap<String, Value>,
40    in_tool_call: bool,
41    function_name_sent: bool,
42    waiting_for_tool_call_end: bool,
43
44    // Token configuration
45    tool_call_start_token: &'static str,
46    tool_call_end_token: &'static str,
47    invoke_end_token: &'static str,
48}
49
50impl MinimaxM2Parser {
51    /// Parse a value from string with consistent logic
52    #[inline]
53    fn parse_value(text: &str) -> Value {
54        // Try parsing as common literals first
55        match text {
56            "true" | "True" => return Value::Bool(true),
57            "false" | "False" => return Value::Bool(false),
58            "null" | "None" => return Value::Null,
59            _ => {}
60        }
61
62        // Try parsing as number
63        if let Ok(num) = text.parse::<i64>() {
64            return Value::Number(num.into());
65        }
66
67        if let Ok(num) = text.parse::<f64>() {
68            if let Some(n) = serde_json::Number::from_f64(num) {
69                return Value::Number(n);
70            }
71        }
72
73        // Default to string
74        Value::String(text.to_string())
75    }
76
77    /// Create a new MiniMax M2 parser
78    #[expect(
79        clippy::expect_used,
80        reason = "regex patterns are compile-time string literals"
81    )]
82    pub fn new() -> Self {
83        // Use (?s) flag for DOTALL mode to handle newlines
84        let tool_call_pattern = r"(?s)<minimax:tool_call>.*?</minimax:tool_call>";
85        let tool_call_extractor = Regex::new(tool_call_pattern).expect("Valid regex pattern");
86
87        let invoke_pattern = r#"(?s)<invoke\s+name="([^"]+)">(.*?)</invoke>"#;
88        let invoke_extractor = Regex::new(invoke_pattern).expect("Valid regex pattern");
89
90        let param_pattern = r#"(?s)<parameter\s+name="([^"]+)">(.*?)</parameter>"#;
91        let param_extractor = Regex::new(param_pattern).expect("Valid regex pattern");
92
93        Self {
94            tool_call_extractor,
95            invoke_extractor,
96            param_extractor,
97            buffer: String::new(),
98            prev_tool_call_arr: Vec::new(),
99            current_tool_id: -1,
100            streamed_args_for_tool: Vec::new(),
101            current_function_name: String::new(),
102            current_parameters: HashMap::new(),
103            in_tool_call: false,
104            function_name_sent: false,
105            waiting_for_tool_call_end: false,
106            tool_call_start_token: "<minimax:tool_call>",
107            tool_call_end_token: "</minimax:tool_call>",
108            invoke_end_token: "</invoke>",
109        }
110    }
111
112    /// Parse parameter tags, coercing each value by its declared schema type when
113    /// known (so a numeric-looking `string` stays a string), else inferring.
114    fn parse_parameters(
115        &self,
116        params_text: &str,
117        param_types: &HashMap<String, String>,
118    ) -> serde_json::Map<String, Value> {
119        let mut parameters = serde_json::Map::new();
120
121        for capture in self.param_extractor.captures_iter(params_text) {
122            let key = capture.get(1).map_or("", |m| m.as_str()).trim();
123            let value_str = capture.get(2).map_or("", |m| m.as_str());
124
125            let decoded_value = Self::decode_xml_entities(value_str);
126            let value = helpers::coerce_by_schema_type(
127                &decoded_value,
128                param_types.get(key).map(String::as_str),
129            )
130            .unwrap_or_else(|| Self::parse_value(&decoded_value));
131
132            parameters.insert(key.to_string(), value);
133        }
134
135        parameters
136    }
137
138    /// Decode common XML entities
139    fn decode_xml_entities(text: &str) -> String {
140        text.replace("&lt;", "<")
141            .replace("&gt;", ">")
142            .replace("&amp;", "&")
143            .replace("&quot;", "\"")
144            .replace("&apos;", "'")
145    }
146
147    /// Parse a single tool call block
148    fn parse_tool_call(&self, block: &str, tools: &[Tool]) -> ParserResult<Option<ToolCall>> {
149        if let Some(captures) = self.invoke_extractor.captures(block) {
150            // Get function name from invoke tag attribute
151            let func_name = captures.get(1).map_or("", |m| m.as_str()).trim();
152
153            // Get parameters text
154            let params_text = captures.get(2).map_or("", |m| m.as_str());
155
156            // Parse parameters, coerced by this function's declared schema.
157            let param_types = helpers::param_types_for_function(tools, func_name);
158            let parameters = self.parse_parameters(params_text, &param_types);
159
160            let arguments_str = serde_json::to_string(&parameters)
161                .map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
162
163            Ok(Some(ToolCall {
164                function: FunctionCall {
165                    name: func_name.to_string(),
166                    arguments: arguments_str,
167                },
168            }))
169        } else {
170            Ok(None)
171        }
172    }
173
174    /// Parse all tool calls from text and return first valid position
175    fn parse_tool_calls_from_text(
176        &self,
177        text: &str,
178        tools: &[Tool],
179    ) -> (Vec<ToolCall>, Option<usize>) {
180        let mut tool_calls = Vec::new();
181        let mut first_valid_pos = None;
182
183        for mat in self.tool_call_extractor.find_iter(text) {
184            match self.parse_tool_call(mat.as_str(), tools) {
185                Ok(Some(tool)) => {
186                    if first_valid_pos.is_none() {
187                        first_valid_pos = Some(mat.start());
188                    }
189                    tool_calls.push(tool);
190                }
191                Ok(None) => continue,
192                Err(e) => {
193                    tracing::debug!("Failed to parse tool call: {}", e);
194                    continue;
195                }
196            }
197        }
198
199        (tool_calls, first_valid_pos)
200    }
201
202    /// Shared non-streaming parse; `tools` empty means infer types from text.
203    fn parse_complete_inner(&self, text: &str, tools: &[Tool]) -> (String, Vec<ToolCall>) {
204        if !self.has_tool_markers(text) {
205            return (text.to_string(), vec![]);
206        }
207        let (tool_calls, first_valid_tool_pos) = self.parse_tool_calls_from_text(text, tools);
208        if tool_calls.is_empty() {
209            return (text.to_string(), vec![]);
210        }
211        let normal_text = match first_valid_tool_pos {
212            Some(pos) => text[..pos].to_string(),
213            None => text.to_string(),
214        };
215        (normal_text, tool_calls)
216    }
217
218    /// Parse and stream parameters incrementally
219    fn parse_and_stream_parameters(&mut self, text: &str, tools: &[Tool]) -> Vec<ToolCallItem> {
220        let mut calls = Vec::new();
221        let param_types = helpers::param_types_for_function(tools, &self.current_function_name);
222
223        // Find all complete parameter patterns in the buffer
224        let param_matches: Vec<_> = self
225            .param_extractor
226            .captures_iter(text)
227            .map(|cap| {
228                let name = cap.get(1).map_or("", |m| m.as_str()).trim().to_string();
229                let value_str = cap.get(2).map_or("", |m| m.as_str());
230                let decoded = Self::decode_xml_entities(value_str);
231
232                // Coerce by declared type when known; otherwise keep the prior
233                // JSON-first-then-infer behavior for nested objects/arrays.
234                let value = helpers::coerce_by_schema_type(
235                    &decoded,
236                    param_types.get(&name).map(String::as_str),
237                )
238                .unwrap_or_else(|| {
239                    if decoded.starts_with('{') || decoded.starts_with('[') {
240                        serde_json::from_str::<Value>(&decoded)
241                            .unwrap_or_else(|_| Self::parse_value(&decoded))
242                    } else {
243                        Self::parse_value(&decoded)
244                    }
245                });
246
247                (name, value)
248            })
249            .collect();
250
251        // Build new parameters map
252        let mut new_params = HashMap::new();
253        for (name, value) in param_matches {
254            new_params.insert(name, value);
255        }
256
257        // If we have new parameters that weren't in current_parameters, stream them
258        if !new_params.is_empty() && new_params != self.current_parameters {
259            let tool_id = self.current_tool_id as usize;
260
261            // Ensure we have enough capacity
262            while self.streamed_args_for_tool.len() <= tool_id {
263                self.streamed_args_for_tool.push(String::new());
264            }
265
266            // Build incremental JSON with single allocation
267            if self.current_parameters.is_empty() {
268                // First parameters - start JSON object but don't close it
269                let mut json_fragment = String::with_capacity(256);
270                json_fragment.push('{');
271
272                let mut first = true;
273                for (key, value) in &new_params {
274                    if !first {
275                        json_fragment.push_str(", ");
276                    }
277                    // serde_json::to_string for String/Value is infallible; write! to String is infallible
278                    let key_json = serde_json::to_string(key).unwrap_or_default();
279                    let value_json = serde_json::to_string(value).unwrap_or_default();
280                    let _ = write!(&mut json_fragment, "{key_json}: {value_json}");
281                    first = false;
282                }
283
284                calls.push(ToolCallItem {
285                    tool_index: tool_id,
286                    name: None,
287                    parameters: json_fragment.clone(),
288                });
289
290                self.streamed_args_for_tool[tool_id] = json_fragment;
291            } else {
292                // Additional parameters - add them incrementally
293                let new_keys: Vec<_> = new_params
294                    .keys()
295                    .filter(|k| !self.current_parameters.contains_key(*k))
296                    .collect();
297
298                if !new_keys.is_empty() {
299                    let mut json_fragment = String::with_capacity(128);
300
301                    for key in new_keys {
302                        let value = &new_params[key];
303                        // serde_json::to_string for String/Value is infallible; write! to String is infallible
304                        let key_json = serde_json::to_string(key).unwrap_or_default();
305                        let value_json = serde_json::to_string(value).unwrap_or_default();
306                        let _ = write!(&mut json_fragment, ", {key_json}: {value_json}");
307                    }
308
309                    calls.push(ToolCallItem {
310                        tool_index: tool_id,
311                        name: None,
312                        parameters: json_fragment.clone(),
313                    });
314
315                    self.streamed_args_for_tool[tool_id].push_str(&json_fragment);
316                }
317            }
318
319            // Update current parameters
320            self.current_parameters = new_params;
321
322            // Update prev_tool_call_arr
323            while self.prev_tool_call_arr.len() <= tool_id {
324                self.prev_tool_call_arr.push(Value::Null);
325            }
326            self.prev_tool_call_arr[tool_id] = serde_json::json!({
327                "name": self.current_function_name,
328                "arguments": self.current_parameters,
329            });
330        }
331
332        calls
333    }
334}
335
336impl Default for MinimaxM2Parser {
337    fn default() -> Self {
338        Self::new()
339    }
340}
341
342#[async_trait]
343impl ToolParser for MinimaxM2Parser {
344    async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
345        Ok(self.parse_complete_inner(text, &[]))
346    }
347
348    async fn parse_complete_with_tools(
349        &self,
350        text: &str,
351        tools: &[Tool],
352    ) -> ParserResult<(String, Vec<ToolCall>)> {
353        Ok(self.parse_complete_inner(text, tools))
354    }
355
356    async fn parse_incremental(
357        &mut self,
358        chunk: &str,
359        tools: &[Tool],
360    ) -> ParserResult<StreamingParseResult> {
361        self.buffer.push_str(chunk);
362        let mut normal_text = String::new();
363        let mut calls = Vec::new();
364
365        loop {
366            // If we're waiting for the tool call end tag, check for it first
367            if self.waiting_for_tool_call_end {
368                if let Some(end_pos) = self.buffer.find(self.tool_call_end_token) {
369                    // Complete tool call found
370                    self.buffer =
371                        self.buffer[end_pos + self.tool_call_end_token.len()..].to_string();
372                    self.in_tool_call = false;
373                    self.waiting_for_tool_call_end = false;
374                    self.function_name_sent = false;
375                    self.current_function_name.clear();
376                    self.current_parameters.clear();
377                    self.current_tool_id += 1;
378                    continue;
379                } else {
380                    // End tag not complete yet, wait for more text
381                    break;
382                }
383            }
384
385            // If we're not in a tool call and don't see a start token, return normal text
386            if !self.in_tool_call && !self.buffer.contains(self.tool_call_start_token) {
387                // Check if buffer might contain a partial start token at the end
388                if let Some(partial_len) =
389                    helpers::ends_with_partial_token(&self.buffer, self.tool_call_start_token)
390                {
391                    // Return everything except the potential partial token
392                    let end = self.buffer.len() - partial_len;
393                    normal_text = self.buffer[..end].to_string();
394                    self.buffer = self.buffer[end..].to_string();
395                } else {
396                    // No partial token, return all as normal text
397                    normal_text.clone_from(&self.buffer);
398                    self.buffer.clear();
399                }
400                break;
401            }
402
403            // Look for tool call start
404            if !self.in_tool_call {
405                if let Some(start) = self.buffer.find(self.tool_call_start_token) {
406                    normal_text = self.buffer[..start].to_string();
407                    self.buffer =
408                        self.buffer[start + self.tool_call_start_token.len()..].to_string();
409
410                    self.in_tool_call = true;
411                    self.function_name_sent = false;
412                    self.current_function_name.clear();
413                    self.current_parameters.clear();
414
415                    continue;
416                } else {
417                    // No start token found
418                    break;
419                }
420            }
421
422            // We're in a tool call, try to parse function name if not sent yet
423            if !self.function_name_sent {
424                // Use regex to extract function name from <invoke name="..."> pattern
425                // Check if we have enough text to match the invoke pattern
426                if let Some(captures) = self.invoke_extractor.captures(&self.buffer) {
427                    let function_name = captures
428                        .get(1)
429                        .map_or("", |m| m.as_str())
430                        .trim()
431                        .to_string();
432
433                    // Forward unknown tool names too — emit a tool_call rather than
434                    // leaking the <invoke> markup into assistant text.
435                    self.current_function_name.clone_from(&function_name);
436                    self.function_name_sent = true;
437
438                    // Initialize tool call tracking
439                    if self.current_tool_id == -1 {
440                        self.current_tool_id = 0;
441                    }
442
443                    // Ensure tracking arrays are large enough
444                    helpers::ensure_capacity(
445                        self.current_tool_id,
446                        &mut self.prev_tool_call_arr,
447                        &mut self.streamed_args_for_tool,
448                    );
449
450                    // Send tool name with empty parameters
451                    calls.push(ToolCallItem {
452                        tool_index: self.current_tool_id as usize,
453                        name: Some(function_name),
454                        parameters: String::new(),
455                    });
456
457                    // Find the position after the opening invoke tag (after the >)
458                    // We only want to remove up to the opening tag, not the full match
459                    if let Some(pos) = self.buffer.find('>') {
460                        self.buffer = self.buffer[pos + 1..].to_string();
461                    }
462                    continue;
463                }
464                // No complete invoke pattern found yet, wait for more text
465                break;
466            }
467
468            // Parse parameters incrementally
469            if self.function_name_sent {
470                // Process parameters and get any calls to emit
471                let buffer_copy = self.buffer.clone(); // TODO: avoid cloning the buffer
472                let parameter_calls = self.parse_and_stream_parameters(&buffer_copy, tools);
473                calls.extend(parameter_calls);
474
475                // Check if tool call is complete (</invoke> found)
476                if let Some(invoke_end) = self.buffer.find(self.invoke_end_token) {
477                    // Add closing brace to complete the JSON object
478                    let tool_id = self.current_tool_id as usize;
479                    if tool_id < self.streamed_args_for_tool.len() {
480                        let current_streamed = &self.streamed_args_for_tool[tool_id];
481                        if !current_streamed.is_empty() && !current_streamed.ends_with('}') {
482                            // Count opening and closing braces to check if JSON is complete
483                            let open_braces = current_streamed.matches('{').count();
484                            let close_braces = current_streamed.matches('}').count();
485                            if open_braces > close_braces {
486                                calls.push(ToolCallItem {
487                                    tool_index: tool_id,
488                                    name: None,
489                                    parameters: "}".to_string(),
490                                });
491                                self.streamed_args_for_tool[tool_id].push('}');
492                            }
493                        }
494                    }
495
496                    // Move buffer past the </invoke>
497                    self.buffer =
498                        self.buffer[invoke_end + self.invoke_end_token.len()..].to_string();
499
500                    // Check if we have the closing </minimax:tool_call>
501                    if let Some(end_pos) = self.buffer.find(self.tool_call_end_token) {
502                        // Complete tool call found
503                        self.buffer =
504                            self.buffer[end_pos + self.tool_call_end_token.len()..].to_string();
505                        self.in_tool_call = false;
506                        self.function_name_sent = false;
507                        self.current_function_name.clear();
508                        self.current_parameters.clear();
509                        self.current_tool_id += 1;
510                        continue;
511                    } else {
512                        // End tag not complete yet, mark that we're waiting for it
513                        self.waiting_for_tool_call_end = true;
514                        break;
515                    }
516                }
517                // Tool call not complete yet, wait for more text
518                break;
519            }
520        }
521
522        Ok(StreamingParseResult { normal_text, calls })
523    }
524
525    fn has_tool_markers(&self, text: &str) -> bool {
526        text.contains(self.tool_call_start_token)
527    }
528
529    fn get_unstreamed_tool_args(&self) -> Option<Vec<ToolCallItem>> {
530        helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
531    }
532
533    fn reset(&mut self) {
534        self.buffer.clear();
535        self.prev_tool_call_arr.clear();
536        self.current_tool_id = -1;
537        self.streamed_args_for_tool.clear();
538        self.current_function_name.clear();
539        self.current_parameters.clear();
540        self.in_tool_call = false;
541        self.function_name_sent = false;
542        self.waiting_for_tool_call_end = false;
543    }
544}