1use std::sync::OnceLock;
2
3use 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#[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 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
48pub struct PythonicParser {
50 buffer: String,
52 emitted: usize,
54}
55
56impl Default for PythonicParser {
57 fn default() -> Self {
58 Self::new()
59 }
60}
61
62impl PythonicParser {
63 pub fn new() -> Self {
65 Self {
66 buffer: String::new(),
67 emitted: 0,
68 }
69 }
70
71 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 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 Ok((text.to_string(), vec![]))
114 } else {
115 Ok((normal_text, calls))
116 }
117 }
118 Err(e) => {
119 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 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 if let Some(end) = find_matching_bracket(&cleaned, start) {
148 let call_text = &cleaned[start..=end];
150
151 match self.parse_complete(call_text).await {
152 Ok((_, calls)) => {
153 let remaining_text = &cleaned[end + 1..];
155 self.buffer = remaining_text.to_string();
156
157 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 self.buffer.clear();
187 return Ok(StreamingParseResult::default());
188 }
189 }
190 } else {
191 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 return Ok(StreamingParseResult::default());
204 }
205 }
206
207 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
225fn 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 }
243
244pub(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}