use async_trait::async_trait;
use crate::traits::{Embedder, Result};
pub struct MockEmbedder {
pub dimension: usize,
}
impl MockEmbedder {
pub fn new(dimension: usize) -> Self {
Self { dimension }
}
}
#[async_trait]
impl Embedder for MockEmbedder {
async fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
let mut results = Vec::with_capacity(texts.len());
for (i, _) in texts.iter().enumerate() {
let mut vec = vec![0.0; self.dimension];
if !vec.is_empty() {
vec[0] = (i as f32) * 0.1;
}
results.push(vec);
}
Ok(results)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_mock_embedder() {
let embedder = MockEmbedder::new(3);
let res = embedder.embed(&["hello", "world"]).await.unwrap();
assert_eq!(res.len(), 2);
assert_eq!(res[0].len(), 3);
}
}