embed_anything/embeddings/cloud/
gemini.rs

1use reqwest::Client;
2use serde::Deserialize;
3use serde_json::json;
4
5use crate::embeddings::embed::EmbeddingResult;
6
7#[derive(Deserialize, Debug, Default)]
8pub struct GeminiEmbedResponse {
9    pub embeddings: Vec<GeminiEmbeddingData>,
10}
11
12#[derive(Deserialize, Debug, Default)]
13pub struct GeminiEmbeddingData {
14    pub embedding: Vec<f32>,
15}
16
17/// Represents a GeminiEmbedder struct that contains the URL and API key for making requests to the Gemini API.
18#[derive(Debug)]
19pub struct GeminiEmbedder {
20    url: String,
21    api_key: String,
22    client: Client,
23}
24
25impl Default for GeminiEmbedder {
26    fn default() -> Self {
27        Self::new(None)
28    }
29}
30
31impl GeminiEmbedder {
32    pub fn new(api_key: Option<String>) -> Self {
33        let api_key =
34            api_key.unwrap_or_else(|| std::env::var("GEMINI_API_KEY").expect("API Key not set"));
35
36        Self {
37            url: "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:embedContent".to_string(),
38            api_key,
39            client: Client::new(),
40        }
41    }
42
43    pub async fn embed(&self, text_batch: &[&str]) -> Result<Vec<EmbeddingResult>, anyhow::Error> {
44        // Convert text_batch to the format expected by Gemini API
45        let contents: Vec<serde_json::Value> = text_batch
46            .iter()
47            .map(|text| {
48                json!({
49                    "parts": [{"text": text}]
50                })
51            })
52            .collect();
53
54        let request_body = json!({
55            "contents": contents,
56            "embedding_config": {
57                "task_type": "SEMANTIC_SIMILARITY"
58            }
59        });
60
61        let response = self
62            .client
63            .post(&self.url)
64            .header("Content-Type", "application/json")
65            .header("x-goog-api-key", &self.api_key)
66            .json(&request_body)
67            .send()
68            .await?;
69
70        let data = response.json::<GeminiEmbedResponse>().await?;
71        let encodings = data
72            .embeddings
73            .iter()
74            .map(|data| EmbeddingResult::DenseVector(data.embedding.clone()))
75            .collect::<Vec<_>>();
76
77        Ok(encodings)
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84
85    #[tokio::test]
86    async fn test_gemini_embed() {
87        let gemini = GeminiEmbedder::default();
88        let contents: Vec<serde_json::Value> = vec!["Hello world"]
89            .iter()
90            .map(|text| {
91                json!({
92                    "parts": [{"text": text}]
93                })
94            })
95            .collect();
96
97        let request_body = json!({
98            "contents": contents,
99            "embedding_config": {
100                "task_type": "SEMANTIC_SIMILARITY"
101            }
102        });
103
104        let response = gemini
105            .client
106            .post(&gemini.url)
107            .header("Content-Type", "application/json")
108            .header("x-goog-api-key", &gemini.api_key)
109            .json(&request_body)
110            .send()
111            .await
112            .unwrap();
113
114        let data = response.json::<GeminiEmbedResponse>().await.unwrap();
115        println!("{:?}", data);
116    }
117}
embed_anything/embeddings/cloud/gemini.rs

embed_anything/embeddings/cloud/
gemini.rs
