swiftide_integrations/ollama/
embed.rs

1use anyhow::{Context as _, Result};
2use async_openai::types::CreateEmbeddingRequestArgs;
3use async_trait::async_trait;
4
5use swiftide_core::{chat_completion::errors::LanguageModelError, EmbeddingModel, Embeddings};
6
7use super::Ollama;
8use crate::openai::openai_error_to_language_model_error;
9
10#[async_trait]
11impl EmbeddingModel for Ollama {
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            .context("Model not set")?;
18
19        let request = CreateEmbeddingRequestArgs::default()
20            .model(model)
21            .input(&input)
22            .build()
23            .map_err(LanguageModelError::permanent)?;
24        tracing::debug!(
25            num_chunks = input.len(),
26            model = &model,
27            "[Embed] Request to openai"
28        );
29        let response = self
30            .client
31            .embeddings()
32            .create(request)
33            .await
34            .map_err(openai_error_to_language_model_error)?;
35
36        let num_embeddings = response.data.len();
37        tracing::debug!(num_embeddings = num_embeddings, "[Embed] Response openai");
38
39        // WARN: Naively assumes that the order is preserved. Might not always be the case.
40        Ok(response.data.into_iter().map(|d| d.embedding).collect())
41    }
42}
swiftide_integrations/ollama/embed.rs

swiftide_integrations/ollama/
embed.rs
