Skip to main content

swiftide_integrations/fastembed/
embedding_model.rs

1use anyhow::Result;
2use async_trait::async_trait;
3use swiftide_core::{EmbeddingModel, Embeddings, chat_completion::errors::LanguageModelError};
4
5use super::{EmbeddingModelType, FastEmbed};
6#[async_trait]
7impl EmbeddingModel for FastEmbed {
8    #[tracing::instrument(skip_all)]
9    async fn embed(&self, input: Vec<String>) -> Result<Embeddings, LanguageModelError> {
10        let mut embedding_model = self.embedding_model.lock().await;
11
12        match &mut *embedding_model {
13            EmbeddingModelType::Dense(model) => model
14                .embed(input, self.batch_size)
15                .map_err(LanguageModelError::permanent),
16            EmbeddingModelType::Sparse(_) => Err(LanguageModelError::PermanentError(
17                "Expected dense model, got sparse".into(),
18            )),
19        }
20    }
21}