Skip to main content

aither_gemini/
embedding.rs

1use aither_core::{EmbeddingModel, Error as AitherError, Result as AitherResult};
2
3use crate::{
4    client::embed_content,
5    config::Gemini,
6    types::{EmbedContentRequest, GeminiContent},
7};
8
9impl EmbeddingModel for Gemini {
10    fn dim(&self) -> usize {
11        self.config().embedding_dimensions
12    }
13
14    fn embed(
15        &self,
16        text: &str,
17    ) -> impl core::future::Future<Output = AitherResult<Vec<f32>>> + Send {
18        let cfg = self.config();
19        let prompt = text.to_owned();
20        async move {
21            let request =
22                EmbedContentRequest::new(&cfg.embedding_model, GeminiContent::text("user", prompt));
23            let response = embed_content(cfg, request)
24                .await
25                .map_err(AitherError::from)?;
26            Ok(response.embedding.values)
27        }
28    }
29}