1#[allow(deprecated)]
2use crate::{
3 batch::{BatchBuilder, BatchHandle},
4 cache::{CacheBuilder, CachedContentHandle},
5 embedding::{
6 BatchContentEmbeddingResponse, BatchEmbedContentsRequest, ContentEmbeddingResponse,
7 EmbedBuilder, EmbedContentRequest,
8 },
9 files::{
10 handle::FileHandle,
11 model::{File, ListFilesResponse},
12 },
13 generation::{ContentBuilder, GenerateContentRequest, GenerationResponse},
14 interactions::{
15 builder::InteractionBuilder,
16 handle::InteractionHandle,
17 model::{CreateInteractionRequest, Interaction},
18 stream::{InteractionEvent, InteractionStream},
19 },
20};
21use eventsource_stream::{EventStreamError, Eventsource};
22use futures::{Stream, StreamExt, TryStreamExt};
23use mime::Mime;
24use reqwest::{
25 header::{HeaderMap, HeaderName, HeaderValue, InvalidHeaderValue},
26 Client, ClientBuilder, RequestBuilder, Response,
27};
28use serde::{Deserialize, Serialize};
29use serde_json::json;
30use snafu::{OptionExt, ResultExt, Snafu};
31use std::{
32 fmt::{self, Formatter},
33 pin::Pin,
34 sync::{Arc, LazyLock},
35};
36use tracing::{instrument, Level, Span};
37use url::Url;
38
39use crate::batch::model::*;
40use crate::cache::model::*;
41
42#[allow(deprecated)]
48pub type GenerationStream = Pin<Box<dyn Stream<Item = Result<GenerationResponse, Error>> + Send>>;
49
50static DEFAULT_BASE_URL: LazyLock<Url> = LazyLock::new(|| {
51 Url::parse("https://generativelanguage.googleapis.com/v1beta/")
52 .expect("unreachable error: failed to parse default base URL")
53});
54
55#[derive(Debug, Default, Clone, PartialEq, Eq, Hash, Deserialize, Serialize)]
61pub enum Model {
62 #[default]
64 #[serde(rename = "models/gemini-3.7-flash")]
65 Gemini37Flash,
66 #[serde(rename = "models/gemini-3.6-flash")]
67 Gemini36Flash,
68 #[serde(rename = "models/gemini-3.5-flash")]
69 Gemini35Flash,
70 #[serde(rename = "models/gemini-3.5-flash-lite")]
71 Gemini35FlashLite,
72 #[serde(rename = "models/gemini-3.1-flash-lite")]
73 Gemini31FlashLite,
74 #[serde(rename = "models/gemini-3.1-pro-preview")]
76 Gemini31Pro,
77 #[serde(rename = "models/gemini-3.1-flash-image")]
79 Gemini31FlashImage,
80 #[serde(rename = "models/gemini-3.1-flash-lite-image")]
82 Gemini31FlashLiteImage,
83 #[serde(rename = "models/gemini-3-pro-image")]
85 Gemini3ProImage,
86 #[serde(rename = "models/gemini-3-flash-preview")]
87 Gemini3Flash,
88 #[serde(rename = "models/gemini-omni-flash")]
91 GeminiOmniFlash,
92 #[serde(rename = "models/gemini-3.1-flash-tts-preview")]
94 Gemini31FlashTts,
95 #[deprecated(
97 since = "2.1.0",
98 note = "gemini-3-pro-preview has been shut down; use Model::Gemini31Pro instead"
99 )]
100 #[serde(rename = "models/gemini-3-pro-preview")]
101 Gemini3Pro,
102 #[serde(rename = "models/gemini-2.5-flash")]
104 Gemini25Flash,
105 #[serde(rename = "models/gemini-2.5-flash-lite")]
106 Gemini25FlashLite,
107 #[serde(rename = "models/gemini-2.5-flash-image")]
110 Gemini25FlashImage,
111 #[serde(rename = "models/gemini-2.5-pro")]
112 Gemini25Pro,
113 #[serde(rename = "models/gemini-2.5-flash-preview-tts")]
115 Gemini25FlashTts,
116 #[serde(rename = "models/gemini-2.5-pro-preview-tts")]
118 Gemini25ProTts,
119 #[serde(rename = "models/gemini-2.5-computer-use-preview-10-2025")]
122 Gemini25ComputerUse,
123 #[serde(rename = "models/gemini-embedding-2")]
126 GeminiEmbedding2,
127 #[serde(rename = "models/gemini-embedding-001")]
129 GeminiEmbedding001,
130 #[deprecated(
131 since = "2.1.0",
132 note = "shut down on January 14, 2026; use Model::GeminiEmbedding2 instead"
133 )]
134 #[serde(rename = "models/text-embedding-004")]
135 TextEmbedding004,
136 #[serde(untagged)]
137 Custom(String),
138}
139
140impl Model {
141 #[allow(deprecated)]
142 pub fn as_str(&self) -> &str {
143 match self {
144 Model::Gemini37Flash => "models/gemini-3.7-flash",
145 Model::Gemini36Flash => "models/gemini-3.6-flash",
146 Model::Gemini35Flash => "models/gemini-3.5-flash",
147 Model::Gemini35FlashLite => "models/gemini-3.5-flash-lite",
148 Model::Gemini31FlashLite => "models/gemini-3.1-flash-lite",
149 Model::Gemini31Pro => "models/gemini-3.1-pro-preview",
150 Model::Gemini31FlashImage => "models/gemini-3.1-flash-image",
151 Model::Gemini31FlashLiteImage => "models/gemini-3.1-flash-lite-image",
152 Model::Gemini3ProImage => "models/gemini-3-pro-image",
153 Model::Gemini3Flash => "models/gemini-3-flash-preview",
154 Model::GeminiOmniFlash => "models/gemini-omni-flash",
155 Model::Gemini31FlashTts => "models/gemini-3.1-flash-tts-preview",
156 Model::Gemini3Pro => "models/gemini-3-pro-preview",
157 Model::Gemini25Flash => "models/gemini-2.5-flash",
158 Model::Gemini25FlashLite => "models/gemini-2.5-flash-lite",
159 Model::Gemini25FlashImage => "models/gemini-2.5-flash-image",
160 Model::Gemini25Pro => "models/gemini-2.5-pro",
161 Model::Gemini25FlashTts => "models/gemini-2.5-flash-preview-tts",
162 Model::Gemini25ProTts => "models/gemini-2.5-pro-preview-tts",
163 Model::Gemini25ComputerUse => "models/gemini-2.5-computer-use-preview-10-2025",
164 Model::GeminiEmbedding2 => "models/gemini-embedding-2",
165 Model::GeminiEmbedding001 => "models/gemini-embedding-001",
166 Model::TextEmbedding004 => "models/text-embedding-004",
167 Model::Custom(model) => model,
168 }
169 }
170}
171
172impl From<String> for Model {
173 fn from(model: String) -> Self {
174 Self::Custom(model)
175 }
176}
177
178impl fmt::Display for Model {
179 #[allow(deprecated)]
180 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
181 match self {
182 Model::Gemini37Flash => write!(f, "models/gemini-3.7-flash"),
183 Model::Gemini36Flash => write!(f, "models/gemini-3.6-flash"),
184 Model::Gemini35Flash => write!(f, "models/gemini-3.5-flash"),
185 Model::Gemini35FlashLite => write!(f, "models/gemini-3.5-flash-lite"),
186 Model::Gemini31FlashLite => write!(f, "models/gemini-3.1-flash-lite"),
187 Model::Gemini31Pro => write!(f, "models/gemini-3.1-pro-preview"),
188 Model::Gemini31FlashImage => write!(f, "models/gemini-3.1-flash-image"),
189 Model::Gemini31FlashLiteImage => write!(f, "models/gemini-3.1-flash-lite-image"),
190 Model::Gemini3ProImage => write!(f, "models/gemini-3-pro-image"),
191 Model::Gemini3Flash => write!(f, "models/gemini-3-flash-preview"),
192 Model::GeminiOmniFlash => write!(f, "models/gemini-omni-flash"),
193 Model::Gemini31FlashTts => write!(f, "models/gemini-3.1-flash-tts-preview"),
194 Model::Gemini3Pro => write!(f, "models/gemini-3-pro-preview"),
195 Model::Gemini25Flash => write!(f, "models/gemini-2.5-flash"),
196 Model::Gemini25FlashLite => write!(f, "models/gemini-2.5-flash-lite"),
197 Model::Gemini25FlashImage => write!(f, "models/gemini-2.5-flash-image"),
198 Model::Gemini25Pro => write!(f, "models/gemini-2.5-pro"),
199 Model::Gemini25FlashTts => write!(f, "models/gemini-2.5-flash-preview-tts"),
200 Model::Gemini25ProTts => write!(f, "models/gemini-2.5-pro-preview-tts"),
201 Model::Gemini25ComputerUse => {
202 write!(f, "models/gemini-2.5-computer-use-preview-10-2025")
203 }
204 Model::GeminiEmbedding2 => write!(f, "models/gemini-embedding-2"),
205 Model::GeminiEmbedding001 => write!(f, "models/gemini-embedding-001"),
206 Model::TextEmbedding004 => write!(f, "models/text-embedding-004"),
207 Model::Custom(model) => write!(f, "{model}"),
208 }
209 }
210}
211
212#[derive(Debug, Snafu)]
213#[snafu(visibility(pub))]
214pub enum Error {
215 #[snafu(display("failed to parse API key"))]
216 InvalidApiKey {
217 source: InvalidHeaderValue,
218 },
219
220 #[snafu(display("failed to construct URL (probably incorrect model name): {suffix}"))]
221 ConstructUrl {
222 source: url::ParseError,
223 suffix: String,
224 },
225
226 PerformRequestNew {
227 source: reqwest::Error,
228 },
229
230 #[snafu(display("failed to perform request to '{url}'"))]
231 PerformRequest {
232 source: reqwest::Error,
233 url: Url,
234 },
235
236 #[snafu(display(
237 "bad response from server; code {code}; description: {}",
238 description.as_deref().unwrap_or("none")
239 ))]
240 BadResponse {
241 code: u16,
243 description: Option<String>,
245 },
246
247 MissingResponseHeader {
248 header: String,
249 },
250
251 #[snafu(display("failed to obtain stream SSE part"))]
252 BadPart {
253 source: EventStreamError<reqwest::Error>,
254 },
255
256 #[snafu(display("failed to deserialize JSON response"))]
257 Deserialize {
258 source: serde_json::Error,
259 },
260
261 #[snafu(display("failed to generate content"))]
262 DecodeResponse {
263 source: reqwest::Error,
264 },
265
266 #[snafu(display("failed to parse URL"))]
267 UrlParse {
268 source: url::ParseError,
269 },
270
271 #[snafu(display("I/O error during file operations"))]
272 Io {
273 source: std::io::Error,
274 },
275
276 #[snafu(display("operation timed out: {name}"))]
277 OperationTimeout {
278 name: String,
279 },
280
281 #[snafu(display("operation failed: {name}, code: {code}, message: {message}"))]
282 OperationFailed {
283 name: String,
284 code: i32,
285 message: String,
286 },
287
288 #[snafu(display("invalid resource name: {name}"))]
289 InvalidResourceName {
290 name: String,
291 },
292}
293
294#[derive(Debug)]
296pub struct GeminiClient {
297 http_client: Client,
298 pub model: Model,
299 base_url: Url,
300}
301
302impl GeminiClient {
303 fn with_base_url<K: AsRef<str>, M: Into<Model>>(
305 client_builder: ClientBuilder,
306 api_key: K,
307 model: M,
308 base_url: Url,
309 ) -> Result<Self, Error> {
310 let headers = HeaderMap::from_iter([(
311 HeaderName::from_static("x-goog-api-key"),
312 HeaderValue::from_str(api_key.as_ref()).context(InvalidApiKeySnafu)?,
313 )]);
314
315 let http_client = client_builder
316 .default_headers(headers)
317 .build()
318 .expect("all parameters must be valid");
319
320 Ok(Self {
321 http_client,
322 model: model.into(),
323 base_url,
324 })
325 }
326
327 #[tracing::instrument(skip_all, err)]
329 async fn check_response(response: Response) -> Result<Response, Error> {
330 let status = response.status();
331 if !status.is_success() {
332 let description = response.text().await.ok();
333 BadResponseSnafu {
334 code: status.as_u16(),
335 description,
336 }
337 .fail()
338 } else {
339 Ok(response)
340 }
341 }
342
343 #[tracing::instrument(skip_all)]
421 #[doc(hidden)]
422 pub async fn perform_request<
423 B: FnOnce(&Client) -> RequestBuilder,
424 D: AsyncFn(Response) -> Result<T, Error>,
425 T,
426 >(
427 &self,
428 builder: B,
429 deserializer: D,
430 ) -> Result<T, Error> {
431 let request = builder(&self.http_client);
432 tracing::debug!("request built successfully");
433 let response = request.send().await.context(PerformRequestNewSnafu)?;
434 tracing::debug!("response received successfully");
435 let response = Self::check_response(response).await?;
436 tracing::debug!("response ok");
437 deserializer(response).await
438 }
439
440 #[tracing::instrument(skip(self), fields(request.type = "get", request.url = %url))]
444 async fn get_json<T: serde::de::DeserializeOwned>(&self, url: Url) -> Result<T, Error> {
445 self.perform_request(
446 |c| c.get(url),
447 async |r| r.json().await.context(DecodeResponseSnafu),
448 )
449 .await
450 }
451
452 #[tracing::instrument(skip(self, body), fields(request.type = "post", request.url = %url))]
456 async fn post_json<Req: serde::Serialize, Res: serde::de::DeserializeOwned>(
457 &self,
458 url: Url,
459 body: &Req,
460 ) -> Result<Res, Error> {
461 self.perform_request(
462 |c| c.post(url).json(body),
463 async |r| r.json().await.context(DecodeResponseSnafu),
464 )
465 .await
466 }
467
468 #[allow(deprecated)]
470 #[instrument(skip_all, fields(
471 model,
472 messages.parts.count = request.contents.len(),
473 tools.present = request.tools.is_some(),
474 system.instruction.present = request.system_instruction.is_some(),
475 cached.content.present = request.cached_content.is_some(),
476 usage.prompt_tokens,
477 usage.candidates_tokens,
478 usage.thoughts_tokens,
479 usage.cached_content_tokens,
480 usage.total_tokens,
481 ), ret(level = Level::TRACE), err)]
482 pub(crate) async fn generate_content_raw(
483 &self,
484 request: GenerateContentRequest,
485 ) -> Result<GenerationResponse, Error> {
486 let url = self.build_url("generateContent")?;
487 let response: GenerationResponse = self.post_json(url, &request).await?;
488
489 if let Some(usage) = &response.usage_metadata {
491 #[rustfmt::skip]
492 Span::current()
493 .record("usage.prompt_tokens", usage.prompt_token_count)
494 .record("usage.candidates_tokens", usage.candidates_token_count)
495 .record("usage.thoughts_tokens", usage.thoughts_token_count)
496 .record("usage.cached_content_tokens", usage.cached_content_token_count)
497 .record("usage.total_tokens", usage.total_token_count);
498
499 tracing::debug!("generation usage evaluated");
500 }
501
502 Ok(response)
503 }
504
505 #[allow(deprecated)]
507 #[instrument(skip_all, fields(
508 model,
509 messages.parts.count = request.contents.len(),
510 tools.present = request.tools.is_some(),
511 system.instruction.present = request.system_instruction.is_some(),
512 cached.content.present = request.cached_content.is_some(),
513 ), err)]
514 pub(crate) async fn generate_content_stream(
515 &self,
516 request: GenerateContentRequest,
517 ) -> Result<GenerationStream, Error> {
518 let mut url = self.build_url("streamGenerateContent")?;
519 url.query_pairs_mut().append_pair("alt", "sse");
520
521 let stream = self
522 .perform_request(
523 |c| c.post(url).json(&request),
524 async |r| Ok(r.bytes_stream()),
525 )
526 .await?;
527
528 Ok(Box::pin(
529 stream
530 .eventsource()
531 .map(|event| event.context(BadPartSnafu))
532 .try_filter(|event| std::future::ready(event.data != "[DONE]"))
533 .and_then(|event| async move {
534 serde_json::from_str::<GenerationResponse>(&event.data)
535 .context(DeserializeSnafu)
536 }),
537 ))
538 }
539
540 #[allow(deprecated)]
542 #[instrument(skip_all, fields(
543 model,
544 messages.parts.count = request.contents.len(),
545 ))]
546 pub(crate) async fn count_tokens(
547 &self,
548 request: GenerateContentRequest,
549 ) -> Result<crate::generation::CountTokensResponse, Error> {
550 let url = self.build_url("countTokens")?;
551 let body = json!({
554 "generateContentRequest": {
555 "model": self.model.as_str(),
556 "contents": request.contents,
557 "generationConfig": request.generation_config,
558 "safetySettings": request.safety_settings,
559 "tools": request.tools,
560 "toolConfig": request.tool_config,
561 "systemInstruction": request.system_instruction,
562 "cachedContent": request.cached_content,
563 }
564 });
565 self.post_json(url, &body).await
566 }
567
568 #[instrument(skip_all, fields(
570 model,
571 task.type = request.task_type.as_ref().map(|t| format!("{t:?}")),
572 task.title = request.title,
573 task.output.dimensionality = request.output_dimensionality,
574 ))]
575 pub(crate) async fn embed_content(
576 &self,
577 request: EmbedContentRequest,
578 ) -> Result<ContentEmbeddingResponse, Error> {
579 let url = self.build_url("embedContent")?;
580 self.post_json(url, &request).await
581 }
582
583 #[instrument(skip_all, fields(batch.size = request.requests.len()))]
585 pub(crate) async fn embed_content_batch(
586 &self,
587 request: BatchEmbedContentsRequest,
588 ) -> Result<BatchContentEmbeddingResponse, Error> {
589 let url = self.build_url("batchEmbedContents")?;
590 self.post_json(url, &request).await
591 }
592
593 #[instrument(skip_all, fields(
595 batch.display_name = request.batch.display_name,
596 batch.size = request.batch.input_config.batch_size(),
597 ))]
598 pub(crate) async fn batch_generate_content(
599 &self,
600 request: BatchGenerateContentRequest,
601 ) -> Result<BatchGenerateContentResponse, Error> {
602 let url = self.build_url("batchGenerateContent")?;
603 self.post_json(url, &request).await
604 }
605
606 #[instrument(skip_all, fields(
608 operation.name = name,
609 ))]
610 pub(crate) async fn get_batch_operation<T: serde::de::DeserializeOwned>(
611 &self,
612 name: &str,
613 ) -> Result<T, Error> {
614 let url = self.build_batch_url(name, None)?;
615 self.get_json(url).await
616 }
617
618 #[instrument(skip_all, fields(
620 page.size = page_size,
621 page.token.present = page_token.is_some(),
622 ))]
623 pub(crate) async fn list_batch_operations(
624 &self,
625 page_size: Option<u32>,
626 page_token: Option<String>,
627 ) -> Result<ListBatchesResponse, Error> {
628 let mut url = self.build_batch_url("batches", None)?;
629
630 if let Some(size) = page_size {
631 url.query_pairs_mut()
632 .append_pair("pageSize", &size.to_string());
633 }
634 if let Some(token) = page_token {
635 url.query_pairs_mut().append_pair("pageToken", &token);
636 }
637
638 self.get_json(url).await
639 }
640
641 #[instrument(skip_all, fields(
643 page.size = page_size,
644 page.token.present = page_token.is_some(),
645 ))]
646 pub(crate) async fn list_files(
647 &self,
648 page_size: Option<u32>,
649 page_token: Option<String>,
650 ) -> Result<ListFilesResponse, Error> {
651 let mut url = self.build_files_url(None)?;
652
653 if let Some(size) = page_size {
654 url.query_pairs_mut()
655 .append_pair("pageSize", &size.to_string());
656 }
657 if let Some(token) = page_token {
658 url.query_pairs_mut().append_pair("pageToken", &token);
659 }
660
661 self.get_json(url).await
662 }
663
664 #[instrument(skip_all, fields(
666 operation.name = name,
667 ))]
668 pub(crate) async fn cancel_batch_operation(&self, name: &str) -> Result<(), Error> {
669 let url = self.build_batch_url(name, Some("cancel"))?;
670 self.perform_request(|c| c.post(url).json(&json!({})), async |_r| Ok(()))
671 .await
672 }
673
674 #[instrument(skip_all, fields(
676 operation.name = name,
677 ))]
678 pub(crate) async fn delete_batch_operation(&self, name: &str) -> Result<(), Error> {
679 let url = self.build_batch_url(name, None)?;
680 self.perform_request(|c| c.delete(url), async |_r| Ok(()))
681 .await
682 }
683
684 async fn create_upload(
685 &self,
686 bytes: usize,
687 display_name: Option<String>,
688 mime_type: Mime,
689 ) -> Result<Url, Error> {
690 let url = self
691 .base_url
692 .join("/upload/v1beta/files")
693 .context(ConstructUrlSnafu {
694 suffix: "/upload/v1beta/files".to_string(),
695 })?;
696
697 self.perform_request(
698 |c| {
699 c.post(url)
700 .header("X-Goog-Upload-Protocol", "resumable")
701 .header("X-Goog-Upload-Command", "start")
702 .header("X-Goog-Upload-Content-Length", bytes.to_string())
703 .header("X-Goog-Upload-Header-Content-Type", mime_type.to_string())
704 .json(&json!({"file": {"displayName": display_name}}))
705 },
706 async |r| {
707 r.headers()
708 .get("X-Goog-Upload-URL")
709 .context(MissingResponseHeaderSnafu {
710 header: "X-Goog-Upload-URL",
711 })
712 .and_then(|upload_url| {
713 upload_url
714 .to_str()
715 .map(str::to_string)
716 .map_err(|_| Error::BadResponse {
717 code: 500,
718 description: Some("Missing upload URL in response".to_string()),
719 })
720 })
721 .and_then(|url| Url::parse(&url).context(UrlParseSnafu))
722 },
723 )
724 .await
725 }
726
727 #[instrument(skip_all, fields(
729 file.size = file_bytes.len(),
730 mime.type = mime_type.to_string(),
731 file.display_name = display_name.as_deref(),
732 ))]
733 pub(crate) async fn upload_file(
734 &self,
735 display_name: Option<String>,
736 file_bytes: Vec<u8>,
737 mime_type: Mime,
738 ) -> Result<File, Error> {
739 let upload_url = self
741 .create_upload(file_bytes.len(), display_name, mime_type)
742 .await?;
743
744 let upload_response = self
746 .http_client
747 .post(upload_url.clone())
748 .header("X-Goog-Upload-Command", "upload, finalize")
749 .header("X-Goog-Upload-Offset", "0")
750 .body(file_bytes)
751 .send()
752 .await
753 .map_err(|e| Error::PerformRequest {
754 source: e,
755 url: upload_url,
756 })?;
757
758 let final_response = Self::check_response(upload_response).await?;
759
760 #[derive(serde::Deserialize)]
761 struct UploadResponse {
762 file: File,
763 }
764
765 let upload_response: UploadResponse =
766 final_response.json().await.context(DecodeResponseSnafu)?;
767 Ok(upload_response.file)
768 }
769
770 #[instrument(skip_all, fields(
772 file.name = name,
773 ))]
774 pub(crate) async fn get_file(&self, name: &str) -> Result<File, Error> {
775 let url = self.build_files_url(Some(name))?;
776 self.get_json(url).await
777 }
778
779 #[instrument(skip_all, fields(
781 file.name = name,
782 ))]
783 pub(crate) async fn delete_file(&self, name: &str) -> Result<(), Error> {
784 let url = self.build_files_url(Some(name))?;
785 self.perform_request(|c| c.delete(url), async |_r| Ok(()))
786 .await
787 }
788
789 #[instrument(skip_all, fields(
791 file.name = name,
792 ))]
793 pub(crate) async fn download_file(&self, name: &str) -> Result<Vec<u8>, Error> {
794 let mut url = self
795 .base_url
796 .join(&format!("/download/v1beta/{name}:download"))
797 .context(ConstructUrlSnafu {
798 suffix: format!("/download/v1beta/{name}:download"),
799 })?;
800 url.query_pairs_mut().append_pair("alt", "media");
801
802 self.perform_request(
803 |c| c.get(url),
804 async |r| {
805 r.bytes()
806 .await
807 .context(DecodeResponseSnafu)
808 .map(|bytes| bytes.to_vec())
809 },
810 )
811 .await
812 }
813
814 pub(crate) async fn create_cached_content(
816 &self,
817 cached_content: CreateCachedContentRequest,
818 ) -> Result<CachedContent, Error> {
819 let url = self.build_cache_url(None)?;
820 self.post_json(url, &cached_content).await
821 }
822
823 pub(crate) async fn get_cached_content(&self, name: &str) -> Result<CachedContent, Error> {
825 let url = self.build_cache_url(Some(name))?;
826 self.get_json(url).await
827 }
828
829 pub(crate) async fn update_cached_content(
831 &self,
832 name: &str,
833 expiration: CacheExpirationRequest,
834 ) -> Result<CachedContent, Error> {
835 let url = self.build_cache_url(Some(name))?;
836
837 let update_payload = match expiration {
839 CacheExpirationRequest::Ttl { ttl } => json!({ "ttl": ttl }),
840 CacheExpirationRequest::ExpireTime { expire_time } => {
841 json!({ "expireTime": expire_time.format(&time::format_description::well_known::Rfc3339).unwrap() })
842 }
843 };
844
845 self.perform_request(
846 |c| c.patch(url.clone()).json(&update_payload),
847 async |r| r.json().await.context(DecodeResponseSnafu),
848 )
849 .await
850 }
851
852 pub(crate) async fn delete_cached_content(&self, name: &str) -> Result<(), Error> {
854 let url = self.build_cache_url(Some(name))?;
855 self.perform_request(|c| c.delete(url.clone()), async |_r| Ok(()))
856 .await
857 }
858
859 pub(crate) async fn list_cached_contents(
861 &self,
862 page_size: Option<i32>,
863 page_token: Option<String>,
864 ) -> Result<ListCachedContentsResponse, Error> {
865 let mut url = self.build_cache_url(None)?;
866
867 if let Some(size) = page_size {
868 url.query_pairs_mut()
869 .append_pair("pageSize", &size.to_string());
870 }
871 if let Some(token) = page_token {
872 url.query_pairs_mut().append_pair("pageToken", &token);
873 }
874
875 self.get_json(url).await
876 }
877
878 #[instrument(skip_all, fields(
882 model = request.model.as_deref().unwrap_or(""),
883 agent = request.agent.as_deref().unwrap_or(""),
884 tools.count = request.tools.len(),
885 background = request.background.unwrap_or(false),
886 previous.interaction.present = request.previous_interaction_id.is_some(),
887 status.code,
888 usage.total_tokens,
889 ))]
890 pub(crate) async fn create_interaction(
891 &self,
892 request: CreateInteractionRequest,
893 ) -> Result<Interaction, Error> {
894 let url = self.build_url_with_suffix("interactions")?;
895 let response: Interaction = self.post_json(url, &request).await?;
896
897 Span::current().record("status.code", response.status.as_ref());
898
899 if let Some(usage) = &response.usage {
900 Span::current().record("usage.total_tokens", usage.total_tokens);
901 }
902
903 Ok(response)
904 }
905
906 #[instrument(skip_all, fields(
908 model = request.model.as_deref().unwrap_or(""),
909 agent = request.agent.as_deref().unwrap_or(""),
910 tools.count = request.tools.len(),
911 ))]
912 pub(crate) async fn create_interaction_stream(
913 &self,
914 mut request: CreateInteractionRequest,
915 ) -> Result<InteractionStream, Error> {
916 let mut url = self.build_url_with_suffix("interactions")?;
917 url.query_pairs_mut().append_pair("alt", "sse");
918 request.stream = Some(true);
919
920 let stream = self
921 .perform_request(
922 |c| c.post(url).json(&request),
923 async |r| Ok(r.bytes_stream()),
924 )
925 .await?;
926
927 Ok(Box::pin(
928 stream
929 .eventsource()
930 .map(|event| event.context(BadPartSnafu))
931 .try_filter(|event| std::future::ready(event.data != "[DONE]"))
932 .and_then(|event| async move {
933 serde_json::from_str::<InteractionEvent>(&event.data).context(DeserializeSnafu)
934 }),
935 ))
936 }
937
938 #[instrument(skip_all, fields(
940 interaction.id = id,
941 ))]
942 pub(crate) async fn get_interaction(&self, id: &str) -> Result<Interaction, Error> {
943 let url = self.build_url_with_suffix(&format!("interactions/{id}"))?;
944 self.get_json(url).await
945 }
946
947 #[instrument(skip_all, fields(
949 interaction.id = id,
950 ))]
951 pub(crate) async fn get_interaction_stream(
952 &self,
953 id: &str,
954 last_event_id: Option<&str>,
955 ) -> Result<InteractionStream, Error> {
956 let mut url = self.build_url_with_suffix(&format!("interactions/{id}"))?;
957 url.query_pairs_mut().append_pair("stream", "true");
958 if let Some(event_id) = last_event_id {
959 url.query_pairs_mut().append_pair("last_event_id", event_id);
960 }
961
962 let stream = self
963 .perform_request(|c| c.get(url), async |r| Ok(r.bytes_stream()))
964 .await?;
965
966 Ok(Box::pin(
967 stream
968 .eventsource()
969 .map(|event| event.context(BadPartSnafu))
970 .try_filter(|event| std::future::ready(event.data != "[DONE]"))
971 .and_then(|event| async move {
972 serde_json::from_str::<InteractionEvent>(&event.data).context(DeserializeSnafu)
973 }),
974 ))
975 }
976
977 #[instrument(skip_all, fields(
979 interaction.id = id,
980 ))]
981 pub(crate) async fn cancel_interaction(&self, id: &str) -> Result<Interaction, Error> {
982 let url = self.build_url_with_suffix(&format!("interactions/{id}/cancel"))?;
983 self.perform_request(
984 |c| c.post(url).json(&json!({})),
985 async |r| r.json().await.context(DecodeResponseSnafu),
986 )
987 .await
988 }
989
990 #[instrument(skip_all, fields(
992 interaction.id = id,
993 ))]
994 pub(crate) async fn delete_interaction(&self, id: &str) -> Result<(), Error> {
995 let url = self.build_url_with_suffix(&format!("interactions/{id}"))?;
996 self.perform_request(|c| c.delete(url), async |_r| Ok(()))
997 .await
998 }
999
1000 #[tracing::instrument(skip(self), ret(level = Level::DEBUG))]
1002 fn build_url_with_suffix(&self, suffix: &str) -> Result<Url, Error> {
1003 self.base_url.join(suffix).context(ConstructUrlSnafu {
1004 suffix: suffix.to_string(),
1005 })
1006 }
1007
1008 #[tracing::instrument(skip(self), ret(level = Level::DEBUG))]
1010 fn build_url(&self, endpoint: &str) -> Result<Url, Error> {
1011 let suffix = format!("{}:{endpoint}", self.model);
1012 self.build_url_with_suffix(&suffix)
1013 }
1014
1015 fn build_batch_url(&self, name: &str, action: Option<&str>) -> Result<Url, Error> {
1017 let suffix = action
1018 .map(|a| format!("{name}:{a}"))
1019 .unwrap_or_else(|| name.to_string());
1020 self.build_url_with_suffix(&suffix)
1021 }
1022
1023 fn build_files_url(&self, name: Option<&str>) -> Result<Url, Error> {
1025 let suffix = name
1026 .map(|n| format!("files/{}", n.strip_prefix("files/").unwrap_or(n)))
1027 .unwrap_or_else(|| "files".to_string());
1028 self.build_url_with_suffix(&suffix)
1029 }
1030
1031 fn build_cache_url(&self, name: Option<&str>) -> Result<Url, Error> {
1033 let suffix = name
1034 .map(|n| {
1035 if n.starts_with("cachedContents/") {
1036 n.to_string()
1037 } else {
1038 format!("cachedContents/{n}")
1039 }
1040 })
1041 .unwrap_or_else(|| "cachedContents".to_string());
1042 self.build_url_with_suffix(&suffix)
1043 }
1044
1045 #[instrument(skip_all, fields(display_name = request.display_name.as_deref()))]
1048 pub async fn create_file_search_store(
1049 &self,
1050 request: crate::file_search::CreateFileSearchStoreRequest,
1051 ) -> Result<crate::file_search::FileSearchStore, Error> {
1052 let url = self.build_url_with_suffix("fileSearchStores")?;
1053 self.post_json(url, &request).await
1054 }
1055
1056 #[instrument(skip_all, fields(store.name = %name))]
1057 pub async fn get_file_search_store(
1058 &self,
1059 name: &str,
1060 ) -> Result<crate::file_search::FileSearchStore, Error> {
1061 let url = self.build_url_with_suffix(name)?;
1062 self.get_json(url).await
1063 }
1064
1065 #[instrument(skip_all, fields(
1066 page.size = page_size,
1067 page.token.present = page_token.is_some(),
1068 ))]
1069 pub async fn list_file_search_stores(
1070 &self,
1071 page_size: Option<u32>,
1072 page_token: Option<&str>,
1073 ) -> Result<crate::file_search::ListFileSearchStoresResponse, Error> {
1074 let mut url = self.build_url_with_suffix("fileSearchStores")?;
1075 if let Some(size) = page_size {
1076 url.query_pairs_mut()
1077 .append_pair("pageSize", &size.to_string());
1078 }
1079 if let Some(token) = page_token {
1080 url.query_pairs_mut().append_pair("pageToken", token);
1081 }
1082 self.get_json(url).await
1083 }
1084
1085 #[instrument(skip_all, fields(store.name = %name, force))]
1086 pub async fn delete_file_search_store(&self, name: &str, force: bool) -> Result<(), Error> {
1087 let mut url = self.build_url_with_suffix(name)?;
1088 if force {
1089 url.query_pairs_mut().append_pair("force", "true");
1090 }
1091 self.perform_request(|c| c.delete(url.clone()), async |_r| Ok(()))
1092 .await
1093 }
1094
1095 #[instrument(skip_all, fields(
1098 store.name = %store_name,
1099 file.size = file_data.len(),
1100 display_name = display_name.as_deref(),
1101 mime.type = mime_type.as_ref().map(|m| m.to_string()),
1102 ))]
1103 pub async fn upload_to_file_search_store(
1104 &self,
1105 store_name: &str,
1106 file_data: Vec<u8>,
1107 display_name: Option<String>,
1108 mime_type: Option<mime::Mime>,
1109 custom_metadata: Option<Vec<crate::file_search::CustomMetadata>>,
1110 chunking_config: Option<crate::file_search::ChunkingConfig>,
1111 ) -> Result<crate::file_search::Operation, Error> {
1112 use crate::file_search::UploadToFileSearchStoreRequest;
1113
1114 let metadata_request = UploadToFileSearchStoreRequest {
1115 display_name,
1116 custom_metadata,
1117 chunking_config,
1118 mime_type: mime_type.clone(),
1119 };
1120
1121 let mime = mime_type.unwrap_or(mime::APPLICATION_OCTET_STREAM);
1122
1123 let init_url = format!("/upload/v1beta/{}:uploadToFileSearchStore", store_name);
1124 let upload_url = self
1125 .initiate_resumable_upload(&init_url, file_data.len(), &mime, Some(&metadata_request))
1126 .await?;
1127
1128 let operation: crate::file_search::Operation =
1129 self.upload_file_data(&upload_url, file_data).await?;
1130 Ok(operation)
1131 }
1132
1133 #[instrument(skip_all, fields(
1136 store.name = %store_name,
1137 file.name = %request.file_name,
1138 ))]
1139 pub async fn import_file_to_search_store(
1140 &self,
1141 store_name: &str,
1142 request: crate::file_search::ImportFileRequest,
1143 ) -> Result<crate::file_search::Operation, Error> {
1144 let url = self.build_url_with_suffix(&format!("{}:importFile", store_name))?;
1145 self.post_json(url, &request).await
1146 }
1147
1148 #[instrument(skip_all, fields(
1151 store.name = %store_name,
1152 document.id = %document_id,
1153 ))]
1154 pub async fn get_document(
1155 &self,
1156 store_name: &str,
1157 document_id: &str,
1158 ) -> Result<crate::file_search::Document, Error> {
1159 let url =
1160 self.build_url_with_suffix(&format!("{}/documents/{}", store_name, document_id))?;
1161 self.get_json(url).await
1162 }
1163
1164 #[instrument(skip_all, fields(
1165 store.name = %store_name,
1166 page.size = page_size,
1167 page.token.present = page_token.is_some(),
1168 ))]
1169 pub async fn list_documents(
1170 &self,
1171 store_name: &str,
1172 page_size: Option<u32>,
1173 page_token: Option<&str>,
1174 ) -> Result<crate::file_search::ListDocumentsResponse, Error> {
1175 let mut url = self.build_url_with_suffix(&format!("{}/documents", store_name))?;
1176 if let Some(size) = page_size {
1177 url.query_pairs_mut()
1178 .append_pair("pageSize", &size.to_string());
1179 }
1180 if let Some(token) = page_token {
1181 url.query_pairs_mut().append_pair("pageToken", token);
1182 }
1183 self.get_json(url).await
1184 }
1185
1186 #[instrument(skip_all, fields(
1187 store.name = %store_name,
1188 document.id = %document_id,
1189 force,
1190 ))]
1191 pub async fn delete_document(
1192 &self,
1193 store_name: &str,
1194 document_id: &str,
1195 force: bool,
1196 ) -> Result<(), Error> {
1197 let mut url =
1198 self.build_url_with_suffix(&format!("{}/documents/{}", store_name, document_id))?;
1199 if force {
1200 url.query_pairs_mut().append_pair("force", "true");
1201 }
1202 self.perform_request(|c| c.delete(url.clone()), async |_r| Ok(()))
1203 .await
1204 }
1205
1206 #[instrument(skip_all, fields(operation.name = %name))]
1209 pub async fn get_operation(&self, name: &str) -> Result<crate::file_search::Operation, Error> {
1210 let url = self.build_url_with_suffix(name)?;
1211 self.get_json(url).await
1212 }
1213
1214 #[instrument(skip(self, metadata))]
1217 async fn initiate_resumable_upload<T: Serialize>(
1218 &self,
1219 path: &str,
1220 total_bytes: usize,
1221 mime_type: &Mime,
1222 metadata: Option<&T>,
1223 ) -> Result<String, Error> {
1224 let url = self.build_url_with_suffix(path)?;
1225
1226 tracing::debug!("initiating resumable upload to {}", url);
1227
1228 let mut request = self
1229 .http_client
1230 .post(url.clone())
1231 .header("X-Goog-Upload-Protocol", "resumable")
1232 .header("X-Goog-Upload-Command", "start")
1233 .header(
1234 "X-Goog-Upload-Header-Content-Length",
1235 total_bytes.to_string(),
1236 )
1237 .header("X-Goog-Upload-Header-Content-Type", mime_type.to_string())
1238 .header("Content-Type", "application/json");
1239
1240 if let Some(metadata) = metadata {
1242 request = request.json(metadata);
1243 } else {
1244 request = request.body("{}");
1245 }
1246
1247 let response = request.send().await.context(PerformRequestNewSnafu)?;
1248
1249 let response = Self::check_response(response).await?;
1251
1252 let upload_url = response
1253 .headers()
1254 .get("x-goog-upload-url")
1255 .and_then(|v| v.to_str().ok())
1256 .ok_or(Error::MissingResponseHeader {
1257 header: "x-goog-upload-url".to_string(),
1258 })?;
1259
1260 tracing::debug!("received upload url: {}", upload_url);
1261 Ok(upload_url.to_string())
1262 }
1263
1264 #[instrument(skip(self, data), fields(data.len = data.len()))]
1265 async fn upload_file_data<T: serde::de::DeserializeOwned>(
1266 &self,
1267 upload_url: &str,
1268 data: Vec<u8>,
1269 ) -> Result<T, Error> {
1270 tracing::debug!("uploading file data to {}", upload_url);
1271
1272 let data_len = data.len();
1273 let response = self
1274 .http_client
1275 .post(upload_url)
1276 .header("Content-Length", data_len.to_string())
1277 .header("X-Goog-Upload-Offset", "0")
1278 .header("X-Goog-Upload-Command", "upload, finalize")
1279 .body(data)
1280 .send()
1281 .await
1282 .context(PerformRequestNewSnafu)?;
1283
1284 tracing::debug!("upload response status: {}", response.status());
1285 let response = Self::check_response(response).await?;
1286
1287 response.json().await.context(DecodeResponseSnafu)
1289 }
1290}
1291
1292pub struct GeminiBuilder {
1326 key: String,
1327 model: Model,
1328 client_builder: ClientBuilder,
1329 base_url: Url,
1330}
1331
1332impl GeminiBuilder {
1333 pub fn new<K: Into<String>>(key: K) -> Self {
1335 Self {
1336 key: key.into(),
1337 model: Model::default(),
1338 client_builder: ClientBuilder::default(),
1339 base_url: DEFAULT_BASE_URL.clone(),
1340 }
1341 }
1342
1343 pub fn with_model<M: Into<Model>>(mut self, model: M) -> Self {
1345 self.model = model.into();
1346 self
1347 }
1348
1349 pub fn with_http_client(mut self, client_builder: ClientBuilder) -> Self {
1351 self.client_builder = client_builder;
1352 self
1353 }
1354
1355 pub fn with_base_url(mut self, base_url: Url) -> Self {
1357 self.base_url = base_url;
1358 self
1359 }
1360
1361 pub fn build(self) -> Result<Gemini, Error> {
1363 Ok(Gemini {
1364 client: Arc::new(GeminiClient::with_base_url(
1365 self.client_builder,
1366 self.key,
1367 self.model,
1368 self.base_url,
1369 )?),
1370 })
1371 }
1372}
1373
1374#[derive(Clone)]
1376pub struct Gemini {
1377 client: Arc<GeminiClient>,
1378}
1379
1380impl Gemini {
1381 pub fn new<K: AsRef<str>>(api_key: K) -> Result<Self, Error> {
1383 Self::with_model(api_key, Model::default())
1384 }
1385
1386 pub fn pro<K: AsRef<str>>(api_key: K) -> Result<Self, Error> {
1388 Self::with_model(api_key, Model::Gemini31Pro)
1389 }
1390
1391 pub fn pro_image<K: AsRef<str>>(api_key: K) -> Result<Self, Error> {
1393 Self::with_model(api_key, Model::Gemini3ProImage)
1394 }
1395
1396 pub fn with_model<K: AsRef<str>, M: Into<Model>>(api_key: K, model: M) -> Result<Self, Error> {
1398 Self::with_model_and_base_url(api_key, model, DEFAULT_BASE_URL.clone())
1399 }
1400
1401 pub fn with_base_url<K: AsRef<str>>(api_key: K, base_url: Url) -> Result<Self, Error> {
1403 Self::with_model_and_base_url(api_key, Model::default(), base_url)
1404 }
1405
1406 pub fn with_model_and_base_url<K: AsRef<str>, M: Into<Model>>(
1408 api_key: K,
1409 model: M,
1410 base_url: Url,
1411 ) -> Result<Self, Error> {
1412 let client =
1413 GeminiClient::with_base_url(Default::default(), api_key, model.into(), base_url)?;
1414 Ok(Self {
1415 client: Arc::new(client),
1416 })
1417 }
1418
1419 #[deprecated(
1421 since = "1.8.0",
1422 note = "Use Gemini::create_interaction() instead. See migration guide: interactions-api/migration-plan.md"
1423 )]
1424 #[allow(deprecated)]
1425 pub fn generate_content(&self) -> ContentBuilder {
1426 ContentBuilder::new(self.client.clone())
1427 }
1428
1429 pub fn create_interaction(&self) -> InteractionBuilder {
1435 InteractionBuilder::new(self.client.clone())
1436 }
1437
1438 pub fn interaction(&self, id: &str) -> InteractionHandle {
1442 InteractionHandle::new(id.to_string(), self.client.clone())
1443 }
1444
1445 pub async fn get_interaction(&self, id: &str) -> Result<Interaction, Error> {
1447 self.client.get_interaction(id).await
1448 }
1449
1450 pub fn embed_content(&self) -> EmbedBuilder {
1452 EmbedBuilder::new(self.client.clone())
1453 }
1454
1455 pub fn batch_generate_content(&self) -> BatchBuilder {
1457 BatchBuilder::new(self.client.clone())
1458 }
1459
1460 pub fn get_batch(&self, name: &str) -> BatchHandle {
1462 BatchHandle::new(name.to_string(), self.client.clone())
1463 }
1464
1465 pub fn list_batches(
1469 &self,
1470 page_size: impl Into<Option<u32>>,
1471 ) -> impl Stream<Item = Result<BatchOperation, Error>> + Send {
1472 let client = self.client.clone();
1473 let page_size = page_size.into();
1474 async_stream::try_stream! {
1475 let mut page_token: Option<String> = None;
1476 loop {
1477 let response = client
1478 .list_batch_operations(page_size, page_token.clone())
1479 .await?;
1480
1481 for operation in response.operations {
1482 yield operation;
1483 }
1484
1485 if let Some(next_page_token) = response.next_page_token {
1486 page_token = Some(next_page_token);
1487 } else {
1488 break;
1489 }
1490 }
1491 }
1492 }
1493
1494 pub fn create_cache(&self) -> CacheBuilder {
1496 CacheBuilder::new(self.client.clone())
1497 }
1498
1499 pub fn get_cached_content(&self, name: &str) -> CachedContentHandle {
1501 CachedContentHandle::new(name.to_string(), self.client.clone())
1502 }
1503
1504 pub fn list_cached_contents(
1508 &self,
1509 page_size: impl Into<Option<i32>>,
1510 ) -> impl Stream<Item = Result<CachedContentSummary, Error>> + Send {
1511 let client = self.client.clone();
1512 let page_size = page_size.into();
1513 async_stream::try_stream! {
1514 let mut page_token: Option<String> = None;
1515 loop {
1516 let response = client
1517 .list_cached_contents(page_size, page_token.clone())
1518 .await?;
1519
1520 for cached_content in response.cached_contents {
1521 yield cached_content;
1522 }
1523
1524 if let Some(next_page_token) = response.next_page_token {
1525 page_token = Some(next_page_token);
1526 } else {
1527 break;
1528 }
1529 }
1530 }
1531 }
1532
1533 pub fn create_file<B: Into<Vec<u8>>>(&self, bytes: B) -> crate::files::builder::FileBuilder {
1535 crate::files::builder::FileBuilder::new(self.client.clone(), bytes)
1536 }
1537
1538 pub async fn get_file(&self, name: &str) -> Result<FileHandle, Error> {
1540 let file = self.client.get_file(name).await?;
1541 Ok(FileHandle::new(self.client.clone(), file))
1542 }
1543
1544 pub fn list_files(
1548 &self,
1549 page_size: impl Into<Option<u32>>,
1550 ) -> impl Stream<Item = Result<FileHandle, Error>> + Send {
1551 let client = self.client.clone();
1552 let page_size = page_size.into();
1553 async_stream::try_stream! {
1554 let mut page_token: Option<String> = None;
1555 loop {
1556 let response = client
1557 .list_files(page_size, page_token.clone())
1558 .await?;
1559
1560 for file in response.files {
1561 yield FileHandle::new(client.clone(), file);
1562 }
1563
1564 if let Some(next_page_token) = response.next_page_token {
1565 page_token = Some(next_page_token);
1566 } else {
1567 break;
1568 }
1569 }
1570 }
1571 }
1572
1573 pub fn create_file_search_store(&self) -> crate::file_search::FileSearchStoreBuilder {
1575 crate::file_search::FileSearchStoreBuilder {
1576 client: self.client.clone(),
1577 display_name: None,
1578 }
1579 }
1580
1581 pub async fn get_file_search_store(
1583 &self,
1584 name: &str,
1585 ) -> Result<crate::file_search::FileSearchStoreHandle, Error> {
1586 let store = self.client.get_file_search_store(name).await?;
1587 Ok(crate::file_search::FileSearchStoreHandle::new(
1588 self.client.clone(),
1589 store,
1590 ))
1591 }
1592
1593 pub fn list_file_search_stores(
1597 &self,
1598 page_size: impl Into<Option<u32>>,
1599 ) -> impl Stream<Item = Result<crate::file_search::FileSearchStoreHandle, Error>> + Send {
1600 let client = self.client.clone();
1601 let page_size = page_size.into();
1602 async_stream::try_stream! {
1603 let mut page_token: Option<String> = None;
1604 loop {
1605 let response = client
1606 .list_file_search_stores(page_size, page_token.as_deref())
1607 .await?;
1608
1609 for store in response.file_search_stores {
1610 yield crate::file_search::FileSearchStoreHandle::new(client.clone(), store);
1611 }
1612
1613 if let Some(next_page_token) = response.next_page_token {
1614 page_token = Some(next_page_token);
1615 } else {
1616 break;
1617 }
1618 }
1619 }
1620 }
1621}