rig/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 crate::embeddings::{self, EmbeddingError};
9
10use super::{Client, client::ApiResponse};
11
12/// `embedding-001` embedding model
13pub const EMBEDDING_001: &str = "embedding-001";
14/// `text-embedding-004` embedding model
15pub const EMBEDDING_004: &str = "text-embedding-004";
16#[derive(Clone)]
17pub struct EmbeddingModel {
18 client: Client,
19 model: String,
20 ndims: Option<usize>,
21}
22
23impl EmbeddingModel {
24 pub fn new(client: Client, model: &str, ndims: Option<usize>) -> Self {
25 Self {
26 client,
27 model: model.to_string(),
28 ndims,
29 }
30 }
31}
32
33impl embeddings::EmbeddingModel for EmbeddingModel {
34 const MAX_DOCUMENTS: usize = 1024;
35
36 fn ndims(&self) -> usize {
37 match self.model.as_str() {
38 EMBEDDING_001 | EMBEDDING_004 => 768,
39 _ => 0, // Default to 0 for unknown models
40 }
41 }
42
43 /// <https://ai.google.dev/api/embeddings#batch_embed_contents-SHELL>
44 #[cfg_attr(feature = "worker", worker::send)]
45 async fn embed_texts(
46 &self,
47 documents: impl IntoIterator<Item = String> + Send,
48 ) -> Result<Vec<embeddings::Embedding>, EmbeddingError> {
49 let documents: Vec<String> = documents.into_iter().collect();
50
51 // Google batch embed requests. See docstrings for API ref link.
52 let requests: Vec<_> = documents
53 .iter()
54 .map(|doc| {
55 json!({
56 "model": format!("models/{}", self.model),
57 "content": json!({
58 "parts": [json!({
59 "text": doc.to_string()
60 })]
61 }),
62 "output_dimensionality": self.ndims,
63 })
64 })
65 .collect();
66
67 let request_body = json!({ "requests": requests });
68
69 tracing::info!("{}", serde_json::to_string_pretty(&request_body).unwrap());
70
71 let response = self
72 .client
73 .post(&format!("/v1beta/models/{}:batchEmbedContents", self.model))
74 .json(&request_body)
75 .send()
76 .await?
77 .error_for_status()?
78 .json::<ApiResponse<gemini_api_types::EmbeddingResponse>>()
79 .await?;
80
81 match response {
82 ApiResponse::Ok(response) => {
83 let docs = documents
84 .into_iter()
85 .zip(response.embeddings)
86 .map(|(document, embedding)| embeddings::Embedding {
87 document,
88 vec: embedding.values,
89 })
90 .collect();
91
92 Ok(docs)
93 }
94 ApiResponse::Err(err) => Err(EmbeddingError::ProviderError(err.message)),
95 }
96 }
97}
98
99// =================================================================
100// Gemini API Types
101// =================================================================
102/// Rust Implementation of the Gemini Types from [Gemini API Reference](https://ai.google.dev/api/embeddings)
103#[allow(dead_code)]
104mod gemini_api_types {
105 use serde::{Deserialize, Serialize};
106 use serde_json::Value;
107
108 use crate::providers::gemini::gemini_api_types::{CodeExecutionResult, ExecutableCode};
109
110 #[derive(Serialize)]
111 #[serde(rename_all = "camelCase")]
112 pub struct EmbedContentRequest {
113 model: String,
114 content: EmbeddingContent,
115 task_type: TaskType,
116 title: String,
117 output_dimensionality: i32,
118 }
119
120 #[derive(Serialize)]
121 pub struct EmbeddingContent {
122 parts: Vec<EmbeddingContentPart>,
123 /// Optional. The producer of the content. Must be either 'user' or 'model'. Useful to set for multi-turn
124 /// conversations, otherwise can be left blank or unset.
125 role: Option<String>,
126 }
127
128 /// A datatype containing media that is part of a multi-part Content message.
129 /// - A Part consists of data which has an associated datatype. A Part can only contain one of the accepted types in Part.data.
130 /// - A Part must have a fixed IANA MIME type identifying the type and subtype of the media if the inlineData field is filled with raw bytes.
131 #[derive(Serialize)]
132 pub struct EmbeddingContentPart {
133 /// Inline text.
134 text: String,
135 /// Inline media bytes.
136 inline_data: Option<Blob>,
137 /// A predicted FunctionCall returned from the model that contains a string representing the [FunctionDeclaration.name]
138 /// with the arguments and their values.
139 function_call: Option<FunctionCall>,
140 /// The result output of a FunctionCall that contains a string representing the [FunctionDeclaration.name] and a structured
141 /// JSON object containing any output from the function is used as context to the model.
142 function_response: Option<FunctionResponse>,
143 /// URI based data.
144 file_data: Option<FileData>,
145 /// Code generated by the model that is meant to be executed.
146 executable_code: Option<ExecutableCode>,
147 /// Result of executing the ExecutableCode.
148 code_execution_result: Option<CodeExecutionResult>,
149 }
150
151 /// Raw media bytes.
152 /// Text should not be sent as raw bytes, use the 'text' field.
153 #[derive(Serialize)]
154 pub struct Blob {
155 /// Raw bytes for media formats.A base64-encoded string.
156 data: String,
157 /// The IANA standard MIME type of the source data. Examples: - image/png - image/jpeg If an unsupported MIME type is
158 /// provided, an error will be returned. For a complete list of supported types, see Supported file formats.
159 mime_type: String,
160 }
161
162 #[derive(Serialize)]
163 pub struct FunctionCall {
164 /// The name of the function to call. Must be a-z, A-Z, 0-9, or contain underscores and dashes, with a maximum length of 63.
165 name: String,
166 /// The function parameters and values in JSON object format.
167 args: Option<Value>,
168 }
169
170 #[derive(Serialize)]
171 pub struct FunctionResponse {
172 /// The name of the function to call. Must be a-z, A-Z, 0-9, or contain underscores and dashes, with a maximum length of 63.
173 name: String,
174 /// The result of the function call in JSON object format.
175 result: Value,
176 }
177
178 #[derive(Serialize)]
179 #[serde(rename_all = "camelCase")]
180 pub struct FileData {
181 /// The URI of the file.
182 file_uri: String,
183 /// The IANA standard MIME type of the source data.
184 mime_type: String,
185 }
186
187 #[derive(Serialize)]
188 #[serde(rename_all = "SCREAMING_SNAKE_CASE")]
189 pub enum TaskType {
190 /// Unset value, which will default to one of the other enum values.
191 Unspecified,
192 /// Specifies the given text is a query in a search/retrieval setting.
193 RetrievalQuery,
194 /// Specifies the given text is a document from the corpus being searched.
195 RetrievalDocument,
196 /// Specifies the given text will be used for STS.
197 SemanticSimilarity,
198 /// Specifies that the given text will be classified.
199 Classification,
200 /// Specifies that the embeddings will be used for clustering.
201 Clustering,
202 /// Specifies that the given text will be used for question answering.
203 QuestionAnswering,
204 /// Specifies that the given text will be used for fact verification.
205 FactVerification,
206 }
207
208 #[derive(Debug, Deserialize)]
209 pub struct EmbeddingResponse {
210 pub embeddings: Vec<EmbeddingValues>,
211 }
212
213 #[derive(Debug, Deserialize)]
214 pub struct EmbeddingValues {
215 pub values: Vec<f64>,
216 }
217}