objectiveai_sdk/functions/check/example_inputs/
array.rs1use rand::Rng;
2use rand::rngs::StdRng;
3use rand::seq::SliceRandom;
4use rand::SeedableRng;
5
6use crate::functions::expression::{ArrayInputSchema, InputValue};
7
8fn length_range(schema: &ArrayInputSchema) -> (usize, usize) {
9 match (schema.min_items, schema.max_items) {
10 (Some(min), Some(max)) => (min as usize, max as usize),
11 (Some(min), None) => (min as usize, min as usize + 10),
12 (None, Some(max)) => ((max as usize).saturating_sub(10), max as usize),
13 (None, None) => (0, 10),
14 }
15}
16
17fn distinct_lengths(schema: &ArrayInputSchema) -> Vec<usize> {
19 let (min, max) = length_range(schema);
20 let mid = (min + max) / 2;
21 let mut lengths = vec![min];
22 if mid != min {
23 lengths.push(mid);
24 }
25 if max != min && max != mid {
26 lengths.push(max);
27 }
28 lengths
29}
30
31pub fn permutations(schema: &ArrayInputSchema) -> usize {
32 let item_perms = super::optional::inner_permutations(schema.items.as_ref());
33 item_perms.saturating_mul(distinct_lengths(schema).len())
34}
35
36pub fn generate<R: Rng>(schema: &ArrayInputSchema, mut rng: R) -> Generator<R> {
37 let item_schema = schema.items.as_ref().clone();
38 let lengths = distinct_lengths(schema);
39
40 let per_length: Vec<Vec<Box<super::multi::Generator>>> = lengths
42 .iter()
43 .map(|&len| {
44 (0..len)
45 .map(|_| Box::new(super::multi::generate(&item_schema, StdRng::seed_from_u64(rng.random::<u64>()))))
46 .collect()
47 })
48 .collect();
49
50 let mut length_order: Vec<usize> = (0..lengths.len()).collect();
51 length_order.shuffle(&mut rng);
52
53 Generator {
54 per_length,
55 length_order,
56 pos: 0,
57 rng,
58 }
59}
60
61pub struct Generator<R: Rng> {
62 per_length: Vec<Vec<Box<super::multi::Generator>>>,
64 length_order: Vec<usize>,
65 pos: usize,
66 rng: R,
67}
68
69impl<R: Rng> Iterator for Generator<R> {
70 type Item = InputValue;
71 fn next(&mut self) -> Option<InputValue> {
72 let li = self.length_order[self.pos];
74 self.pos += 1;
75 if self.pos >= self.length_order.len() {
76 self.pos = 0;
77 self.length_order.shuffle(&mut self.rng);
78 }
79
80 let arr: Vec<InputValue> = self.per_length[li]
82 .iter_mut()
83 .map(|g| g.next().unwrap())
84 .collect();
85
86 Some(InputValue::Array(arr))
87 }
88}