Skip to main content

dynamo_parsers/tool_calling/pythonic/
pythonic_parser.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use 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
14/// Get the compiled regex pattern for pythonic tool calls
15/// Initialize the regex pattern once, no need to compile it everytime
16fn get_pythonic_regex() -> &'static Regex {
17    PYTHONIC_REGEX.get_or_init(|| {
18        // Format Structure: [tool1(arg1=val1, arg2=val2), tool2(arg1=val3)]
19        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    // Remove unexpected python tags if any
26    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                // Safety: `Value::Object` is always valid JSON, so serialization cannot fail
85                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                // Ruff stores large integers as their source text, so the
99                // existing i64/u64-or-string behavior does not need a bigint
100                // dependency.
101                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        // Handle Python lists as expressions, not constants
114        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        // Handle Python dictionaries as expressions, not constants
120        Expr::Dict(expr_dict) => {
121            let mut dict_map = std::collections::HashMap::new();
122            for item in &expr_dict.items {
123                // Keys should be strings for JSON compatibility
124                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    // Early exit if no content
151    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    // normal text is everything before the first match
163    let normal_text = stripped
164        .split(&matches[0])
165        .next()
166        .unwrap() // Safety: `split()` always returns at least one element (the string before the first delimiter, or the entire string if delimiter not found)
167        .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    // Early return for empty input
176    if trimmed.is_empty() {
177        return false;
178    }
179    // Heuristic: Pythonic tool calls always start with a '[' somewhere in the chunk
180    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] // helper
193    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] // helper
212    fn test_get_regex_matches_simple_case() {
213        // Simple Case
214        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] // helper
221    fn test_get_regex_matches_text_before_and_after() {
222        // Spacing in arg and value and text before and after
223        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] // helper, TOOLCALLING.batch.7
230    fn test_get_regex_matches_new_line_in_arg_and_value() {
231        // New Line in Arg and value
232        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] // helper
239    fn test_get_regex_matches_no_call() {
240        // No Call
241        let message = "Hey yo !";
242        let matches = get_regex_matches(message);
243        assert_eq!(matches.len(), 0);
244    }
245
246    // DEPRECATED(parser-fixture-duplicate): Duplicate of YAML fixture coverage: TOOLCALLING.batch.2.a in tests/parity/toolcalling/fixtures/pythonic/TOOLCALLING.batch.2.yaml.
247    #[test] // TOOLCALLING.batch.2
248    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()); // TODO: Add support for normal text
255        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    // DEPRECATED(parser-fixture-duplicate): Duplicate of YAML fixture coverage: TOOLCALLING.batch.2.c, TOOLCALLING.batch.8.c in tests/parity/toolcalling/fixtures/pythonic/TOOLCALLING.batch.2.yaml, tests/parity/toolcalling/fixtures/pythonic/TOOLCALLING.batch.8.yaml.
264    #[test] // TOOLCALLING.batch.2, TOOLCALLING.batch.8
265    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    // DEPRECATED(parser-fixture-duplicate): Duplicate of YAML fixture coverage: TOOLCALLING.batch.2.c, TOOLCALLING.batch.8.c in tests/parity/toolcalling/fixtures/pythonic/TOOLCALLING.batch.2.yaml, tests/parity/toolcalling/fixtures/pythonic/TOOLCALLING.batch.8.yaml.
281    #[test] // TOOLCALLING.batch.2, TOOLCALLING.batch.8, TOOLCALLING.fmt.2
282    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] // TOOLCALLING.batch.3
298    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    // DEPRECATED(parser-fixture-duplicate): Duplicate of YAML fixture coverage: TOOLCALLING.batch.2.a in tests/parity/toolcalling/fixtures/pythonic/TOOLCALLING.batch.2.yaml.
307    #[test] // TOOLCALLING.batch.2, TOOLCALLING.fmt.3
308    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    // DEPRECATED(parser-fixture-duplicate): Duplicate of YAML fixture coverage: TOOLCALLING.batch.7.a in tests/parity/toolcalling/fixtures/pythonic/TOOLCALLING.batch.7.yaml.
324    #[test] // TOOLCALLING.batch.7
325    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    // DEPRECATED(parser-fixture-duplicate): Duplicate of YAML fixture coverage: TOOLCALLING.batch.7.d in tests/parity/toolcalling/fixtures/pythonic/TOOLCALLING.batch.7.yaml.
340    #[test] // TOOLCALLING.batch.7
341    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] // helper
387    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] // helper
394    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] // helper
401    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] // helper
408    fn test_detect_tool_call_start_pythonic_false_positive() {
409        // Since we detect just "[" as tool call start token, this will be a false positive
410        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}