Skip to main content

tool_parser/parsers/
pythonic.rs

1use std::sync::OnceLock;
2
3/// Pythonic format parser for tool calls
4///
5/// Handles Python function call syntax within square brackets:
6/// ```text
7/// [tool1(arg1=val1, arg2=val2), tool2(arg1=val3)]
8/// ```
9///
10/// This format is used by Llama models and uses Python literals
11/// rather than JSON for arguments.
12/// Reference: https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct?chat_template=default
13use async_trait::async_trait;
14use num_traits::ToPrimitive;
15use openai_protocol::common::Tool;
16use regex::Regex;
17use rustpython_parser::{
18    ast::{Constant, Expr, Mod, UnaryOp},
19    parse, Mode,
20};
21use serde_json::{Map, Number, Value};
22
23use crate::{
24    errors::{ParserError, ParserResult},
25    parsers::helpers,
26    traits::ToolParser,
27    types::{FunctionCall, StreamingParseResult, ToolCall, ToolCallItem},
28};
29
30static PYTHONIC_BLOCK_REGEX: OnceLock<Regex> = OnceLock::new();
31
32/// Lazily compiled regex that locates pythonic tool call blocks.
33#[expect(
34    clippy::expect_used,
35    reason = "regex pattern is a compile-time string literal"
36)]
37fn pythonic_block_regex() -> &'static Regex {
38    PYTHONIC_BLOCK_REGEX.get_or_init(|| {
39        // Matches one or more function calls inside a list. The `(?s)` flag allows
40        // newlines inside argument lists while keeping the pattern anchored to
41        // identifiers followed by parentheses, preventing plain lists like
42        // `[1, 2, 3]` from matching.
43        Regex::new(r"(?s)\[\s*[A-Za-z_]\w*\s*\(.*?\)\s*(?:,\s*[A-Za-z_]\w*\s*\(.*?\)\s*)*\]")
44            .expect("pythonic tool call regex must compile")
45    })
46}
47
48/// Parser for Pythonic tool call format
49pub struct PythonicParser {
50    /// Buffer for accumulating chunks
51    buffer: String,
52    /// Calls handed out so far: the next accepted call's index.
53    emitted: usize,
54}
55
56impl Default for PythonicParser {
57    fn default() -> Self {
58        Self::new()
59    }
60}
61
62impl PythonicParser {
63    /// Create a new Pythonic parser
64    pub fn new() -> Self {
65        Self {
66            buffer: String::new(),
67            emitted: 0,
68        }
69    }
70
71    /// Extract the first pythonic tool call block and return it along with the
72    /// surrounding "normal" content.
73    fn extract_tool_calls(text: &str) -> Option<(String, String)> {
74        pythonic_block_regex().find(text).map(|mat| {
75            let block = mat.as_str().to_string();
76            let normal = format!("{}{}", &text[..mat.start()], &text[mat.end()..]);
77            (block, normal)
78        })
79    }
80
81    /// Strip special tokens that Llama models might output
82    fn strip_special_tokens(text: &str) -> String {
83        text.replace("<|python_start|>", "")
84            .replace("<|python_end|>", "")
85    }
86
87    fn parse_tool_call_block(block: &str) -> ParserResult<Vec<ToolCall>> {
88        let expr = parse_python_expression(block)?;
89        match expr {
90            Expr::List(list_expr) => list_expr
91                .elts
92                .into_iter()
93                .enumerate()
94                .map(|(idx, call_expr)| build_tool_call(call_expr, idx))
95                .collect(),
96            _ => Err(ParserError::ParsingFailed(
97                "Expected a list of function calls in pythonic tool call".to_string(),
98            )),
99        }
100    }
101}
102
103#[async_trait]
104impl ToolParser for PythonicParser {
105    async fn parse_complete(&self, text: &str) -> ParserResult<(String, Vec<ToolCall>)> {
106        let cleaned = Self::strip_special_tokens(text);
107
108        if let Some((tool_calls_text, normal_text)) = Self::extract_tool_calls(&cleaned) {
109            match Self::parse_tool_call_block(&tool_calls_text) {
110                Ok(calls) => {
111                    if calls.is_empty() {
112                        // No tools successfully parsed despite having markers
113                        Ok((text.to_string(), vec![]))
114                    } else {
115                        Ok((normal_text, calls))
116                    }
117                }
118                Err(e) => {
119                    // Log and return entire text as fallback
120                    tracing::debug!("Failed to parse pythonic tool calls: {}", e);
121                    Ok((text.to_string(), vec![]))
122                }
123            }
124        } else {
125            Ok((text.to_string(), vec![]))
126        }
127    }
128
129    async fn parse_incremental(
130        &mut self,
131        chunk: &str,
132        tools: &[Tool],
133    ) -> ParserResult<StreamingParseResult> {
134        self.buffer.push_str(chunk);
135
136        let cleaned = Self::strip_special_tokens(&self.buffer);
137
138        // Look for opening bracket
139        if let Some(start) = cleaned.find('[') {
140            let normal_text = if start > 0 {
141                cleaned[..start].to_string()
142            } else {
143                String::new()
144            };
145
146            // Look for matching closing bracket
147            if let Some(end) = find_matching_bracket(&cleaned, start) {
148                // Found complete tool call - extract it and parse using parse_complete
149                let call_text = &cleaned[start..=end];
150
151                match self.parse_complete(call_text).await {
152                    Ok((_, calls)) => {
153                        // Update buffer with remaining text after tool call
154                        let remaining_text = &cleaned[end + 1..];
155                        self.buffer = remaining_text.to_string();
156
157                        // Validate tool names and convert ToolCall to ToolCallItem.
158                        // Indices count the accepted calls across the stream, so an
159                        // unknown name shifts nothing and a later batch continues.
160                        let tool_indices = helpers::get_tool_indices(tools);
161                        let mut items: Vec<ToolCallItem> = Vec::with_capacity(calls.len());
162                        for tool in calls {
163                            if !tool_indices.contains_key(&tool.function.name) {
164                                tracing::debug!(
165                                    "Invalid tool name '{}' - skipping",
166                                    tool.function.name
167                                );
168                                continue;
169                            }
170                            items.push(ToolCallItem {
171                                tool_index: self.emitted + items.len(),
172                                name: Some(tool.function.name),
173                                parameters: tool.function.arguments,
174                            });
175                        }
176                        self.emitted += items.len();
177
178                        return Ok(StreamingParseResult {
179                            normal_text,
180                            calls: items,
181                        });
182                    }
183                    Err(e) => {
184                        tracing::debug!("Failed to parse pythonic tool call: {}", e);
185                        // Clear buffer on error
186                        self.buffer.clear();
187                        return Ok(StreamingParseResult::default());
188                    }
189                }
190            } else {
191                // We have an opening bracket but no closing bracket yet
192                // Put back everything from the bracket onwards
193                self.buffer = cleaned[start..].to_string();
194
195                if !normal_text.is_empty() {
196                    return Ok(StreamingParseResult {
197                        normal_text,
198                        calls: vec![],
199                    });
200                }
201
202                // Still accumulating a potential tool call
203                return Ok(StreamingParseResult::default());
204            }
205        }
206
207        // No tool call bracket found
208        self.buffer.clear();
209        Ok(StreamingParseResult {
210            normal_text: cleaned,
211            calls: vec![],
212        })
213    }
214
215    fn has_tool_markers(&self, text: &str) -> bool {
216        let cleaned = Self::strip_special_tokens(text);
217        if pythonic_block_regex().is_match(&cleaned) {
218            return true;
219        }
220
221        false
222    }
223}
224
225/// Find the matching closing bracket for the opening bracket at start position.
226/// Properly handles nested brackets.
227fn find_matching_bracket(buffer: &str, start: usize) -> Option<usize> {
228    let mut bracket_count = 0;
229    let chars: Vec<char> = buffer.chars().collect();
230
231    for (i, &ch) in chars.iter().enumerate().skip(start) {
232        if ch == '[' {
233            bracket_count += 1;
234        } else if ch == ']' {
235            bracket_count -= 1;
236            if bracket_count == 0 {
237                return Some(i);
238            }
239        }
240    }
241    None // No matching bracket found
242}
243
244// Also used by the Sarashina parser to parse its Python-literal payload;
245// keep the (source -> JSON) behavior stable when editing.
246pub(crate) fn parse_python_expression(source: &str) -> ParserResult<Expr> {
247    let module = parse(source, Mode::Expression, "<pythonic_tool_call>")
248        .map_err(|err| ParserError::ParsingFailed(err.to_string()))?;
249
250    match module {
251        Mod::Expression(expr_mod) => Ok(*expr_mod.body),
252        _ => Err(ParserError::ParsingFailed(
253            "Expected a Python expression".to_string(),
254        )),
255    }
256}
257
258fn build_tool_call(expr: Expr, _index: usize) -> ParserResult<ToolCall> {
259    match expr {
260        Expr::Call(call_expr) => {
261            if !call_expr.args.is_empty() {
262                return Err(ParserError::ParsingFailed(
263                    "Positional arguments are not supported in pythonic tool calls".to_string(),
264                ));
265            }
266
267            let function_name = match *call_expr.func {
268                Expr::Name(name_expr) => name_expr.id.to_string(),
269                _ => {
270                    return Err(ParserError::ParsingFailed(
271                        "Unsupported function reference in pythonic tool call".to_string(),
272                    ))
273                }
274            };
275
276            let mut arguments_map = Map::with_capacity(call_expr.keywords.len());
277            for keyword in call_expr.keywords {
278                let arg_name = keyword.arg.ok_or_else(|| {
279                    ParserError::ParsingFailed(
280                        "pythonic tool calls do not support **kwargs".to_string(),
281                    )
282                })?;
283                let value_json = expression_to_json(&keyword.value)?;
284                arguments_map.insert(arg_name.to_string(), value_json);
285            }
286
287            let arguments_json = Value::Object(arguments_map);
288            let arguments_string = serde_json::to_string(&arguments_json)?;
289
290            Ok(ToolCall {
291                function: FunctionCall {
292                    name: function_name,
293                    arguments: arguments_string,
294                },
295            })
296        }
297        _ => Err(ParserError::ParsingFailed(
298            "Expected function calls inside pythonic tool call list".to_string(),
299        )),
300    }
301}
302
303pub(crate) fn expression_to_json(expr: &Expr) -> ParserResult<Value> {
304    match expr {
305        Expr::Constant(expr_constant) => constant_to_json(&expr_constant.value),
306        Expr::List(list_expr) => collect_sequence(&list_expr.elts).map(Value::Array),
307        Expr::Tuple(tuple_expr) => collect_sequence(&tuple_expr.elts).map(Value::Array),
308        Expr::Dict(dict_expr) => {
309            collect_dict(&dict_expr.keys, &dict_expr.values).map(Value::Object)
310        }
311        Expr::UnaryOp(unary_expr) => match unary_expr.op {
312            UnaryOp::USub => match unary_expr.operand.as_ref() {
313                Expr::Constant(const_expr) => negate_constant(&const_expr.value),
314                _ => Err(ParserError::ParsingFailed(
315                    "Unsupported unary operand in pythonic tool call".to_string(),
316                )),
317            },
318            UnaryOp::UAdd => expression_to_json(unary_expr.operand.as_ref()),
319            _ => Err(ParserError::ParsingFailed(format!(
320                "Unsupported unary operator in pythonic tool call: {:?}",
321                unary_expr.op
322            ))),
323        },
324        Expr::Name(name_expr) => Ok(Value::String(name_expr.id.to_string())),
325        _ => Err(ParserError::ParsingFailed(format!(
326            "Unsupported expression in pythonic tool call: {expr:?}"
327        ))),
328    }
329}
330
331fn constant_to_json(constant: &Constant) -> ParserResult<Value> {
332    match constant {
333        Constant::None => Ok(Value::Null),
334        Constant::Bool(b) => Ok(Value::Bool(*b)),
335        Constant::Int(value) => Ok(integer_constant_to_value(value, false)),
336        Constant::Float(f) => Number::from_f64(*f).map(Value::Number).ok_or_else(|| {
337            ParserError::ParsingFailed("Invalid float literal in pythonic tool call".to_string())
338        }),
339        Constant::Str(s) => Ok(Value::String(s.clone())),
340        Constant::Bytes(bytes) => Ok(Value::String(String::from_utf8_lossy(bytes).into_owned())),
341        Constant::Tuple(values) => constant_tuple_to_array(values).map(Value::Array),
342        Constant::Ellipsis | Constant::Complex { .. } => Err(ParserError::ParsingFailed(
343            "Unsupported literal in pythonic tool call".to_string(),
344        )),
345    }
346}
347
348fn negate_constant(constant: &Constant) -> ParserResult<Value> {
349    match constant {
350        Constant::Int(value) => Ok(integer_constant_to_value(value, true)),
351        Constant::Float(f) => Number::from_f64(-f).map(Value::Number).ok_or_else(|| {
352            ParserError::ParsingFailed("Invalid float literal in pythonic tool call".to_string())
353        }),
354        _ => Err(ParserError::ParsingFailed(
355            "Unsupported unary operand in pythonic tool call".to_string(),
356        )),
357    }
358}
359
360fn value_to_key_string(value: Value) -> ParserResult<String> {
361    match value {
362        Value::String(s) => Ok(s),
363        Value::Number(num) => Ok(num.to_string()),
364        Value::Bool(b) => Ok(b.to_string()),
365        Value::Null => Ok("null".to_string()),
366        other => Err(ParserError::ParsingFailed(format!(
367            "Unsupported key type in pythonic tool call: {other:?}"
368        ))),
369    }
370}
371
372fn collect_sequence(elements: &[Expr]) -> ParserResult<Vec<Value>> {
373    elements.iter().map(expression_to_json).collect()
374}
375
376fn collect_dict(keys: &[Option<Expr>], values: &[Expr]) -> ParserResult<Map<String, Value>> {
377    let mut map = Map::with_capacity(keys.len());
378    for (key_expr, value_expr) in keys.iter().zip(values.iter()) {
379        let key_expr = key_expr.as_ref().ok_or_else(|| {
380            ParserError::ParsingFailed("pythonic tool calls do not support **kwargs".to_string())
381        })?;
382        let key_value = expression_to_json(key_expr)?;
383        let key = value_to_key_string(key_value)?;
384        let value_json = expression_to_json(value_expr)?;
385        map.insert(key, value_json);
386    }
387    Ok(map)
388}
389
390fn constant_tuple_to_array(values: &[Constant]) -> ParserResult<Vec<Value>> {
391    values.iter().map(constant_to_json).collect()
392}
393
394fn integer_constant_to_value<T>(value: &T, negate: bool) -> Value
395where
396    T: ToPrimitive + std::fmt::Display,
397{
398    if let Some(mut i) = value.to_i64() {
399        if negate {
400            i = -i;
401        }
402        return Value::Number(Number::from(i));
403    }
404
405    if negate {
406        if let Some(u) = value.to_u64() {
407            if u <= i64::MAX as u64 {
408                return Value::Number(Number::from(-(u as i64)));
409            }
410            return Value::String(format!("-{value}"));
411        }
412        Value::String(format!("-{value}"))
413    } else if let Some(u) = value.to_u64() {
414        Value::Number(Number::from(u))
415    } else {
416        Value::String(value.to_string())
417    }
418}