rig_core/providers/gemini/
embedding.rs

1// ================================================================
2//! Google Gemini Embeddings Integration
3//! From [Gemini API Reference](https://ai.google.dev/api/embeddings)
4// ================================================================
5
6use serde_json::json;
7
8use super::{Client, client::ApiResponse};
9use crate::{
10    embeddings::{self, EmbeddingError},
11    http_client::HttpClientExt,
12    wasm_compat::WasmCompatSend,
13};
14
15/// `gemini-embedding-001` embedding model (3072 dimensions by default)
16pub const EMBEDDING_001: &str = "gemini-embedding-001";
17/// `text-embedding-004` embedding model (768 dimensions by default)
18pub const EMBEDDING_004: &str = "text-embedding-004";
19
20/// Returns the default output dimensionality for known Gemini embedding models.
21///
22/// See <https://ai.google.dev/gemini-api/docs/models#gemini-embedding>
23fn model_default_ndims(model: &str) -> Option<usize> {
24    match model {
25        EMBEDDING_001 => Some(3072),
26        EMBEDDING_004 => Some(768),
27        _ => None,
28    }
29}
30
31#[derive(Clone)]
32pub struct EmbeddingModel<T = reqwest::Client> {
33    client: Client<T>,
34    model: String,
35    ndims: usize,
36}
37
38impl<T> EmbeddingModel<T> {
39    pub fn new(client: Client<T>, model: impl Into<String>, ndims: usize) -> Self {
40        Self {
41            client,
42            model: model.into(),
43            ndims,
44        }
45    }
46
47    pub fn with_model(client: Client<T>, model: &str, ndims: usize) -> Self {
48        Self {
49            client,
50            model: model.to_string(),
51            ndims,
52        }
53    }
54}
55
56impl<T> embeddings::EmbeddingModel for EmbeddingModel<T>
57where
58    T: Clone + HttpClientExt + 'static,
59{
60    type Client = Client<T>;
61
62    const MAX_DOCUMENTS: usize = 1024;
63
64    fn make(client: &Self::Client, model: impl Into<String>, dims: Option<usize>) -> Self {
65        let model = model.into();
66        let ndims = dims.or_else(|| model_default_ndims(&model)).unwrap_or(768);
67        Self::new(client.clone(), model, ndims)
68    }
69
70    fn ndims(&self) -> usize {
71        self.ndims
72    }
73
74    /// <https://ai.google.dev/api/embeddings#batch_embed_contents-SHELL>
75    async fn embed_texts(
76        &self,
77        documents: impl IntoIterator<Item = String> + WasmCompatSend,
78    ) -> Result<Vec<embeddings::Embedding>, EmbeddingError> {
79        let documents: Vec<String> = documents.into_iter().collect();
80
81        // Google batch embed requests. See docstrings for API ref link.
82        let requests: Vec<_> = documents
83            .iter()
84            .map(|doc| {
85                json!({
86                    "model": format!("models/{}", self.model),
87                    "content": json!({
88                        "parts": [json!({
89                            "text": doc.to_string()
90                        })]
91                    }),
92                    "output_dimensionality": self.ndims,
93                })
94            })
95            .collect();
96
97        let request_body = json!({ "requests": requests  });
98
99        if let Ok(pretty_body) = serde_json::to_string_pretty(&request_body) {
100            tracing::trace!(
101                target: "rig::embedding",
102                "Sending embedding request to Gemini API {pretty_body}"
103            );
104        }
105
106        let request_body = serde_json::to_vec(&request_body)?;
107        let path = format!("/v1beta/models/{}:batchEmbedContents", self.model);
108        let req = self
109            .client
110            .post(path.as_str())?
111            .body(request_body)
112            .map_err(|e| EmbeddingError::HttpError(e.into()))?;
113        let response = self.client.send::<_, Vec<u8>>(req).await?;
114
115        let status = response.status();
116        let body = response.into_body().await?;
117
118        // A non-success status may carry an error body that does not match the
119        // `ApiResponse` shape (Gemini nests its error under `error`, with no
120        // top-level `message`), so preserve it verbatim before trying to
121        // deserialize the success envelope — otherwise it would surface as a
122        // `JsonError` and the provider_response_* helpers would lose it.
123        if !status.is_success() {
124            return Err(EmbeddingError::from_http_response(
125                status,
126                String::from_utf8_lossy(&body),
127            ));
128        }
129
130        match serde_json::from_slice::<ApiResponse<gemini_api_types::EmbeddingResponse>>(&body)? {
131            ApiResponse::Ok(response) => {
132                let docs = documents
133                    .into_iter()
134                    .zip(response.embeddings)
135                    .map(|(document, embedding)| embeddings::Embedding {
136                        document,
137                        vec: embedding
138                            .values
139                            .into_iter()
140                            .filter_map(|n| n.as_f64())
141                            .collect(),
142                    })
143                    .collect();
144
145                Ok(docs)
146            }
147            ApiResponse::Err(err) => {
148                tracing::warn!(message = %err.message, "provider returned an error response");
149                Err(EmbeddingError::from_http_response(
150                    status,
151                    String::from_utf8_lossy(&body),
152                ))
153            }
154        }
155    }
156}
157
158// =================================================================
159// Gemini API Types
160// =================================================================
161/// Rust Implementation of the Gemini Types from [Gemini API Reference](https://ai.google.dev/api/embeddings)
162mod gemini_api_types {
163    use serde::Deserialize;
164
165    #[derive(Debug, Deserialize)]
166    pub struct EmbeddingResponse {
167        pub embeddings: Vec<EmbeddingValues>,
168    }
169
170    #[derive(Debug, Deserialize)]
171    pub struct EmbeddingValues {
172        #[serde(default)]
173        pub values: Vec<serde_json::Number>,
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180
181    #[test]
182    fn test_embedding_values_deserializes_without_empty_values_field() {
183        let values: gemini_api_types::EmbeddingValues =
184            serde_json::from_str("{}").expect("empty embedding values should deserialize");
185        assert!(values.values.is_empty());
186    }
187
188    #[test]
189    fn test_model_default_ndims_lookup() {
190        assert_eq!(model_default_ndims(EMBEDDING_001), Some(3072));
191        assert_eq!(model_default_ndims(EMBEDDING_004), Some(768));
192        assert_eq!(model_default_ndims("unknown-model"), None);
193    }
194
195    #[test]
196    fn test_make_resolves_default_dims() {
197        let client = Client::new("test_key").unwrap();
198
199        // EMBEDDING_001 defaults to 3072
200        let model =
201            <EmbeddingModel as embeddings::EmbeddingModel>::make(&client, EMBEDDING_001, None);
202        assert_eq!(embeddings::EmbeddingModel::ndims(&model), 3072);
203
204        // EMBEDDING_004 defaults to 768
205        let model =
206            <EmbeddingModel as embeddings::EmbeddingModel>::make(&client, EMBEDDING_004, None);
207        assert_eq!(embeddings::EmbeddingModel::ndims(&model), 768);
208
209        // Unknown model falls back to 768
210        let model = <EmbeddingModel as embeddings::EmbeddingModel>::make(
211            &client,
212            "some-future-model",
213            None,
214        );
215        assert_eq!(embeddings::EmbeddingModel::ndims(&model), 768);
216    }
217
218    #[test]
219    fn test_make_respects_explicit_dims() {
220        let client = Client::new("test_key").unwrap();
221
222        let model =
223            <EmbeddingModel as embeddings::EmbeddingModel>::make(&client, EMBEDDING_001, Some(256));
224        assert_eq!(embeddings::EmbeddingModel::ndims(&model), 256);
225    }
226
227    #[test]
228    fn test_new_uses_provided_ndims() {
229        let client = Client::new("test_key").unwrap();
230
231        let model = EmbeddingModel::new(client, EMBEDDING_001, 512);
232        assert_eq!(embeddings::EmbeddingModel::ndims(&model), 512);
233    }
234
235    #[tokio::test]
236    async fn embedding_non_success_preserves_status_and_body() {
237        use crate::client::embeddings::EmbeddingsClient;
238        use crate::embeddings::EmbeddingModel as _;
239        use crate::test_utils::RecordingHttpClient;
240
241        // A realistic Gemini error body: the error is nested under `error` with no
242        // top-level `message`, so it does NOT match the `ApiResponse` envelope. The
243        // non-success status guard must preserve it verbatim (otherwise it would
244        // surface as a `JsonError`).
245        let body =
246            r#"{"error":{"code":503,"message":"service unavailable","status":"UNAVAILABLE"}}"#;
247        let http_client =
248            RecordingHttpClient::with_error_response(http::StatusCode::SERVICE_UNAVAILABLE, body);
249        let client = Client::builder()
250            .api_key("test-key")
251            .http_client(http_client)
252            .build()
253            .expect("build client");
254        let model = client.embedding_model(EMBEDDING_001);
255
256        let error = model
257            .embed_texts(vec!["hello".to_string()])
258            .await
259            .expect_err("should fail with non-success status");
260
261        assert!(matches!(error, EmbeddingError::HttpError(_)));
262        assert_eq!(
263            error.provider_response_status(),
264            Some(http::StatusCode::SERVICE_UNAVAILABLE)
265        );
266        assert_eq!(error.provider_response_body(), Some(body));
267    }
268
269    #[tokio::test]
270    async fn embedding_2xx_error_envelope_preserves_status_and_body() {
271        use crate::client::embeddings::EmbeddingsClient;
272        use crate::embeddings::EmbeddingModel as _;
273        use crate::test_utils::RecordingHttpClient;
274
275        // 200 OK whose body deserializes to `ApiResponse::Err` (requires a top-level
276        // `message`; absence of `embeddings` keeps it off the success variant).
277        let body =
278            r#"{"message":"boom","error":{"code":503,"message":"boom","status":"UNAVAILABLE"}}"#;
279        let http_client = RecordingHttpClient::new(body); // 200 OK
280        let client = Client::builder()
281            .api_key("test-key")
282            .http_client(http_client)
283            .build()
284            .expect("build client");
285        let model = client.embedding_model(EMBEDDING_001);
286
287        let error = model
288            .embed_texts(vec!["hello".to_string()])
289            .await
290            .expect_err("should fail with provider error envelope");
291
292        match &error {
293            EmbeddingError::ProviderResponse(stored) => {
294                assert_eq!(stored.body, body);
295                assert_eq!(stored.status, Some(http::StatusCode::OK));
296            }
297            other => panic!("expected ProviderResponse, got {other:?}"),
298        }
299    }
300}
rig_core/providers/gemini/embedding.rs

rig_core/providers/gemini/
embedding.rs
