llm_embedding/
service.rs

1//! High-level embedding service wrapping a provider.
2
3use crate::provider::{EmbeddingProvider, EmbeddingResult};
4
5/// Convenience wrapper that adds count-based batch handling on top of a provider.
6#[derive(Debug, Clone)]
7pub struct EmbeddingService<P: EmbeddingProvider> {
8    provider: P,
9    batch_size: usize,
10}
11
12impl<P: EmbeddingProvider> EmbeddingService<P> {
13    /// Creates a service with the given provider and batch size.
14    pub fn new(provider: P, batch_size: usize) -> Self {
15        Self {
16            provider,
17            batch_size: batch_size.max(1),
18        }
19    }
20
21    /// Returns the wrapped provider.
22    pub fn provider(&self) -> &P {
23        &self.provider
24    }
25
26    /// Embeds texts in batches, concatenating results in input order.
27    pub async fn embed_batch(&self, texts: &[String]) -> crate::error::Result<EmbeddingResult> {
28        let mut all = Vec::new();
29        let mut prompt_tokens = 0;
30        let mut total_tokens = 0;
31
32        for chunk in texts.chunks(self.batch_size) {
33            let result = self.provider.embed(chunk).await?;
34            all.extend(result.embeddings);
35            prompt_tokens += result.prompt_tokens;
36            total_tokens += result.total_tokens;
37        }
38
39        Ok(EmbeddingResult {
40            embeddings: all,
41            prompt_tokens,
42            total_tokens,
43        })
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50    use crate::error::EmbeddingError;
51    use async_trait::async_trait;
52
53    struct MockProvider;
54
55    #[async_trait]
56    impl EmbeddingProvider for MockProvider {
57        async fn embed(&self, texts: &[String]) -> crate::error::Result<EmbeddingResult> {
58            if texts.iter().any(|text| text == "boom") {
59                return Err(EmbeddingError::Transport("mock failure".into()));
60            }
61            Ok(EmbeddingResult {
62                embeddings: texts.iter().map(|text| vec![text.len() as f32]).collect(),
63                prompt_tokens: texts.len() as u64,
64                total_tokens: texts.len() as u64,
65            })
66        }
67
68        fn dimension(&self) -> usize {
69            1
70        }
71
72        fn model_name(&self) -> &str {
73            "mock"
74        }
75    }
76
77    #[tokio::test]
78    async fn batches_preserve_order_and_sum_usage() {
79        let service = EmbeddingService::new(MockProvider, 2);
80        let texts: Vec<String> = ["a", "bb", "ccc", "dddd"]
81            .iter()
82            .map(ToString::to_string)
83            .collect();
84        let result = service.embed_batch(&texts).await.expect("batch embed");
85        assert_eq!(
86            result.embeddings,
87            vec![vec![1.0], vec![2.0], vec![3.0], vec![4.0]]
88        );
89        assert_eq!(result.prompt_tokens, 4);
90        assert_eq!(result.total_tokens, 4);
91    }
92
93    #[tokio::test]
94    async fn empty_input_returns_empty_result() {
95        let service = EmbeddingService::new(MockProvider, 2);
96        let result = service.embed_batch(&[]).await.expect("empty batch");
97        assert!(result.embeddings.is_empty());
98    }
99
100    #[tokio::test]
101    async fn provider_error_propagates() {
102        let service = EmbeddingService::new(MockProvider, 2);
103        let result = service.embed_batch(&["boom".to_string()]).await;
104        assert!(result.is_err());
105    }
106}
llm_embedding/service.rs

llm_embedding/
service.rs
