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        if let EmbeddingModelType::Dense(embedding_model) = &*self.embedding_model {
11            embedding_model
12                .embed(input, self.batch_size)
13                .map_err(LanguageModelError::permanent)
14        } else {
15            Err(LanguageModelError::PermanentError(
16                "Expected dense model, got sparse".into(),
17            ))
18        }
19    }
20}