Skip to main content

directory_indexer/embedding/
provider.rs

1use crate::error::Result;
2use async_trait::async_trait;
3
4#[derive(Debug, Clone)]
5pub struct EmbeddingResponse {
6    pub embeddings: Vec<Vec<f32>>,
7    pub model: String,
8    pub usage: Option<EmbeddingUsage>,
9}
10
11#[derive(Debug, Clone)]
12pub struct EmbeddingUsage {
13    pub prompt_tokens: Option<u32>,
14    pub total_tokens: Option<u32>,
15}
16
17#[async_trait]
18pub trait EmbeddingProvider: Send + Sync {
19    fn model_name(&self) -> &str;
20    fn embedding_dimension(&self) -> usize;
21
22    async fn generate_embeddings(&self, texts: Vec<String>) -> Result<EmbeddingResponse>;
23    async fn generate_embedding(&self, text: String) -> Result<Vec<f32>> {
24        let response = self.generate_embeddings(vec![text]).await?;
25        response.embeddings.into_iter().next().ok_or_else(|| {
26            crate::error::IndexerError::embedding("No embedding returned".to_string())
27        })
28    }
29
30    async fn health_check(&self) -> Result<bool>;
31}