objectiveai_sdk/functions/check/example_inputs/
string.rs1use 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}