rig_core/providers/gemini/
extension.rs

1//! Typed request options and reply extras for the Gemini API, keyed by
2//! [`PROVIDER_NAME`](super::PROVIDER_NAME). One [`GeminiOptions`] serves both
3//! routes: its `"*"` section goes to GenerateContent and Interactions alike,
4//! and each route section only to its own route.
5//!
6//! [`GeminiExtras`] reads the reply document of either route. A field the
7//! route taken does not return reads `None`.
8//!
9//! ```
10//! use rig_core::completion::CompletionRequest;
11//! use rig_core::providers::gemini::extension::GeminiOptions;
12//!
13//! let options = GeminiOptions::new().top_k(40);
14//! let request = CompletionRequest::new("hi").provider_option(options);
15//! # let _ = request;
16//! ```
17
18use std::collections::BTreeMap;
19
20use serde::de::Error as _;
21use serde::ser::SerializeMap;
22use serde::{Deserialize, Serialize, Serializer};
23use serde_json::Value;
24
25use crate::completion::provider_options::{ExtensionOptions, ProviderExtension, ReplyExtras};
26use crate::message::Api;
27
28use super::ImageSize;
29use super::interactions_api::InteractionSpeechConfig;
30
31/// The API name of the GenerateContent route.
32const GENERATE_CONTENT: &str = "gemini.generate_content";
33
34/// The API name of the Interactions route.
35const INTERACTIONS: &str = "gemini.interactions";
36
37/// The Gemini API's extension: [`GeminiOptions`] and [`GeminiExtras`].
38#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
39pub struct GeminiExt;
40
41impl ProviderExtension for GeminiExt {
42    const PROVIDER: &'static str = super::PROVIDER_NAME;
43    type Options = GeminiOptions;
44    type Extras = GeminiExtras;
45}
46
47/// Whether `value` is its type's default, so a section or config that sets
48/// nothing writes no key.
49fn is_default<T: Default + PartialEq>(value: &T) -> bool {
50    *value == T::default()
51}
52
53/// The Gemini API's request options: a section both routes read, and one
54/// per route.
55///
56/// The `generationConfig` setters here (`top_k`, `include_thoughts`, ...)
57/// write the same entry as the nested [`GenerateContentOptions`] form, so
58/// only GenerateContent sends them. [`generate_content`](Self::generate_content)
59/// replaces the whole section, including entries set before it. Safety
60/// settings stay per section, because each route spells them differently.
61#[non_exhaustive]
62#[derive(Clone, Debug, Default, PartialEq, Serialize)]
63pub struct GeminiOptions {
64    /// Fields both routes spell the same.
65    #[serde(rename = "*")]
66    pub shared: GeminiShared,
67    /// Fields only GenerateContent reads.
68    #[serde(rename = "gemini.generate_content")]
69    pub generate_content: GenerateContentOptions,
70    /// Fields only Interactions reads.
71    #[serde(rename = "gemini.interactions")]
72    pub interactions: InteractionsOptions,
73}
74
75impl GeminiOptions {
76    /// No field set.
77    pub fn new() -> Self {
78        Self::default()
79    }
80
81    /// Whether the API stores the request and its reply (`store`).
82    pub fn store(mut self, store: bool) -> Self {
83        self.shared.store = Some(store);
84        self
85    }
86
87    /// Add the label `key` = `value` (`labels`), user metadata the API
88    /// reports billing by.
89    pub fn label(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
90        self.shared.labels.insert(key.into(), value.into());
91        self
92    }
93
94    /// The GenerateContent section.
95    pub fn generate_content(mut self, section: GenerateContentOptions) -> Self {
96        self.generate_content = section;
97        self
98    }
99
100    /// The Interactions section.
101    pub fn interactions(mut self, section: InteractionsOptions) -> Self {
102        self.interactions = section;
103        self
104    }
105
106    /// Apply `set` to the GenerateContent `generationConfig` entries.
107    fn generation_config_entry(
108        mut self,
109        set: impl FnOnce(GenerationConfig) -> GenerationConfig,
110    ) -> Self {
111        let common = std::mem::take(&mut self.generate_content.generation_config.common);
112        self.generate_content.generation_config.common = set(common);
113        self
114    }
115
116    /// GenerateContent `generationConfig.thinkingConfig.includeThoughts`,
117    /// as [`GenerationConfig::include_thoughts`] in the GenerateContent
118    /// section. Interactions reads its own `thinking_summaries` instead.
119    pub fn include_thoughts(self, include: bool) -> Self {
120        self.generation_config_entry(|config| config.include_thoughts(include))
121    }
122
123    /// GenerateContent `generationConfig.topK`, as
124    /// [`GenerationConfig::top_k`] in the GenerateContent section.
125    pub fn top_k(self, top_k: u32) -> Self {
126        self.generation_config_entry(|config| config.top_k(top_k))
127    }
128
129    /// GenerateContent `generationConfig.presencePenalty`, as
130    /// [`GenerationConfig::presence_penalty`] in the GenerateContent section.
131    pub fn presence_penalty(self, penalty: f64) -> Self {
132        self.generation_config_entry(|config| config.presence_penalty(penalty))
133    }
134
135    /// GenerateContent `generationConfig.frequencyPenalty`, as
136    /// [`GenerationConfig::frequency_penalty`] in the GenerateContent
137    /// section.
138    pub fn frequency_penalty(self, penalty: f64) -> Self {
139        self.generation_config_entry(|config| config.frequency_penalty(penalty))
140    }
141
142    /// GenerateContent `generationConfig.responseLogprobs`, as
143    /// [`GenerationConfig::response_logprobs`] in the GenerateContent
144    /// section.
145    pub fn response_logprobs(self, enable: bool) -> Self {
146        self.generation_config_entry(|config| config.response_logprobs(enable))
147    }
148
149    /// GenerateContent `generationConfig.logprobs`, as
150    /// [`GenerationConfig::logprobs`] in the GenerateContent section.
151    pub fn logprobs(self, top: u32) -> Self {
152        self.generation_config_entry(|config| config.logprobs(top))
153    }
154
155    /// GenerateContent `generationConfig.candidateCount`, as
156    /// [`GenerationConfig::candidate_count`] in the GenerateContent section.
157    pub fn candidate_count(self, count: CandidateCount) -> Self {
158        self.generation_config_entry(|config| config.candidate_count(count))
159    }
160
161    /// GenerateContent `generationConfig.responseModalities`, as
162    /// [`GenerationConfig::response_modalities`] in the GenerateContent
163    /// section.
164    pub fn response_modalities(
165        self,
166        modalities: impl IntoIterator<Item = ResponseModality>,
167    ) -> Self {
168        self.generation_config_entry(|config| config.response_modalities(modalities))
169    }
170
171    /// GenerateContent `generationConfig.mediaResolution`, as
172    /// [`GenerationConfig::media_resolution`] in the GenerateContent
173    /// section.
174    pub fn media_resolution(self, resolution: MediaResolution) -> Self {
175        self.generation_config_entry(|config| config.media_resolution(resolution))
176    }
177}
178
179impl ExtensionOptions for GeminiOptions {
180    type Ext = GeminiExt;
181}
182
183/// The fields both Gemini routes read, at the top level of the body.
184#[non_exhaustive]
185#[derive(Clone, Debug, Default, PartialEq, Serialize)]
186pub struct GeminiShared {
187    /// `store`: whether the API stores the request and its reply.
188    #[serde(skip_serializing_if = "Option::is_none")]
189    pub store: Option<bool>,
190    /// `labels`: user metadata, each key and value at most 63 lowercase
191    /// letters, digits, underscores and dashes.
192    #[serde(skip_serializing_if = "BTreeMap::is_empty")]
193    pub labels: BTreeMap<String, String>,
194}
195
196/// The GenerateContent-only fields: `generationConfig` entries and
197/// `safetySettings`.
198#[non_exhaustive]
199#[derive(Clone, Debug, Default, PartialEq, Serialize)]
200pub struct GenerateContentOptions {
201    /// `generationConfig` entries, merged key by key with the ones the
202    /// request and its generation options write.
203    #[serde(rename = "generationConfig", skip_serializing_if = "is_default")]
204    pub generation_config: GeminiGenerationConfig,
205    /// `safetySettings`, which take the place of the `null` sent otherwise.
206    #[serde(rename = "safetySettings", skip_serializing_if = "Vec::is_empty")]
207    pub safety_settings: Vec<SafetySetting>,
208}
209
210impl GenerateContentOptions {
211    /// No field set.
212    pub fn new() -> Self {
213        Self::default()
214    }
215
216    /// The `generationConfig` entries every GenerateContent route takes.
217    pub fn generation_config(mut self, config: GenerationConfig) -> Self {
218        self.generation_config.common = config;
219        self
220    }
221
222    /// `generationConfig.enableEnhancedCivicAnswers`.
223    pub fn enable_enhanced_civic_answers(mut self, enable: bool) -> Self {
224        self.generation_config.enable_enhanced_civic_answers = Some(enable);
225        self
226    }
227
228    /// Add a safety setting: block `category` at `threshold`.
229    pub fn safety_setting(mut self, category: HarmCategory, threshold: HarmBlockThreshold) -> Self {
230        self.safety_settings
231            .push(SafetySetting::new(category, threshold));
232        self
233    }
234}
235
236/// The Gemini API's `generationConfig` entries: the ones every
237/// GenerateContent route takes and the API's own.
238#[non_exhaustive]
239#[derive(Clone, Debug, Default, PartialEq, Serialize)]
240pub struct GeminiGenerationConfig {
241    /// The entries every GenerateContent route takes.
242    #[serde(flatten)]
243    pub common: GenerationConfig,
244    /// `enableEnhancedCivicAnswers`.
245    #[serde(
246        rename = "enableEnhancedCivicAnswers",
247        skip_serializing_if = "Option::is_none"
248    )]
249    pub enable_enhanced_civic_answers: Option<bool>,
250}
251
252/// The `generationConfig` entries the Gemini API, Vertex AI and the gRPC
253/// API all take. None of them is one a request or its generation options
254/// write: `temperature`, `maxOutputTokens`, `topP`, `seed`, `stopSequences`
255/// and the thinking level or budget are set there.
256#[non_exhaustive]
257#[derive(Clone, Debug, Default, PartialEq, Serialize)]
258#[serde(rename_all = "camelCase")]
259pub struct GenerationConfig {
260    /// `thinkingConfig.includeThoughts`, next to the thinking level or
261    /// budget the request's reasoning sets.
262    #[serde(skip_serializing_if = "is_default")]
263    pub thinking_config: ThinkingConfig,
264    /// `topK`.
265    #[serde(skip_serializing_if = "Option::is_none")]
266    pub top_k: Option<u32>,
267    /// `presencePenalty`.
268    #[serde(skip_serializing_if = "Option::is_none")]
269    pub presence_penalty: Option<f64>,
270    /// `frequencyPenalty`.
271    #[serde(skip_serializing_if = "Option::is_none")]
272    pub frequency_penalty: Option<f64>,
273    /// `responseLogprobs`: return the chosen tokens' log probabilities.
274    #[serde(skip_serializing_if = "Option::is_none")]
275    pub response_logprobs: Option<bool>,
276    /// `logprobs`: how many top candidates to return per token.
277    #[serde(skip_serializing_if = "Option::is_none")]
278    pub logprobs: Option<u32>,
279    /// `candidateCount`.
280    #[serde(skip_serializing_if = "Option::is_none")]
281    pub candidate_count: Option<CandidateCount>,
282    /// `responseModalities`.
283    #[serde(skip_serializing_if = "Vec::is_empty")]
284    pub response_modalities: Vec<ResponseModality>,
285    /// `imageConfig`.
286    #[serde(skip_serializing_if = "Option::is_none")]
287    pub image_config: Option<ImageConfig>,
288    /// `speechConfig`.
289    #[serde(skip_serializing_if = "Option::is_none")]
290    pub speech_config: Option<SpeechConfig>,
291    /// `mediaResolution`.
292    #[serde(skip_serializing_if = "Option::is_none")]
293    pub media_resolution: Option<MediaResolution>,
294}
295
296impl GenerationConfig {
297    /// No entry set.
298    pub fn new() -> Self {
299        Self::default()
300    }
301
302    /// `thinkingConfig.includeThoughts`: return thought summaries.
303    pub fn include_thoughts(mut self, include: bool) -> Self {
304        self.thinking_config.include_thoughts = Some(include);
305        self
306    }
307
308    /// `topK`.
309    pub fn top_k(mut self, top_k: u32) -> Self {
310        self.top_k = Some(top_k);
311        self
312    }
313
314    /// `presencePenalty`.
315    pub fn presence_penalty(mut self, penalty: f64) -> Self {
316        self.presence_penalty = Some(penalty);
317        self
318    }
319
320    /// `frequencyPenalty`.
321    pub fn frequency_penalty(mut self, penalty: f64) -> Self {
322        self.frequency_penalty = Some(penalty);
323        self
324    }
325
326    /// `responseLogprobs`.
327    pub fn response_logprobs(mut self, enable: bool) -> Self {
328        self.response_logprobs = Some(enable);
329        self
330    }
331
332    /// `logprobs`.
333    pub fn logprobs(mut self, top: u32) -> Self {
334        self.logprobs = Some(top);
335        self
336    }
337
338    /// `candidateCount`.
339    pub fn candidate_count(mut self, count: CandidateCount) -> Self {
340        self.candidate_count = Some(count);
341        self
342    }
343
344    /// `responseModalities`.
345    pub fn response_modalities(
346        mut self,
347        modalities: impl IntoIterator<Item = ResponseModality>,
348    ) -> Self {
349        self.response_modalities = modalities.into_iter().collect();
350        self
351    }
352
353    /// `imageConfig`.
354    pub fn image_config(mut self, config: ImageConfig) -> Self {
355        self.image_config = Some(config);
356        self
357    }
358
359    /// `speechConfig`.
360    pub fn speech_config(mut self, config: SpeechConfig) -> Self {
361        self.speech_config = Some(config);
362        self
363    }
364
365    /// `mediaResolution`.
366    pub fn media_resolution(mut self, resolution: MediaResolution) -> Self {
367        self.media_resolution = Some(resolution);
368        self
369    }
370}
371
372/// The provider half of `thinkingConfig`.
373#[non_exhaustive]
374#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
375#[serde(rename_all = "camelCase")]
376pub struct ThinkingConfig {
377    /// `includeThoughts`.
378    #[serde(skip_serializing_if = "Option::is_none")]
379    pub include_thoughts: Option<bool>,
380}
381
382/// How many candidates to generate. Rig reads only the first candidate of
383/// a reply, so one is the only count it can ask for.
384#[non_exhaustive]
385#[derive(Clone, Copy, Debug, PartialEq, Eq)]
386pub enum CandidateCount {
387    /// One candidate, sent as `1`.
388    One,
389}
390
391impl Serialize for CandidateCount {
392    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
393        match self {
394            Self::One => serializer.serialize_u32(1),
395        }
396    }
397}
398
399/// A modality of the reply.
400#[non_exhaustive]
401#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
402#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
403pub enum ResponseModality {
404    /// Text.
405    Text,
406    /// Images.
407    Image,
408    /// Audio.
409    Audio,
410}
411
412/// The resolution media inputs are read at.
413#[non_exhaustive]
414#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
415pub enum MediaResolution {
416    /// `MEDIA_RESOLUTION_LOW`.
417    #[serde(rename = "MEDIA_RESOLUTION_LOW")]
418    Low,
419    /// `MEDIA_RESOLUTION_MEDIUM`.
420    #[serde(rename = "MEDIA_RESOLUTION_MEDIUM")]
421    Medium,
422    /// `MEDIA_RESOLUTION_HIGH`.
423    #[serde(rename = "MEDIA_RESOLUTION_HIGH")]
424    High,
425}
426
427/// How generated images are shaped.
428#[non_exhaustive]
429#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
430#[serde(rename_all = "camelCase")]
431pub struct ImageConfig {
432    /// `aspectRatio`, such as `"16:9"`.
433    #[serde(skip_serializing_if = "Option::is_none")]
434    pub aspect_ratio: Option<String>,
435    /// `imageSize`.
436    #[serde(skip_serializing_if = "Option::is_none")]
437    pub image_size: Option<ImageSize>,
438}
439
440impl ImageConfig {
441    /// No entry set.
442    pub fn new() -> Self {
443        Self::default()
444    }
445
446    /// `aspectRatio`.
447    pub fn aspect_ratio(mut self, ratio: impl Into<String>) -> Self {
448        self.aspect_ratio = Some(ratio.into());
449        self
450    }
451
452    /// `imageSize`.
453    pub fn image_size(mut self, size: ImageSize) -> Self {
454        self.image_size = Some(size);
455        self
456    }
457}
458
459/// How generated speech sounds: one voice, or one per speaker.
460#[non_exhaustive]
461#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
462#[serde(rename_all = "camelCase")]
463pub struct SpeechConfig {
464    /// `voiceConfig`: the one voice.
465    #[serde(skip_serializing_if = "Option::is_none")]
466    pub voice_config: Option<VoiceConfig>,
467    /// `multiSpeakerVoiceConfig`: a voice per speaker.
468    #[serde(skip_serializing_if = "Option::is_none")]
469    pub multi_speaker_voice_config: Option<MultiSpeakerVoiceConfig>,
470    /// `languageCode`, such as `"en-US"`.
471    #[serde(skip_serializing_if = "Option::is_none")]
472    pub language_code: Option<String>,
473}
474
475impl SpeechConfig {
476    /// Speak with the prebuilt voice `voice_name`.
477    pub fn voice(voice_name: impl Into<String>) -> Self {
478        Self {
479            voice_config: Some(VoiceConfig::prebuilt(voice_name)),
480            multi_speaker_voice_config: None,
481            language_code: None,
482        }
483    }
484
485    /// Speak each `(speaker, voice_name)` pair's lines with its prebuilt
486    /// voice.
487    pub fn speakers<S: Into<String>, V: Into<String>>(
488        speakers: impl IntoIterator<Item = (S, V)>,
489    ) -> Self {
490        let speaker_voice_configs = speakers
491            .into_iter()
492            .map(|(speaker, voice)| SpeakerVoiceConfig {
493                speaker: speaker.into(),
494                voice_config: VoiceConfig::prebuilt(voice),
495            })
496            .collect();
497        Self {
498            voice_config: None,
499            multi_speaker_voice_config: Some(MultiSpeakerVoiceConfig {
500                speaker_voice_configs,
501            }),
502            language_code: None,
503        }
504    }
505
506    /// `languageCode`.
507    pub fn language_code(mut self, code: impl Into<String>) -> Self {
508        self.language_code = Some(code.into());
509        self
510    }
511}
512
513/// A voice.
514#[non_exhaustive]
515#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
516#[serde(rename_all = "camelCase")]
517pub struct VoiceConfig {
518    /// `prebuiltVoiceConfig`.
519    pub prebuilt_voice_config: PrebuiltVoiceConfig,
520}
521
522impl VoiceConfig {
523    /// The prebuilt voice `voice_name`, such as `"Kore"`.
524    pub fn prebuilt(voice_name: impl Into<String>) -> Self {
525        Self {
526            prebuilt_voice_config: PrebuiltVoiceConfig {
527                voice_name: voice_name.into(),
528            },
529        }
530    }
531}
532
533/// A prebuilt voice, by name.
534#[non_exhaustive]
535#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
536#[serde(rename_all = "camelCase")]
537pub struct PrebuiltVoiceConfig {
538    /// `voiceName`.
539    pub voice_name: String,
540}
541
542/// A voice per speaker.
543#[non_exhaustive]
544#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
545#[serde(rename_all = "camelCase")]
546pub struct MultiSpeakerVoiceConfig {
547    /// `speakerVoiceConfigs`.
548    pub speaker_voice_configs: Vec<SpeakerVoiceConfig>,
549}
550
551/// One speaker's voice.
552#[non_exhaustive]
553#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
554#[serde(rename_all = "camelCase")]
555pub struct SpeakerVoiceConfig {
556    /// `speaker`: the name the prompt gives the speaker.
557    pub speaker: String,
558    /// `voiceConfig`.
559    pub voice_config: VoiceConfig,
560}
561
562/// A harm category a safety setting applies to.
563#[non_exhaustive]
564#[derive(Clone, Copy, Debug, PartialEq, Eq)]
565pub enum HarmCategory {
566    /// Hate speech.
567    HateSpeech,
568    /// Dangerous content.
569    DangerousContent,
570    /// Harassment.
571    Harassment,
572    /// Sexually explicit content.
573    SexuallyExplicit,
574    /// Content that may be used to harm civic integrity.
575    CivicIntegrity,
576    /// Jailbreak attempts.
577    Jailbreak,
578}
579
580impl HarmCategory {
581    /// The GenerateContent spelling, such as `HARM_CATEGORY_HATE_SPEECH`.
582    pub fn as_str(self) -> &'static str {
583        match self {
584            Self::HateSpeech => "HARM_CATEGORY_HATE_SPEECH",
585            Self::DangerousContent => "HARM_CATEGORY_DANGEROUS_CONTENT",
586            Self::Harassment => "HARM_CATEGORY_HARASSMENT",
587            Self::SexuallyExplicit => "HARM_CATEGORY_SEXUALLY_EXPLICIT",
588            Self::CivicIntegrity => "HARM_CATEGORY_CIVIC_INTEGRITY",
589            Self::Jailbreak => "HARM_CATEGORY_JAILBREAK",
590        }
591    }
592
593    /// The Interactions spelling, such as `hate_speech`.
594    fn interactions(self) -> &'static str {
595        match self {
596            Self::HateSpeech => "hate_speech",
597            Self::DangerousContent => "dangerous_content",
598            Self::Harassment => "harassment",
599            Self::SexuallyExplicit => "sexually_explicit",
600            Self::CivicIntegrity => "civic_integrity",
601            Self::Jailbreak => "jailbreak",
602        }
603    }
604}
605
606/// The probability at and above which content is blocked.
607#[non_exhaustive]
608#[derive(Clone, Copy, Debug, PartialEq, Eq)]
609pub enum HarmBlockThreshold {
610    /// Block low probability and above.
611    BlockLowAndAbove,
612    /// Block medium probability and above.
613    BlockMediumAndAbove,
614    /// Block only high probability.
615    BlockOnlyHigh,
616    /// Block nothing.
617    BlockNone,
618    /// Turn the safety filter off.
619    Off,
620}
621
622impl HarmBlockThreshold {
623    /// The GenerateContent spelling, such as `BLOCK_ONLY_HIGH`.
624    pub fn as_str(self) -> &'static str {
625        match self {
626            Self::BlockLowAndAbove => "BLOCK_LOW_AND_ABOVE",
627            Self::BlockMediumAndAbove => "BLOCK_MEDIUM_AND_ABOVE",
628            Self::BlockOnlyHigh => "BLOCK_ONLY_HIGH",
629            Self::BlockNone => "BLOCK_NONE",
630            Self::Off => "OFF",
631        }
632    }
633
634    /// The Interactions spelling, such as `block_only_high`.
635    fn interactions(self) -> &'static str {
636        match self {
637            Self::BlockLowAndAbove => "block_low_and_above",
638            Self::BlockMediumAndAbove => "block_medium_and_above",
639            Self::BlockOnlyHigh => "block_only_high",
640            Self::BlockNone => "block_none",
641            Self::Off => "off",
642        }
643    }
644}
645
646/// Block `category` at `threshold`. Sent as
647/// `{"category": "HARM_CATEGORY_…", "threshold": "BLOCK_…"}` on
648/// GenerateContent and `{"type": "…", "threshold": "block_…"}` on
649/// Interactions.
650#[non_exhaustive]
651#[derive(Clone, Copy, Debug, PartialEq, Eq)]
652pub struct SafetySetting {
653    /// The category.
654    pub category: HarmCategory,
655    /// The threshold.
656    pub threshold: HarmBlockThreshold,
657}
658
659impl SafetySetting {
660    /// Block `category` at `threshold`.
661    pub fn new(category: HarmCategory, threshold: HarmBlockThreshold) -> Self {
662        Self {
663            category,
664            threshold,
665        }
666    }
667}
668
669impl Serialize for SafetySetting {
670    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
671        let mut map = serializer.serialize_map(Some(2))?;
672        map.serialize_entry("category", self.category.as_str())?;
673        map.serialize_entry("threshold", self.threshold.as_str())?;
674        map.end()
675    }
676}
677
678/// `settings` in the Interactions spelling.
679fn interactions_safety<S: Serializer>(
680    settings: &[SafetySetting],
681    serializer: S,
682) -> Result<S::Ok, S::Error> {
683    serializer.collect_seq(settings.iter().map(|setting| {
684        BTreeMap::from([
685            ("type", setting.category.interactions()),
686            ("threshold", setting.threshold.interactions()),
687        ])
688    }))
689}
690
691/// The Interactions-only fields, at the top level of the create body but
692/// for `generation_config`.
693#[non_exhaustive]
694#[derive(Clone, Debug, Default, PartialEq, Serialize)]
695pub struct InteractionsOptions {
696    /// `agent`: the agent to run in place of the model, which the body then
697    /// leaves out.
698    #[serde(skip_serializing_if = "Option::is_none")]
699    pub agent: Option<String>,
700    /// `agent_config`.
701    #[serde(skip_serializing_if = "Option::is_none")]
702    pub agent_config: Option<AgentConfig>,
703    /// `background`: run the interaction in the background.
704    #[serde(skip_serializing_if = "Option::is_none")]
705    pub background: Option<bool>,
706    /// `previous_interaction_id`: continue a stored interaction.
707    #[serde(skip_serializing_if = "Option::is_none")]
708    pub previous_interaction_id: Option<String>,
709    /// `safety_settings`.
710    #[serde(
711        skip_serializing_if = "Vec::is_empty",
712        serialize_with = "interactions_safety"
713    )]
714    pub safety_settings: Vec<SafetySetting>,
715    /// `generation_config` entries, merged key by key with the ones the
716    /// request and its generation options write.
717    #[serde(skip_serializing_if = "is_default")]
718    pub generation_config: InteractionsGenerationConfig,
719}
720
721impl InteractionsOptions {
722    /// No field set.
723    pub fn new() -> Self {
724        Self::default()
725    }
726
727    /// `agent`.
728    pub fn agent(mut self, agent: impl Into<String>) -> Self {
729        self.agent = Some(agent.into());
730        self
731    }
732
733    /// `agent_config`.
734    pub fn agent_config(mut self, config: AgentConfig) -> Self {
735        self.agent_config = Some(config);
736        self
737    }
738
739    /// `background`.
740    pub fn background(mut self, background: bool) -> Self {
741        self.background = Some(background);
742        self
743    }
744
745    /// `previous_interaction_id`.
746    pub fn previous_interaction_id(mut self, id: impl Into<String>) -> Self {
747        self.previous_interaction_id = Some(id.into());
748        self
749    }
750
751    /// Add a safety setting: block `category` at `threshold`.
752    pub fn safety_setting(mut self, category: HarmCategory, threshold: HarmBlockThreshold) -> Self {
753        self.safety_settings
754            .push(SafetySetting::new(category, threshold));
755        self
756    }
757
758    /// `generation_config.thinking_summaries`.
759    pub fn thinking_summaries(mut self, summaries: ThinkingSummaries) -> Self {
760        self.generation_config.thinking_summaries = Some(summaries);
761        self
762    }
763
764    /// `generation_config.speech_config`.
765    pub fn speech_config(mut self, config: InteractionSpeechConfig) -> Self {
766        self.generation_config.speech_config = Some(config);
767        self
768    }
769}
770
771/// The provider entries of the Interactions `generation_config`.
772#[non_exhaustive]
773#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
774pub struct InteractionsGenerationConfig {
775    /// `thinking_summaries`.
776    #[serde(skip_serializing_if = "Option::is_none")]
777    pub thinking_summaries: Option<ThinkingSummaries>,
778    /// `speech_config`.
779    #[serde(skip_serializing_if = "Option::is_none")]
780    pub speech_config: Option<InteractionSpeechConfig>,
781}
782
783/// Whether the reply carries summaries of the model's thoughts.
784#[non_exhaustive]
785#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
786#[serde(rename_all = "snake_case")]
787pub enum ThinkingSummaries {
788    /// `auto`.
789    Auto,
790    /// `none`.
791    None,
792}
793
794/// An agent's configuration, tagged by `type`.
795#[non_exhaustive]
796#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
797#[serde(tag = "type", rename_all = "kebab-case")]
798pub enum AgentConfig {
799    /// `dynamic`.
800    Dynamic,
801    /// `deep-research`.
802    DeepResearch(DeepResearchConfig),
803}
804
805/// The Deep Research agent's `agent_config` fields. A field left unset
806/// takes the API's default.
807#[non_exhaustive]
808#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
809pub struct DeepResearchConfig {
810    /// `thinking_summaries`: stream the agent's intermediate thoughts.
811    #[serde(skip_serializing_if = "Option::is_none")]
812    pub thinking_summaries: Option<ThinkingSummaries>,
813    /// `visualization`: whether the agent may draw charts and images.
814    #[serde(skip_serializing_if = "Option::is_none")]
815    pub visualization: Option<Visualization>,
816    /// `collaborative_planning`: answer with a research plan to review
817    /// instead of the report. Send `false` on a later turn to approve it.
818    #[serde(skip_serializing_if = "Option::is_none")]
819    pub collaborative_planning: Option<bool>,
820}
821
822impl DeepResearchConfig {
823    /// No field set.
824    pub fn new() -> Self {
825        Self::default()
826    }
827
828    /// `thinking_summaries`.
829    pub fn thinking_summaries(mut self, summaries: ThinkingSummaries) -> Self {
830        self.thinking_summaries = Some(summaries);
831        self
832    }
833
834    /// `visualization`.
835    pub fn visualization(mut self, visualization: Visualization) -> Self {
836        self.visualization = Some(visualization);
837        self
838    }
839
840    /// `collaborative_planning`.
841    pub fn collaborative_planning(mut self, plan: bool) -> Self {
842        self.collaborative_planning = Some(plan);
843        self
844    }
845}
846
847impl From<DeepResearchConfig> for AgentConfig {
848    fn from(config: DeepResearchConfig) -> Self {
849        Self::DeepResearch(config)
850    }
851}
852
853/// Whether an agent may generate charts and images.
854#[non_exhaustive]
855#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
856#[serde(rename_all = "snake_case")]
857pub enum Visualization {
858    /// `auto`: when the prompt asks for visuals.
859    Auto,
860    /// `off`.
861    Off,
862}
863
864/// The Gemini API's reply fields rig does not normalize, from the reply
865/// document of either route. Each field says which route returns it; on the
866/// other route it reads `None`, as does a field the reply left out. A
867/// candidate field is the first candidate's, the one rig reads.
868#[non_exhaustive]
869#[derive(Clone, Debug, Default, PartialEq)]
870pub struct GeminiExtras {
871    /// GenerateContent `modelVersion`.
872    pub model_version: Option<String>,
873    /// GenerateContent `responseId`.
874    pub response_id: Option<String>,
875    /// The tier that served the request: GenerateContent
876    /// `usageMetadata.serviceTier`, Interactions `service_tier`.
877    pub service_tier: Option<String>,
878    /// GenerateContent `usageMetadata.promptTokensDetails`.
879    pub prompt_tokens_details: Option<Vec<ModalityTokenCount>>,
880    /// GenerateContent `usageMetadata.cacheTokensDetails`.
881    pub cache_tokens_details: Option<Vec<ModalityTokenCount>>,
882    /// GenerateContent `usageMetadata.candidatesTokensDetails`.
883    pub candidates_tokens_details: Option<Vec<ModalityTokenCount>>,
884    /// GenerateContent `usageMetadata.toolUsePromptTokensDetails`.
885    pub tool_use_prompt_tokens_details: Option<Vec<ModalityTokenCount>>,
886    /// GenerateContent `promptFeedback`.
887    pub prompt_feedback: Option<PromptFeedback>,
888    /// GenerateContent `safetyRatings` of the candidate.
889    pub safety_ratings: Option<Vec<SafetyRating>>,
890    /// GenerateContent `finishMessage` of the candidate.
891    pub finish_message: Option<String>,
892    /// GenerateContent `citationMetadata` of the candidate.
893    pub citation_metadata: Option<CitationMetadata>,
894    /// GenerateContent `groundingMetadata` of the candidate.
895    pub grounding_metadata: Option<GroundingMetadata>,
896    /// GenerateContent `urlContextMetadata` of the candidate.
897    pub url_context_metadata: Option<UrlContextMetadata>,
898    /// GenerateContent `avgLogprobs` of the candidate.
899    pub avg_logprobs: Option<f64>,
900    /// GenerateContent `logprobsResult` of the candidate.
901    pub logprobs_result: Option<LogprobsResult>,
902    /// Interactions `id`.
903    pub id: Option<String>,
904    /// Interactions `status`.
905    pub status: Option<InteractionStatus>,
906    /// Interactions `created`, an RFC 3339 timestamp.
907    pub created: Option<String>,
908    /// Interactions `updated`, an RFC 3339 timestamp.
909    pub updated: Option<String>,
910    /// Interactions `usage.input_tokens_by_modality`.
911    pub input_tokens_by_modality: Option<Vec<ModalityTokens>>,
912    /// Interactions `usage.output_tokens_by_modality`.
913    pub output_tokens_by_modality: Option<Vec<ModalityTokens>>,
914    /// Interactions `usage.cached_tokens_by_modality`.
915    pub cached_tokens_by_modality: Option<Vec<ModalityTokens>>,
916    /// Interactions `usage.grounding_tool_count`.
917    pub grounding_tool_count: Option<Vec<GroundingToolCount>>,
918}
919
920impl ReplyExtras for GeminiExtras {
921    fn from_reply(api: &Api, raw: &Value) -> Result<Self, serde_json::Error> {
922        match api.as_str() {
923            GENERATE_CONTENT => Ok(GenerateContentReply::deserialize(raw)?.into()),
924            INTERACTIONS => Ok(InteractionReply::deserialize(raw)?.into()),
925            other => Err(serde_json::Error::custom(format!(
926                "the Gemini API returns no `{other}` reply"
927            ))),
928        }
929    }
930}
931
932/// The parts of a `generateContent` reply [`GeminiExtras`] reads.
933#[derive(Deserialize)]
934#[serde(rename_all = "camelCase")]
935struct GenerateContentReply {
936    model_version: Option<String>,
937    response_id: Option<String>,
938    usage_metadata: Option<UsageDetails>,
939    prompt_feedback: Option<PromptFeedback>,
940    #[serde(default)]
941    candidates: Vec<CandidateDetails>,
942}
943
944#[derive(Deserialize)]
945#[serde(rename_all = "camelCase")]
946struct UsageDetails {
947    service_tier: Option<String>,
948    prompt_tokens_details: Option<Vec<ModalityTokenCount>>,
949    cache_tokens_details: Option<Vec<ModalityTokenCount>>,
950    candidates_tokens_details: Option<Vec<ModalityTokenCount>>,
951    tool_use_prompt_tokens_details: Option<Vec<ModalityTokenCount>>,
952}
953
954#[derive(Deserialize)]
955#[serde(rename_all = "camelCase")]
956struct CandidateDetails {
957    safety_ratings: Option<Vec<SafetyRating>>,
958    finish_message: Option<String>,
959    citation_metadata: Option<CitationMetadata>,
960    grounding_metadata: Option<GroundingMetadata>,
961    url_context_metadata: Option<UrlContextMetadata>,
962    avg_logprobs: Option<f64>,
963    logprobs_result: Option<LogprobsResult>,
964}
965
966impl From<GenerateContentReply> for GeminiExtras {
967    fn from(reply: GenerateContentReply) -> Self {
968        let mut extras = Self {
969            model_version: reply.model_version,
970            response_id: reply.response_id,
971            prompt_feedback: reply.prompt_feedback,
972            ..Self::default()
973        };
974        if let Some(usage) = reply.usage_metadata {
975            extras.service_tier = usage.service_tier;
976            extras.prompt_tokens_details = usage.prompt_tokens_details;
977            extras.cache_tokens_details = usage.cache_tokens_details;
978            extras.candidates_tokens_details = usage.candidates_tokens_details;
979            extras.tool_use_prompt_tokens_details = usage.tool_use_prompt_tokens_details;
980        }
981        if let Some(candidate) = reply.candidates.into_iter().next() {
982            extras.safety_ratings = candidate.safety_ratings;
983            extras.finish_message = candidate.finish_message;
984            extras.citation_metadata = candidate.citation_metadata;
985            extras.grounding_metadata = candidate.grounding_metadata;
986            extras.url_context_metadata = candidate.url_context_metadata;
987            extras.avg_logprobs = candidate.avg_logprobs;
988            extras.logprobs_result = candidate.logprobs_result;
989        }
990        extras
991    }
992}
993
994/// The parts of an interaction resource [`GeminiExtras`] reads.
995#[derive(Deserialize)]
996struct InteractionReply {
997    id: Option<String>,
998    status: Option<InteractionStatus>,
999    service_tier: Option<String>,
1000    created: Option<String>,
1001    updated: Option<String>,
1002    usage: Option<InteractionUsage>,
1003}
1004
1005#[derive(Deserialize)]
1006struct InteractionUsage {
1007    input_tokens_by_modality: Option<Vec<ModalityTokens>>,
1008    output_tokens_by_modality: Option<Vec<ModalityTokens>>,
1009    cached_tokens_by_modality: Option<Vec<ModalityTokens>>,
1010    grounding_tool_count: Option<Vec<GroundingToolCount>>,
1011}
1012
1013impl From<InteractionReply> for GeminiExtras {
1014    fn from(reply: InteractionReply) -> Self {
1015        let mut extras = Self {
1016            id: reply.id,
1017            status: reply.status,
1018            service_tier: reply.service_tier,
1019            created: reply.created,
1020            updated: reply.updated,
1021            ..Self::default()
1022        };
1023        if let Some(usage) = reply.usage {
1024            extras.input_tokens_by_modality = usage.input_tokens_by_modality;
1025            extras.output_tokens_by_modality = usage.output_tokens_by_modality;
1026            extras.cached_tokens_by_modality = usage.cached_tokens_by_modality;
1027            extras.grounding_tool_count = usage.grounding_tool_count;
1028        }
1029        extras
1030    }
1031}
1032
1033/// A token count of one modality, as GenerateContent reports it.
1034#[non_exhaustive]
1035#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
1036#[serde(rename_all = "camelCase")]
1037pub struct ModalityTokenCount {
1038    /// `modality`, such as `TEXT`.
1039    pub modality: Option<String>,
1040    /// `tokenCount`.
1041    pub token_count: Option<u64>,
1042}
1043
1044/// Why the prompt was blocked, and how it rated.
1045#[non_exhaustive]
1046#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
1047#[serde(rename_all = "camelCase")]
1048pub struct PromptFeedback {
1049    /// `blockReason`, such as `SAFETY`, when the prompt was blocked.
1050    pub block_reason: Option<String>,
1051    /// `safetyRatings`.
1052    pub safety_ratings: Option<Vec<SafetyRating>>,
1053}
1054
1055/// How a prompt or candidate rated in one harm category.
1056#[non_exhaustive]
1057#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
1058#[serde(rename_all = "camelCase")]
1059pub struct SafetyRating {
1060    /// `category`, such as `HARM_CATEGORY_HARASSMENT`.
1061    pub category: Option<String>,
1062    /// `probability`, such as `NEGLIGIBLE`.
1063    pub probability: Option<String>,
1064    /// `blocked`: whether this rating blocked the content.
1065    pub blocked: Option<bool>,
1066}
1067
1068/// The sources a candidate recites.
1069#[non_exhaustive]
1070#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
1071#[serde(rename_all = "camelCase")]
1072pub struct CitationMetadata {
1073    /// `citationSources`.
1074    pub citation_sources: Option<Vec<CitationSource>>,
1075}
1076
1077/// One recited source.
1078#[non_exhaustive]
1079#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
1080#[serde(rename_all = "camelCase")]
1081pub struct CitationSource {
1082    /// `startIndex` of the reciting text.
1083    pub start_index: Option<u64>,
1084    /// `endIndex` of the reciting text.
1085    pub end_index: Option<u64>,
1086    /// `uri` of the source.
1087    pub uri: Option<String>,
1088    /// `license` of the source.
1089    pub license: Option<String>,
1090}
1091
1092/// What grounded a candidate in search or retrieval.
1093#[non_exhaustive]
1094#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
1095#[serde(rename_all = "camelCase")]
1096pub struct GroundingMetadata {
1097    /// `webSearchQueries`.
1098    pub web_search_queries: Option<Vec<String>>,
1099    /// `searchEntryPoint`.
1100    pub search_entry_point: Option<SearchEntryPoint>,
1101    /// `groundingChunks`.
1102    pub grounding_chunks: Option<Vec<GroundingChunk>>,
1103    /// `groundingSupports`.
1104    pub grounding_supports: Option<Vec<GroundingSupport>>,
1105}
1106
1107/// The search suggestion to show with a grounded answer.
1108#[non_exhaustive]
1109#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
1110#[serde(rename_all = "camelCase")]
1111pub struct SearchEntryPoint {
1112    /// `renderedContent`: HTML and CSS to embed.
1113    pub rendered_content: Option<String>,
1114}
1115
1116/// One grounding source.
1117#[non_exhaustive]
1118#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
1119#[serde(rename_all = "camelCase")]
1120pub struct GroundingChunk {
1121    /// `web`: a web page.
1122    pub web: Option<WebChunk>,
1123    /// `retrievedContext`: a retrieved document.
1124    pub retrieved_context: Option<RetrievedContext>,
1125}
1126
1127/// A web page a candidate is grounded in.
1128#[non_exhaustive]
1129#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
1130pub struct WebChunk {
1131    /// `uri`.
1132    pub uri: Option<String>,
1133    /// `title`.
1134    pub title: Option<String>,
1135}
1136
1137/// A retrieved document a candidate is grounded in.
1138#[non_exhaustive]
1139#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
1140pub struct RetrievedContext {
1141    /// `uri`.
1142    pub uri: Option<String>,
1143    /// `title`.
1144    pub title: Option<String>,
1145    /// `text`.
1146    pub text: Option<String>,
1147}
1148
1149/// A span of the candidate and the chunks that support it.
1150#[non_exhaustive]
1151#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
1152#[serde(rename_all = "camelCase")]
1153pub struct GroundingSupport {
1154    /// `segment`.
1155    pub segment: Option<Segment>,
1156    /// `groundingChunkIndices`, into `groundingChunks`.
1157    pub grounding_chunk_indices: Option<Vec<u32>>,
1158    /// `confidenceScores`, one per index.
1159    pub confidence_scores: Option<Vec<f64>>,
1160}
1161
1162/// A span of a candidate's content.
1163#[non_exhaustive]
1164#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
1165#[serde(rename_all = "camelCase")]
1166pub struct Segment {
1167    /// `partIndex`.
1168    pub part_index: Option<u32>,
1169    /// `startIndex`, in bytes.
1170    pub start_index: Option<u64>,
1171    /// `endIndex`, in bytes.
1172    pub end_index: Option<u64>,
1173    /// `text`.
1174    pub text: Option<String>,
1175}
1176
1177/// The URLs the URL context tool retrieved.
1178#[non_exhaustive]
1179#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
1180#[serde(rename_all = "camelCase")]
1181pub struct UrlContextMetadata {
1182    /// `urlMetadata`.
1183    pub url_metadata: Option<Vec<UrlMetadata>>,
1184}
1185
1186/// One URL the URL context tool retrieved.
1187#[non_exhaustive]
1188#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
1189#[serde(rename_all = "camelCase")]
1190pub struct UrlMetadata {
1191    /// `retrievedUrl`.
1192    pub retrieved_url: Option<String>,
1193    /// `urlRetrievalStatus`, such as `URL_RETRIEVAL_STATUS_SUCCESS`.
1194    pub url_retrieval_status: Option<String>,
1195}
1196
1197/// The log probabilities of a candidate's tokens.
1198#[non_exhaustive]
1199#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
1200#[serde(rename_all = "camelCase")]
1201pub struct LogprobsResult {
1202    /// `topCandidates`, one entry per decoding step.
1203    pub top_candidates: Option<Vec<TopCandidates>>,
1204    /// `chosenCandidates`, one per decoding step.
1205    pub chosen_candidates: Option<Vec<LogprobsCandidate>>,
1206    /// `logProbabilitySum`.
1207    pub log_probability_sum: Option<f64>,
1208}
1209
1210/// The most likely tokens of one decoding step.
1211#[non_exhaustive]
1212#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
1213pub struct TopCandidates {
1214    /// `candidates`, most likely first.
1215    pub candidates: Option<Vec<LogprobsCandidate>>,
1216}
1217
1218/// A token and its log probability.
1219#[non_exhaustive]
1220#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
1221#[serde(rename_all = "camelCase")]
1222pub struct LogprobsCandidate {
1223    /// `token`.
1224    pub token: Option<String>,
1225    /// `tokenId`.
1226    pub token_id: Option<i64>,
1227    /// `logProbability`.
1228    pub log_probability: Option<f64>,
1229}
1230
1231/// Where an interaction stands.
1232#[non_exhaustive]
1233#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
1234#[serde(rename_all = "snake_case")]
1235pub enum InteractionStatus {
1236    /// `in_progress`.
1237    InProgress,
1238    /// `requires_action`.
1239    RequiresAction,
1240    /// `incomplete`.
1241    Incomplete,
1242    /// `budget_exceeded`.
1243    BudgetExceeded,
1244    /// `completed`.
1245    Completed,
1246    /// `failed`.
1247    Failed,
1248    /// `cancelled`.
1249    Cancelled,
1250    /// `queued`: waiting for capacity, such as the flex tier's.
1251    Queued,
1252    /// A status this crate does not know yet, as the API spelled it.
1253    #[serde(untagged)]
1254    Unknown(String),
1255}
1256
1257/// A token count of one modality, as Interactions reports it.
1258#[non_exhaustive]
1259#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
1260pub struct ModalityTokens {
1261    /// `modality`, such as `text`.
1262    pub modality: Option<String>,
1263    /// `tokens`.
1264    pub tokens: Option<u64>,
1265}
1266
1267/// How often a grounding tool ran.
1268#[non_exhaustive]
1269#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
1270pub struct GroundingToolCount {
1271    /// `type`, such as `google_search`.
1272    #[serde(rename = "type")]
1273    pub kind: Option<String>,
1274    /// `count`.
1275    pub count: Option<u64>,
1276    /// `search_query_count`.
1277    pub search_query_count: Option<u64>,
1278}
1279
1280#[cfg(test)]
1281mod tests;
rig_core/providers/gemini/extension.rs

rig_core/providers/gemini/
extension.rs
