tool_parser/parsers/
glm4_moe.rs

1use std::collections::HashMap;
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/// GLM-4 MoE format parser for tool calls
16///
17/// Handles both GLM-4 MoE and GLM-4.7 MoE formats:
18/// - GLM-4: `<tool_call>{name}\n<arg_key>{key}</arg_key>\n<arg_value>{value}</arg_value>\n</tool_call>`
19/// - GLM-4.7: `<tool_call>{name}<arg_key>{key}</arg_key><arg_value>{value}</arg_value></tool_call>`
20///
21/// Features:
22/// - XML-style tags for tool calls
23/// - Key-value pairs for arguments
24/// - Support for multiple sequential tool calls
25pub struct Glm4MoeParser {
26    /// Regex for extracting complete tool calls
27    tool_call_extractor: Regex,
28    /// Regex for extracting function details
29    func_detail_extractor: Regex,
30    /// Regex for extracting argument key-value pairs
31    arg_extractor: Regex,
32
33    /// Buffer for accumulating incomplete patterns across chunks
34    buffer: String,
35
36    /// Stores complete tool call info (name and arguments) for each tool being parsed
37    prev_tool_call_arr: Vec<Value>,
38
39    /// Index of currently streaming tool call (-1 means no active tool)
40    current_tool_id: i32,
41
42    /// Tracks raw JSON string content streamed to client for each tool's arguments
43    streamed_args_for_tool: Vec<String>,
44
45    /// Token configuration
46    bot_token: &'static str,
47    eot_token: &'static str,
48}
49
50impl Glm4MoeParser {
51    /// Create a new generic GLM MoE parser with a custom func_detail_extractor pattern
52    ///
53    /// # Arguments
54    /// - `func_detail_pattern`: Regex pattern for extracting function name and arguments
55    ///   - For GLM-4: `r"(?s)<tool_call>([^\n]*)\n(.*)</tool_call>"`
56    ///   - For GLM-4.7: `r"(?s)<tool_call>\s*([^<\s]+)\s*(.*?)</tool_call>"`
57    #[expect(
58        clippy::expect_used,
59        reason = "regex patterns are compile-time string literals"
60    )]
61    pub(crate) fn new(func_detail_pattern: &str) -> Self {
62        // Use (?s) flag for DOTALL mode to handle newlines
63        let tool_call_pattern = r"(?s)<tool_call>.*?</tool_call>";
64        let tool_call_extractor = Regex::new(tool_call_pattern).expect("Valid regex pattern");
65
66        let func_detail_extractor = Regex::new(func_detail_pattern).expect("Valid regex pattern");
67
68        let arg_pattern = r"(?s)<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)</arg_value>";
69        let arg_extractor = Regex::new(arg_pattern).expect("Valid regex pattern");
70
71        Self {
72            tool_call_extractor,
73            func_detail_extractor,
74            arg_extractor,
75            buffer: String::new(),
76            prev_tool_call_arr: Vec::new(),
77            current_tool_id: -1,
78            streamed_args_for_tool: Vec::new(),
79            bot_token: "<tool_call>",
80            eot_token: "</tool_call>",
81        }
82    }
83
84    /// Create a new GLM-4.5/4.6 MoE parser (with newline-based format)
85    pub fn glm45() -> Self {
86        Self::new(r"(?s)<tool_call>([^\n]*)\n(.*)</tool_call>")
87    }
88
89    /// Create a new GLM-4.7 MoE parser (with whitespace-based format)
90    pub fn glm47() -> Self {
91        Self::new(r"(?s)<tool_call>\s*([^<\s]+)\s*(.*?)</tool_call>")
92    }
93
94    /// Parse arguments, coercing each value by its declared schema type and
95    /// falling back to [`infer_value`] when the type is unknown.
96    fn parse_arguments(
97        &self,
98        args_text: &str,
99        param_types: &HashMap<String, String>,
100    ) -> serde_json::Map<String, Value> {
101        let mut arguments = serde_json::Map::new();
102
103        for capture in self.arg_extractor.captures_iter(args_text) {
104            let key = capture.get(1).map_or("", |m| m.as_str()).trim();
105            let value_str = capture.get(2).map_or("", |m| m.as_str()).trim();
106
107            let value =
108                helpers::coerce_by_schema_type(value_str, param_types.get(key).map(String::as_str))
109                    .unwrap_or_else(|| infer_value(value_str));
110
111            arguments.insert(key.to_string(), value);
112        }
113
114        arguments
115    }
116
117    /// Parse a single tool call block
118    fn parse_tool_call(&self, block: &str, tools: &[Tool]) -> ParserResult<Option<ToolCall>> {
119        if let Some(captures) = self.func_detail_extractor.captures(block) {
120            // Get function name
121            let func_name = captures.get(1).map_or("", |m| m.as_str()).trim();
122
123            // Get arguments text
124            let args_text = captures.get(2).map_or("", |m| m.as_str());
125
126            let param_types = helpers::param_types_for_function(tools, func_name);
127            let arguments = self.parse_arguments(args_text, &param_types);
128
129            let arguments_str = serde_json::to_string(&arguments)
130                .map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
131
132            Ok(Some(ToolCall {
133                function: FunctionCall {
134                    name: func_name.to_string(),
135                    arguments: arguments_str,
136                },
137            }))
138        } else {
139            Ok(None)
140        }
141    }
142
143    /// Parse all tool calls from text (shared logic for complete and incremental parsing)
144    fn parse_tool_calls_from_text(&self, text: &str, tools: &[Tool]) -> Vec<ToolCall> {
145        let mut parsed = Vec::new();
146
147        for mat in self.tool_call_extractor.find_iter(text) {
148            match self.parse_tool_call(mat.as_str(), tools) {
149                Ok(Some(tool)) => parsed.push(tool),
150                Ok(None) => continue,
151                Err(e) => {
152                    tracing::debug!("Failed to parse tool call: {}", e);
153                    continue;
154                }
155            }
156        }
157
158        parsed
159    }
160}
161
162impl Glm4MoeParser {
163    /// Shared non-streaming parse, schema-aware when `tools` are provided.
164    fn parse_complete_inner(
165        &self,
166        text: &str,
167        tools: &[Tool],
168    ) -> ParserResult<(String, Vec<ToolCall>)> {
169        if !self.has_tool_markers(text) {
170            return Ok((text.to_string(), vec![]));
171        }
172
173        // Find where tool calls begin
174        // Safe: has_tool_markers() already confirmed the marker exists
175        let idx = text
176            .find("<tool_call>")
177            .ok_or_else(|| ParserError::ParsingFailed("tool call marker not found".to_string()))?;
178        let normal_text = text[..idx].to_string();
179
180        let parsed = self.parse_tool_calls_from_text(text, tools);
181
182        // If no tools were successfully parsed despite having markers, return entire text as fallback
183        if parsed.is_empty() {
184            return Ok((text.to_string(), vec![]));
185        }
186
187        Ok((normal_text, parsed))
188    }
189}
190
191/// Infer a JSON value from raw text when the schema type is unknown: JSON
192/// (numbers/bools/null/objects/arrays), then Python-style literals, then string.
193fn infer_value(value_str: &str) -> Value {
194    if let Ok(json_val) = serde_json::from_str::<Value>(value_str) {
195        return json_val;
196    }
197    match value_str {
198        "true" | "True" => Value::Bool(true),
199        "false" | "False" => Value::Bool(false),
200        "null" | "None" => Value::Null,
201        _ => {
202            if let Ok(num) = value_str.parse::<i64>() {
203                Value::Number(num.into())
204            } else if let Ok(num) = value_str.parse::<f64>() {
205                serde_json::Number::from_f64(num)
206                    .map_or_else(|| Value::String(value_str.to_string()), Value::Number)
207            } else {
208                Value::String(value_str.to_string())
209            }
210        }
211    }
212}
213
214impl Default for Glm4MoeParser {
215    fn default() -> Self {
216        Self::glm45()
217    }
218}
219
220#[async_trait]
221impl ToolParser for Glm4MoeParser {
222    async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
223        self.parse_complete_inner(text, &[])
224    }
225
226    async fn parse_complete_with_tools(
227        &self,
228        text: &str,
229        tools: &[Tool],
230    ) -> ParserResult<(String, Vec<ToolCall>)> {
231        self.parse_complete_inner(text, tools)
232    }
233
234    async fn parse_incremental(
235        &mut self,
236        chunk: &str,
237        tools: &[Tool],
238    ) -> ParserResult<StreamingParseResult> {
239        // Python logic: Wait for complete tool call, then parse it all at once
240        self.buffer.push_str(chunk);
241        let current_text = &self.buffer.clone();
242
243        // Check if we have bot_token
244        let start = current_text.find(self.bot_token);
245        if start.is_none() {
246            self.buffer.clear();
247            // If we're in the middle of streaming (current_tool_id > 0), don't return text
248            let normal_text = if self.current_tool_id > 0 {
249                String::new()
250            } else {
251                current_text.clone()
252            };
253            return Ok(StreamingParseResult {
254                normal_text,
255                calls: vec![],
256            });
257        }
258
259        // Check if we have eot_token (end of tool call)
260        let end = current_text.find(self.eot_token);
261        if let Some(end_pos) = end {
262            // We have a complete tool call!
263
264            // Initialize state if this is the first tool call
265            if self.current_tool_id == -1 {
266                self.current_tool_id = 0;
267                self.prev_tool_call_arr = Vec::new();
268                self.streamed_args_for_tool = vec![String::new()];
269            }
270
271            // Ensure we have enough entries in our tracking arrays
272            helpers::ensure_capacity(
273                self.current_tool_id,
274                &mut self.prev_tool_call_arr,
275                &mut self.streamed_args_for_tool,
276            );
277
278            // Parse the complete block using shared helper
279            let block_end = end_pos + self.eot_token.len();
280            let parsed_tools = self.parse_tool_calls_from_text(&current_text[..block_end], tools);
281
282            // Extract normal text before tool calls
283            let idx = current_text.find(self.bot_token);
284            let normal_text = if let Some(pos) = idx {
285                current_text[..pos].trim().to_string()
286            } else {
287                String::new()
288            };
289
290            // Build tool indices for validation
291            let tool_indices = helpers::get_tool_indices(tools);
292
293            let mut calls = Vec::new();
294
295            if !parsed_tools.is_empty() {
296                // Take the first tool and convert to ToolCallItem
297                let tool_call = &parsed_tools[0];
298                let tool_id = self.current_tool_id as usize;
299
300                // Validate tool name
301                if !tool_indices.contains_key(&tool_call.function.name) {
302                    // Invalid tool name - skip this tool, preserve indexing for next tool
303                    tracing::debug!("Invalid tool name '{}' - skipping", tool_call.function.name);
304                    helpers::reset_current_tool_state(
305                        &mut self.buffer,
306                        &mut false, // glm45_moe/glm47_moe doesn't track name_sent per tool
307                        &mut self.streamed_args_for_tool,
308                        &self.prev_tool_call_arr,
309                    );
310                    return Ok(StreamingParseResult::default());
311                }
312
313                calls.push(ToolCallItem {
314                    tool_index: tool_id,
315                    name: Some(tool_call.function.name.clone()),
316                    parameters: tool_call.function.arguments.clone(),
317                });
318
319                // Store in tracking arrays
320                if self.prev_tool_call_arr.len() <= tool_id {
321                    self.prev_tool_call_arr
322                        .resize_with(tool_id + 1, || Value::Null);
323                }
324
325                // Parse parameters as JSON and store
326                if let Ok(args) = serde_json::from_str::<Value>(&tool_call.function.arguments) {
327                    self.prev_tool_call_arr[tool_id] = serde_json::json!({
328                        "name": tool_call.function.name,
329                        "arguments": args,
330                    });
331                }
332
333                if self.streamed_args_for_tool.len() <= tool_id {
334                    self.streamed_args_for_tool
335                        .resize_with(tool_id + 1, String::new);
336                }
337                self.streamed_args_for_tool[tool_id].clone_from(&tool_call.function.arguments);
338
339                self.current_tool_id += 1;
340            }
341
342            // Remove processed portion from buffer
343            self.buffer = current_text[block_end..].to_string();
344            return Ok(StreamingParseResult { normal_text, calls });
345        }
346
347        // No complete tool call yet - return normal text before start token
348        // Safe: start.is_none() case was handled above (early return)
349        let Some(start_pos) = start else {
350            return Ok(StreamingParseResult::default());
351        };
352        let normal_text = current_text[..start_pos].to_string();
353        self.buffer = current_text[start_pos..].to_string();
354
355        Ok(StreamingParseResult {
356            normal_text,
357            calls: vec![],
358        })
359    }
360
361    fn has_tool_markers(&self, text: &str) -> bool {
362        text.contains(self.bot_token)
363    }
364
365    fn get_unstreamed_tool_args(&self) -> Option<Vec<ToolCallItem>> {
366        helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
367    }
368
369    fn reset(&mut self) {
370        self.buffer.clear();
371        self.prev_tool_call_arr.clear();
372        self.current_tool_id = -1;
373        self.streamed_args_for_tool.clear();
374    }
375}
376
377#[cfg(test)]
378mod tests {
379    use openai_protocol::common::Function;
380
381    use super::*;
382
383    fn tool_with_props(props: Value) -> Vec<Tool> {
384        vec![Tool {
385            tool_type: "function".to_string(),
386            function: Function {
387                name: "f".to_string(),
388                description: None,
389                parameters: serde_json::json!({"type": "object", "properties": props}),
390                strict: None,
391            },
392        }]
393    }
394
395    // String-typed params stay strings even when they look numeric/bool/array.
396    #[tokio::test]
397    async fn test_schema_aware_coercion_keeps_strings() {
398        let tools = tool_with_props(serde_json::json!({
399            "limit": {"type": "string"},
400            "flag": {"type": "string"},
401            "coords": {"type": "string"},
402            "count": {"type": "integer"},
403        }));
404        let text = "<tool_call>f\n\
405            <arg_key>limit</arg_key>\n<arg_value>4</arg_value>\n\
406            <arg_key>flag</arg_key>\n<arg_value>true</arg_value>\n\
407            <arg_key>coords</arg_key>\n<arg_value>[60,30]</arg_value>\n\
408            <arg_key>count</arg_key>\n<arg_value>5</arg_value>\n\
409            </tool_call>";
410        let (_, calls) = Glm4MoeParser::glm45()
411            .parse_complete_with_tools(text, &tools)
412            .await
413            .unwrap();
414        assert_eq!(calls.len(), 1);
415        let args: Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
416        assert_eq!(args["limit"], Value::String("4".to_string()));
417        assert_eq!(args["flag"], Value::String("true".to_string()));
418        assert_eq!(args["coords"], Value::String("[60,30]".to_string()));
419        assert_eq!(args["count"], Value::Number(5.into()));
420    }
421
422    // The streaming path threads `tools` separately, so cover it too.
423    #[tokio::test]
424    async fn test_streaming_schema_aware_coercion() {
425        let tools = tool_with_props(serde_json::json!({
426            "limit": {"type": "string"},
427            "count": {"type": "integer"},
428        }));
429        let text = "<tool_call>f\n\
430            <arg_key>limit</arg_key>\n<arg_value>4</arg_value>\n\
431            <arg_key>count</arg_key>\n<arg_value>5</arg_value>\n\
432            </tool_call>";
433        let result = Glm4MoeParser::glm45()
434            .parse_incremental(text, &tools)
435            .await
436            .unwrap();
437        let args: Value = serde_json::from_str(&result.calls[0].parameters).unwrap();
438        assert_eq!(args["limit"], Value::String("4".to_string()));
439        assert_eq!(args["count"], Value::Number(5.into()));
440    }
441}
tool_parser/parsers/glm4_moe.rs

tool_parser/parsers/
glm4_moe.rs
