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
15pub struct QwenXmlParser {
29 extractor: Regex,
31
32 buffer: String,
34
35 prev_tool_call_arr: Vec<Value>,
37
38 current_tool_id: i32,
40
41 current_tool_name_sent: bool,
43
44 streamed_args_for_tool: Vec<String>,
46
47 tool_call_start_token: &'static str,
49 tool_call_end_token: &'static str,
50
51 in_tool_call: bool,
53 current_function_name: String,
54 current_parameters: serde_json::Map<String, Value>,
55
56 xml_function_pattern: Regex,
58 xml_param_pattern: Regex,
59 xml_param_open_pattern: Regex,
60
61 open_parameter: Option<OpenParameter>,
63}
64
65const CLOSING_TAGS: [&str; 3] = ["</parameter>", "</function>", "</tool_call>"];
68
69struct OpenParameter {
74 key: String,
75 value_start: usize,
77 streamed_end: Option<usize>,
81 sent: String,
84 literal: bool,
87}
88
89fn json_string_body(text: &str) -> String {
91 let quoted = serde_json::to_string(text).unwrap_or_default();
92 quoted
93 .strip_prefix('"')
94 .and_then(|body| body.strip_suffix('"'))
95 .unwrap_or_default()
96 .to_string()
97}
98
99fn safe_val(raw: &str) -> Value {
113 let trimmed = raw.trim();
114
115 if let Ok(v) = serde_json::from_str::<Value>(trimmed) {
117 return v;
118 }
119
120 match trimmed {
122 "True" => return Value::Bool(true),
123 "False" => return Value::Bool(false),
124 "None" => return Value::Null,
125 _ => {}
126 }
127
128 Value::String(trimmed.to_string())
130}
131
132fn coerce_value(raw: &str, declared_type: Option<&str>) -> Value {
138 let trimmed = raw.trim();
139 helpers::coerce_by_schema_type(trimmed, declared_type).unwrap_or_else(|| safe_val(raw))
140}
141
142impl QwenXmlParser {
143 #[expect(
145 clippy::expect_used,
146 reason = "regex patterns are compile-time string literals"
147 )]
148 pub fn new() -> Self {
149 let pattern = r"(?s)<tool_call>\s*(.*?)\s*</tool_call>";
151 let extractor = Regex::new(pattern).expect("Valid regex pattern");
152
153 let xml_function_pattern =
155 Regex::new(r"<function=([^>]+)>").expect("Valid XML function pattern");
156 let xml_param_pattern = Regex::new(r"(?s)<parameter=([^>]+)>(.*?)</parameter>")
157 .expect("Valid XML parameter pattern");
158 let xml_param_open_pattern =
159 Regex::new(r"<parameter=([^>]+)>").expect("Valid XML parameter start pattern");
160
161 Self {
162 extractor,
163 buffer: String::new(),
164 prev_tool_call_arr: Vec::new(),
165 current_tool_id: -1,
166 current_tool_name_sent: false,
167 streamed_args_for_tool: Vec::new(),
168 tool_call_start_token: "<tool_call>",
169 tool_call_end_token: "</tool_call>",
170 in_tool_call: false,
171 current_function_name: String::new(),
172 current_parameters: serde_json::Map::new(),
173 xml_function_pattern,
174 xml_param_pattern,
175 xml_param_open_pattern,
176 open_parameter: None,
177 }
178 }
179
180 fn parse_xml_format(&self, content: &str, tools: &[Tool]) -> ParserResult<Option<ToolCall>> {
182 let function_captures = self
183 .xml_function_pattern
184 .captures(content)
185 .ok_or_else(|| ParserError::ParsingFailed("No function name found".to_string()))?;
186
187 let function_name = function_captures
188 .get(1)
189 .ok_or_else(|| ParserError::ParsingFailed("Function name capture failed".to_string()))?
190 .as_str()
191 .trim()
192 .to_string();
193
194 if function_name.is_empty() {
195 return Ok(None);
196 }
197
198 let param_types = helpers::param_types_for_function(tools, &function_name);
199 let mut parameters = serde_json::Map::new();
200
201 for cap in self.xml_param_pattern.captures_iter(content) {
202 if let (Some(key_match), Some(value_match)) = (cap.get(1), cap.get(2)) {
203 let key = key_match.as_str().trim().to_string();
204 let value = value_match.as_str();
205 let json_value = coerce_value(value, param_types.get(&key).map(String::as_str));
206 parameters.insert(key, json_value);
207 }
208 }
209
210 let arguments = serde_json::to_string(¶meters)
211 .map_err(|e| ParserError::ParsingFailed(e.to_string()))?;
212
213 Ok(Some(ToolCall {
214 function: FunctionCall {
215 name: function_name,
216 arguments,
217 },
218 }))
219 }
220
221 fn parse_and_stream_parameters(
224 &mut self,
225 tools: &[Tool],
226 parameter_end: usize,
227 ) -> Vec<ToolCallItem> {
228 let mut calls: Vec<ToolCallItem> = vec![];
229 let param_types = helpers::param_types_for_function(tools, &self.current_function_name);
230
231 let mut new_params = serde_json::Map::new();
233 for cap in self
234 .xml_param_pattern
235 .captures_iter(&self.buffer[..parameter_end])
236 {
237 if let (Some(key_match), Some(value_match)) = (cap.get(1), cap.get(2)) {
238 let key = key_match.as_str().trim().to_string();
239 let value = value_match.as_str();
240 let json_value = coerce_value(value, param_types.get(&key).map(String::as_str));
241 new_params.insert(key, json_value);
242 }
243 }
244
245 if new_params != self.current_parameters {
247 let current_args = &mut self.streamed_args_for_tool[self.current_tool_id as usize];
248
249 if self.current_parameters.is_empty() {
250 let mut items = Vec::new();
252 for (key, value) in &new_params {
253 let key_json =
254 serde_json::to_string(key).unwrap_or_else(|_| format!("\"{key}\""));
255 let value_json = serde_json::to_string(value).unwrap_or_default();
256 items.push(format!("{key_json}: {value_json}"));
257 }
258 let json_fragment = format!("{{{}", items.join(", "));
259
260 calls.push(ToolCallItem {
261 tool_index: self.current_tool_id as usize,
262 name: None,
263 parameters: json_fragment.clone(),
264 });
265 *current_args = json_fragment;
266 } else {
267 let new_keys: Vec<_> = new_params
269 .keys()
270 .filter(|k| !self.current_parameters.contains_key(*k))
271 .collect();
272
273 if !new_keys.is_empty() {
274 let mut continuation_parts = Vec::new();
275 for key in new_keys {
276 if let Some(value) = new_params.get(key) {
277 let key_json =
278 serde_json::to_string(key).unwrap_or_else(|_| format!("\"{key}\""));
279 let value_json = serde_json::to_string(value).unwrap_or_default();
280 continuation_parts.push(format!("{key_json}: {value_json}"));
281 }
282 }
283
284 let json_fragment = format!(", {}", continuation_parts.join(", "));
285
286 calls.push(ToolCallItem {
287 tool_index: self.current_tool_id as usize,
288 name: None,
289 parameters: json_fragment.clone(),
290 });
291 current_args.push_str(&json_fragment);
292 }
293 }
294
295 self.current_parameters.clone_from(&new_params);
297 if let Some(tool_obj) =
298 self.prev_tool_call_arr[self.current_tool_id as usize].as_object_mut()
299 {
300 tool_obj.insert("arguments".to_string(), Value::Object(new_params));
301 }
302 }
303
304 calls
305 }
306
307 fn finish_closed_parameter(
311 &mut self,
312 tools: &[Tool],
313 parameter_end: usize,
314 ) -> Vec<ToolCallItem> {
315 let closed_at = self.open_parameter.as_ref().and_then(|open| {
316 self.buffer[open.value_start..parameter_end]
317 .find("</parameter>")
318 .map(|close| open.value_start + close)
319 });
320 let Some(value_end) = closed_at else {
321 return Vec::new();
322 };
323 let param_types = helpers::param_types_for_function(tools, &self.current_function_name);
324 self.finish_open_parameter(¶m_types, value_end)
325 }
326
327 fn stream_open_parameter(&mut self, tools: &[Tool], parameter_end: usize) -> Vec<ToolCallItem> {
332 let mut calls = Vec::new();
333 let param_types = helpers::param_types_for_function(tools, &self.current_function_name);
334
335 let region = &self.buffer[..parameter_end];
337 let Some(open_tag) = self.xml_param_open_pattern.captures_iter(region).last() else {
338 return calls;
339 };
340 let value_start = open_tag.get(0).map_or(parameter_end, |m| m.end());
341 let value_region = ®ion[value_start..];
342 if value_region.contains("</parameter>") {
343 return calls;
344 }
345 let key = open_tag.get(1).map_or("", |m| m.as_str()).trim();
346 if param_types.get(key).map(String::as_str) != Some("string") {
347 return calls;
348 }
349 if self
350 .open_parameter
351 .as_ref()
352 .is_none_or(|open| open.key != key || open.value_start != value_start)
353 {
354 self.open_parameter = Some(OpenParameter {
355 key: key.to_string(),
356 value_start,
357 streamed_end: None,
358 sent: String::new(),
359 literal: false,
360 });
361 }
362
363 let region_end = match value_region.find("</function>") {
366 Some(end) => value_start + end,
367 None => {
368 let held = CLOSING_TAGS
369 .iter()
370 .filter_map(|tag| helpers::ends_with_partial_token(value_region, tag))
371 .max()
372 .unwrap_or(0);
373 parameter_end - held
374 }
375 };
376
377 let tool_index = self.current_tool_id as usize;
378 let Some(open) = self.open_parameter.as_mut() else {
379 return calls;
380 };
381 if open.literal {
382 return calls;
383 }
384 let from = match open.streamed_end {
385 Some(end) => end,
386 None => {
387 let raw = &self.buffer[value_start..region_end];
388 let from = value_start + (raw.len() - raw.trim_start().len());
389 if from >= region_end {
390 return calls;
391 }
392 if self.buffer[from..].starts_with('"') {
393 open.literal = true;
394 return calls;
395 }
396 from
397 }
398 };
399 if region_end <= from {
400 return calls;
401 }
402 let fresh = self.buffer[from..region_end].trim_end();
403 if fresh.is_empty() {
404 return calls;
405 }
406 let body = json_string_body(fresh);
407 let mut delta = String::new();
408 if open.streamed_end.is_none() {
409 let key_json =
410 serde_json::to_string(&open.key).unwrap_or_else(|_| format!("\"{}\"", open.key));
411 delta.push_str(if self.streamed_args_for_tool[tool_index].is_empty() {
412 "{"
413 } else {
414 ", "
415 });
416 delta.push_str(&key_json);
417 delta.push_str(": \"");
418 }
419 delta.push_str(&body);
420 open.sent.push_str(&body);
421 open.streamed_end = Some(from + fresh.len());
422 self.streamed_args_for_tool[tool_index].push_str(&delta);
423 calls.push(ToolCallItem {
424 tool_index,
425 name: None,
426 parameters: delta,
427 });
428 calls
429 }
430
431 fn finish_open_parameter(
435 &mut self,
436 param_types: &HashMap<String, String>,
437 value_end: usize,
438 ) -> Vec<ToolCallItem> {
439 let Some(open) = self.open_parameter.take() else {
440 return Vec::new();
441 };
442 if open.streamed_end.is_none() {
443 return Vec::new();
444 }
445 let tool_index = self.current_tool_id as usize;
446 let value = coerce_value(
447 &self.buffer[open.value_start..value_end],
448 param_types.get(&open.key).map(String::as_str),
449 );
450 let full = serde_json::to_string(&value).unwrap_or_default();
451 let body = full
452 .strip_prefix('"')
453 .and_then(|body| body.strip_suffix('"'))
454 .unwrap_or_default();
455 let mut tail = body
458 .strip_prefix(open.sent.as_str())
459 .unwrap_or_default()
460 .to_string();
461 tail.push('"');
462 self.streamed_args_for_tool[tool_index].push_str(&tail);
463 if let Some(arguments) = self.prev_tool_call_arr[tool_index]
464 .get_mut("arguments")
465 .and_then(Value::as_object_mut)
466 {
467 arguments.insert(open.key.clone(), value.clone());
468 }
469 self.current_parameters.insert(open.key, value);
470 vec![ToolCallItem {
471 tool_index,
472 name: None,
473 parameters: tail,
474 }]
475 }
476
477 fn parse_complete_inner(
479 &self,
480 text: &str,
481 tools: &[Tool],
482 ) -> ParserResult<(String, Vec<ToolCall>)> {
483 if !self.has_tool_markers(text) {
485 return Ok((text.to_string(), vec![]));
486 }
487
488 let idx = text
491 .find(self.tool_call_start_token)
492 .ok_or_else(|| ParserError::ParsingFailed("tool call marker not found".to_string()))?;
493 let normal_text = text[..idx].to_string();
494
495 let mut parsed = Vec::new();
497 for captures in self.extractor.captures_iter(text) {
498 if let Some(content_str) = captures.get(1) {
499 let content = content_str.as_str().trim();
500
501 match self.parse_xml_format(content, tools) {
502 Ok(Some(tool)) => parsed.push(tool),
503 Ok(None) => continue,
504 Err(e) => {
505 tracing::warn!("Failed to parse XML tool call: {:?}", e);
506 continue;
507 }
508 }
509 }
510 }
511
512 if parsed.is_empty() {
514 return Ok((text.to_string(), vec![]));
515 }
516
517 Ok((normal_text, parsed))
518 }
519
520 fn reset_streaming_state(&mut self) {
522 self.in_tool_call = false;
523 self.current_tool_name_sent = false;
524 self.current_function_name.clear();
525 self.current_parameters.clear();
526 self.open_parameter = None;
527 }
528}
529
530impl Default for QwenXmlParser {
531 fn default() -> Self {
532 Self::new()
533 }
534}
535
536#[async_trait]
537impl ToolParser for QwenXmlParser {
538 async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
539 self.parse_complete_inner(text, &[])
540 }
541
542 async fn parse_complete_with_tools(
543 &self,
544 text: &str,
545 tools: &[Tool],
546 ) -> ParserResult<(String, Vec<ToolCall>)> {
547 self.parse_complete_inner(text, tools)
548 }
549
550 async fn parse_incremental(
551 &mut self,
552 chunk: &str,
553 tools: &[Tool],
554 ) -> ParserResult<StreamingParseResult> {
555 self.buffer.push_str(chunk);
556
557 let mut normal_text = String::new();
558 let mut calls: Vec<ToolCallItem> = vec![];
559
560 let tool_indices = helpers::get_tool_indices(tools);
562
563 loop {
564 if !self.in_tool_call && !self.buffer.contains(self.tool_call_start_token) {
566 if helpers::ends_with_partial_token(&self.buffer, self.tool_call_start_token)
568 .is_none()
569 {
570 normal_text.push_str(&self.buffer);
571 self.buffer.clear();
572 }
573 break;
574 }
575
576 if !self.in_tool_call {
578 if let Some(s) = self.buffer.find(self.tool_call_start_token) {
579 normal_text.push_str(&self.buffer[..s]);
580 self.buffer = self.buffer[s + self.tool_call_start_token.len()..].to_string();
581 self.in_tool_call = true;
582 self.current_tool_name_sent = false;
583 self.current_function_name.clear();
584 self.current_parameters.clear();
585 continue;
586 } else {
587 break;
588 }
589 }
590
591 if !self.current_tool_name_sent {
593 if let Some(captures) = self.xml_function_pattern.captures(&self.buffer) {
594 if let Some(name_match) = captures.get(1) {
595 let function_name = name_match.as_str().trim().to_string();
596
597 if tool_indices.contains_key(&function_name) {
599 self.current_function_name.clone_from(&function_name);
600 self.current_tool_name_sent = true;
601
602 if self.current_tool_id == -1 {
604 self.current_tool_id = 0;
605 }
606
607 helpers::ensure_capacity(
609 self.current_tool_id,
610 &mut self.prev_tool_call_arr,
611 &mut self.streamed_args_for_tool,
612 );
613
614 self.prev_tool_call_arr[self.current_tool_id as usize] = serde_json::json!({
616 "name": function_name,
617 "arguments": {}
618 });
619
620 calls.push(ToolCallItem {
622 tool_index: self.current_tool_id as usize,
623 name: Some(function_name),
624 parameters: String::new(),
625 });
626
627 self.buffer =
630 self.buffer[captures.get(0).map_or(0, |m| m.end())..].to_string();
631 continue;
632 } else {
633 tracing::warn!("Invalid function name: {}", function_name);
635 self.reset_streaming_state();
636 normal_text.push_str(&self.buffer);
637 self.buffer.clear();
638 break;
639 }
640 }
641 } else {
642 break;
644 }
645 }
646
647 if self.current_tool_name_sent {
649 let end_pos = self.buffer.find(self.tool_call_end_token);
650 let parameter_end = end_pos.unwrap_or(self.buffer.len());
651 calls.extend(self.finish_closed_parameter(tools, parameter_end));
654 calls.extend(self.parse_and_stream_parameters(tools, parameter_end));
655 calls.extend(self.stream_open_parameter(tools, parameter_end));
656
657 if let Some(end_pos) = end_pos {
659 let unclosed_string = self
662 .open_parameter
663 .take()
664 .is_some_and(|open| open.streamed_end.is_some());
665 let current_args =
666 &mut self.streamed_args_for_tool[self.current_tool_id as usize];
667 let mut closing = String::new();
668 if unclosed_string {
669 closing.push('"');
670 }
671 closing.push_str(if current_args.is_empty() { "{}" } else { "}" });
672 calls.push(ToolCallItem {
673 tool_index: self.current_tool_id as usize,
674 name: None,
675 parameters: closing.clone(),
676 });
677 current_args.push_str(&closing);
678
679 self.buffer =
681 self.buffer[end_pos + self.tool_call_end_token.len()..].to_string();
682 self.reset_streaming_state();
683 self.current_tool_id += 1;
684 continue;
685 } else {
686 break;
688 }
689 }
690
691 break;
692 }
693
694 Ok(StreamingParseResult { normal_text, calls })
695 }
696
697 fn has_tool_markers(&self, text: &str) -> bool {
698 text.contains(self.tool_call_start_token)
699 }
700
701 fn get_unstreamed_tool_args(&self) -> Option<Vec<ToolCallItem>> {
702 let tool_index = self.prev_tool_call_arr.len().checked_sub(1)?;
703 if self
705 .open_parameter
706 .as_ref()
707 .is_some_and(|open| open.streamed_end.is_some())
708 {
709 return Some(vec![ToolCallItem {
710 tool_index,
711 name: None,
712 parameters: "\"}".to_string(),
713 }]);
714 }
715 let actual = self.streamed_args_for_tool.get(tool_index)?;
716 let expected = self.prev_tool_call_arr[tool_index].get("arguments")?;
717 if !actual.is_empty()
720 && serde_json::from_str::<Value>(&format!("{actual}}}"))
721 .ok()
722 .as_ref()
723 == Some(expected)
724 {
725 return Some(vec![ToolCallItem {
726 tool_index,
727 name: None,
728 parameters: "}".to_string(),
729 }]);
730 }
731 helpers::get_unstreamed_args(&self.prev_tool_call_arr, &self.streamed_args_for_tool)
732 }
733
734 fn reset(&mut self) {
735 helpers::reset_parser_state(
736 &mut self.buffer,
737 &mut self.prev_tool_call_arr,
738 &mut self.current_tool_id,
739 &mut self.current_tool_name_sent,
740 &mut self.streamed_args_for_tool,
741 );
742 self.reset_streaming_state();
743 }
744}
745
746#[cfg(test)]
747mod tests {
748 use openai_protocol::common::Function;
749
750 use super::*;
751
752 #[test]
753 fn test_safe_val_json() {
754 assert_eq!(safe_val("42"), Value::Number(42.into()));
755 assert_eq!(safe_val("1.5"), serde_json::json!(1.5));
756 assert_eq!(safe_val("true"), Value::Bool(true));
757 assert_eq!(safe_val("false"), Value::Bool(false));
758 assert_eq!(safe_val("null"), Value::Null);
759 assert_eq!(
760 safe_val(r#"{"key": "value"}"#),
761 serde_json::json!({"key": "value"})
762 );
763 assert_eq!(safe_val(r"[1, 2, 3]"), serde_json::json!([1, 2, 3]));
764 }
765
766 #[test]
767 fn test_safe_val_python_literals() {
768 assert_eq!(safe_val("True"), Value::Bool(true));
769 assert_eq!(safe_val("False"), Value::Bool(false));
770 assert_eq!(safe_val("None"), Value::Null);
771 }
772
773 #[test]
774 fn test_safe_val_string_fallback() {
775 assert_eq!(
776 safe_val("hello world"),
777 Value::String("hello world".to_string())
778 );
779 assert_eq!(safe_val(" spaces "), Value::String("spaces".to_string()));
780 }
781
782 #[test]
785 fn test_safe_val_preserves_html_entities() {
786 assert_eq!(
787 safe_val("<div>"),
788 Value::String("<div>".to_string())
789 );
790 assert_eq!(
791 safe_val("Tom & Jerry"),
792 Value::String("Tom & Jerry".to_string())
793 );
794 assert_eq!(safe_val("it's"), Value::String("it's".to_string()));
796 assert_eq!(safe_val("<"), Value::String("<".to_string()));
797 }
798
799 fn tool_with_props(props: Value) -> Vec<Tool> {
800 vec![Tool {
801 tool_type: "function".to_string(),
802 function: Function {
803 name: "f".to_string(),
804 description: None,
805 parameters: serde_json::json!({"type": "object", "properties": props}),
806 strict: None,
807 extra: Default::default(),
808 },
809 }]
810 }
811
812 #[tokio::test]
814 async fn test_schema_aware_coercion_keeps_strings() {
815 let tools = tool_with_props(serde_json::json!({
816 "limit": {"type": "string"},
817 "flag": {"type": "string"},
818 "coords": {"type": "string"},
819 "cfg": {"type": "string"},
820 "count": {"type": "integer"},
821 }));
822 let text = "<tool_call>\n<function=f>\n\
823 <parameter=limit>4</parameter>\n\
824 <parameter=flag>true</parameter>\n\
825 <parameter=coords>[60,30]</parameter>\n\
826 <parameter=cfg>{\"a\": 1}</parameter>\n\
827 <parameter=count>5</parameter>\n\
828 </function>\n</tool_call>";
829 let (_, calls) = QwenXmlParser::new()
830 .parse_complete_with_tools(text, &tools)
831 .await
832 .unwrap();
833 assert_eq!(calls.len(), 1);
834 let args: Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
835 assert_eq!(args["limit"], Value::String("4".to_string()));
836 assert_eq!(args["flag"], Value::String("true".to_string()));
837 assert_eq!(args["coords"], Value::String("[60,30]".to_string()));
838 assert_eq!(args["cfg"], Value::String("{\"a\": 1}".to_string()));
839 assert_eq!(args["count"], Value::Number(5.into()));
840 }
841
842 #[tokio::test]
844 async fn test_streaming_schema_aware_coercion() {
845 let tools = tool_with_props(serde_json::json!({
846 "limit": {"type": "string"},
847 "count": {"type": "integer"},
848 }));
849 let text = "<tool_call>\n<function=f>\n\
850 <parameter=limit>4</parameter>\n\
851 <parameter=count>5</parameter>\n\
852 </function>\n</tool_call>";
853 let result = QwenXmlParser::new()
854 .parse_incremental(text, &tools)
855 .await
856 .unwrap();
857 let args: String = result.calls.iter().map(|c| c.parameters.as_str()).collect();
858 assert!(
859 args.contains(r#""limit": "4""#),
860 "string param must stay string: {args}"
861 );
862 assert!(
863 args.contains(r#""count": 5"#),
864 "int param must coerce: {args}"
865 );
866 }
867
868 #[tokio::test]
873 async fn test_arg_values_with_entities_roundtrip_unchanged() {
874 let literal = "<a>Tom & Jerry</a> <x> it's";
875
876 let text = format!(
879 "<tool_call>\n<function=f>\n\
880 <parameter=content>\n{literal}\n</parameter>\n\
881 </function>\n</tool_call>"
882 );
883
884 let (_, calls) = QwenXmlParser::new().parse_complete(&text).await.unwrap();
887 assert_eq!(calls.len(), 1);
888 let args: Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
889 assert_eq!(args["content"], Value::String(literal.to_string()));
890
891 let tools = tool_with_props(serde_json::json!({
893 "content": {"type": "string"},
894 }));
895 let (_, calls) = QwenXmlParser::new()
896 .parse_complete_with_tools(&text, &tools)
897 .await
898 .unwrap();
899 let args: Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
900 assert_eq!(args["content"], Value::String(literal.to_string()));
901 }
902
903 fn bash_tool() -> Vec<Tool> {
904 vec![Tool {
905 tool_type: "function".to_string(),
906 function: Function {
907 name: "bash".to_string(),
908 description: None,
909 parameters: serde_json::json!({
910 "type": "object",
911 "properties": {
912 "command": {"type": "string"},
913 "timeout": {"type": "integer"}
914 }
915 }),
916 strict: None,
917 extra: Default::default(),
918 },
919 }]
920 }
921
922 async fn argument_deltas(
925 parser: &mut QwenXmlParser,
926 chunks: &[&str],
927 tools: &[Tool],
928 ) -> Vec<(usize, String)> {
929 let mut deltas = Vec::new();
930 for (index, chunk) in chunks.iter().enumerate() {
931 for call in parser.parse_incremental(chunk, tools).await.unwrap().calls {
932 if call.name.is_none() && !call.parameters.is_empty() {
933 deltas.push((index, call.parameters));
934 }
935 }
936 }
937 deltas
938 }
939
940 #[tokio::test]
944 async fn test_string_values_stream_before_the_parameter_closes() {
945 let tools = bash_tool();
946 let chunks = [
947 "<tool_call>",
948 "\n",
949 "<function=",
950 "bash",
951 ">\n",
952 "<parameter=",
953 "command",
954 ">\n",
955 "ls",
956 " -",
957 "la",
958 " /",
959 "tmp",
960 "\n",
961 "</parameter>",
962 "\n</function>",
963 "\n</tool_call>",
964 ];
965 let close_at = chunks.iter().position(|c| *c == "</parameter>").unwrap();
966 let mut parser = QwenXmlParser::new();
967 let deltas = argument_deltas(&mut parser, &chunks, &tools).await;
968
969 assert!(deltas.len() > 2, "arguments arrived as {deltas:?}");
970 assert_eq!(deltas[0], (8, "{\"command\": \"ls".to_string()));
971 let before_close = deltas.iter().filter(|(at, _)| *at < close_at).count();
972 assert_eq!(before_close, 5, "{deltas:?}");
973 let joined: String = deltas.iter().map(|(_, d)| d.as_str()).collect();
974 assert_eq!(joined, "{\"command\": \"ls -la /tmp\"}");
975 assert!(parser.get_unstreamed_tool_args().is_none());
976 assert!(parser.take_unstreamed_normal_text().is_empty());
977 }
978
979 #[tokio::test]
982 async fn test_streamed_string_values_match_the_complete_parse_at_every_split() {
983 let tools = bash_tool();
984 let command = "<parameter=command>\n echo \"a\\b\"\t<tag> \n\n</parameter>\n";
985 let timeout = "<parameter=timeout>30</parameter>\n";
986 for text in [
989 format!("<tool_call>\n<function=bash>\n{command}{timeout}</function>\n</tool_call>"),
990 format!("<tool_call>\n<function=bash>\n{timeout}{command}</function>\n</tool_call>"),
991 ] {
992 let (_, complete) = QwenXmlParser::new()
993 .parse_complete_with_tools(&text, &tools)
994 .await
995 .unwrap();
996 let expected: Value = serde_json::from_str(&complete[0].function.arguments).unwrap();
997 assert_eq!(expected["command"], "echo \"a\\b\"\t<tag>");
998 assert_eq!(expected["timeout"], 30);
999
1000 let mut feeds: Vec<Vec<String>> = vec![text.chars().map(|c| c.to_string()).collect()];
1001 for (split, _) in text.char_indices().skip(1) {
1002 feeds.push(vec![text[..split].to_string(), text[split..].to_string()]);
1003 }
1004 for chunks in feeds {
1005 let mut parser = QwenXmlParser::new();
1006 let mut args = String::new();
1007 for chunk in &chunks {
1008 for call in parser.parse_incremental(chunk, &tools).await.unwrap().calls {
1009 args.push_str(&call.parameters);
1010 }
1011 }
1012 assert!(parser.get_unstreamed_tool_args().is_none(), "{chunks:?}");
1013 assert_eq!(
1014 serde_json::from_str::<Value>(&args).ok().as_ref(),
1015 Some(&expected),
1016 "args {args}; chunks {chunks:?}"
1017 );
1018 }
1019 }
1020 }
1021
1022 #[tokio::test]
1025 async fn test_values_that_need_the_complete_text_arrive_whole() {
1026 let tools = bash_tool();
1027 for (text, expected) in [
1028 (
1029 "<tool_call><function=bash><parameter=timeout>30</parameter></function></tool_call>",
1030 ["{\"timeout\": 30", "}"],
1031 ),
1032 (
1033 "<tool_call><function=bash><parameter=command>\"ls\"</parameter></function></tool_call>",
1034 ["{\"command\": \"ls\"", "}"],
1035 ),
1036 ] {
1037 let chunks: Vec<String> = text.chars().map(|c| c.to_string()).collect();
1038 let chunk_refs: Vec<&str> = chunks.iter().map(String::as_str).collect();
1039 let mut parser = QwenXmlParser::new();
1040 let deltas: Vec<String> = argument_deltas(&mut parser, &chunk_refs, &tools)
1041 .await
1042 .into_iter()
1043 .map(|(_, delta)| delta)
1044 .collect();
1045 assert_eq!(deltas, expected);
1046 }
1047 }
1048
1049 #[tokio::test]
1052 async fn test_end_of_stream_closes_a_partially_streamed_value() {
1053 let tools = bash_tool();
1054 let mut parser = QwenXmlParser::new();
1055 let deltas = argument_deltas(
1056 &mut parser,
1057 &[
1058 "<tool_call><function=bash><parameter=command>ls -l",
1059 "a /tm",
1060 ],
1061 &tools,
1062 )
1063 .await;
1064 let mut args: String = deltas.into_iter().map(|(_, d)| d).collect();
1065 assert_eq!(args, "{\"command\": \"ls -la /tm");
1066 for item in parser.get_unstreamed_tool_args().unwrap() {
1067 args.push_str(&item.parameters);
1068 }
1069 assert_eq!(args, "{\"command\": \"ls -la /tm\"}");
1070 assert!(parser.take_unstreamed_normal_text().is_empty());
1071
1072 let mut parser = QwenXmlParser::new();
1073 let deltas = argument_deltas(
1074 &mut parser,
1075 &["<tool_call><function=bash><parameter=command>\n"],
1076 &tools,
1077 )
1078 .await;
1079 assert!(deltas.is_empty(), "{deltas:?}");
1080 let pending: String = parser
1081 .get_unstreamed_tool_args()
1082 .unwrap()
1083 .into_iter()
1084 .map(|item| item.parameters)
1085 .collect();
1086 assert_eq!(pending, "{}");
1087 }
1088}