use async_trait::async_trait;
#[derive(Debug, Clone, PartialEq)]
pub struct Embeddings {
pub vectors: Vec<Vec<f32>>,
pub model: String,
pub input_tokens: u64,
}
impl Embeddings {
#[must_use]
pub fn new(model: impl Into<String>) -> Self {
Self {
vectors: Vec::new(),
model: model.into(),
input_tokens: 0,
}
}
#[must_use]
pub fn vector(mut self, vector: Vec<f32>) -> Self {
self.vectors.push(vector);
self
}
#[must_use]
pub fn usage(mut self, input_tokens: u64) -> Self {
self.input_tokens = input_tokens;
self
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EmbedError {
NotConfigured,
InvalidInput(String),
Provider(String),
}
impl std::fmt::Display for EmbedError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let scrub = crate::logging::scrub_text;
match self {
Self::NotConfigured => f.write_str("no embedder is wired"),
Self::InvalidInput(message) => {
write!(f, "embed request rejected: {}", scrub(message))
}
Self::Provider(message) => write!(f, "embed failed: {}", scrub(message)),
}
}
}
impl std::error::Error for EmbedError {}
#[async_trait]
pub trait Embedder: Send + Sync {
async fn embed(&self, texts: &[String]) -> Result<Embeddings, EmbedError>;
}