MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-03-10 22:21:22 -07:00
Code Issues Packages Projects Releases Wiki Activity

chore: clean up unused models and use consts (#16246)

Browse Source
This commit is contained in:
Sehoon Shon
2026-01-09 10:03:46 -05:00
committed by GitHub
parent 41e627a7ee
commit aa480e5fbb
2 changed files with 27 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
packages/core/src/core/tokenLimits.test.ts
View File

@@ -6,14 +6,24 @@
import { describe, it, expect } from 'vitest';
import { tokenLimit, DEFAULT_TOKEN_LIMIT } from './tokenLimits.js';
import {
DEFAULT_GEMINI_FLASH_LITE_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_MODEL,
PREVIEW_GEMINI_FLASH_MODEL,
PREVIEW_GEMINI_MODEL,
} from '../config/models.js';
describe('tokenLimit', () => {
it('should return the correct token limit for gemini-1.5-pro', () => {
expect(tokenLimit('gemini-1.5-pro')).toBe(2_097_152);
it('should return the correct token limit for default models', () => {
expect(tokenLimit(DEFAULT_GEMINI_MODEL)).toBe(1_048_576);
expect(tokenLimit(DEFAULT_GEMINI_FLASH_MODEL)).toBe(1_048_576);
expect(tokenLimit(DEFAULT_GEMINI_FLASH_LITE_MODEL)).toBe(1_048_576);
});
it('should return the correct token limit for gemini-1.5-flash', () => {
expect(tokenLimit('gemini-1.5-flash')).toBe(1_048_576);
it('should return the correct token limit for preview models', () => {
expect(tokenLimit(PREVIEW_GEMINI_MODEL)).toBe(1_048_576);
expect(tokenLimit(PREVIEW_GEMINI_FLASH_MODEL)).toBe(1_048_576);
});
it('should return the default token limit for an unknown model', () => {

25
packages/core/src/core/tokenLimits.ts
View File

@@ -4,6 +4,14 @@
* SPDX-License-Identifier: Apache-2.0
*/
import {
DEFAULT_GEMINI_FLASH_LITE_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_MODEL,
PREVIEW_GEMINI_FLASH_MODEL,
PREVIEW_GEMINI_MODEL,
} from '../config/models.js';
type Model = string;
type TokenCount = number;
@@ -13,19 +21,12 @@ export function tokenLimit(model: Model): TokenCount {
// Add other models as they become relevant or if specified by config
// Pulled from https://ai.google.dev/gemini-api/docs/models
switch (model) {
case 'gemini-1.5-pro':
return 2_097_152;
case 'gemini-1.5-flash':
case 'gemini-2.5-pro-preview-05-06':
case 'gemini-2.5-pro-preview-06-05':
case 'gemini-2.5-pro':
case 'gemini-2.5-flash-preview-05-20':
case 'gemini-2.5-flash':
case 'gemini-2.5-flash-lite':
case 'gemini-2.0-flash':
case PREVIEW_GEMINI_MODEL:
case PREVIEW_GEMINI_FLASH_MODEL:
case DEFAULT_GEMINI_MODEL:
case DEFAULT_GEMINI_FLASH_MODEL:
case DEFAULT_GEMINI_FLASH_LITE_MODEL:
return 1_048_576;
case 'gemini-2.0-flash-preview-image-generation':
return 32_000;
default:
return DEFAULT_TOKEN_LIMIT;
}