Skip to content

Commit 2beefc8

Browse files
authored
Fix gemini model token limits (#29584)
Release Notes: - N/A
1 parent 5092f0f commit 2beefc8

File tree

1 file changed

+11
-9
lines changed

1 file changed

+11
-9
lines changed

crates/google_ai/src/google_ai.rs

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -451,16 +451,18 @@ impl Model {
451451
}
452452

453453
pub fn max_token_count(&self) -> usize {
454+
const ONE_MILLION: usize = 1_048_576;
455+
const TWO_MILLION: usize = 2_097_152;
454456
match self {
455-
Model::Gemini15Pro => 2_000_000,
456-
Model::Gemini15Flash => 1_000_000,
457-
Model::Gemini20Pro => 2_000_000,
458-
Model::Gemini20Flash => 1_000_000,
459-
Model::Gemini20FlashThinking => 1_000_000,
460-
Model::Gemini20FlashLite => 1_000_000,
461-
Model::Gemini25ProExp0325 => 1_000_000,
462-
Model::Gemini25ProPreview0325 => 1_000_000,
463-
Model::Gemini25FlashPreview0417 => 1_000_000,
457+
Model::Gemini15Pro => TWO_MILLION,
458+
Model::Gemini15Flash => ONE_MILLION,
459+
Model::Gemini20Pro => TWO_MILLION,
460+
Model::Gemini20Flash => ONE_MILLION,
461+
Model::Gemini20FlashThinking => ONE_MILLION,
462+
Model::Gemini20FlashLite => ONE_MILLION,
463+
Model::Gemini25ProExp0325 => ONE_MILLION,
464+
Model::Gemini25ProPreview0325 => ONE_MILLION,
465+
Model::Gemini25FlashPreview0417 => ONE_MILLION,
464466
Model::Custom { max_tokens, .. } => *max_tokens,
465467
}
466468
}

0 commit comments

Comments
 (0)