llm-embedding 0.1.0

Embedding client for LLM providers
Documentation 
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
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

//! High-level embedding service wrapping a provider.

use crate::provider::{EmbeddingProvider, EmbeddingResult};

/// Convenience wrapper that adds count-based batch handling on top of a provider.
#[derive(Debug, Clone)]
pub struct EmbeddingService<P: EmbeddingProvider> {
    provider: P,
    batch_size: usize,
}

impl<P: EmbeddingProvider> EmbeddingService<P> {
    /// Creates a service with the given provider and batch size.
    pub fn new(provider: P, batch_size: usize) -> Self {
        Self {
            provider,
            batch_size: batch_size.max(1),
        }
    }

    /// Returns the wrapped provider.
    pub fn provider(&self) -> &P {
        &self.provider
    }

    /// Embeds texts in batches, concatenating results in input order.
    pub async fn embed_batch(&self, texts: &[String]) -> crate::error::Result<EmbeddingResult> {
        let mut all = Vec::new();
        let mut prompt_tokens = 0;
        let mut total_tokens = 0;

        for chunk in texts.chunks(self.batch_size) {
            let result = self.provider.embed(chunk).await?;
            all.extend(result.embeddings);
            prompt_tokens += result.prompt_tokens;
            total_tokens += result.total_tokens;
        }

        Ok(EmbeddingResult {
            embeddings: all,
            prompt_tokens,
            total_tokens,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::error::EmbeddingError;
    use async_trait::async_trait;

    struct MockProvider;

    #[async_trait]
    impl EmbeddingProvider for MockProvider {
        async fn embed(&self, texts: &[String]) -> crate::error::Result<EmbeddingResult> {
            if texts.iter().any(|text| text == "boom") {
                return Err(EmbeddingError::Transport("mock failure".into()));
            }
            Ok(EmbeddingResult {
                embeddings: texts.iter().map(|text| vec![text.len() as f32]).collect(),
                prompt_tokens: texts.len() as u64,
                total_tokens: texts.len() as u64,
            })
        }

        fn dimension(&self) -> usize {
            1
        }

        fn model_name(&self) -> &str {
            "mock"
        }
    }

    #[tokio::test]
    async fn batches_preserve_order_and_sum_usage() {
        let service = EmbeddingService::new(MockProvider, 2);
        let texts: Vec<String> = ["a", "bb", "ccc", "dddd"]
            .iter()
            .map(ToString::to_string)
            .collect();
        let result = service.embed_batch(&texts).await.expect("batch embed");
        assert_eq!(
            result.embeddings,
            vec![vec![1.0], vec![2.0], vec![3.0], vec![4.0]]
        );
        assert_eq!(result.prompt_tokens, 4);
        assert_eq!(result.total_tokens, 4);
    }

    #[tokio::test]
    async fn empty_input_returns_empty_result() {
        let service = EmbeddingService::new(MockProvider, 2);
        let result = service.embed_batch(&[]).await.expect("empty batch");
        assert!(result.embeddings.is_empty());
    }

    #[tokio::test]
    async fn provider_error_propagates() {
        let service = EmbeddingService::new(MockProvider, 2);
        let result = service.embed_batch(&["boom".to_string()]).await;
        assert!(result.is_err());
    }
}