skm_embed/
provider.rs

1//! EmbeddingProvider trait definition.
2
3use async_trait::async_trait;
4
5use crate::embedding::Embedding;
6use crate::error::EmbedError;
7
8/// Core embedding provider trait.
9///
10/// All implementations must be Send + Sync for use across async boundaries.
11/// `embed` is async to support both local inference (via spawn_blocking)
12/// and API-based providers (network IO).
13#[async_trait]
14pub trait EmbeddingProvider: Send + Sync {
15    /// Generate embeddings for a batch of texts.
16    ///
17    /// Returns one embedding per input text, in the same order.
18    async fn embed(&self, texts: &[&str]) -> Result<Vec<Embedding>, EmbedError>;
19
20    /// Vector dimensionality of this provider.
21    fn dimensions(&self) -> usize;
22
23    /// Human-readable model identifier.
24    fn model_id(&self) -> &str;
25
26    /// Maximum batch size supported.
27    fn max_batch_size(&self) -> usize {
28        128 // Default, can be overridden
29    }
30
31    /// Embed a single text (convenience method).
32    async fn embed_one(&self, text: &str) -> Result<Embedding, EmbedError> {
33        if text.is_empty() {
34            return Err(EmbedError::EmptyInput);
35        }
36        let mut embeddings = self.embed(&[text]).await?;
37        embeddings.pop().ok_or_else(|| EmbedError::Embedding("No embedding returned".into()))
38    }
39}
40
41#[cfg(test)]
42pub mod tests {
43    use super::*;
44    use std::sync::atomic::{AtomicUsize, Ordering};
45
46    /// Mock embedding provider for testing.
47    pub struct MockProvider {
48        pub dimensions: usize,
49        pub call_count: AtomicUsize,
50    }
51
52    impl MockProvider {
53        pub fn new(dimensions: usize) -> Self {
54            Self {
55                dimensions,
56                call_count: AtomicUsize::new(0),
57            }
58        }
59    }
60
61    #[async_trait]
62    impl EmbeddingProvider for MockProvider {
63        async fn embed(&self, texts: &[&str]) -> Result<Vec<Embedding>, EmbedError> {
64            self.call_count.fetch_add(1, Ordering::SeqCst);
65
66            let mut results = Vec::with_capacity(texts.len());
67            for text in texts {
68                if text.is_empty() {
69                    return Err(EmbedError::EmptyInput);
70                }
71
72                // Create a deterministic "embedding" based on text hash
73                let hash = xxhash_rust::xxh64::xxh64(text.as_bytes(), 0);
74                let mut vector = vec![0.0f32; self.dimensions];
75
76                // Fill with deterministic values
77                for (i, v) in vector.iter_mut().enumerate() {
78                    *v = ((hash.wrapping_add(i as u64)) % 1000) as f32 / 1000.0;
79                }
80
81                results.push(Embedding::new(vector, hash));
82            }
83
84            Ok(results)
85        }
86
87        fn dimensions(&self) -> usize {
88            self.dimensions
89        }
90
91        fn model_id(&self) -> &str {
92            "mock-embedding-model"
93        }
94    }
95
96    #[tokio::test]
97    async fn test_mock_provider_embed() {
98        let provider = MockProvider::new(384);
99        let texts = vec!["hello", "world"];
100        let embeddings = provider.embed(&texts).await.unwrap();
101
102        assert_eq!(embeddings.len(), 2);
103        assert_eq!(embeddings[0].dimensions(), 384);
104        assert!(embeddings[0].is_normalized());
105    }
106
107    #[tokio::test]
108    async fn test_mock_provider_embed_one() {
109        let provider = MockProvider::new(384);
110        let embedding = provider.embed_one("test").await.unwrap();
111
112        assert_eq!(embedding.dimensions(), 384);
113    }
114
115    #[tokio::test]
116    async fn test_mock_provider_empty_input() {
117        let provider = MockProvider::new(384);
118        let result = provider.embed_one("").await;
119
120        assert!(result.is_err());
121        assert!(matches!(result.unwrap_err(), EmbedError::EmptyInput));
122    }
123
124    #[tokio::test]
125    async fn test_mock_provider_deterministic() {
126        let provider = MockProvider::new(384);
127
128        let e1 = provider.embed_one("test").await.unwrap();
129        let e2 = provider.embed_one("test").await.unwrap();
130
131        // Same input should produce same embedding
132        assert_eq!(e1.text_hash, e2.text_hash);
133        assert_eq!(e1.vector, e2.vector);
134    }
135}
skm_embed/provider.rs

skm_embed/
provider.rs
