gemini_adk_rs/utils/
model_name.rs1use 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
8static GEMINI_VERSION_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"^gemini-(\d+)").unwrap());
11
12pub 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
37pub fn is_gemini_model(model_string: &str) -> bool {
41 extract_model_name(model_string).starts_with("gemini-")
42}
43
44pub fn is_gemini1_model(model_string: &str) -> bool {
46 extract_model_name(model_string).starts_with("gemini-1")
47}
48
49pub 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 #[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 #[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 #[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 #[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}