mixa/zed
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

google_ai: Refactor token count methods in Google AI (#45184)

Browse Source 
The change simplifies the `max_token_count` and `max_output_tokens`
methods by grouping Gemini models with identical token limits.

Release Notes:

- N/A
This commit is contained in:
Xiaobo Liu
2025-12-18 10:12:40 +08:00
committed by GitHub
parent ea37057814
commit 225a2a8a20
1 changed files with 10 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
crates/google_ai/src/google_ai.rs
View File

@@ -566,22 +566,22 @@ impl Model {
pub fn max_token_count(&self) -> u64 {
match self {
Self::Gemini25FlashLite => 1_048_576,
Self::Gemini25Flash => 1_048_576,
Self::Gemini25Pro => 1_048_576,
Self::Gemini3Pro => 1_048_576,
Self::Gemini3Flash => 1_048_576,
Self::Gemini25FlashLite
| Self::Gemini25Flash
| Self::Gemini25Pro
| Self::Gemini3Pro
| Self::Gemini3Flash => 1_048_576,
Self::Custom { max_tokens, .. } => *max_tokens,
}
}
pub fn max_output_tokens(&self) -> Option<u64> {
match self {
Model::Gemini25FlashLite => Some(65_536),
Model::Gemini25Flash => Some(65_536),
Model::Gemini25Pro => Some(65_536),
Model::Gemini3Pro => Some(65_536),
Model::Gemini3Flash => Some(65_536),
Model::Gemini25FlashLite
| Model::Gemini25Flash
| Model::Gemini25Pro
| Model::Gemini3Pro
| Model::Gemini3Flash => Some(65_536),
Model::Custom { .. } => None,
}
}