cerebro 0.1.8

Blazing-fast, storage-agnostic semantic memory engine for AI Agents — written in pure Rust
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

use async_trait::async_trait;
use crate::traits::{Embedder, Result};

/// A mock embedder for tests or local development without dependencies.
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);
    }
}