zoo_embedding/
mock_generator.rs

1use crate::model_type::EmbeddingModelType;
2use crate::zoo_embedding_errors::ZooEmbeddingError;
3use async_trait::async_trait;
4use crate::embedding_generator::EmbeddingGenerator;
5
6#[derive(Clone)]
7pub struct MockGenerator {
8    model_type: EmbeddingModelType,
9    num_embeddings: usize,
10}
11
12impl MockGenerator {
13    pub fn new(model_type: EmbeddingModelType, num_embeddings: usize) -> Self {
14        MockGenerator {
15            model_type,
16            num_embeddings,
17        }
18    }
19}
20
21#[async_trait]
22impl EmbeddingGenerator for MockGenerator {
23    fn model_type(&self) -> EmbeddingModelType {
24        self.model_type.clone()
25    }
26
27    fn set_model_type(&mut self, model_type: EmbeddingModelType) {
28        self.model_type = model_type;
29    }
30
31    fn box_clone(&self) -> Box<dyn EmbeddingGenerator> {
32        Box::new((*self).clone())
33    }
34
35    fn generate_embedding_blocking(&self, _input_string: &str) -> Result<Vec<f32>, ZooEmbeddingError> {
36        Ok(vec![0.0; self.num_embeddings])
37    }
38
39    fn generate_embeddings_blocking(&self, input_strings: &Vec<String>) -> Result<Vec<Vec<f32>>, ZooEmbeddingError> {
40        Ok(input_strings.iter().map(|_| vec![0.0; self.num_embeddings]).collect())
41    }
42
43    async fn generate_embedding(&self, _input_string: &str) -> Result<Vec<f32>, ZooEmbeddingError> {
44        Ok(vec![0.0; self.num_embeddings])
45    }
46
47    async fn generate_embeddings(&self, input_strings: &Vec<String>) -> Result<Vec<Vec<f32>>, ZooEmbeddingError> {
48        Ok(input_strings.iter().map(|_| vec![0.0; self.num_embeddings]).collect())
49    }
50}
zoo_embedding/mock_generator.rs

zoo_embedding/
mock_generator.rs
