google_ai: Add support for Gemini 2.0 models (#24448)

IaVashik created 11 months ago

Add support for the newly released Gemini 2.0 models from Google announced this new family of models earlier this week (2025-02-05).

Release Notes:

- Added support for Google's new Gemini 2.0 models.

Change summary

crates/google_ai/src/google_ai.rs              | 19 +++++++++++++++++--
crates/language_model/src/model/cloud_model.rs |  3 +++
2 files changed, 20 insertions(+), 2 deletions(-)

Detailed changes

crates/google_ai/src/google_ai.rs

@@ -305,8 +305,14 @@ pub enum Model {
     Gemini15Pro,
     #[serde(rename = "gemini-1.5-flash")]
     Gemini15Flash,
-    #[serde(rename = "gemini-2.0-flash-exp")]
+    #[serde(rename = "gemini-2.0-pro-exp")]
+    Gemini20Pro,
+    #[serde(rename = "gemini-2.0-flash")]
     Gemini20Flash,
+    #[serde(rename = "gemini-2.0-flash-thinking-exp")]
+    Gemini20FlashThinking,
+    #[serde(rename = "gemini-2.0-flash-lite-preview")]
+    Gemini20FlashLite,
     #[serde(rename = "custom")]
     Custom {
         name: String,
@@ -321,7 +327,10 @@ impl Model {
         match self {
             Model::Gemini15Pro => "gemini-1.5-pro",
             Model::Gemini15Flash => "gemini-1.5-flash",
-            Model::Gemini20Flash => "gemini-2.0-flash-exp",
+            Model::Gemini20Pro => "gemini-2.0-pro-exp",
+            Model::Gemini20Flash => "gemini-2.0-flash",
+            Model::Gemini20FlashThinking => "gemini-2.0-flash-thinking-exp",
+            Model::Gemini20FlashLite => "gemini-2.0-flash-lite-preview",
             Model::Custom { name, .. } => name,
         }
     }
@@ -330,7 +339,10 @@ impl Model {
         match self {
             Model::Gemini15Pro => "Gemini 1.5 Pro",
             Model::Gemini15Flash => "Gemini 1.5 Flash",
+            Model::Gemini20Pro => "Gemini 2.0 Pro",
             Model::Gemini20Flash => "Gemini 2.0 Flash",
+            Model::Gemini20FlashThinking => "Gemini 2.0 Flash Thinking",
+            Model::Gemini20FlashLite => "Gemini 2.0 Flash Lite",
             Self::Custom {
                 name, display_name, ..
             } => display_name.as_ref().unwrap_or(name),
@@ -341,7 +353,10 @@ impl Model {
         match self {
             Model::Gemini15Pro => 2_000_000,
             Model::Gemini15Flash => 1_000_000,
+            Model::Gemini20Pro => 2_000_000,
             Model::Gemini20Flash => 1_000_000,
+            Model::Gemini20FlashThinking => 1_000_000,
+            Model::Gemini20FlashLite => 1_000_000,
             Model::Custom { max_tokens, .. } => *max_tokens,
         }
     }

crates/language_model/src/model/cloud_model.rs

@@ -90,7 +90,10 @@ impl CloudModel {
             Self::Google(model) => match model {
                 google_ai::Model::Gemini15Pro
                 | google_ai::Model::Gemini15Flash
+                | google_ai::Model::Gemini20Pro
                 | google_ai::Model::Gemini20Flash
+                | google_ai::Model::Gemini20FlashThinking
+                | google_ai::Model::Gemini20FlashLite
                 | google_ai::Model::Custom { .. } => {
                     LanguageModelAvailability::RequiresPlan(Plan::ZedPro)
                 }