objectiveai_sdk/functions/check/example_inputs/
array.rs1use rand::Rng;
2use rand::SeedableRng;
3use rand::rngs::StdRng;
4use rand::seq::SliceRandom;
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(|_| {
46 Box::new(super::multi::generate(
47 &item_schema,
48 StdRng::seed_from_u64(rng.random::<u64>()),
49 ))
50 })
51 .collect()
52 })
53 .collect();
54
55 let mut length_order: Vec<usize> = (0..lengths.len()).collect();
56 length_order.shuffle(&mut rng);
57
58 Generator {
59 per_length,
60 length_order,
61 pos: 0,
62 rng,
63 }
64}
65
66pub struct Generator<R: Rng> {
67 per_length: Vec<Vec<Box<super::multi::Generator>>>,
69 length_order: Vec<usize>,
70 pos: usize,
71 rng: R,
72}
73
74impl<R: Rng> Iterator for Generator<R> {
75 type Item = InputValue;
76 fn next(&mut self) -> Option<InputValue> {
77 let li = self.length_order[self.pos];
79 self.pos += 1;
80 if self.pos >= self.length_order.len() {
81 self.pos = 0;
82 self.length_order.shuffle(&mut self.rng);
83 }
84
85 let arr: Vec<InputValue> = self.per_length[li]
87 .iter_mut()
88 .map(|g| g.next().unwrap())
89 .collect();
90
91 Some(InputValue::Array(arr))
92 }
93}