sc/embeddings/
provider.rs

1//! Embedding provider trait.
2//!
3//! Defines the interface that all embedding providers must implement.
4//! Uses async methods for HTTP-based providers.
5
6use crate::error::Result;
7use super::types::ProviderInfo;
8
9/// Trait for embedding providers.
10///
11/// Implemented by Ollama and HuggingFace providers.
12/// The trait is object-safe to allow runtime provider selection.
13pub trait EmbeddingProvider: Send + Sync {
14    /// Get provider metadata.
15    fn info(&self) -> ProviderInfo;
16
17    /// Check if the provider is available.
18    ///
19    /// For Ollama, this checks if the server is running and the model is available.
20    /// For HuggingFace, this checks if the API token is valid.
21    fn is_available(&self) -> impl std::future::Future<Output = bool> + Send;
22
23    /// Generate embedding for a single text.
24    fn generate_embedding(&self, text: &str) -> impl std::future::Future<Output = Result<Vec<f32>>> + Send;
25
26    /// Generate embeddings for multiple texts (batch).
27    ///
28    /// Default implementation calls `generate_embedding` for each text.
29    fn generate_embeddings(&self, texts: &[&str]) -> impl std::future::Future<Output = Result<Vec<Vec<f32>>>> + Send {
30        async move {
31            let mut results = Vec::with_capacity(texts.len());
32            for text in texts {
33                results.push(self.generate_embedding(text).await?);
34            }
35            Ok(results)
36        }
37    }
38}
39
40/// Boxed provider for dynamic dispatch.
41///
42/// Since the trait has async methods with `impl Future`, we need this wrapper
43/// for runtime polymorphism.
44pub struct BoxedProvider {
45    inner: Box<dyn EmbeddingProviderBoxed + Send + Sync>,
46}
47
48/// Object-safe version of EmbeddingProvider for boxing.
49pub trait EmbeddingProviderBoxed: Send + Sync {
50    fn info(&self) -> ProviderInfo;
51    fn is_available_boxed(&self) -> std::pin::Pin<Box<dyn std::future::Future<Output = bool> + Send + '_>>;
52    fn generate_embedding_boxed(&self, text: &str) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<f32>>> + Send + '_>>;
53    fn generate_embeddings_boxed(&self, texts: &[&str]) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<Vec<f32>>>> + Send + '_>>;
54}
55
56impl BoxedProvider {
57    /// Create a new boxed provider.
58    pub fn new<P: EmbeddingProvider + 'static>(provider: P) -> Self {
59        Self {
60            inner: Box::new(BoxedProviderWrapper(provider)),
61        }
62    }
63
64    /// Get provider metadata.
65    pub fn info(&self) -> ProviderInfo {
66        self.inner.info()
67    }
68
69    /// Check if the provider is available.
70    pub async fn is_available(&self) -> bool {
71        self.inner.is_available_boxed().await
72    }
73
74    /// Generate embedding for a single text.
75    pub async fn generate_embedding(&self, text: &str) -> Result<Vec<f32>> {
76        self.inner.generate_embedding_boxed(text).await
77    }
78
79    /// Generate embeddings for multiple texts (batch).
80    pub async fn generate_embeddings(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
81        self.inner.generate_embeddings_boxed(texts).await
82    }
83}
84
85/// Wrapper to implement EmbeddingProviderBoxed for any EmbeddingProvider.
86struct BoxedProviderWrapper<P: EmbeddingProvider + 'static>(P);
87
88impl<P: EmbeddingProvider + 'static> EmbeddingProviderBoxed for BoxedProviderWrapper<P> {
89    fn info(&self) -> ProviderInfo {
90        self.0.info()
91    }
92
93    fn is_available_boxed(&self) -> std::pin::Pin<Box<dyn std::future::Future<Output = bool> + Send + '_>> {
94        Box::pin(self.0.is_available())
95    }
96
97    fn generate_embedding_boxed(&self, text: &str) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<f32>>> + Send + '_>> {
98        // Convert to owned String to avoid lifetime issues with the future.
99        // The underlying provider's generate_embedding takes &str, so we pass a reference
100        // to the owned string in the async block.
101        let text_owned = text.to_string();
102        Box::pin(async move {
103            // We need a reference that lives in this async block
104            self.0.generate_embedding(&text_owned).await
105        })
106    }
107
108    fn generate_embeddings_boxed(&self, texts: &[&str]) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<Vec<f32>>>> + Send + '_>> {
109        // Convert to owned Strings to avoid lifetime issues.
110        let texts_owned: Vec<String> = texts.iter().map(|s| (*s).to_string()).collect();
111        Box::pin(async move {
112            // Create references to the owned strings for the provider call
113            let refs: Vec<&str> = texts_owned.iter().map(String::as_str).collect();
114            self.0.generate_embeddings(&refs).await
115        })
116    }
117}
sc/embeddings/provider.rs

sc/embeddings/
provider.rs
