Skip to main content

adk_rag/
embedding.rs

1//! Embedding provider trait for generating vector embeddings from text.
2
3use async_trait::async_trait;
4
5use crate::error::Result;
6
7/// A provider that generates vector embeddings from text input.
8///
9/// Implementations wrap specific embedding backends (Gemini, OpenAI, etc.)
10/// behind a unified async interface. The default [`embed_batch`](EmbeddingProvider::embed_batch)
11/// implementation calls [`embed`](EmbeddingProvider::embed) sequentially;
12/// backends that support native batching should override it.
13///
14/// # Example
15///
16/// ```rust,ignore
17/// use adk_rag::EmbeddingProvider;
18///
19/// let provider = MyEmbeddingProvider::new();
20/// let embedding = provider.embed("hello world").await?;
21/// assert_eq!(embedding.len(), provider.dimensions());
22/// ```
23#[async_trait]
24pub trait EmbeddingProvider: Send + Sync {
25    /// Generate an embedding vector for a single text input.
26    async fn embed(&self, text: &str) -> Result<Vec<f32>>;
27
28    /// Generate embedding vectors for a batch of text inputs.
29    ///
30    /// The default implementation calls [`embed`](EmbeddingProvider::embed)
31    /// sequentially for each input. Override this method if the backend
32    /// supports native batch embedding for better throughput.
33    async fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
34        let mut results = Vec::with_capacity(texts.len());
35        for text in texts {
36            results.push(self.embed(text).await?);
37        }
38        Ok(results)
39    }
40
41    /// Return the dimensionality of embeddings produced by this provider.
42    fn dimensions(&self) -> usize;
43}