gemini_adk_rs/utils/
model_name.rs

1use once_cell::sync::Lazy;
2use regex::Regex;
3
4static MODEL_PATH_RE: Lazy<Regex> = Lazy::new(|| {
5    Regex::new(r"^projects/[^/]+/locations/[^/]+/publishers/[^/]+/models/(.+)$").unwrap()
6});
7
8/// Regex to extract the major version digit from a gemini model name.
9/// Matches "gemini-" followed by one or more digits (the major version).
10static GEMINI_VERSION_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"^gemini-(\d+)").unwrap());
11
12/// Extract simple model name from a fully-qualified resource path.
13///
14/// If `model_string` matches the pattern
15/// `projects/{project}/locations/{location}/publishers/{publisher}/models/{model}`,
16/// returns just the `{model}` portion. Otherwise returns `model_string` as-is.
17///
18/// # Examples
19///
20/// ```
21/// use gemini_adk_rs::utils::model_name::extract_model_name;
22///
23/// assert_eq!(
24///     extract_model_name("projects/my-proj/locations/us-central1/publishers/google/models/gemini-2.5-flash"),
25///     "gemini-2.5-flash"
26/// );
27/// assert_eq!(extract_model_name("gemini-2.5-flash"), "gemini-2.5-flash");
28/// ```
29pub fn extract_model_name(model_string: &str) -> &str {
30    MODEL_PATH_RE
31        .captures(model_string)
32        .and_then(|caps| caps.get(1))
33        .map(|m| m.as_str())
34        .unwrap_or(model_string)
35}
36
37/// Returns `true` if the model name starts with `"gemini-"`.
38///
39/// The model name is first extracted from any fully-qualified resource path.
40pub fn is_gemini_model(model_string: &str) -> bool {
41    extract_model_name(model_string).starts_with("gemini-")
42}
43
44/// Returns `true` if the model is a Gemini 1.x model (name starts with `"gemini-1"`).
45pub fn is_gemini1_model(model_string: &str) -> bool {
46    extract_model_name(model_string).starts_with("gemini-1")
47}
48
49/// Returns `true` if the model is Gemini 2.x or above.
50///
51/// Parses the major version number from the model name. If the name does not
52/// match the expected `gemini-{major}...` pattern, returns `false`.
53pub fn is_gemini2_or_above(model_string: &str) -> bool {
54    let name = extract_model_name(model_string);
55    GEMINI_VERSION_RE
56        .captures(name)
57        .and_then(|caps| caps.get(1))
58        .and_then(|m| m.as_str().parse::<u32>().ok())
59        .map(|major| major >= 2)
60        .unwrap_or(false)
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    // ── extract_model_name ──────────────────────────────────────────
68
69    #[test]
70    fn extract_from_full_path() {
71        assert_eq!(
72            extract_model_name(
73                "projects/my-proj/locations/us-central1/publishers/google/models/gemini-2.5-flash"
74            ),
75            "gemini-2.5-flash"
76        );
77    }
78
79    #[test]
80    fn extract_from_simple_name() {
81        assert_eq!(extract_model_name("gemini-2.5-flash"), "gemini-2.5-flash");
82    }
83
84    #[test]
85    fn extract_preserves_suffixes() {
86        assert_eq!(
87            extract_model_name("projects/p/locations/l/publishers/pub/models/gemini-1.5-pro-002"),
88            "gemini-1.5-pro-002"
89        );
90    }
91
92    #[test]
93    fn extract_non_gemini_full_path() {
94        assert_eq!(
95            extract_model_name("projects/p/locations/l/publishers/pub/models/custom-model"),
96            "custom-model"
97        );
98    }
99
100    #[test]
101    fn extract_empty_string() {
102        assert_eq!(extract_model_name(""), "");
103    }
104
105    // ── is_gemini_model ─────────────────────────────────────────────
106
107    #[test]
108    fn gemini_model_simple() {
109        assert!(is_gemini_model("gemini-2.5-flash"));
110        assert!(is_gemini_model("gemini-1.5-pro"));
111    }
112
113    #[test]
114    fn non_gemini_model() {
115        assert!(!is_gemini_model("claude-3-opus"));
116        assert!(!is_gemini_model("gpt-4"));
117        assert!(!is_gemini_model("custom-model"));
118    }
119
120    #[test]
121    fn gemini_model_full_path() {
122        assert!(is_gemini_model(
123            "projects/p/locations/l/publishers/pub/models/gemini-2.5-flash"
124        ));
125    }
126
127    // ── is_gemini1_model ────────────────────────────────────────────
128
129    #[test]
130    fn gemini1_models() {
131        assert!(is_gemini1_model("gemini-1.5-pro"));
132        assert!(is_gemini1_model("gemini-1.5-flash"));
133        assert!(is_gemini1_model("gemini-1.0-pro"));
134    }
135
136    #[test]
137    fn gemini1_full_path() {
138        assert!(is_gemini1_model(
139            "projects/p/locations/l/publishers/pub/models/gemini-1.5-pro-002"
140        ));
141    }
142
143    #[test]
144    fn gemini2_not_gemini1() {
145        assert!(!is_gemini1_model("gemini-2.5-flash"));
146        assert!(!is_gemini1_model("gemini-2.0-flash"));
147    }
148
149    #[test]
150    fn non_gemini_not_gemini1() {
151        assert!(!is_gemini1_model("gpt-4"));
152    }
153
154    // ── is_gemini2_or_above ─────────────────────────────────────────
155
156    #[test]
157    fn gemini2_or_above_positive() {
158        assert!(is_gemini2_or_above("gemini-2.5-flash"));
159        assert!(is_gemini2_or_above("gemini-2.0-flash"));
160        assert!(is_gemini2_or_above("gemini-3.0-ultra"));
161    }
162
163    #[test]
164    fn gemini2_or_above_negative() {
165        assert!(!is_gemini2_or_above("gemini-1.5-pro"));
166        assert!(!is_gemini2_or_above("gemini-1.0-pro"));
167    }
168
169    #[test]
170    fn gemini2_or_above_full_path() {
171        assert!(is_gemini2_or_above(
172            "projects/p/locations/l/publishers/pub/models/gemini-2.5-flash"
173        ));
174    }
175
176    #[test]
177    fn gemini2_or_above_non_gemini() {
178        assert!(!is_gemini2_or_above("custom-model"));
179        assert!(!is_gemini2_or_above("gpt-4"));
180    }
181
182    #[test]
183    fn gemini2_or_above_edge_cases() {
184        assert!(!is_gemini2_or_above(""));
185        assert!(!is_gemini2_or_above("gemini-"));
186    }
187}