Skip to main content

tool_parser/parsers/
qwen_xml.rs

1use async_trait::async_trait;
2use openai_protocol::common::Tool;
3use regex::Regex;
4use serde_json::Value;
5
6use crate::{
7    errors::{ParserError, ParserResult},
8    parsers::helpers,
9    traits::ToolParser,
10    types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
11};
12
13/// Qwen XML format parser for tool calls
14///
15/// Handles the Qwen XML specific XML format:
16/// `<tool_call>\n<function=name>\n<parameter=key>value</parameter>\n</function>\n</tool_call>`
17///
18/// Features:
19/// - Tool Call Tags: `<tool_call>` and `</tool_call>` wrap each individual call
20/// - XML-style function declaration: `<function=name>`
21/// - XML-style parameters: `<parameter=key>value</parameter>`
22///
23/// Reference: https://huggingface.co/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8?chat_template=default
24pub struct QwenXmlParser {
25    /// Regex for extracting tool calls in parse_complete
26    extractor: Regex,
27
28    /// Buffer for accumulating incomplete patterns across chunks
29    buffer: String,
30
31    /// Stores complete tool call info (name and arguments) for each tool being parsed
32    prev_tool_call_arr: Vec<Value>,
33
34    /// Index of currently streaming tool call (-1 means no active tool)
35    current_tool_id: i32,
36
37    /// Flag for whether current tool's name has been sent to client
38    current_tool_name_sent: bool,
39
40    /// Tracks raw JSON string content streamed to client for each tool's arguments
41    streamed_args_for_tool: Vec<String>,
42
43    /// Token configuration
44    tool_call_start_token: &'static str,
45    tool_call_end_token: &'static str,
46
47    /// XML format streaming state
48    in_tool_call: bool,
49    current_function_name: String,
50    current_parameters: serde_json::Map<String, Value>,
51
52    /// Precompiled regex patterns for XML format parsing
53    xml_function_pattern: Regex,
54    xml_param_pattern: Regex,
55}
56
57/// Decode HTML entities in a string (equivalent to Python's html.unescape)
58///
59/// Handles common HTML entities like &amp; &lt; &gt; &quot; &#39; and numeric entities
60fn html_unescape(s: &str) -> String {
61    let mut result = String::with_capacity(s.len());
62    let mut chars = s.chars().peekable();
63
64    while let Some(c) = chars.next() {
65        if c == '&' {
66            let mut entity = String::new();
67            let mut consumed_semicolon = false;
68            while let Some(&next) = chars.peek() {
69                if next == ';' {
70                    chars.next();
71                    consumed_semicolon = true;
72                    break;
73                }
74                if next.is_alphanumeric() || next == '#' {
75                    // Safe: peek() returned Some, so next() will too
76                    if let Some(ch) = chars.next() {
77                        entity.push(ch);
78                    }
79                } else {
80                    break;
81                }
82            }
83
84            let decoded = match entity.as_str() {
85                "amp" => "&",
86                "lt" => "<",
87                "gt" => ">",
88                "quot" => "\"",
89                "apos" => "'",
90                "nbsp" => "\u{00A0}",
91                s if s.starts_with('#') => {
92                    let num_str = &s[1..];
93                    let code_point = if num_str.starts_with('x') || num_str.starts_with('X') {
94                        u32::from_str_radix(&num_str[1..], 16).ok()
95                    } else {
96                        num_str.parse::<u32>().ok()
97                    };
98                    if let Some(cp) = code_point {
99                        if let Some(ch) = char::from_u32(cp) {
100                            result.push(ch);
101                            continue;
102                        }
103                    }
104                    // Invalid numeric entity, reconstruct original
105                    result.push('&');
106                    result.push_str(&entity);
107                    if consumed_semicolon {
108                        result.push(';');
109                    }
110                    continue;
111                }
112                _ => {
113                    // Unknown entity, reconstruct original
114                    result.push('&');
115                    result.push_str(&entity);
116                    if consumed_semicolon {
117                        result.push(';');
118                    }
119                    continue;
120                }
121            };
122            result.push_str(decoded);
123        } else {
124            result.push(c);
125        }
126    }
127
128    result
129}
130
131/// Parse a raw parameter value, similar to Python's _safe_val
132///
133/// 1. Decode HTML entities
134/// 2. Try to parse as JSON (numbers, booleans, null, objects, arrays)
135/// 3. Fall back to string if JSON parsing fails
136fn safe_val(raw: &str) -> Value {
137    let unescaped = html_unescape(raw.trim());
138
139    // Try JSON parsing first
140    if let Ok(v) = serde_json::from_str::<Value>(&unescaped) {
141        return v;
142    }
143
144    // Handle Python-style literals (True, False, None)
145    match unescaped.as_str() {
146        "True" => return Value::Bool(true),
147        "False" => return Value::Bool(false),
148        "None" => return Value::Null,
149        _ => {}
150    }
151
152    // Fall back to string
153    Value::String(unescaped)
154}
155
156/// Coerce an XML parameter value by its declared schema type, falling back to
157/// [`safe_val`] inference when the type is unknown.
158fn coerce_value(raw: &str, declared_type: Option<&str>) -> Value {
159    let decoded = html_unescape(raw.trim());
160    helpers::coerce_by_schema_type(&decoded, declared_type).unwrap_or_else(|| safe_val(raw))
161}
162
163impl QwenXmlParser {
164    /// Create a new Qwen XML parser
165    #[expect(
166        clippy::expect_used,
167        reason = "regex patterns are compile-time string literals"
168    )]
169    pub fn new() -> Self {
170        // Support XML format: <tool_call>\n<function=name>\n<parameter=key>value</parameter>\n</function>\n</tool_call>
171        let pattern = r"(?s)<tool_call>\s*(.*?)\s*</tool_call>";
172        let extractor = Regex::new(pattern).expect("Valid regex pattern");
173
174        // Precompile XML format regex patterns for performance
175        let xml_function_pattern =
176            Regex::new(r"<function=([^>]+)>").expect("Valid XML function pattern");
177        let xml_param_pattern = Regex::new(r"(?s)<parameter=([^>]+)>(.*?)</parameter>")
178            .expect("Valid XML parameter pattern");
179
180        Self {
181            extractor,
182            buffer: String::new(),
183            prev_tool_call_arr: Vec::new(),
184            current_tool_id: -1,
185            current_tool_name_sent: false,
186            streamed_args_for_tool: Vec::new(),
187            tool_call_start_token: "<tool_call>",
188            tool_call_end_token: "</tool_call>",
189            in_tool_call: false,
190            current_function_name: String::new(),
191            current_parameters: serde_json::Map::new(),
192            xml_function_pattern,
193            xml_param_pattern,
194        }
195    }
196
197    /// Parse XML format tool call: <function=name><parameter=key>value</parameter></function>
198    fn parse_xml_format(&self, content: &str, tools: &[Tool]) -> ParserResult<Option<ToolCall>> {
199        let function_captures = self
200            .xml_function_pattern
201            .captures(content)
202            .ok_or_else(|| ParserError::ParsingFailed("No function name found".to_string()))?;
203
204        let function_name = function_captures
205            .get(1)
206            .ok_or_else(|| ParserError::ParsingFailed("Function name capture failed".to_string()))?
207            .as_str()
208            .trim()
209            .to_string();
210
211        if function_name.is_empty() {
212            return Ok(None);
213        }
214
215        let param_types = helpers::param_types_for_function(tools, &function_name);
216        let mut parameters = serde_json::Map::new();
217
218        for cap in self.xml_param_pattern.captures_iter(content) {
219            if let (Some(key_match), Some(value_match)) = (cap.get(1), cap.get(2)) {
220                let key = key_match.as_str().trim().to_string();
221                let value = value_match.as_str();
222                let json_value = coerce_value(value, param_types.get(&key).map(String::as_str));
223                parameters.insert(key, json_value);
224            }
225        }
226
227        let arguments = serde_json::to_string(&parameters)
228            .map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
229
230        Ok(Some(ToolCall {
231            function: FunctionCall {
232                name: function_name,
233                arguments,
234            },
235        }))
236    }
237
238    /// Parse and stream complete parameters from buffer
239    /// Returns tool call items to emit (similar to Python's _parse_and_stream_parameters)
240    fn parse_and_stream_parameters(&mut self, tools: &[Tool]) -> Vec<ToolCallItem> {
241        let mut calls: Vec<ToolCallItem> = vec![];
242        let param_types = helpers::param_types_for_function(tools, &self.current_function_name);
243
244        // Find all complete parameter patterns in buffer
245        let mut new_params = serde_json::Map::new();
246        for cap in self.xml_param_pattern.captures_iter(&self.buffer) {
247            if let (Some(key_match), Some(value_match)) = (cap.get(1), cap.get(2)) {
248                let key = key_match.as_str().trim().to_string();
249                let value = value_match.as_str();
250                let json_value = coerce_value(value, param_types.get(&key).map(String::as_str));
251                new_params.insert(key, json_value);
252            }
253        }
254
255        // Calculate parameter diff and stream updates
256        if new_params != self.current_parameters {
257            let current_args = &mut self.streamed_args_for_tool[self.current_tool_id as usize];
258
259            if self.current_parameters.is_empty() {
260                // First parameter(s) - build JSON fragment (without closing brace)
261                let mut items = Vec::new();
262                for (key, value) in &new_params {
263                    let key_json =
264                        serde_json::to_string(key).unwrap_or_else(|_| format!("\"{key}\""));
265                    let value_json = serde_json::to_string(value).unwrap_or_default();
266                    items.push(format!("{key_json}: {value_json}"));
267                }
268                let json_fragment = format!("{{{}", items.join(", "));
269
270                calls.push(ToolCallItem {
271                    tool_index: self.current_tool_id as usize,
272                    name: None,
273                    parameters: json_fragment.clone(),
274                });
275                *current_args = json_fragment;
276            } else {
277                // Additional parameters - add them incrementally
278                let new_keys: Vec<_> = new_params
279                    .keys()
280                    .filter(|k| !self.current_parameters.contains_key(*k))
281                    .collect();
282
283                if !new_keys.is_empty() {
284                    let mut continuation_parts = Vec::new();
285                    for key in new_keys {
286                        if let Some(value) = new_params.get(key) {
287                            let key_json =
288                                serde_json::to_string(key).unwrap_or_else(|_| format!("\"{key}\""));
289                            let value_json = serde_json::to_string(value).unwrap_or_default();
290                            continuation_parts.push(format!("{key_json}: {value_json}"));
291                        }
292                    }
293
294                    let json_fragment = format!(", {}", continuation_parts.join(", "));
295
296                    calls.push(ToolCallItem {
297                        tool_index: self.current_tool_id as usize,
298                        name: None,
299                        parameters: json_fragment.clone(),
300                    });
301                    current_args.push_str(&json_fragment);
302                }
303            }
304
305            // Update current state
306            self.current_parameters.clone_from(&new_params);
307            if let Some(tool_obj) =
308                self.prev_tool_call_arr[self.current_tool_id as usize].as_object_mut()
309            {
310                tool_obj.insert("arguments".to_string(), Value::Object(new_params));
311            }
312        }
313
314        calls
315    }
316
317    /// Shared non-streaming parse, schema-aware when `tools` are provided.
318    fn parse_complete_inner(
319        &self,
320        text: &str,
321        tools: &[Tool],
322    ) -> ParserResult<(String, Vec<ToolCall>)> {
323        // Check if text contains Qwen XML format
324        if !self.has_tool_markers(text) {
325            return Ok((text.to_string(), vec![]));
326        }
327
328        // Find where the first tool call begins
329        // Safe: has_tool_markers() already confirmed the marker exists
330        let idx = text
331            .find(self.tool_call_start_token)
332            .ok_or_else(|| ParserError::ParsingFailed("tool call marker not found".to_string()))?;
333        let normal_text = text[..idx].to_string();
334
335        // Extract tool calls
336        let mut parsed = Vec::new();
337        for captures in self.extractor.captures_iter(text) {
338            if let Some(content_str) = captures.get(1) {
339                let content = content_str.as_str().trim();
340
341                match self.parse_xml_format(content, tools) {
342                    Ok(Some(tool)) => parsed.push(tool),
343                    Ok(None) => continue,
344                    Err(e) => {
345                        tracing::warn!("Failed to parse XML tool call: {:?}", e);
346                        continue;
347                    }
348                }
349            }
350        }
351
352        // If no tools were successfully parsed despite having markers, return entire text
353        if parsed.is_empty() {
354            return Ok((text.to_string(), vec![]));
355        }
356
357        Ok((normal_text, parsed))
358    }
359
360    /// Reset streaming state for next tool call
361    fn reset_streaming_state(&mut self) {
362        self.in_tool_call = false;
363        self.current_tool_name_sent = false;
364        self.current_function_name.clear();
365        self.current_parameters.clear();
366    }
367}
368
369impl Default for QwenXmlParser {
370    fn default() -> Self {
371        Self::new()
372    }
373}
374
375#[async_trait]
376impl ToolParser for QwenXmlParser {
377    async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
378        self.parse_complete_inner(text, &[])
379    }
380
381    async fn parse_complete_with_tools(
382        &self,
383        text: &str,
384        tools: &[Tool],
385    ) -> ParserResult<(String, Vec<ToolCall>)> {
386        self.parse_complete_inner(text, tools)
387    }
388
389    async fn parse_incremental(
390        &mut self,
391        chunk: &str,
392        tools: &[Tool],
393    ) -> ParserResult<StreamingParseResult> {
394        self.buffer.push_str(chunk);
395
396        let mut normal_text = String::new();
397        let mut calls: Vec<ToolCallItem> = vec![];
398
399        // Build tool indices for validation
400        let tool_indices = helpers::get_tool_indices(tools);
401
402        loop {
403            // If we're not in a tool call and don't see a start token, return normal text
404            if !self.in_tool_call && !self.buffer.contains(self.tool_call_start_token) {
405                // Check for partial start token
406                if helpers::ends_with_partial_token(&self.buffer, self.tool_call_start_token)
407                    .is_none()
408                {
409                    normal_text.push_str(&self.buffer);
410                    self.buffer.clear();
411                }
412                break;
413            }
414
415            // Look for tool call start
416            if !self.in_tool_call {
417                if let Some(s) = self.buffer.find(self.tool_call_start_token) {
418                    normal_text.push_str(&self.buffer[..s]);
419                    self.buffer = self.buffer[s + self.tool_call_start_token.len()..].to_string();
420                    self.in_tool_call = true;
421                    self.current_tool_name_sent = false;
422                    self.current_function_name.clear();
423                    self.current_parameters.clear();
424                    continue;
425                } else {
426                    break;
427                }
428            }
429
430            // We're in a tool call, try to parse function name if not sent yet
431            if !self.current_tool_name_sent {
432                if let Some(captures) = self.xml_function_pattern.captures(&self.buffer) {
433                    if let Some(name_match) = captures.get(1) {
434                        let function_name = name_match.as_str().trim().to_string();
435
436                        // Validate function name
437                        if tool_indices.contains_key(&function_name) {
438                            self.current_function_name.clone_from(&function_name);
439                            self.current_tool_name_sent = true;
440
441                            // Initialize tool call tracking
442                            if self.current_tool_id == -1 {
443                                self.current_tool_id = 0;
444                            }
445
446                            // Ensure tracking arrays are large enough
447                            helpers::ensure_capacity(
448                                self.current_tool_id,
449                                &mut self.prev_tool_call_arr,
450                                &mut self.streamed_args_for_tool,
451                            );
452
453                            // Store tool call info
454                            self.prev_tool_call_arr[self.current_tool_id as usize] = serde_json::json!({
455                                "name": function_name,
456                                "arguments": {}
457                            });
458
459                            // Send tool name
460                            calls.push(ToolCallItem {
461                                tool_index: self.current_tool_id as usize,
462                                name: Some(function_name),
463                                parameters: String::new(),
464                            });
465
466                            // Remove processed function declaration from buffer
467                            // Safe: captures.get(0) always returns Some (group 0 is the entire match)
468                            self.buffer =
469                                self.buffer[captures.get(0).map_or(0, |m| m.end())..].to_string();
470                            continue;
471                        } else {
472                            // Invalid function name, reset state
473                            tracing::warn!("Invalid function name: {}", function_name);
474                            self.reset_streaming_state();
475                            normal_text.push_str(&self.buffer);
476                            self.buffer.clear();
477                            break;
478                        }
479                    }
480                } else {
481                    // Function name not complete yet, wait for more text
482                    break;
483                }
484            }
485
486            // Parse parameters (only complete ones)
487            if self.current_tool_name_sent {
488                let param_calls = self.parse_and_stream_parameters(tools);
489                calls.extend(param_calls);
490
491                // Check if tool call is complete
492                if let Some(end_pos) = self.buffer.find(self.tool_call_end_token) {
493                    // Close JSON object if we have parameters
494                    let current_args = &self.streamed_args_for_tool[self.current_tool_id as usize];
495                    if !current_args.is_empty() {
496                        // Count braces to check if JSON is complete
497                        let open_braces = current_args.matches('{').count();
498                        let close_braces = current_args.matches('}').count();
499                        if open_braces > close_braces {
500                            calls.push(ToolCallItem {
501                                tool_index: self.current_tool_id as usize,
502                                name: None,
503                                parameters: "}".to_string(),
504                            });
505                            self.streamed_args_for_tool[self.current_tool_id as usize].push('}');
506                        }
507                    }
508
509                    // Complete the tool call
510                    self.buffer =
511                        self.buffer[end_pos + self.tool_call_end_token.len()..].to_string();
512                    self.reset_streaming_state();
513                    self.current_tool_id += 1;
514                    continue;
515                } else {
516                    // Tool call not complete yet, wait for more text
517                    break;
518                }
519            }
520
521            break;
522        }
523
524        Ok(StreamingParseResult { normal_text, calls })
525    }
526
527    fn has_tool_markers(&self, text: &str) -> bool {
528        text.contains(self.tool_call_start_token)
529    }
530
531    fn get_unstreamed_tool_args(&self) -> Option<Vec<ToolCallItem>> {
532        helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
533    }
534
535    fn reset(&mut self) {
536        helpers::reset_parser_state(
537            &mut self.buffer,
538            &mut self.prev_tool_call_arr,
539            &mut self.current_tool_id,
540            &mut self.current_tool_name_sent,
541            &mut self.streamed_args_for_tool,
542        );
543        self.reset_streaming_state();
544    }
545}
546
547#[cfg(test)]
548mod tests {
549    use openai_protocol::common::Function;
550
551    use super::*;
552
553    #[test]
554    fn test_html_unescape_basic() {
555        assert_eq!(html_unescape("&amp;"), "&");
556        assert_eq!(html_unescape("&lt;"), "<");
557        assert_eq!(html_unescape("&gt;"), ">");
558        assert_eq!(html_unescape("&quot;"), "\"");
559        assert_eq!(html_unescape("&apos;"), "'");
560    }
561
562    #[test]
563    fn test_html_unescape_numeric() {
564        assert_eq!(html_unescape("&#60;"), "<");
565        assert_eq!(html_unescape("&#x3C;"), "<");
566        assert_eq!(html_unescape("&#x3c;"), "<");
567    }
568
569    #[test]
570    fn test_html_unescape_mixed() {
571        assert_eq!(
572            html_unescape("Hello &amp; World &lt;tag&gt;"),
573            "Hello & World <tag>"
574        );
575    }
576
577    #[test]
578    fn test_html_unescape_unknown() {
579        // Unknown entities with semicolon should be preserved as-is
580        assert_eq!(html_unescape("&unknown;"), "&unknown;");
581        // Unterminated entities should NOT have semicolon added
582        assert_eq!(html_unescape("&foo bar"), "&foo bar");
583        assert_eq!(html_unescape("&"), "&");
584        assert_eq!(html_unescape("& "), "& ");
585    }
586
587    #[test]
588    fn test_safe_val_json() {
589        assert_eq!(safe_val("42"), Value::Number(42.into()));
590        assert_eq!(safe_val("1.5"), serde_json::json!(1.5));
591        assert_eq!(safe_val("true"), Value::Bool(true));
592        assert_eq!(safe_val("false"), Value::Bool(false));
593        assert_eq!(safe_val("null"), Value::Null);
594        assert_eq!(
595            safe_val(r#"{"key": "value"}"#),
596            serde_json::json!({"key": "value"})
597        );
598        assert_eq!(safe_val(r"[1, 2, 3]"), serde_json::json!([1, 2, 3]));
599    }
600
601    #[test]
602    fn test_safe_val_python_literals() {
603        assert_eq!(safe_val("True"), Value::Bool(true));
604        assert_eq!(safe_val("False"), Value::Bool(false));
605        assert_eq!(safe_val("None"), Value::Null);
606    }
607
608    #[test]
609    fn test_safe_val_string_fallback() {
610        assert_eq!(
611            safe_val("hello world"),
612            Value::String("hello world".to_string())
613        );
614        assert_eq!(safe_val("  spaces  "), Value::String("spaces".to_string()));
615    }
616
617    #[test]
618    fn test_safe_val_html_entities() {
619        assert_eq!(safe_val("&lt;div&gt;"), Value::String("<div>".to_string()));
620        assert_eq!(
621            safe_val("Tom &amp; Jerry"),
622            Value::String("Tom & Jerry".to_string())
623        );
624    }
625
626    fn tool_with_props(props: Value) -> Vec<Tool> {
627        vec![Tool {
628            tool_type: "function".to_string(),
629            function: Function {
630                name: "f".to_string(),
631                description: None,
632                parameters: serde_json::json!({"type": "object", "properties": props}),
633                strict: None,
634            },
635        }]
636    }
637
638    // String-typed params stay strings even when they look numeric/bool/array/object.
639    #[tokio::test]
640    async fn test_schema_aware_coercion_keeps_strings() {
641        let tools = tool_with_props(serde_json::json!({
642            "limit": {"type": "string"},
643            "flag": {"type": "string"},
644            "coords": {"type": "string"},
645            "cfg": {"type": "string"},
646            "count": {"type": "integer"},
647        }));
648        let text = "<tool_call>\n<function=f>\n\
649            <parameter=limit>4</parameter>\n\
650            <parameter=flag>true</parameter>\n\
651            <parameter=coords>[60,30]</parameter>\n\
652            <parameter=cfg>{\"a\": 1}</parameter>\n\
653            <parameter=count>5</parameter>\n\
654            </function>\n</tool_call>";
655        let (_, calls) = QwenXmlParser::new()
656            .parse_complete_with_tools(text, &tools)
657            .await
658            .unwrap();
659        assert_eq!(calls.len(), 1);
660        let args: Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
661        assert_eq!(args["limit"], Value::String("4".to_string()));
662        assert_eq!(args["flag"], Value::String("true".to_string()));
663        assert_eq!(args["coords"], Value::String("[60,30]".to_string()));
664        assert_eq!(args["cfg"], Value::String("{\"a\": 1}".to_string()));
665        assert_eq!(args["count"], Value::Number(5.into()));
666    }
667
668    // The streaming path threads `tools` separately, so cover it too.
669    #[tokio::test]
670    async fn test_streaming_schema_aware_coercion() {
671        let tools = tool_with_props(serde_json::json!({
672            "limit": {"type": "string"},
673            "count": {"type": "integer"},
674        }));
675        let text = "<tool_call>\n<function=f>\n\
676            <parameter=limit>4</parameter>\n\
677            <parameter=count>5</parameter>\n\
678            </function>\n</tool_call>";
679        let result = QwenXmlParser::new()
680            .parse_incremental(text, &tools)
681            .await
682            .unwrap();
683        let args: String = result.calls.iter().map(|c| c.parameters.as_str()).collect();
684        assert!(
685            args.contains(r#""limit": "4""#),
686            "string param must stay string: {args}"
687        );
688        assert!(
689            args.contains(r#""count": 5"#),
690            "int param must coerce: {args}"
691        );
692    }
693}