google_gemini_rs/google/
mod.rs

1//! Wrapper types for supported Google AI Models
2
3use std::fmt::Display;
4
5use thiserror::Error;
6
7use crate::google::common::Modality;
8
9pub mod common;
10pub mod request;
11pub mod response;
12
13#[derive(Error, Debug)]
14pub enum Error {
15    #[error("{0}")]
16    NotFound(String),
17}
18
19const GEMINI_2_0_FLASH_EXP_IMAGE_GEN: &str = "gemini-2.0-flash-exp-image-generation";
20const GEMINI_2_0_FLASH: &str = "gemini-2.0-flash";
21const GEMINI_2_5_FLASH: &str = "gemini-2.5-flash";
22const GEMINI_2_5_FLASH_LITE: &str = "gemini-2.5-flash-lite";
23const GEMINI_2_5_PRO: &str = "gemini-2.5-pro";
24
25/// Supported Google AI models.  Some models have different capabilities than others, so this
26/// enum may be used to branch the different capabilities.
27#[derive(Debug, Clone, Hash, Eq, PartialEq)]
28pub enum GoogleModelVariant {
29    Gemini20FlashExpImageGen,
30    Gemini20Flash,
31    Gemini25Flash,
32    Gemini25Pro,
33    Gemini25FlashLight,
34}
35
36impl GoogleModelVariant {
37    fn name(&self) -> String {
38        match self {
39            GoogleModelVariant::Gemini20FlashExpImageGen => GEMINI_2_0_FLASH_EXP_IMAGE_GEN,
40            GoogleModelVariant::Gemini20Flash => GEMINI_2_0_FLASH,
41            GoogleModelVariant::Gemini25Flash => GEMINI_2_5_FLASH,
42            GoogleModelVariant::Gemini25Pro => GEMINI_2_5_PRO,
43            GoogleModelVariant::Gemini25FlashLight => GEMINI_2_5_FLASH_LITE,
44        }
45        .to_string()
46    }
47
48    fn inputs(&self) -> Vec<Modality> {
49        match self {
50            GoogleModelVariant::Gemini20FlashExpImageGen => vec![
51                Modality::Text,
52                Modality::Video,
53                Modality::Image,
54                Modality::Audio,
55            ],
56            GoogleModelVariant::Gemini20Flash => vec![
57                Modality::Text,
58                Modality::Video,
59                Modality::Image,
60                Modality::Audio,
61            ],
62            GoogleModelVariant::Gemini25Flash => vec![
63                Modality::Text,
64                Modality::Video,
65                Modality::Image,
66                Modality::Audio,
67            ],
68            GoogleModelVariant::Gemini25Pro => vec![
69                Modality::Text,
70                Modality::Video,
71                Modality::Image,
72                Modality::Audio,
73            ],
74            GoogleModelVariant::Gemini25FlashLight => vec![
75                Modality::Text,
76                Modality::Video,
77                Modality::Image,
78                Modality::Audio,
79            ],
80        }
81    }
82
83    fn outputs(&self) -> Vec<Modality> {
84        match self {
85            GoogleModelVariant::Gemini20FlashExpImageGen => {
86                vec![Modality::Text, Modality::Image]
87            }
88            GoogleModelVariant::Gemini20Flash => vec![Modality::Text],
89            GoogleModelVariant::Gemini25Flash => vec![Modality::Text],
90            GoogleModelVariant::Gemini25Pro => vec![Modality::Text],
91            GoogleModelVariant::Gemini25FlashLight => vec![Modality::Text],
92        }
93    }
94}
95
96#[derive(Debug, Clone)]
97pub struct GoogleModel {
98    pub variant: GoogleModelVariant,
99    pub name: String,
100    pub input: Vec<Modality>,
101    pub output: Vec<Modality>,
102}
103
104impl GoogleModel {
105    pub fn new(variant: GoogleModelVariant, suffix: Option<String>) -> Self {
106        let name = if let Some(suffix) = suffix {
107            format!("{}-{suffix}", variant.name())
108        } else {
109            variant.name()
110        };
111
112        let input = variant.inputs();
113        let output = variant.outputs();
114
115        Self {
116            variant,
117            name,
118            input,
119            output,
120        }
121    }
122}
123
124impl TryFrom<&str> for GoogleModel {
125    type Error = Error;
126
127    fn try_from(value: &str) -> Result<Self, Error> {
128        let (model, preview) = if let Some((model, preview)) = value.split_once("-preview") {
129            (model, Some(format!("preview{preview}")))
130        } else {
131            (value, None)
132        };
133
134        println!("Model: {model} preview: {preview:?}");
135
136        let variant = match model {
137            GEMINI_2_5_PRO => Ok(GoogleModelVariant::Gemini25Pro),
138            GEMINI_2_5_FLASH => Ok(GoogleModelVariant::Gemini25Flash),
139            GEMINI_2_5_FLASH_LITE => Ok(GoogleModelVariant::Gemini25FlashLight),
140            GEMINI_2_0_FLASH => Ok(GoogleModelVariant::Gemini20Flash),
141            GEMINI_2_0_FLASH_EXP_IMAGE_GEN => Ok(GoogleModelVariant::Gemini20FlashExpImageGen),
142            _ => Err(Error::NotFound(format!("No such model: {value}"))),
143        }?;
144
145        Ok(GoogleModel::new(variant, preview))
146    }
147}
148
149impl Display for GoogleModel {
150    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
151        write!(f, "{}", self.name)
152    }
153}
google_gemini_rs/google/mod.rs

google_gemini_rs/google/
mod.rs
