docling_rag/embed/
gemini.rs

1//! Google Gemini embedding provider. Uses the Generative Language API's
2//! `embedContent` endpoint with `outputDimensionality` so `gemini-embedding-001`
3//! (natively 3072-dim, Matryoshka-truncatable) returns the configured dimension.
4
5use super::Embedder;
6use crate::{RagConfig, RagError, Result};
7use async_trait::async_trait;
8use serde::{Deserialize, Serialize};
9
10const API_BASE: &str = "https://generativelanguage.googleapis.com/v1beta";
11
12/// Embedder backed by the Gemini API.
13#[derive(Debug, Clone)]
14pub struct GeminiEmbedder {
15    client: reqwest::Client,
16    api_key: String,
17    model: String,
18    dim: usize,
19    id: String,
20}
21
22#[derive(Serialize)]
23struct EmbedContentReq<'a> {
24    model: String,
25    content: Content<'a>,
26    #[serde(rename = "outputDimensionality")]
27    output_dim: usize,
28}
29
30#[derive(Serialize)]
31struct Content<'a> {
32    parts: Vec<Part<'a>>,
33}
34
35#[derive(Serialize)]
36struct Part<'a> {
37    text: &'a str,
38}
39
40#[derive(Deserialize)]
41struct EmbedContentResp {
42    embedding: Embedding,
43}
44
45#[derive(Deserialize)]
46struct Embedding {
47    #[serde(default)]
48    values: Vec<f32>,
49}
50
51impl GeminiEmbedder {
52    /// Build from config; errors if `GEMINI_API_KEY` is unset.
53    pub fn from_config(cfg: &RagConfig) -> Result<Self> {
54        let api_key = cfg.gemini_api_key.clone().ok_or_else(|| {
55            RagError::config("GEMINI_API_KEY is required for the gemini provider")
56        })?;
57        Ok(GeminiEmbedder {
58            client: reqwest::Client::new(),
59            api_key,
60            model: cfg.gemini_model.clone(),
61            dim: cfg.embed_dim,
62            id: format!("gemini:{}", cfg.gemini_model),
63        })
64    }
65}
66
67#[async_trait]
68impl Embedder for GeminiEmbedder {
69    async fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
70        // The public API embeds one content per call; issue them sequentially to
71        // stay well within rate limits (batch endpoints exist but add complexity).
72        let url = format!("{API_BASE}/models/{}:embedContent", self.model);
73        let mut out = Vec::with_capacity(texts.len());
74        for text in texts {
75            let req = EmbedContentReq {
76                model: format!("models/{}", self.model),
77                content: Content {
78                    parts: vec![Part { text }],
79                },
80                output_dim: self.dim,
81            };
82            let resp = self
83                .client
84                .post(&url)
85                .header("x-goog-api-key", &self.api_key)
86                .json(&req)
87                .send()
88                .await?
89                .error_for_status()?;
90            let body: EmbedContentResp = resp.json().await?;
91            if body.embedding.values.is_empty() {
92                return Err(RagError::Embedding(
93                    "gemini returned an empty embedding".into(),
94                ));
95            }
96            out.push(body.embedding.values);
97        }
98        Ok(out)
99    }
100
101    fn dim(&self) -> usize {
102        self.dim
103    }
104
105    fn id(&self) -> &str {
106        &self.id
107    }
108}
docling_rag/embed/gemini.rs

docling_rag/embed/
gemini.rs
