objectiveai_sdk/functions/expression/
runtime.rs

1//! JMESPath runtime with custom functions for expression evaluation.
2//!
3//! Extends the standard JMESPath runtime with additional functions:
4//! - `add(a, b)` - Addition
5//! - `subtract(a, b)` - Subtraction
6//! - `multiply(a, b)` - Multiplication
7//! - `divide(a, b)` - Division (returns null if dividing by zero)
8//! - `mod(a, b)` - Modulo (returns null if dividing by zero)
9//! - `json_parse(s)` - Parse a JSON string
10//! - `is_null(v)` - Check if a value is null
11//! - `if(cond, then, else)` - Conditional expression
12
13use std::sync::LazyLock;
14
15/// Global JMESPath runtime instance with custom functions.
16pub static JMESPATH_RUNTIME: LazyLock<jmespath::Runtime> =
17    LazyLock::new(|| {
18        use jmespath::{
19            Context, ErrorReason, JmespathError, Rcvar, Runtime, RuntimeError,
20            Variable,
21            functions::{ArgumentType, CustomFunction, Signature},
22        };
23        use serde_json::Number;
24        use std::rc::Rc;
25
26        // convert arg
27        fn arg_as_number(
28            arg: &Rcvar,
29            ctx: &Context,
30            position: usize,
31        ) -> Result<f64, JmespathError> {
32            arg.as_number().ok_or_else(|| {
33                JmespathError::new(
34                    ctx.expression,
35                    ctx.offset,
36                    ErrorReason::Runtime(RuntimeError::InvalidType {
37                        expected: "number".to_string(),
38                        actual: arg.get_type().to_string(),
39                        position,
40                    }),
41                )
42            })
43        }
44        #[allow(dead_code)]
45        fn arg_as_string(
46            arg: &Rcvar,
47            ctx: &Context,
48            position: usize,
49        ) -> Result<String, JmespathError> {
50            arg.as_string().cloned().ok_or_else(|| {
51                JmespathError::new(
52                    ctx.expression,
53                    ctx.offset,
54                    ErrorReason::Runtime(RuntimeError::InvalidType {
55                        expected: "string".to_string(),
56                        actual: arg.get_type().to_string(),
57                        position,
58                    }),
59                )
60            })
61        }
62        #[allow(dead_code)]
63        fn arg_as_bool(
64            arg: &Rcvar,
65            ctx: &Context,
66            position: usize,
67        ) -> Result<bool, JmespathError> {
68            arg.as_boolean().ok_or_else(|| {
69                JmespathError::new(
70                    ctx.expression,
71                    ctx.offset,
72                    ErrorReason::Runtime(RuntimeError::InvalidType {
73                        expected: "boolean".to_string(),
74                        actual: arg.get_type().to_string(),
75                        position,
76                    }),
77                )
78            })
79        }
80        fn arg_as_array<'var>(
81            arg: &'var Rcvar,
82            ctx: &Context,
83            position: usize,
84        ) -> Result<&'var Vec<Rcvar>, JmespathError> {
85            arg.as_array().ok_or_else(|| {
86                JmespathError::new(
87                    ctx.expression,
88                    ctx.offset,
89                    ErrorReason::Runtime(RuntimeError::InvalidType {
90                        expected: "array".to_string(),
91                        actual: arg.get_type().to_string(),
92                        position,
93                    }),
94                )
95            })
96        }
97
98        // extract arg
99        fn any_arg(
100            args: &[Rcvar],
101            ctx: &Context,
102            position: usize,
103            expect_args_len: usize,
104        ) -> Result<Rcvar, JmespathError> {
105            args.get(position)
106                .ok_or_else(|| {
107                    JmespathError::new(
108                        ctx.expression,
109                        ctx.offset,
110                        ErrorReason::Runtime(
111                            RuntimeError::NotEnoughArguments {
112                                expected: expect_args_len,
113                                actual: args.len(),
114                            },
115                        ),
116                    )
117                })
118                .cloned()
119        }
120        fn number_arg(
121            args: &[Rcvar],
122            ctx: &Context,
123            position: usize,
124            expect_args_len: usize,
125        ) -> Result<f64, JmespathError> {
126            let arg = any_arg(args, ctx, position, expect_args_len)?;
127            arg_as_number(&arg, ctx, position)
128        }
129        #[allow(dead_code)]
130        fn nullable_number_arg(
131            args: &[Rcvar],
132            ctx: &Context,
133            position: usize,
134            expect_args_len: usize,
135        ) -> Result<Option<f64>, JmespathError> {
136            let arg = any_arg(args, ctx, position, expect_args_len)?;
137            if arg.is_null() {
138                Ok(None)
139            } else {
140                Ok(Some(arg_as_number(&arg, ctx, position)?))
141            }
142        }
143        #[allow(dead_code)]
144        fn string_arg(
145            args: &[Rcvar],
146            ctx: &Context,
147            position: usize,
148            expect_args_len: usize,
149        ) -> Result<String, JmespathError> {
150            let arg = any_arg(args, ctx, position, expect_args_len)?;
151            arg_as_string(&arg, ctx, position)
152        }
153        #[allow(dead_code)]
154        fn nullable_string_arg(
155            args: &[Rcvar],
156            ctx: &Context,
157            position: usize,
158            expect_args_len: usize,
159        ) -> Result<Option<String>, JmespathError> {
160            let arg = any_arg(args, ctx, position, expect_args_len)?;
161            if arg.is_null() {
162                Ok(None)
163            } else {
164                Ok(Some(arg_as_string(&arg, ctx, position)?))
165            }
166        }
167        #[allow(dead_code)]
168        fn bool_arg(
169            args: &[Rcvar],
170            ctx: &Context,
171            position: usize,
172            expect_args_len: usize,
173        ) -> Result<bool, JmespathError> {
174            let arg = any_arg(args, ctx, position, expect_args_len)?;
175            arg_as_bool(&arg, ctx, position)
176        }
177        #[allow(dead_code)]
178        fn nullable_bool_arg(
179            args: &[Rcvar],
180            ctx: &Context,
181            position: usize,
182            expect_args_len: usize,
183        ) -> Result<Option<bool>, JmespathError> {
184            let arg = any_arg(args, ctx, position, expect_args_len)?;
185            if arg.is_null() {
186                Ok(None)
187            } else {
188                Ok(Some(arg_as_bool(&arg, ctx, position)?))
189            }
190        }
191        #[allow(dead_code)]
192        fn array_arg(
193            args: &[Rcvar],
194            ctx: &Context,
195            position: usize,
196            expect_args_len: usize,
197        ) -> Result<Vec<Rcvar>, JmespathError> {
198            let arg = any_arg(args, ctx, position, expect_args_len)?;
199            let array = arg_as_array(&arg, ctx, position)?;
200            Ok(array.clone())
201        }
202        fn number_array_arg(
203            args: &[Rcvar],
204            ctx: &Context,
205            position: usize,
206            expect_args_len: usize,
207        ) -> Result<Vec<f64>, JmespathError> {
208            let arg = any_arg(args, ctx, position, expect_args_len)?;
209            let array = arg_as_array(&arg, ctx, position)?;
210            let mut numbers = Vec::with_capacity(array.len());
211            for item in array.iter() {
212                let number = arg_as_number(item, ctx, position)?;
213                numbers.push(number);
214            }
215            Ok(numbers)
216        }
217
218        // return value
219        fn rcvar_f64(n: f64) -> Rcvar {
220            Rc::new(Variable::Number(
221                Number::from_f64(n).unwrap_or(Number::from_f64(0.0).unwrap()),
222            ))
223        }
224        #[allow(dead_code)]
225        fn rcvar_f64_u64(n: f64) -> Rcvar {
226            Rc::new(Variable::Number(Number::from(n.round() as u64)))
227        }
228
229        let mut runtime = Runtime::new();
230
231        // https://jmespath.org/specification.html
232        runtime.register_builtin_functions();
233
234        // basic math
235        runtime.register_function(
236            "add",
237            Box::new(CustomFunction::new(
238                Signature::new(
239                    vec![ArgumentType::Number, ArgumentType::Number],
240                    None,
241                ),
242                Box::new(|args: &[Rcvar], ctx: &mut Context| {
243                    let a = number_arg(args, ctx, 0, 2)?;
244                    let b = number_arg(args, ctx, 1, 2)?;
245                    Ok(rcvar_f64(a + b))
246                }),
247            )),
248        );
249        runtime.register_function(
250            "subtract",
251            Box::new(CustomFunction::new(
252                Signature::new(
253                    vec![ArgumentType::Number, ArgumentType::Number],
254                    None,
255                ),
256                Box::new(|args: &[Rcvar], ctx: &mut Context| {
257                    let a = number_arg(args, ctx, 0, 2)?;
258                    let b = number_arg(args, ctx, 1, 2)?;
259                    Ok(rcvar_f64(a - b))
260                }),
261            )),
262        );
263        runtime.register_function(
264            "multiply",
265            Box::new(CustomFunction::new(
266                Signature::new(
267                    vec![ArgumentType::Number, ArgumentType::Number],
268                    None,
269                ),
270                Box::new(|args: &[Rcvar], ctx: &mut Context| {
271                    let a = number_arg(args, ctx, 0, 2)?;
272                    let b = number_arg(args, ctx, 1, 2)?;
273                    Ok(rcvar_f64(a * b))
274                }),
275            )),
276        );
277        runtime.register_function(
278            "divide",
279            Box::new(CustomFunction::new(
280                Signature::new(
281                    vec![ArgumentType::Number, ArgumentType::Number],
282                    None,
283                ),
284                Box::new(|args: &[Rcvar], ctx: &mut Context| {
285                    let a = number_arg(args, ctx, 0, 2)?;
286                    let b = number_arg(args, ctx, 1, 2)?;
287                    if b == 0.0 {
288                        Ok(Rc::new(Variable::Null))
289                    } else {
290                        Ok(rcvar_f64(a / b))
291                    }
292                }),
293            )),
294        );
295        runtime.register_function(
296            "mod",
297            Box::new(CustomFunction::new(
298                Signature::new(
299                    vec![ArgumentType::Number, ArgumentType::Number],
300                    None,
301                ),
302                Box::new(|args: &[Rcvar], ctx: &mut Context| {
303                    let a = number_arg(args, ctx, 0, 2)?;
304                    let b = number_arg(args, ctx, 1, 2)?;
305                    if b == 0.0 {
306                        Ok(Rc::new(Variable::Null))
307                    } else {
308                        Ok(rcvar_f64(a % b))
309                    }
310                }),
311            )),
312        );
313
314        // zips a 2D array and maps each column with an expref
315        // if sub-arrays are of different lengths, fills missing values with null
316        runtime.register_function(
317            "zip_map",
318            Box::new(CustomFunction::new(
319                Signature::new(
320                    vec![
321                        ArgumentType::Expref,
322                        ArgumentType::TypedArray(Box::new(ArgumentType::Array)),
323                    ],
324                    None,
325                ),
326                Box::new(|args: &[Rcvar], ctx: &mut Context| {
327                    let expref = args[0].as_expref().unwrap();
328                    let input_array = args[1].as_array().unwrap();
329                    let mut output_array = Vec::with_capacity(
330                        input_array
331                            .iter()
332                            .map(|v| v.as_array().unwrap().len())
333                            .max()
334                            .unwrap_or_default(),
335                    );
336                    for i in 0..output_array.capacity() {
337                        let mut column = Vec::with_capacity(input_array.len());
338                        for row in input_array.iter() {
339                            let row_array = row.as_array().unwrap();
340                            if let Some(value) = row_array.get(i) {
341                                column.push(value.clone());
342                            } else {
343                                column.push(Rc::new(Variable::Null));
344                            }
345                        }
346                        output_array.push(jmespath::interpret(
347                            &Rc::new(Variable::Array(column)),
348                            &expref,
349                            ctx,
350                        )?);
351                    }
352                    Ok(Rc::new(Variable::Array(output_array)))
353                }),
354            )),
355        );
356
357        // L1 normalization of a number array
358        runtime.register_function(
359            "l1_normalize",
360            Box::new(CustomFunction::new(
361                Signature::new(
362                    vec![ArgumentType::TypedArray(Box::new(
363                        ArgumentType::Number,
364                    ))],
365                    None,
366                ),
367                Box::new(|args: &[Rcvar], ctx: &mut Context| {
368                    let numbers = number_array_arg(args, ctx, 0, 1)?;
369                    let sum: f64 = numbers.iter().map(|n| n.abs()).sum();
370                    if numbers.len() == 0 {
371                        Ok(Rc::new(Variable::Array(Vec::new())))
372                    } else if sum == 0.0 {
373                        Ok(Rc::new(Variable::Array(
374                            numbers
375                                .iter()
376                                .map(|_| {
377                                    Rc::new(Variable::Number(
378                                        Number::from_f64(
379                                            1.0 / numbers.len() as f64,
380                                        )
381                                        .unwrap(),
382                                    ))
383                                })
384                                .collect(),
385                        )))
386                    } else {
387                        Ok(Rc::new(Variable::Array(
388                            numbers
389                                .iter()
390                                .map(|n| rcvar_f64(n / sum))
391                                .collect(),
392                        )))
393                    }
394                }),
395            )),
396        );
397
398        runtime
399    });
objectiveai_sdk/functions/expression/runtime.rs

objectiveai_sdk/functions/expression/
runtime.rs
