gemini_client_api/gemini/
embed.rs

1use super::ask::BASE_URL;
2use super::error::GeminiResponseError;
3use super::types::embedding::*;
4use super::types::request::Part;
5use reqwest::Client;
6use std::time::Duration;
7
8/// Client for generating embeddings using Gemini embedding models.
9///
10/// # Example
11/// ```no_run
12/// use gemini_client_api::gemini::embed::GeminiEmbedding;
13/// use gemini_client_api::gemini::types::embedding::TaskType;
14///
15/// # async fn run() {
16/// let embedder = GeminiEmbedding::new("YOUR_API_KEY", "gemini-embedding-001")
17///     .set_task_type(TaskType::RetrievalDocument);
18///
19/// let response = embedder.embed_text("Hello, world!").await.unwrap();
20/// println!("Embedding dimension: {}", response.embedding().values().len());
21/// # }
22/// ```
23#[derive(Clone, Debug)]
24pub struct GeminiEmbedding {
25    client: Client,
26    api_key: String,
27    model: String,
28    config: Option<EmbedContentConfig>,
29}
30
31impl GeminiEmbedding {
32    /// Creates a new `GeminiEmbedding` client.
33    ///
34    /// # Arguments
35    /// * `api_key` - Your Gemini API key. Get one from [Google AI studio](https://aistudio.google.com/app/apikey).
36    /// * `model` - The embedding model to use (e.g., `"gemini-embedding-001"`).
37    ///   See [embedding models](https://ai.google.dev/gemini-api/docs/models#gemini-embedding).
38    pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Self {
39        Self {
40            client: Client::builder()
41                .timeout(Duration::from_secs(60))
42                .build()
43                .unwrap(),
44            api_key: api_key.into(),
45            model: model.into(),
46            config: None,
47        }
48    }
49    /// Creates a new `GeminiEmbedding` client with a custom `reqwest::Client`.
50    ///
51    /// # Arguments
52    /// * `api_key` - Your Gemini API key.
53    /// * `model` - The embedding model to use.
54    /// * `client` - A custom `reqwest::Client` for making requests.
55    pub fn new_with_client(
56        api_key: impl Into<String>,
57        model: impl Into<String>,
58        client: Client,
59    ) -> Self {
60        Self {
61            client,
62            api_key: api_key.into(),
63            model: model.into(),
64            config: None,
65        }
66    }
67    /// Sets the task type for the embedding.
68    ///
69    /// The task type helps the model produce better embeddings tailored for the specific use case.
70    pub fn set_task_type(mut self, task_type: TaskType) -> Self {
71        let output_dimensionality = self
72            .config
73            .as_ref()
74            .and_then(|c| c.output_dimensionality().clone());
75        self.config = Some(EmbedContentConfig::new(
76            Some(task_type),
77            output_dimensionality,
78        ));
79        self
80    }
81    /// Sets the output dimensionality for the embedding.
82    ///
83    /// Allows reducing the embedding dimension for storage/performance optimization
84    /// via [Matryoshka Representation Learning](https://ai.google.dev/gemini-api/docs/embeddings#matryoshka).
85    pub fn set_output_dimensionality(mut self, output_dimensionality: u32) -> Self {
86        let task_type = self.config.as_ref().and_then(|c| c.task_type().clone());
87        self.config = Some(EmbedContentConfig::new(
88            task_type,
89            Some(output_dimensionality),
90        ));
91        self
92    }
93    pub fn set_api_key(mut self, api_key: impl Into<String>) -> Self {
94        self.api_key = api_key.into();
95        self
96    }
97    pub fn set_model(mut self, model: impl Into<String>) -> Self {
98        self.model = model.into();
99        self
100    }
101    /// Sets the full embedding configuration, replacing any previously set config.
102    pub fn set_config(mut self, config: EmbedContentConfig) -> Self {
103        self.config = Some(config);
104        self
105    }
106
107    /// Generates an embedding for the given content parts.
108    ///
109    /// # Arguments
110    /// * `content` - The content parts to embed (e.g., text, inline data).
111    ///
112    /// # Errors
113    /// Returns `GeminiResponseError` on network failure or API error.
114    pub async fn embed(
115        &self,
116        content: Vec<Part>,
117    ) -> Result<EmbedContentResponse, GeminiResponseError> {
118        let req_url = format!(
119            "{BASE_URL}/{}:embedContent?key={}",
120            self.model, self.api_key
121        );
122
123        let request_body = EmbedContentRequest {
124            model: format!("models/{}", self.model),
125            content: Content::new(content),
126            task_type: self.config.as_ref().and_then(|c| c.task_type().clone()),
127            output_dimensionality: self
128                .config
129                .as_ref()
130                .and_then(|c| c.output_dimensionality().clone()),
131        };
132
133        let response = self
134            .client
135            .post(req_url)
136            .json(&request_body)
137            .send()
138            .await
139            .map_err(|e| GeminiResponseError::ReqwestError(e))?;
140
141        if !response.status().is_success() {
142            let error = response
143                .json()
144                .await
145                .map_err(|e| GeminiResponseError::ReqwestError(e))?;
146            return Err(GeminiResponseError::StatusNotOk(error));
147        }
148
149        let embed_response: EmbedContentResponse = response
150            .json()
151            .await
152            .map_err(|e| GeminiResponseError::ReqwestError(e))?;
153        Ok(embed_response)
154    }
155
156    /// Convenience method to generate an embedding for a single text string.
157    ///
158    /// # Arguments
159    /// * `text` - The text to embed.
160    ///
161    /// # Example
162    /// ```no_run
163    /// # use gemini_client_api::gemini::embed::GeminiEmbedding;
164    /// # async fn run() {
165    /// let embedder = GeminiEmbedding::new("YOUR_API_KEY", "gemini-embedding-001");
166    /// let response = embedder.embed_text("What is the meaning of life?").await.unwrap();
167    /// println!("Got {} dimensions", response.embedding().values().len());
168    /// # }
169    /// ```
170    pub async fn embed_text(
171        &self,
172        text: impl Into<String>,
173    ) -> Result<EmbedContentResponse, GeminiResponseError> {
174        let part: Part = text.into().into();
175        self.embed(vec![part]).await
176    }
177}
gemini_client_api/gemini/embed.rs

gemini_client_api/gemini/
embed.rs
