swiftide_integrations/openai/
embed.rs

1use async_trait::async_trait;
2
3use swiftide_core::{EmbeddingModel, Embeddings, chat_completion::errors::LanguageModelError};
4
5use super::GenericOpenAI;
6use crate::openai::openai_error_to_language_model_error;
7
8#[async_trait]
9impl<C: async_openai::config::Config + std::default::Default + Sync + Send + std::fmt::Debug>
10    EmbeddingModel for GenericOpenAI<C>
11{
12    async fn embed(&self, input: Vec<String>) -> Result<Embeddings, LanguageModelError> {
13        let model = self
14            .default_options
15            .embed_model
16            .as_ref()
17            .ok_or(LanguageModelError::PermanentError("Model not set".into()))?;
18
19        let request = self
20            .embed_request_defaults()
21            .model(model)
22            .input(&input)
23            .build()
24            .map_err(LanguageModelError::permanent)?;
25
26        tracing::debug!(
27            num_chunks = input.len(),
28            model = &model,
29            "[Embed] Request to openai"
30        );
31        let response = self
32            .client
33            .embeddings()
34            .create(request)
35            .await
36            .map_err(openai_error_to_language_model_error)?;
37
38        let num_embeddings = response.data.len();
39        tracing::debug!(num_embeddings = num_embeddings, "[Embed] Response openai");
40
41        // WARN: Naively assumes that the order is preserved. Might not always be the case.
42        Ok(response.data.into_iter().map(|d| d.embedding).collect())
43    }
44}