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

gemini_client_api/gemini/
embed.rs
