Skip to main content

llm_embedding/
provider.rs

1//! Provider trait and result types for embeddings.
2
3use async_trait::async_trait;
4use serde::{Deserialize, Serialize};
5
6use crate::error::{EmbeddingError, Result};
7
8/// Result of an embedding operation.
9#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
10pub struct EmbeddingResult {
11    /// One embedding vector per input text, in input order.
12    pub embeddings: Vec<Vec<f32>>,
13    /// Number of prompt tokens reported by the provider.
14    pub prompt_tokens: u64,
15    /// Total number of tokens reported by the provider.
16    pub total_tokens: u64,
17}
18
19/// Port for text-to-vector embedding capability.
20#[async_trait]
21pub trait EmbeddingProvider: Send + Sync {
22    /// Embeds a batch of texts, returning vectors in input order.
23    async fn embed(&self, texts: &[String]) -> Result<EmbeddingResult>;
24
25    /// Embeds a single text.
26    async fn embed_one(&self, text: &str) -> Result<Vec<f32>> {
27        let result = self.embed(&[text.to_string()]).await?;
28        result
29            .embeddings
30            .into_iter()
31            .next()
32            .ok_or_else(|| EmbeddingError::Decode("provider returned no embedding".into()))
33    }
34
35    /// Expected vector dimension of this provider.
36    fn dimension(&self) -> usize;
37
38    /// Model name served by this provider.
39    fn model_name(&self) -> &str;
40}