Skip to main content

objectiveai_sdk/functions/check/example_inputs/
string.rs

1use rand::Rng;
2use rand::seq::SliceRandom;
3
4use crate::functions::expression::{InputValue, StringInputSchema};
5
6pub fn permutations(schema: &StringInputSchema) -> usize {
7    if let Some(ref e) = schema.r#enum {
8        e.len()
9    } else {
10        2
11    }
12}
13
14pub fn generate<R: Rng>(
15    schema: &StringInputSchema,
16    mut rng: R,
17) -> Generator<R> {
18    let count = permutations(schema);
19    let mut indices: Vec<usize> = (0..count).collect();
20    indices.shuffle(&mut rng);
21    let variants = if let Some(ref e) = schema.r#enum {
22        Some(e.clone())
23    } else {
24        None
25    };
26    Generator {
27        variants,
28        indices,
29        pos: 0,
30        rng,
31    }
32}
33
34pub struct Generator<R: Rng> {
35    variants: Option<Vec<String>>,
36    indices: Vec<usize>,
37    pos: usize,
38    rng: R,
39}
40
41impl<R: Rng> Iterator for Generator<R> {
42    type Item = InputValue;
43    fn next(&mut self) -> Option<InputValue> {
44        if self.indices.is_empty() {
45            return None;
46        }
47        if self.pos >= self.indices.len() {
48            self.indices.shuffle(&mut self.rng);
49            self.pos = 0;
50        }
51        let index = self.indices[self.pos];
52        self.pos += 1;
53        Some(if let Some(ref e) = self.variants {
54            InputValue::String(e[index].clone())
55        } else if index == 0 {
56            InputValue::String(String::new())
57        } else {
58            InputValue::String(random_string(&mut self.rng))
59        })
60    }
61}
62
63pub fn random_string(rng: &mut impl Rng) -> String {
64    let len = rng.random_range(1..=32);
65    (0..len)
66        .map(|_| rng.random_range(b'a'..=b'z') as char)
67        .collect()
68}