gemini_genai_rs/embed/
mod.rs

1//! Embedding API — embedContent.
2//!
3//! Feature-gated behind `embed`.
4
5use serde::{Deserialize, Serialize};
6
7use crate::client::Client;
8use crate::client::http::HttpError;
9use crate::protocol::types::{Content, ModelId};
10use crate::transport::auth::ServiceEndpoint;
11
12/// Configuration for embed requests.
13#[derive(Debug, Clone)]
14pub struct EmbedContentConfig {
15    /// Content to embed.
16    pub content: Content,
17    /// Optional task type for better embeddings.
18    pub task_type: Option<TaskType>,
19    /// Optional title (for RETRIEVAL_DOCUMENT task type).
20    pub title: Option<String>,
21    /// Optional output dimensionality.
22    pub output_dimensionality: Option<u32>,
23}
24
25impl EmbedContentConfig {
26    /// Create an embed config from text.
27    pub fn from_text(text: impl Into<String>) -> Self {
28        Self {
29            content: Content::user(text),
30            task_type: None,
31            title: None,
32            output_dimensionality: None,
33        }
34    }
35
36    /// Set the task type.
37    pub fn task_type(mut self, task_type: TaskType) -> Self {
38        self.task_type = Some(task_type);
39        self
40    }
41
42    /// Set the title (for document retrieval).
43    pub fn title(mut self, title: impl Into<String>) -> Self {
44        self.title = Some(title.into());
45        self
46    }
47
48    /// Set the output dimensionality.
49    pub fn output_dimensionality(mut self, dim: u32) -> Self {
50        self.output_dimensionality = Some(dim);
51        self
52    }
53}
54
55/// Task type for embedding optimization.
56#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
57#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
58pub enum TaskType {
59    /// Embedding for a search query.
60    RetrievalQuery,
61    /// Embedding for a document to be retrieved.
62    RetrievalDocument,
63    /// Embedding for similarity comparison.
64    SemanticSimilarity,
65    /// Embedding used as classifier input.
66    Classification,
67    /// Embedding used for clustering.
68    Clustering,
69    /// Embedding for question answering.
70    QuestionAnswering,
71    /// Embedding for fact verification.
72    FactVerification,
73}
74
75/// Response from embedContent.
76#[derive(Debug, Clone, Serialize, Deserialize)]
77#[serde(rename_all = "camelCase")]
78pub struct EmbedContentResponse {
79    /// The embedding values.
80    pub embedding: ContentEmbedding,
81}
82
83/// Embedding vector.
84#[derive(Debug, Clone, Serialize, Deserialize)]
85#[serde(rename_all = "camelCase")]
86pub struct ContentEmbedding {
87    /// The embedding values (float vector).
88    pub values: Vec<f32>,
89}
90
91/// Errors from the Embed API.
92#[derive(Debug, thiserror::Error)]
93pub enum EmbedError {
94    #[error(transparent)]
95    /// Transport-level HTTP failure.
96    Http(#[from] HttpError),
97    #[error("Failed to parse response: {0}")]
98    /// Response body failed to parse.
99    Parse(#[from] serde_json::Error),
100    #[error("Auth error: {0}")]
101    /// Authentication/authorization failure.
102    Auth(#[from] crate::session::AuthError),
103}
104
105impl Client {
106    /// Embed text content using the default model.
107    pub async fn embed_content(
108        &self,
109        text: impl Into<String>,
110    ) -> Result<EmbedContentResponse, EmbedError> {
111        self.embed_content_with(EmbedContentConfig::from_text(text), None)
112            .await
113    }
114
115    /// Embed content with full configuration.
116    pub async fn embed_content_with(
117        &self,
118        config: EmbedContentConfig,
119        model: Option<&ModelId>,
120    ) -> Result<EmbedContentResponse, EmbedError> {
121        let model = model.unwrap_or(self.default_model());
122        let url = self.rest_url_for(ServiceEndpoint::EmbedContent, model);
123        let headers = self.auth_headers().await?;
124
125        let mut body = serde_json::json!({
126            "content": config.content,
127        });
128
129        if let Some(task_type) = config.task_type {
130            body["taskType"] = serde_json::to_value(task_type).unwrap();
131        }
132        if let Some(title) = config.title {
133            body["title"] = serde_json::Value::String(title);
134        }
135        if let Some(dim) = config.output_dimensionality {
136            body["outputDimensionality"] = serde_json::json!(dim);
137        }
138
139        let json = self.http_client().post_json(&url, headers, &body).await?;
140        Ok(serde_json::from_value(json)?)
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147
148    #[test]
149    fn parse_embed_response() {
150        let json = serde_json::json!({
151            "embedding": {
152                "values": [0.1, 0.2, 0.3, 0.4]
153            }
154        });
155        let resp: EmbedContentResponse = serde_json::from_value(json).unwrap();
156        assert_eq!(resp.embedding.values.len(), 4);
157        assert!((resp.embedding.values[0] - 0.1).abs() < f32::EPSILON);
158    }
159
160    #[test]
161    fn embed_config_builder() {
162        let config = EmbedContentConfig::from_text("Hello")
163            .task_type(TaskType::RetrievalQuery)
164            .output_dimensionality(256);
165        assert_eq!(config.task_type, Some(TaskType::RetrievalQuery));
166        assert_eq!(config.output_dimensionality, Some(256));
167    }
168}
gemini_genai_rs/embed/mod.rs

gemini_genai_rs/embed/
mod.rs
