From aa480e5fbbbe4acd99b9f3f66de67511fa7aab6c Mon Sep 17 00:00:00 2001 From: Sehoon Shon Date: Fri, 9 Jan 2026 10:03:46 -0500 Subject: [PATCH] chore: clean up unused models and use consts (#16246) --- packages/core/src/core/tokenLimits.test.ts | 18 ++++++++++++---- packages/core/src/core/tokenLimits.ts | 25 +++++++++++----------- 2 files changed, 27 insertions(+), 16 deletions(-) diff --git a/packages/core/src/core/tokenLimits.test.ts b/packages/core/src/core/tokenLimits.test.ts index 1bff09d315..5a5092d7ea 100644 --- a/packages/core/src/core/tokenLimits.test.ts +++ b/packages/core/src/core/tokenLimits.test.ts @@ -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', () => { diff --git a/packages/core/src/core/tokenLimits.ts b/packages/core/src/core/tokenLimits.ts index d238cdb3a0..39a3443e36 100644 --- a/packages/core/src/core/tokenLimits.ts +++ b/packages/core/src/core/tokenLimits.ts @@ -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; }