1use super::super::ToolDefinition;
5use super::response::{CalledFunction, ToolCallResponse, ToolCallType};
6use regex::Regex;
7use ruff_python_ast::{Expr, Number as PythonNumber};
8use ruff_python_parser::parse_expression;
9use serde_json::{Number as JsonNumber, Value, json};
10use std::sync::OnceLock;
11
12static PYTHONIC_REGEX: OnceLock<Regex> = OnceLock::new();
13
14fn get_pythonic_regex() -> &'static Regex {
17 PYTHONIC_REGEX.get_or_init(|| {
18 let pattern = r"\[([a-zA-Z]+\w*\(([a-zA-Z]+\w*=.*?,\s*)*([a-zA-Z]+\w*=.*?\s?)?\),\s*)*([a-zA-Z]+\w*\(([a-zA-Z]+\w*=.*?,\s*)*([a-zA-Z]+\w*=.*?\s*)?\)\s*)+\]";
20 Regex::new(pattern).expect("Failed to compile pythonic regex pattern")
21 })
22}
23
24fn strip_text(message: &str) -> String {
25 message
27 .replace("<|python_start|>", "")
28 .replace("<|python_end|>", "")
29}
30
31fn get_regex_matches(message: &str) -> Vec<String> {
32 let re = get_pythonic_regex();
33 let mut matches = Vec::new();
34 for cap in re.find_iter(message) {
35 matches.push(cap.as_str().to_string());
36 }
37 matches
38}
39
40pub fn parse_tool_calls(src: &str) -> anyhow::Result<Vec<ToolCallResponse>> {
41 let parsed = parse_expression(src)?;
42 let elts = match parsed.expr() {
43 Expr::List(expr_list) => &expr_list.elts,
44 _ => return Ok(vec![]),
45 };
46
47 let mut res = Vec::with_capacity(elts.len());
48 for (idx, elt) in elts.iter().enumerate() {
49 let (func, keywords) = match elt {
50 Expr::Call(call) => (&call.func, &call.arguments.keywords),
51 _ => continue,
52 };
53
54 let name = match func.as_ref() {
55 Expr::Name(name) => name.id.clone(),
56 _ => continue,
57 };
58
59 let mut obj = serde_json::Map::new();
60 for keyword in keywords.iter() {
61 let Some(arg_ident) = keyword.arg.as_ref() else {
62 tracing::debug!(
63 "Skipping **kwargs in pythonic tool call for function {}",
64 name
65 );
66 continue;
67 };
68
69 match const_expr(&keyword.value) {
70 Ok(value) => {
71 obj.insert(arg_ident.to_string(), value);
72 }
73 Err(e) => {
74 tracing::debug!("Skipping non-constant argument {}: {}", arg_ident, e);
75 }
76 }
77 }
78
79 res.push(ToolCallResponse {
80 id: format!("call-{}", idx + 1),
81 tp: ToolCallType::Function,
82 function: CalledFunction {
83 name: name.to_string(),
84 arguments: serde_json::to_string(&Value::Object(obj))?,
86 },
87 });
88 }
89 Ok(res)
90}
91
92fn const_expr(e: &Expr) -> Result<Value, Box<dyn std::error::Error>> {
93 match e {
94 Expr::BooleanLiteral(literal) => Ok(json!(literal.value)),
95 Expr::NoneLiteral(_) => Ok(Value::Null),
96 Expr::NumberLiteral(literal) => Ok(match &literal.value {
97 PythonNumber::Int(i) => {
98 if let Some(v) = i.as_i64() {
102 Value::Number(JsonNumber::from(v))
103 } else if let Some(v) = i.as_u64() {
104 Value::Number(JsonNumber::from(v))
105 } else {
106 Value::String(i.to_string())
107 }
108 }
109 PythonNumber::Float(f) => json!(f),
110 PythonNumber::Complex { .. } => return Err("unsupported constant type".into()),
111 }),
112 Expr::StringLiteral(literal) => Ok(json!(literal.value.to_str())),
113 Expr::List(expr_list) => {
115 let list_values: Result<Vec<Value>, Box<dyn std::error::Error>> =
116 expr_list.elts.iter().map(|e| const_expr(e)).collect();
117 Ok(json!(list_values?))
118 }
119 Expr::Dict(expr_dict) => {
121 let mut dict_map = std::collections::HashMap::new();
122 for item in &expr_dict.items {
123 let key = match &item.key {
125 Some(k) => match const_expr(k)? {
126 Value::String(s) => s,
127 other => other.to_string(),
128 },
129 None => {
130 return Err(
131 "dictionary unpacking (**kwargs) not supported in constants".into()
132 );
133 }
134 };
135 let value = const_expr(&item.value)?;
136 dict_map.insert(key, value);
137 }
138 Ok(json!(dict_map))
139 }
140 _ => Err("only constant values, lists, and dicts are allowed".into()),
141 }
142}
143
144pub fn try_tool_call_parse_pythonic(
145 message: &str,
146 _tools: Option<&[ToolDefinition]>,
147) -> anyhow::Result<(Vec<ToolCallResponse>, Option<String>)> {
148 let stripped = strip_text(message).trim().to_string();
149
150 if stripped.is_empty() {
152 return Ok((vec![], Some(String::new())));
153 }
154
155 let matches = get_regex_matches(&stripped);
156 if matches.is_empty() {
157 return Ok((vec![], Some(stripped)));
158 }
159
160 let tool_response = parse_tool_calls(&matches[0]);
161
162 let normal_text = stripped
164 .split(&matches[0])
165 .next()
166 .unwrap() .trim()
168 .to_string();
169
170 Ok((tool_response?, Some(normal_text)))
171}
172
173pub fn detect_tool_call_start_pythonic(chunk: &str) -> bool {
174 let trimmed = chunk.trim();
175 if trimmed.is_empty() {
177 return false;
178 }
179 trimmed.contains('[')
181}
182
183#[cfg(test)]
184mod tests {
185 use super::*;
186
187 fn extract_name_and_args(call: ToolCallResponse) -> (String, serde_json::Value) {
188 let args: serde_json::Value = serde_json::from_str(&call.function.arguments).unwrap();
189 (call.function.name, args)
190 }
191
192 #[test] fn test_strip_text() {
194 let message = "Hello, world!";
195 let stripped = strip_text(message);
196 assert_eq!(stripped, "Hello, world!");
197
198 let message = "<|python_start|>foo(a=1, b=2)<|python_end|>";
199 let stripped = strip_text(message);
200 assert_eq!(stripped, "foo(a=1, b=2)");
201
202 let message = "<|python_start|>foo(a=1, b=2)";
203 let stripped = strip_text(message);
204 assert_eq!(stripped, "foo(a=1, b=2)");
205
206 let message = "foo(a=1, b=2)<|python_end|>";
207 let stripped = strip_text(message);
208 assert_eq!(stripped, "foo(a=1, b=2)");
209 }
210
211 #[test] fn test_get_regex_matches_simple_case() {
213 let message = "[foo(a=1, b=2), bar(x=3)]";
215 let matches = get_regex_matches(message);
216 assert_eq!(matches.len(), 1);
217 assert_eq!(matches[0], "[foo(a=1, b=2), bar(x=3)]");
218 }
219
220 #[test] fn test_get_regex_matches_text_before_and_after() {
222 let message = "Hey yo ! [foo(a=1, b=2), bar(x= 3)] Hey yo";
224 let matches = get_regex_matches(message);
225 assert_eq!(matches.len(), 1);
226 assert_eq!(matches[0], "[foo(a=1, b=2), bar(x= 3)]");
227 }
228
229 #[test] fn test_get_regex_matches_new_line_in_arg_and_value() {
231 let message = "Hey \n yo ! [foo(a=1,b=2), \n bar(x=3)] Hey yo";
233 let matches = get_regex_matches(message);
234 assert_eq!(matches.len(), 1);
235 assert_eq!(matches[0], "[foo(a=1,b=2), \n bar(x=3)]");
236 }
237
238 #[test] fn test_get_regex_matches_no_call() {
240 let message = "Hey yo !";
242 let matches = get_regex_matches(message);
243 assert_eq!(matches.len(), 0);
244 }
245
246 #[test] fn test_parse_tool_call_parse_pythonic_basic() {
249 let message = "[foo(a=1, b=2), bar(x=3)]";
250 let (result, content) = try_tool_call_parse_pythonic(message, None).unwrap();
251 assert_eq!(content, Some("".to_string()));
252 assert!(!result.is_empty());
253 assert_eq!(result.len(), 2);
254 let (name, args) = extract_name_and_args(result[0].clone()); assert_eq!(name, "foo");
256 assert_eq!(args["a"], 1);
257 assert_eq!(args["b"], 2);
258 let (name, args) = extract_name_and_args(result[1].clone());
259 assert_eq!(name, "bar");
260 assert_eq!(args["x"], 3);
261 }
262
263 #[test] fn test_parse_tool_call_parse_pythonic_with_text() {
266 let message = "Hey yo ! [foo(a=1, b=2), bar(x=3)] Hey yo";
267 let (result, content) = try_tool_call_parse_pythonic(message, None).unwrap();
268 assert_eq!(content, Some("Hey yo !".to_string()));
269 assert!(!result.is_empty());
270 assert_eq!(result.len(), 2);
271 let (name, args) = extract_name_and_args(result[0].clone());
272 assert_eq!(name, "foo");
273 assert_eq!(args["a"], 1);
274 assert_eq!(args["b"], 2);
275 let (name, args) = extract_name_and_args(result[1].clone());
276 assert_eq!(name, "bar");
277 assert_eq!(args["x"], 3);
278 }
279
280 #[test] fn test_parse_tool_call_parse_pythonic_with_text_and_new_line() {
283 let message = "Hey \n yo ! [foo(a=1, b=2), bar(x=3)] Hey yo";
284 let (result, content) = try_tool_call_parse_pythonic(message, None).unwrap();
285 assert_eq!(content, Some("Hey \n yo !".to_string()));
286 assert!(!result.is_empty());
287 assert_eq!(result.len(), 2);
288 let (name, args) = extract_name_and_args(result[0].clone());
289 assert_eq!(name, "foo");
290 assert_eq!(args["a"], 1);
291 assert_eq!(args["b"], 2);
292 let (name, args) = extract_name_and_args(result[1].clone());
293 assert_eq!(name, "bar");
294 assert_eq!(args["x"], 3);
295 }
296
297 #[test] fn test_parse_tool_call_parse_pythonic_with_no_calls() {
299 let message = "Hey \n yo !";
300 let (result, content) = try_tool_call_parse_pythonic(message, None).unwrap();
301 assert_eq!(content, Some("Hey \n yo !".to_string()));
302 assert!(result.is_empty());
303 assert_eq!(result.len(), 0)
304 }
305
306 #[test] fn test_parse_tool_call_parse_pythonic_with_python_tags() {
309 let message = "<|python_start|>[foo(a=1, b=2), bar(x=3)]<|python_end|>";
310 let (result, content) = try_tool_call_parse_pythonic(message, None).unwrap();
311 assert_eq!(content, Some("".to_string()));
312 assert!(!result.is_empty());
313 assert_eq!(result.len(), 2);
314 let (name, args) = extract_name_and_args(result[0].clone());
315 assert_eq!(name, "foo");
316 assert_eq!(args["a"], 1);
317 assert_eq!(args["b"], 2);
318 let (name, args) = extract_name_and_args(result[1].clone());
319 assert_eq!(name, "bar");
320 assert_eq!(args["x"], 3);
321 }
322
323 #[test] fn test_parse_tool_call_parse_pythonic_with_list_arg_values() {
326 let message = "[foo(a=[1, 2, 3], b=2), bar(x=[3, 4, 5])]";
327 let (result, _) = try_tool_call_parse_pythonic(message, None).unwrap();
328 assert!(!result.is_empty());
329 assert_eq!(result.len(), 2);
330 let (name, args) = extract_name_and_args(result[0].clone());
331 assert_eq!(name, "foo");
332 assert_eq!(args["a"], json!([1, 2, 3]));
333 assert_eq!(args["b"], 2);
334 let (name, args) = extract_name_and_args(result[1].clone());
335 assert_eq!(name, "bar");
336 assert_eq!(args["x"], json!([3, 4, 5]));
337 }
338
339 #[test] fn test_parse_tool_call_parse_pythonic_with_dict_arg_values() {
342 let message = "[foo(a={'a': 1, 'b': 2}, b=2), bar(x={'x': 3, 'y': {'e': 'f'}})]";
343 let (result, _) = try_tool_call_parse_pythonic(message, None).unwrap();
344 assert!(!result.is_empty());
345 assert_eq!(result.len(), 2);
346 let (name, args) = extract_name_and_args(result[0].clone());
347 assert_eq!(name, "foo");
348 assert_eq!(args["a"], json!({"a": 1, "b": 2}));
349 assert_eq!(args["b"], 2);
350 let (name, args) = extract_name_and_args(result[1].clone());
351 assert_eq!(name, "bar");
352 assert_eq!(args["x"], json!({"x": 3, "y": {"e": "f"}}));
353 }
354
355 #[test]
356 fn test_parse_tool_call_parse_pythonic_scalar_literals() {
357 let message = concat!(
358 "[foo(flag=True, empty=None, ratio=1.5, text='hello', ",
359 "huge=18446744073709551616)]"
360 );
361 let (result, _) = try_tool_call_parse_pythonic(message, None).unwrap();
362 let (name, args) = extract_name_and_args(result[0].clone());
363
364 assert_eq!(name, "foo");
365 assert_eq!(args["flag"], true);
366 assert_eq!(args["empty"], Value::Null);
367 assert_eq!(args["ratio"], 1.5);
368 assert_eq!(args["text"], "hello");
369 assert_eq!(args["huge"], "18446744073709551616");
370 }
371
372 #[test]
373 fn test_parse_tool_calls_skips_unsupported_arguments() {
374 let result = parse_tool_calls("[foo(valid=1, computed=1 + 2, **extra)]").unwrap();
375 let (name, args) = extract_name_and_args(result[0].clone());
376
377 assert_eq!(name, "foo");
378 assert_eq!(args, json!({"valid": 1}));
379 }
380}
381
382#[cfg(test)]
383mod detect_parser_tests {
384 use super::*;
385
386 #[test] fn test_detect_tool_call_start_pythonic_chunk_with_tool_call_start_token() {
388 let text = r#"[foo(a=1, b=2), bar(x=3)]"#;
389 let result = detect_tool_call_start_pythonic(text);
390 assert!(result);
391 }
392
393 #[test] fn test_detect_tool_call_start_pythonic_chunk_without_tool_call_start_token() {
395 let text = r#"foo(a=1, b=2)"#;
396 let result = detect_tool_call_start_pythonic(text);
397 assert!(!result);
398 }
399
400 #[test] fn test_detect_tool_call_start_pythonic_chunk_with_tool_call_start_token_in_middle() {
402 let text = r#"information: [foo(a=1, b=2), bar(x=3)]"#;
403 let result = detect_tool_call_start_pythonic(text);
404 assert!(result);
405 }
406
407 #[test] fn test_detect_tool_call_start_pythonic_false_positive() {
409 let text = r#"Hey [ There is one tool call here . foo(a=1, b=2)"#;
411 let result = detect_tool_call_start_pythonic(text);
412 assert!(result);
413 }
414}