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 factory::GrammarStart,
11 parsers::helpers,
12 traits::ToolParser,
13 types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
14};
15
16pub struct Glm4MoeParser {
27 tool_call_extractor: Regex,
29 func_detail_extractor: Regex,
31 arg_extractor: Regex,
33
34 buffer: String,
36
37 prev_tool_call_arr: Vec<Value>,
39
40 current_tool_id: i32,
42
43 streamed_args_for_tool: Vec<String>,
45
46 bot_token: &'static str,
48 eot_token: &'static str,
49}
50
51impl Glm4MoeParser {
52 #[expect(
59 clippy::expect_used,
60 reason = "regex patterns are compile-time string literals"
61 )]
62 pub(crate) fn new(func_detail_pattern: &str) -> Self {
63 let tool_call_pattern = r"(?s)<tool_call>.*?</tool_call>";
65 let tool_call_extractor = Regex::new(tool_call_pattern).expect("Valid regex pattern");
66
67 let func_detail_extractor = Regex::new(func_detail_pattern).expect("Valid regex pattern");
68
69 let arg_pattern = r"(?s)<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)</arg_value>";
70 let arg_extractor = Regex::new(arg_pattern).expect("Valid regex pattern");
71
72 Self {
73 tool_call_extractor,
74 func_detail_extractor,
75 arg_extractor,
76 buffer: String::new(),
77 prev_tool_call_arr: Vec::new(),
78 current_tool_id: -1,
79 streamed_args_for_tool: Vec::new(),
80 bot_token: "<tool_call>",
81 eot_token: "</tool_call>",
82 }
83 }
84
85 pub fn glm45() -> Self {
87 Self::new(r"(?s)<tool_call>([^\n]*)\n(.*)</tool_call>")
88 }
89
90 pub fn glm47() -> Self {
92 Self::new(r"(?s)<tool_call>\s*([^<\s]+)\s*(.*?)</tool_call>")
93 }
94
95 pub fn build_structural_tag(tools: &[Tool], at_least_one: bool) -> Value {
123 let tags: Vec<Value> = tools
124 .iter()
125 .filter(|tool| !tool.function.name.is_empty())
126 .map(|tool| {
127 serde_json::json!({
128 "type": "tag",
129 "begin": format!("<tool_call>{}", tool.function.name),
130 "content": {
131 "type": "json_schema",
132 "json_schema": tool.function.parameters_schema(),
133 "style": "glm_xml",
134 },
135 "end": "</tool_call>",
136 })
137 })
138 .collect();
139 serde_json::json!({
140 "format": {
141 "type": "triggered_tags",
142 "triggers": ["<tool_call>"],
143 "tags": tags,
144 "at_least_one": at_least_one,
145 }
146 })
147 }
148
149 pub fn reasoning_prefix(start: GrammarStart) -> Option<Value> {
167 const CALL_MARKERS: [&str; 6] = [
168 "<tool_call>",
169 "</tool_call>",
170 "<arg_key>",
171 "</arg_key>",
172 "<arg_value>",
173 "</arg_value>",
174 ];
175 match start {
176 GrammarStart::AfterReasoning => None,
177 GrammarStart::Unknown => Some(serde_json::json!({
178 "type": "any_text",
179 "excludes": CALL_MARKERS,
180 })),
181 GrammarStart::FirstToken => {
182 let excludes = [&["<think>", "</think>"][..], &CALL_MARKERS[..]].concat();
183 Some(serde_json::json!({
184 "type": "tag",
185 "begin": "",
186 "content": { "type": "any_text", "excludes": excludes },
187 "end": "</think>",
188 }))
189 }
190 }
191 }
192
193 fn parse_arguments(
196 &self,
197 args_text: &str,
198 param_types: &HashMap<String, String>,
199 ) -> serde_json::Map<String, Value> {
200 let mut arguments = serde_json::Map::new();
201
202 for capture in self.arg_extractor.captures_iter(args_text) {
203 let key = capture.get(1).map_or("", |m| m.as_str()).trim();
204 let value_str = capture.get(2).map_or("", |m| m.as_str());
205
206 let value =
207 helpers::coerce_by_schema_type(value_str, param_types.get(key).map(String::as_str))
208 .unwrap_or_else(|| infer_value(value_str.trim()));
209
210 arguments.insert(key.to_string(), value);
211 }
212
213 arguments
214 }
215
216 fn parse_tool_call(&self, block: &str, tools: &[Tool]) -> ParserResult<Option<ToolCall>> {
218 if let Some(captures) = self.func_detail_extractor.captures(block) {
219 let func_name = captures.get(1).map_or("", |m| m.as_str()).trim();
221
222 let args_text = captures.get(2).map_or("", |m| m.as_str());
224
225 let param_types = helpers::param_types_for_function(tools, func_name);
226 let arguments = self.parse_arguments(args_text, ¶m_types);
227
228 let arguments_str = serde_json::to_string(&arguments)
229 .map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
230
231 Ok(Some(ToolCall {
232 function: FunctionCall {
233 name: func_name.to_string(),
234 arguments: arguments_str,
235 },
236 }))
237 } else {
238 Ok(None)
239 }
240 }
241
242 fn parse_tool_calls_from_text(&self, text: &str, tools: &[Tool]) -> Vec<ToolCall> {
244 let mut parsed = Vec::new();
245
246 for mat in self.tool_call_extractor.find_iter(text) {
247 match self.parse_tool_call(mat.as_str(), tools) {
248 Ok(Some(tool)) => parsed.push(tool),
249 Ok(None) => continue,
250 Err(e) => {
251 tracing::debug!("Failed to parse tool call: {}", e);
252 continue;
253 }
254 }
255 }
256
257 parsed
258 }
259}
260
261impl Glm4MoeParser {
262 fn parse_complete_inner(
264 &self,
265 text: &str,
266 tools: &[Tool],
267 ) -> ParserResult<(String, Vec<ToolCall>)> {
268 if !self.has_tool_markers(text) {
269 return Ok((text.to_string(), vec![]));
270 }
271
272 let idx = text
275 .find("<tool_call>")
276 .ok_or_else(|| ParserError::ParsingFailed("tool call marker not found".to_string()))?;
277 let normal_text = text[..idx].to_string();
278
279 let parsed = self.parse_tool_calls_from_text(text, tools);
280
281 if parsed.is_empty() {
283 return Ok((text.to_string(), vec![]));
284 }
285
286 Ok((normal_text, parsed))
287 }
288}
289
290fn infer_value(value_str: &str) -> Value {
293 if let Ok(json_val) = serde_json::from_str::<Value>(value_str) {
294 return json_val;
295 }
296 match value_str {
297 "true" | "True" => Value::Bool(true),
298 "false" | "False" => Value::Bool(false),
299 "null" | "None" => Value::Null,
300 _ => {
301 if let Ok(num) = value_str.parse::<i64>() {
302 Value::Number(num.into())
303 } else if let Ok(num) = value_str.parse::<f64>() {
304 serde_json::Number::from_f64(num)
305 .map_or_else(|| Value::String(value_str.to_string()), Value::Number)
306 } else {
307 Value::String(value_str.to_string())
308 }
309 }
310 }
311}
312
313impl Default for Glm4MoeParser {
314 fn default() -> Self {
315 Self::glm45()
316 }
317}
318
319#[async_trait]
320impl ToolParser for Glm4MoeParser {
321 async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
322 self.parse_complete_inner(text, &[])
323 }
324
325 async fn parse_complete_with_tools(
326 &self,
327 text: &str,
328 tools: &[Tool],
329 ) -> ParserResult<(String, Vec<ToolCall>)> {
330 self.parse_complete_inner(text, tools)
331 }
332
333 async fn parse_incremental(
334 &mut self,
335 chunk: &str,
336 tools: &[Tool],
337 ) -> ParserResult<StreamingParseResult> {
338 self.buffer.push_str(chunk);
340 let current_text = &self.buffer.clone();
341
342 let start = current_text.find(self.bot_token);
344 if start.is_none() {
345 self.buffer.clear();
346 let normal_text = if self.current_tool_id > 0 {
348 String::new()
349 } else {
350 current_text.clone()
351 };
352 return Ok(StreamingParseResult {
353 normal_text,
354 calls: vec![],
355 });
356 }
357
358 let end = current_text.find(self.eot_token);
360 if let Some(end_pos) = end {
361 if self.current_tool_id == -1 {
365 self.current_tool_id = 0;
366 self.prev_tool_call_arr = Vec::new();
367 self.streamed_args_for_tool = vec![String::new()];
368 }
369
370 helpers::ensure_capacity(
372 self.current_tool_id,
373 &mut self.prev_tool_call_arr,
374 &mut self.streamed_args_for_tool,
375 );
376
377 let block_end = end_pos + self.eot_token.len();
379 let parsed_tools = self.parse_tool_calls_from_text(¤t_text[..block_end], tools);
380
381 let idx = current_text.find(self.bot_token);
383 let normal_text = if let Some(pos) = idx {
384 current_text[..pos].trim().to_string()
385 } else {
386 String::new()
387 };
388
389 let tool_indices = helpers::get_tool_indices(tools);
391
392 let mut calls = Vec::new();
393
394 if !parsed_tools.is_empty() {
395 let tool_call = &parsed_tools[0];
397 let tool_id = self.current_tool_id as usize;
398
399 if !tool_indices.contains_key(&tool_call.function.name) {
401 tracing::debug!("Invalid tool name '{}' - skipping", tool_call.function.name);
403 helpers::reset_current_tool_state(
404 &mut self.buffer,
405 &mut false, &mut self.streamed_args_for_tool,
407 &self.prev_tool_call_arr,
408 );
409 return Ok(StreamingParseResult::default());
410 }
411
412 calls.push(ToolCallItem {
413 tool_index: tool_id,
414 name: Some(tool_call.function.name.clone()),
415 parameters: tool_call.function.arguments.clone(),
416 });
417
418 if self.prev_tool_call_arr.len() <= tool_id {
420 self.prev_tool_call_arr
421 .resize_with(tool_id + 1, || Value::Null);
422 }
423
424 if let Ok(args) = serde_json::from_str::<Value>(&tool_call.function.arguments) {
426 self.prev_tool_call_arr[tool_id] = serde_json::json!({
427 "name": tool_call.function.name,
428 "arguments": args,
429 });
430 }
431
432 if self.streamed_args_for_tool.len() <= tool_id {
433 self.streamed_args_for_tool
434 .resize_with(tool_id + 1, String::new);
435 }
436 self.streamed_args_for_tool[tool_id].clone_from(&tool_call.function.arguments);
437
438 self.current_tool_id += 1;
439 }
440
441 self.buffer = current_text[block_end..].to_string();
443 return Ok(StreamingParseResult { normal_text, calls });
444 }
445
446 let Some(start_pos) = start else {
449 return Ok(StreamingParseResult::default());
450 };
451 let normal_text = current_text[..start_pos].to_string();
452 self.buffer = current_text[start_pos..].to_string();
453
454 Ok(StreamingParseResult {
455 normal_text,
456 calls: vec![],
457 })
458 }
459
460 fn has_tool_markers(&self, text: &str) -> bool {
461 text.contains(self.bot_token)
462 }
463
464 fn get_unstreamed_tool_args(&self) -> Option<Vec<ToolCallItem>> {
465 helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
466 }
467
468 fn reset(&mut self) {
469 self.buffer.clear();
470 self.prev_tool_call_arr.clear();
471 self.current_tool_id = -1;
472 self.streamed_args_for_tool.clear();
473 }
474}
475
476#[cfg(test)]
477mod tests {
478 use openai_protocol::common::Function;
479
480 use super::*;
481
482 fn tool(name: &str, strict: Option<bool>, parameters: Value) -> Tool {
483 Tool {
484 tool_type: "function".to_string(),
485 function: Function {
486 name: name.to_string(),
487 description: None,
488 parameters: Some(parameters),
489 strict,
490 extra: Default::default(),
491 },
492 }
493 }
494
495 #[test]
499 fn structural_tag_mirrors_the_xgrammar_glm_4_7_model() {
500 let schema =
501 serde_json::json!({"type": "object", "properties": {"city": {"type": "string"}}});
502 let tools = vec![
503 tool("get_weather", None, schema.clone()),
504 tool("lookup", Some(false), schema.clone()),
505 tool("ping", None, serde_json::json!({})),
506 tool("", None, schema.clone()),
507 ];
508 let tag = Glm4MoeParser::build_structural_tag(&tools, true);
509 let format = &tag["format"];
510 assert_eq!(format["type"], "triggered_tags");
511 assert_eq!(format["triggers"], serde_json::json!(["<tool_call>"]));
512 assert_eq!(format["at_least_one"], true);
513 let tags = format["tags"].as_array().unwrap();
514 assert_eq!(tags.len(), 3, "a nameless tool has no tag");
515 assert_eq!(tags[0]["begin"], "<tool_call>get_weather");
516 assert_eq!(tags[0]["end"], "</tool_call>");
517 assert_eq!(tags[0]["content"]["type"], "json_schema");
518 assert_eq!(tags[0]["content"]["style"], "glm_xml");
519 assert_eq!(tags[0]["content"]["json_schema"], schema);
520 assert_eq!(
521 tags[1]["content"]["json_schema"], schema,
522 "strict is not a constraint switch here, as in the other builders"
523 );
524 assert_eq!(
525 tags[2]["content"]["json_schema"],
526 serde_json::json!({}),
527 "a no-argument tool keeps its empty schema"
528 );
529
530 let tag = Glm4MoeParser::build_structural_tag(&tools, false);
531 assert_eq!(tag["format"]["at_least_one"], false);
532 }
533
534 #[test]
535 fn reasoning_prefix_is_free_text_that_owes_nothing_at_its_end() {
536 let prefix = Glm4MoeParser::reasoning_prefix(GrammarStart::Unknown).unwrap();
537 assert_eq!(prefix["type"], "any_text");
538 assert!(
539 prefix.get("end").is_none(),
540 "nothing the model must write before the calls: {prefix}"
541 );
542 let excludes = prefix["excludes"].as_array().unwrap();
543 for token in [
544 "<tool_call>",
545 "</tool_call>",
546 "<arg_key>",
547 "</arg_key>",
548 "<arg_value>",
549 "</arg_value>",
550 ] {
551 assert!(
552 excludes.contains(&Value::String(token.to_string())),
553 "{token} must not appear before the forced call"
554 );
555 }
556 assert!(
557 !excludes.contains(&Value::String("</think>".to_string())),
558 "the model closes its thought here when the engine applies the grammar from the first token"
559 );
560 }
561
562 #[test]
567 fn reasoning_prefix_follows_where_the_engine_starts_the_grammar() {
568 assert_eq!(
569 Glm4MoeParser::reasoning_prefix(GrammarStart::AfterReasoning),
570 None,
571 "the grammar starts at the engine's gate: the call is forced right there"
572 );
573 let block = Glm4MoeParser::reasoning_prefix(GrammarStart::FirstToken).unwrap();
574 assert_eq!(block["type"], "tag");
575 assert_eq!(block["begin"], "");
576 assert_eq!(block["end"], "</think>");
577 assert_eq!(block["content"]["type"], "any_text");
578 let excludes = block["content"]["excludes"].as_array().unwrap();
579 for token in ["<think>", "</think>", "<tool_call>", "</arg_value>"] {
580 assert!(
581 excludes.contains(&Value::String(token.to_string())),
582 "{token} must not appear inside the thought block"
583 );
584 }
585 }
586
587 fn tool_with_props(props: Value) -> Vec<Tool> {
588 vec![Tool {
589 tool_type: "function".to_string(),
590 function: Function {
591 name: "f".to_string(),
592 description: None,
593 parameters: Some(serde_json::json!({"type": "object", "properties": props})),
594 strict: None,
595 extra: Default::default(),
596 },
597 }]
598 }
599
600 #[tokio::test]
602 async fn test_schema_aware_coercion_keeps_strings() {
603 let tools = tool_with_props(serde_json::json!({
604 "limit": {"type": "string"},
605 "flag": {"type": "string"},
606 "coords": {"type": "string"},
607 "count": {"type": "integer"},
608 }));
609 let text = "<tool_call>f\n\
610 <arg_key>limit</arg_key>\n<arg_value>4</arg_value>\n\
611 <arg_key>flag</arg_key>\n<arg_value>true</arg_value>\n\
612 <arg_key>coords</arg_key>\n<arg_value>[60,30]</arg_value>\n\
613 <arg_key>count</arg_key>\n<arg_value>5</arg_value>\n\
614 </tool_call>";
615 let (_, calls) = Glm4MoeParser::glm45()
616 .parse_complete_with_tools(text, &tools)
617 .await
618 .unwrap();
619 assert_eq!(calls.len(), 1);
620 let args: Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
621 assert_eq!(args["limit"], Value::String("4".to_string()));
622 assert_eq!(args["flag"], Value::String("true".to_string()));
623 assert_eq!(args["coords"], Value::String("[60,30]".to_string()));
624 assert_eq!(args["count"], Value::Number(5.into()));
625 }
626
627 #[tokio::test]
629 async fn test_streaming_schema_aware_coercion() {
630 let tools = tool_with_props(serde_json::json!({
631 "limit": {"type": "string"},
632 "count": {"type": "integer"},
633 }));
634 let text = "<tool_call>f\n\
635 <arg_key>limit</arg_key>\n<arg_value>4</arg_value>\n\
636 <arg_key>count</arg_key>\n<arg_value>5</arg_value>\n\
637 </tool_call>";
638 let result = Glm4MoeParser::glm45()
639 .parse_incremental(text, &tools)
640 .await
641 .unwrap();
642 let args: Value = serde_json::from_str(&result.calls[0].parameters).unwrap();
643 assert_eq!(args["limit"], Value::String("4".to_string()));
644 assert_eq!(args["count"], Value::Number(5.into()));
645 }
646
647 #[tokio::test]
648 async fn test_declared_strings_preserve_whitespace() {
649 let tools = tool_with_props(serde_json::json!({
650 "code": {"type": "string"},
651 "needle": {"type": "string"},
652 "blank": {"type": "string"},
653 "quoted": {"type": "string"},
654 "count": {"type": "integer"},
655 }));
656 let args_text = concat!(
657 "<arg_key>code</arg_key><arg_value> return 1\n</arg_value>",
658 "<arg_key>needle</arg_key><arg_value> total = 1\n</arg_value>",
659 "<arg_key>blank</arg_key><arg_value> \t\n</arg_value>",
660 "<arg_key>quoted</arg_key><arg_value> \" return 1\\n\" </arg_value>",
661 "<arg_key>count</arg_key><arg_value> 42\n</arg_value>",
662 "<arg_key>unknown</arg_key><arg_value> legacy text\n</arg_value>",
663 );
664 let expected = serde_json::json!({
665 "code": " return 1\n",
666 "needle": " total = 1\n",
667 "blank": " \t\n",
668 "quoted": " return 1\n",
669 "count": 42,
670 "unknown": "legacy text",
671 });
672
673 for (make_parser, name_separator) in [
674 (Glm4MoeParser::glm45 as fn() -> Glm4MoeParser, "\n"),
675 (Glm4MoeParser::glm47 as fn() -> Glm4MoeParser, ""),
676 ] {
677 let text = format!("<tool_call>f{name_separator}{args_text}</tool_call>");
678 let (_, calls) = make_parser()
679 .parse_complete_with_tools(&text, &tools)
680 .await
681 .unwrap();
682 assert_eq!(calls.len(), 1);
683 let args: Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
684 assert_eq!(args, expected);
685
686 let mut parser = make_parser();
689 let start = "<tool_call>".len();
690 let result = parser
691 .parse_incremental(&text[..start], &tools)
692 .await
693 .unwrap();
694 assert!(result.calls.is_empty());
695 let mut arguments = String::new();
696 for chunk in text.as_bytes()[start..].chunks(7) {
697 let result = parser
698 .parse_incremental(std::str::from_utf8(chunk).unwrap(), &tools)
699 .await
700 .unwrap();
701 for call in result.calls {
702 arguments.push_str(&call.parameters);
703 }
704 }
705 let args: Value = serde_json::from_str(&arguments).unwrap();
706 assert_eq!(args, expected);
707 }
708 }
709}