From 8acd8a273de04430e3160a1c037fc8ed12152c21 Mon Sep 17 00:00:00 2001 From: Michael Sloan Date: Mon, 28 Apr 2025 21:12:59 -0600 Subject: [PATCH] Fix gemini model token limits (#29584) Release Notes: - N/A --- crates/google_ai/src/google_ai.rs | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/crates/google_ai/src/google_ai.rs b/crates/google_ai/src/google_ai.rs index 66a3e6e9c68719b17f5293a9f48ba9d44ada8e8e..dd86a7388dde597ad2e711b17aab91829f35027c 100644 --- a/crates/google_ai/src/google_ai.rs +++ b/crates/google_ai/src/google_ai.rs @@ -455,16 +455,18 @@ impl Model { } pub fn max_token_count(&self) -> usize { + const ONE_MILLION: usize = 1_048_576; + const TWO_MILLION: usize = 2_097_152; 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::Gemini25ProExp0325 => 1_000_000, - Model::Gemini25ProPreview0325 => 1_000_000, - Model::Gemini25FlashPreview0417 => 1_000_000, + Model::Gemini15Pro => TWO_MILLION, + Model::Gemini15Flash => ONE_MILLION, + Model::Gemini20Pro => TWO_MILLION, + Model::Gemini20Flash => ONE_MILLION, + Model::Gemini20FlashThinking => ONE_MILLION, + Model::Gemini20FlashLite => ONE_MILLION, + Model::Gemini25ProExp0325 => ONE_MILLION, + Model::Gemini25ProPreview0325 => ONE_MILLION, + Model::Gemini25FlashPreview0417 => ONE_MILLION, Model::Custom { max_tokens, .. } => *max_tokens, } }