swiftide_integrations/dashscope/
embed.rs

1use super::Dashscope;
2use anyhow::{Context as _, Result};
3use async_openai::types::CreateEmbeddingRequestArgs;
4use async_trait::async_trait;
5use swiftide_core::{chat_completion::errors::LanguageModelError, EmbeddingModel, Embeddings};
6
7use crate::openai::openai_error_to_language_model_error;
8
9#[async_trait]
10impl EmbeddingModel for Dashscope {
11    async fn embed(&self, input: Vec<String>) -> Result<Embeddings, LanguageModelError> {
12        let model = self
13            .default_options
14            .embed_model
15            .as_ref()
16            .context("Model not set")
17            .map_err(LanguageModelError::permanent)?;
18
19        let dimensions = self.default_options.dimensions;
20        let request = CreateEmbeddingRequestArgs::default()
21            .model(model)
22            .dimensions(dimensions)
23            .input(&input)
24            .build()
25            .map_err(LanguageModelError::permanent)?;
26        tracing::debug!(
27            num_chunks = input.len(),
28            model = &model,
29            "[Embed] Request to qwen"
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}