mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-23 11:34:44 -07:00
feat(models): support Gemini 3.1 Pro Preview and fixes (#19676)
This commit is contained in:
@@ -44,7 +44,10 @@ export function resolvePolicyChain(
|
||||
const configuredModel = config.getModel();
|
||||
|
||||
let chain;
|
||||
const resolvedModel = resolveModel(modelFromConfig);
|
||||
const resolvedModel = resolveModel(
|
||||
modelFromConfig,
|
||||
config.getGemini31LaunchedSync?.() ?? false,
|
||||
);
|
||||
const isAutoPreferred = preferredModel ? isAutoModel(preferredModel) : false;
|
||||
const isAutoConfigured = isAutoModel(configuredModel);
|
||||
const hasAccessToPreview = config.getHasAccessToPreviewModel?.() ?? true;
|
||||
|
||||
@@ -16,6 +16,7 @@ export const ExperimentFlags = {
|
||||
MASKING_PROTECTION_THRESHOLD: 45758817,
|
||||
MASKING_PRUNABLE_THRESHOLD: 45758818,
|
||||
MASKING_PROTECT_LATEST_TURN: 45758819,
|
||||
GEMINI_3_1_PRO_LAUNCHED: 45760185,
|
||||
} as const;
|
||||
|
||||
export type ExperimentFlagName =
|
||||
|
||||
@@ -1071,6 +1071,12 @@ export class Config {
|
||||
// Reset availability status when switching auth (e.g. from limited key to OAuth)
|
||||
this.modelAvailabilityService.reset();
|
||||
|
||||
// Clear stale authType to ensure getGemini31LaunchedSync doesn't return stale results
|
||||
// during the transition.
|
||||
if (this.contentGeneratorConfig) {
|
||||
this.contentGeneratorConfig.authType = undefined;
|
||||
}
|
||||
|
||||
const newContentGeneratorConfig = await createContentGeneratorConfig(
|
||||
this,
|
||||
authMethod,
|
||||
@@ -1350,7 +1356,10 @@ export class Config {
|
||||
if (pooled.remaining !== undefined) {
|
||||
return pooled.remaining;
|
||||
}
|
||||
const primaryModel = resolveModel(this.getModel());
|
||||
const primaryModel = resolveModel(
|
||||
this.getModel(),
|
||||
this.getGemini31LaunchedSync(),
|
||||
);
|
||||
return this.modelQuotas.get(primaryModel)?.remaining;
|
||||
}
|
||||
|
||||
@@ -1359,7 +1368,10 @@ export class Config {
|
||||
if (pooled.limit !== undefined) {
|
||||
return pooled.limit;
|
||||
}
|
||||
const primaryModel = resolveModel(this.getModel());
|
||||
const primaryModel = resolveModel(
|
||||
this.getModel(),
|
||||
this.getGemini31LaunchedSync(),
|
||||
);
|
||||
return this.modelQuotas.get(primaryModel)?.limit;
|
||||
}
|
||||
|
||||
@@ -1368,7 +1380,10 @@ export class Config {
|
||||
if (pooled.resetTime !== undefined) {
|
||||
return pooled.resetTime;
|
||||
}
|
||||
const primaryModel = resolveModel(this.getModel());
|
||||
const primaryModel = resolveModel(
|
||||
this.getModel(),
|
||||
this.getGemini31LaunchedSync(),
|
||||
);
|
||||
return this.modelQuotas.get(primaryModel)?.resetTime;
|
||||
}
|
||||
|
||||
@@ -2253,6 +2268,36 @@ export class Config {
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns whether Gemini 3.1 has been launched.
|
||||
* This method is async and ensures that experiments are loaded before returning the result.
|
||||
*/
|
||||
async getGemini31Launched(): Promise<boolean> {
|
||||
await this.ensureExperimentsLoaded();
|
||||
return this.getGemini31LaunchedSync();
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns whether Gemini 3.1 has been launched.
|
||||
*
|
||||
* Note: This method should only be called after startup, once experiments have been loaded.
|
||||
* If you need to call this during startup or from an async context, use
|
||||
* getGemini31Launched instead.
|
||||
*/
|
||||
getGemini31LaunchedSync(): boolean {
|
||||
const authType = this.contentGeneratorConfig?.authType;
|
||||
if (
|
||||
authType === AuthType.USE_GEMINI ||
|
||||
authType === AuthType.USE_VERTEX_AI
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
return (
|
||||
this.experiments?.flags[ExperimentFlags.GEMINI_3_1_PRO_LAUNCHED]
|
||||
?.boolValue ?? false
|
||||
);
|
||||
}
|
||||
|
||||
private async ensureExperimentsLoaded(): Promise<void> {
|
||||
if (!this.experimentsPromise) {
|
||||
return;
|
||||
|
||||
@@ -25,8 +25,41 @@ import {
|
||||
PREVIEW_GEMINI_FLASH_MODEL,
|
||||
PREVIEW_GEMINI_MODEL_AUTO,
|
||||
DEFAULT_GEMINI_MODEL_AUTO,
|
||||
isActiveModel,
|
||||
PREVIEW_GEMINI_3_1_MODEL,
|
||||
PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
|
||||
isPreviewModel,
|
||||
isProModel,
|
||||
} from './models.js';
|
||||
|
||||
describe('isPreviewModel', () => {
|
||||
it('should return true for preview models', () => {
|
||||
expect(isPreviewModel(PREVIEW_GEMINI_MODEL)).toBe(true);
|
||||
expect(isPreviewModel(PREVIEW_GEMINI_3_1_MODEL)).toBe(true);
|
||||
expect(isPreviewModel(PREVIEW_GEMINI_FLASH_MODEL)).toBe(true);
|
||||
expect(isPreviewModel(PREVIEW_GEMINI_MODEL_AUTO)).toBe(true);
|
||||
});
|
||||
|
||||
it('should return false for non-preview models', () => {
|
||||
expect(isPreviewModel(DEFAULT_GEMINI_MODEL)).toBe(false);
|
||||
expect(isPreviewModel('gemini-1.5-pro')).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('isProModel', () => {
|
||||
it('should return true for models containing "pro"', () => {
|
||||
expect(isProModel('gemini-3-pro-preview')).toBe(true);
|
||||
expect(isProModel('gemini-2.5-pro')).toBe(true);
|
||||
expect(isProModel('pro')).toBe(true);
|
||||
});
|
||||
|
||||
it('should return false for models without "pro"', () => {
|
||||
expect(isProModel('gemini-3-flash-preview')).toBe(false);
|
||||
expect(isProModel('gemini-2.5-flash')).toBe(false);
|
||||
expect(isProModel('auto')).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('isCustomModel', () => {
|
||||
it('should return true for models not starting with gemini-', () => {
|
||||
expect(isCustomModel('testing')).toBe(true);
|
||||
@@ -115,6 +148,12 @@ describe('getDisplayString', () => {
|
||||
);
|
||||
});
|
||||
|
||||
it('should return PREVIEW_GEMINI_3_1_MODEL for PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL', () => {
|
||||
expect(getDisplayString(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL)).toBe(
|
||||
PREVIEW_GEMINI_3_1_MODEL,
|
||||
);
|
||||
});
|
||||
|
||||
it('should return the model name as is for other models', () => {
|
||||
expect(getDisplayString('custom-model')).toBe('custom-model');
|
||||
expect(getDisplayString(DEFAULT_GEMINI_FLASH_LITE_MODEL)).toBe(
|
||||
@@ -146,6 +185,16 @@ describe('resolveModel', () => {
|
||||
expect(model).toBe(PREVIEW_GEMINI_MODEL);
|
||||
});
|
||||
|
||||
it('should return Gemini 3.1 Pro when auto-gemini-3 is requested and useGemini3_1 is true', () => {
|
||||
const model = resolveModel(PREVIEW_GEMINI_MODEL_AUTO, true);
|
||||
expect(model).toBe(PREVIEW_GEMINI_3_1_MODEL);
|
||||
});
|
||||
|
||||
it('should return Gemini 3.1 Pro Custom Tools when auto-gemini-3 is requested, useGemini3_1 is true, and useCustomToolModel is true', () => {
|
||||
const model = resolveModel(PREVIEW_GEMINI_MODEL_AUTO, true, true);
|
||||
expect(model).toBe(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL);
|
||||
});
|
||||
|
||||
it('should return the Default Pro model when auto-gemini-2.5 is requested', () => {
|
||||
const model = resolveModel(DEFAULT_GEMINI_MODEL_AUTO);
|
||||
expect(model).toBe(DEFAULT_GEMINI_MODEL);
|
||||
@@ -239,4 +288,71 @@ describe('resolveClassifierModel', () => {
|
||||
resolveClassifierModel(PREVIEW_GEMINI_MODEL_AUTO, GEMINI_MODEL_ALIAS_PRO),
|
||||
).toBe(PREVIEW_GEMINI_MODEL);
|
||||
});
|
||||
|
||||
it('should return Gemini 3.1 Pro when alias is pro and useGemini3_1 is true', () => {
|
||||
expect(
|
||||
resolveClassifierModel(
|
||||
PREVIEW_GEMINI_MODEL_AUTO,
|
||||
GEMINI_MODEL_ALIAS_PRO,
|
||||
true,
|
||||
),
|
||||
).toBe(PREVIEW_GEMINI_3_1_MODEL);
|
||||
});
|
||||
|
||||
it('should return Gemini 3.1 Pro Custom Tools when alias is pro, useGemini3_1 is true, and useCustomToolModel is true', () => {
|
||||
expect(
|
||||
resolveClassifierModel(
|
||||
PREVIEW_GEMINI_MODEL_AUTO,
|
||||
GEMINI_MODEL_ALIAS_PRO,
|
||||
true,
|
||||
true,
|
||||
),
|
||||
).toBe(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL);
|
||||
});
|
||||
});
|
||||
|
||||
describe('isActiveModel', () => {
|
||||
it('should return true for valid models when useGemini3_1 is false', () => {
|
||||
expect(isActiveModel(DEFAULT_GEMINI_MODEL)).toBe(true);
|
||||
expect(isActiveModel(PREVIEW_GEMINI_MODEL)).toBe(true);
|
||||
expect(isActiveModel(DEFAULT_GEMINI_FLASH_MODEL)).toBe(true);
|
||||
});
|
||||
|
||||
it('should return true for unknown models and aliases', () => {
|
||||
expect(isActiveModel('invalid-model')).toBe(false);
|
||||
expect(isActiveModel(GEMINI_MODEL_ALIAS_AUTO)).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false for PREVIEW_GEMINI_MODEL when useGemini3_1 is true', () => {
|
||||
expect(isActiveModel(PREVIEW_GEMINI_MODEL, true)).toBe(false);
|
||||
});
|
||||
|
||||
it('should return true for other valid models when useGemini3_1 is true', () => {
|
||||
expect(isActiveModel(DEFAULT_GEMINI_MODEL, true)).toBe(true);
|
||||
});
|
||||
|
||||
it('should correctly filter Gemini 3.1 models based on useCustomToolModel when useGemini3_1 is true', () => {
|
||||
// When custom tools are preferred, standard 3.1 should be inactive
|
||||
expect(isActiveModel(PREVIEW_GEMINI_3_1_MODEL, true, true)).toBe(false);
|
||||
expect(
|
||||
isActiveModel(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL, true, true),
|
||||
).toBe(true);
|
||||
|
||||
// When custom tools are NOT preferred, custom tools 3.1 should be inactive
|
||||
expect(isActiveModel(PREVIEW_GEMINI_3_1_MODEL, true, false)).toBe(true);
|
||||
expect(
|
||||
isActiveModel(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL, true, false),
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false for both Gemini 3.1 models when useGemini3_1 is false', () => {
|
||||
expect(isActiveModel(PREVIEW_GEMINI_3_1_MODEL, false, true)).toBe(false);
|
||||
expect(isActiveModel(PREVIEW_GEMINI_3_1_MODEL, false, false)).toBe(false);
|
||||
expect(
|
||||
isActiveModel(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL, false, true),
|
||||
).toBe(false);
|
||||
expect(
|
||||
isActiveModel(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL, false, false),
|
||||
).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -5,6 +5,9 @@
|
||||
*/
|
||||
|
||||
export const PREVIEW_GEMINI_MODEL = 'gemini-3-pro-preview';
|
||||
export const PREVIEW_GEMINI_3_1_MODEL = 'gemini-3.1-pro-preview';
|
||||
export const PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL =
|
||||
'gemini-3.1-pro-preview-customtools';
|
||||
export const PREVIEW_GEMINI_FLASH_MODEL = 'gemini-3-flash-preview';
|
||||
export const DEFAULT_GEMINI_MODEL = 'gemini-2.5-pro';
|
||||
export const DEFAULT_GEMINI_FLASH_MODEL = 'gemini-2.5-flash';
|
||||
@@ -12,6 +15,8 @@ export const DEFAULT_GEMINI_FLASH_LITE_MODEL = 'gemini-2.5-flash-lite';
|
||||
|
||||
export const VALID_GEMINI_MODELS = new Set([
|
||||
PREVIEW_GEMINI_MODEL,
|
||||
PREVIEW_GEMINI_3_1_MODEL,
|
||||
PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
|
||||
PREVIEW_GEMINI_FLASH_MODEL,
|
||||
DEFAULT_GEMINI_MODEL,
|
||||
DEFAULT_GEMINI_FLASH_MODEL,
|
||||
@@ -37,20 +42,29 @@ export const DEFAULT_THINKING_MODE = 8192;
|
||||
* to a concrete model name.
|
||||
*
|
||||
* @param requestedModel The model alias or concrete model name requested by the user.
|
||||
* @param useGemini3_1 Whether to use Gemini 3.1 Pro Preview for auto/pro aliases.
|
||||
* @returns The resolved concrete model name.
|
||||
*/
|
||||
export function resolveModel(requestedModel: string): string {
|
||||
export function resolveModel(
|
||||
requestedModel: string,
|
||||
useGemini3_1: boolean = false,
|
||||
useCustomToolModel: boolean = false,
|
||||
): string {
|
||||
switch (requestedModel) {
|
||||
case PREVIEW_GEMINI_MODEL_AUTO: {
|
||||
case PREVIEW_GEMINI_MODEL:
|
||||
case PREVIEW_GEMINI_MODEL_AUTO:
|
||||
case GEMINI_MODEL_ALIAS_AUTO:
|
||||
case GEMINI_MODEL_ALIAS_PRO: {
|
||||
if (useGemini3_1) {
|
||||
return useCustomToolModel
|
||||
? PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL
|
||||
: PREVIEW_GEMINI_3_1_MODEL;
|
||||
}
|
||||
return PREVIEW_GEMINI_MODEL;
|
||||
}
|
||||
case DEFAULT_GEMINI_MODEL_AUTO: {
|
||||
return DEFAULT_GEMINI_MODEL;
|
||||
}
|
||||
case GEMINI_MODEL_ALIAS_AUTO:
|
||||
case GEMINI_MODEL_ALIAS_PRO: {
|
||||
return PREVIEW_GEMINI_MODEL;
|
||||
}
|
||||
case GEMINI_MODEL_ALIAS_FLASH: {
|
||||
return PREVIEW_GEMINI_FLASH_MODEL;
|
||||
}
|
||||
@@ -73,6 +87,8 @@ export function resolveModel(requestedModel: string): string {
|
||||
export function resolveClassifierModel(
|
||||
requestedModel: string,
|
||||
modelAlias: string,
|
||||
useGemini3_1: boolean = false,
|
||||
useCustomToolModel: boolean = false,
|
||||
): string {
|
||||
if (modelAlias === GEMINI_MODEL_ALIAS_FLASH) {
|
||||
if (
|
||||
@@ -89,7 +105,7 @@ export function resolveClassifierModel(
|
||||
}
|
||||
return resolveModel(GEMINI_MODEL_ALIAS_FLASH);
|
||||
}
|
||||
return resolveModel(requestedModel);
|
||||
return resolveModel(requestedModel, useGemini3_1, useCustomToolModel);
|
||||
}
|
||||
export function getDisplayString(model: string) {
|
||||
switch (model) {
|
||||
@@ -101,6 +117,8 @@ export function getDisplayString(model: string) {
|
||||
return PREVIEW_GEMINI_MODEL;
|
||||
case GEMINI_MODEL_ALIAS_FLASH:
|
||||
return PREVIEW_GEMINI_FLASH_MODEL;
|
||||
case PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL:
|
||||
return PREVIEW_GEMINI_3_1_MODEL;
|
||||
default:
|
||||
return model;
|
||||
}
|
||||
@@ -115,11 +133,22 @@ export function getDisplayString(model: string) {
|
||||
export function isPreviewModel(model: string): boolean {
|
||||
return (
|
||||
model === PREVIEW_GEMINI_MODEL ||
|
||||
model === PREVIEW_GEMINI_3_1_MODEL ||
|
||||
model === PREVIEW_GEMINI_FLASH_MODEL ||
|
||||
model === PREVIEW_GEMINI_MODEL_AUTO
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Checks if the model is a Pro model.
|
||||
*
|
||||
* @param model The model name to check.
|
||||
* @returns True if the model is a Pro model.
|
||||
*/
|
||||
export function isProModel(model: string): boolean {
|
||||
return model.toLowerCase().includes('pro');
|
||||
}
|
||||
|
||||
/**
|
||||
* Checks if the model is a Gemini 3 model.
|
||||
*
|
||||
@@ -188,3 +217,35 @@ export function isAutoModel(model: string): boolean {
|
||||
export function supportsMultimodalFunctionResponse(model: string): boolean {
|
||||
return model.startsWith('gemini-3-');
|
||||
}
|
||||
|
||||
/**
|
||||
* Checks if the given model is considered active based on the current configuration.
|
||||
*
|
||||
* @param model The model name to check.
|
||||
* @param useGemini3_1 Whether Gemini 3.1 Pro Preview is enabled.
|
||||
* @returns True if the model is active.
|
||||
*/
|
||||
export function isActiveModel(
|
||||
model: string,
|
||||
useGemini3_1: boolean = false,
|
||||
useCustomToolModel: boolean = false,
|
||||
): boolean {
|
||||
if (!VALID_GEMINI_MODELS.has(model)) {
|
||||
return false;
|
||||
}
|
||||
if (useGemini3_1) {
|
||||
if (model === PREVIEW_GEMINI_MODEL) {
|
||||
return false;
|
||||
}
|
||||
if (useCustomToolModel) {
|
||||
return model !== PREVIEW_GEMINI_3_1_MODEL;
|
||||
} else {
|
||||
return model !== PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL;
|
||||
}
|
||||
} else {
|
||||
return (
|
||||
model !== PREVIEW_GEMINI_3_1_MODEL &&
|
||||
model !== PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -542,7 +542,10 @@ export class GeminiClient {
|
||||
|
||||
// Availability logic: The configured model is the source of truth,
|
||||
// including any permanent fallbacks (config.setModel) or manual overrides.
|
||||
return resolveModel(this.config.getActiveModel());
|
||||
return resolveModel(
|
||||
this.config.getActiveModel(),
|
||||
this.config.getGemini31LaunchedSync?.() ?? false,
|
||||
);
|
||||
}
|
||||
|
||||
private async *processTurn(
|
||||
|
||||
@@ -146,7 +146,12 @@ export async function createContentGenerator(
|
||||
return new LoggingContentGenerator(fakeGenerator, gcConfig);
|
||||
}
|
||||
const version = await getVersion();
|
||||
const model = resolveModel(gcConfig.getModel());
|
||||
const model = resolveModel(
|
||||
gcConfig.getModel(),
|
||||
config.authType === AuthType.USE_GEMINI ||
|
||||
config.authType === AuthType.USE_VERTEX_AI ||
|
||||
((await gcConfig.getGemini31Launched?.()) ?? false),
|
||||
);
|
||||
const customHeadersEnv =
|
||||
process.env['GEMINI_CLI_CUSTOM_HEADERS'] || undefined;
|
||||
const userAgent = `GeminiCLI/${version}/${model} (${process.platform}; ${process.arch})`;
|
||||
|
||||
@@ -496,13 +496,14 @@ export class GeminiChat {
|
||||
const initialActiveModel = this.config.getActiveModel();
|
||||
|
||||
const apiCall = async () => {
|
||||
const useGemini3_1 = (await this.config.getGemini31Launched?.()) ?? false;
|
||||
// Default to the last used model (which respects arguments/availability selection)
|
||||
let modelToUse = resolveModel(lastModelToUse);
|
||||
let modelToUse = resolveModel(lastModelToUse, useGemini3_1);
|
||||
|
||||
// If the active model has changed (e.g. due to a fallback updating the config),
|
||||
// we switch to the new active model.
|
||||
if (this.config.getActiveModel() !== initialActiveModel) {
|
||||
modelToUse = resolveModel(this.config.getActiveModel());
|
||||
modelToUse = resolveModel(this.config.getActiveModel(), useGemini3_1);
|
||||
}
|
||||
|
||||
if (modelToUse !== lastModelToUse) {
|
||||
|
||||
@@ -58,7 +58,10 @@ export class PromptProvider {
|
||||
const enabledToolNames = new Set(toolNames);
|
||||
const approvedPlanPath = config.getApprovedPlanPath();
|
||||
|
||||
const desiredModel = resolveModel(config.getActiveModel());
|
||||
const desiredModel = resolveModel(
|
||||
config.getActiveModel(),
|
||||
config.getGemini31LaunchedSync?.() ?? false,
|
||||
);
|
||||
const isModernModel = supportsModernFeatures(desiredModel);
|
||||
const activeSnippets = isModernModel ? snippets : legacySnippets;
|
||||
const contextFilenames = getAllGeminiMdFilenames();
|
||||
@@ -231,7 +234,10 @@ export class PromptProvider {
|
||||
}
|
||||
|
||||
getCompressionPrompt(config: Config): string {
|
||||
const desiredModel = resolveModel(config.getActiveModel());
|
||||
const desiredModel = resolveModel(
|
||||
config.getActiveModel(),
|
||||
config.getGemini31LaunchedSync?.() ?? false,
|
||||
);
|
||||
const isModernModel = supportsModernFeatures(desiredModel);
|
||||
const activeSnippets = isModernModel ? snippets : legacySnippets;
|
||||
return activeSnippets.getCompressionPrompt();
|
||||
|
||||
@@ -18,11 +18,14 @@ import {
|
||||
DEFAULT_GEMINI_MODEL,
|
||||
DEFAULT_GEMINI_MODEL_AUTO,
|
||||
PREVIEW_GEMINI_MODEL_AUTO,
|
||||
PREVIEW_GEMINI_3_1_MODEL,
|
||||
PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
|
||||
} from '../../config/models.js';
|
||||
import { promptIdContext } from '../../utils/promptIdContext.js';
|
||||
import type { Content } from '@google/genai';
|
||||
import type { ResolvedModelConfig } from '../../services/modelConfigService.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
import { AuthType } from '../../core/contentGenerator.js';
|
||||
|
||||
vi.mock('../../core/baseLlmClient.js');
|
||||
|
||||
@@ -53,6 +56,10 @@ describe('ClassifierStrategy', () => {
|
||||
},
|
||||
getModel: vi.fn().mockReturnValue(DEFAULT_GEMINI_MODEL_AUTO),
|
||||
getNumericalRoutingEnabled: vi.fn().mockResolvedValue(false),
|
||||
getGemini31Launched: vi.fn().mockResolvedValue(false),
|
||||
getContentGeneratorConfig: vi.fn().mockReturnValue({
|
||||
authType: AuthType.LOGIN_WITH_GOOGLE,
|
||||
}),
|
||||
} as unknown as Config;
|
||||
mockBaseLlmClient = {
|
||||
generateJson: vi.fn(),
|
||||
@@ -339,4 +346,49 @@ describe('ClassifierStrategy', () => {
|
||||
// Since requestedModel is Pro, and choice is flash, it should resolve to Flash
|
||||
expect(decision?.model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
});
|
||||
|
||||
describe('Gemini 3.1 and Custom Tools Routing', () => {
|
||||
it('should route to PREVIEW_GEMINI_3_1_MODEL when Gemini 3.1 is launched', async () => {
|
||||
vi.mocked(mockConfig.getGemini31Launched).mockResolvedValue(true);
|
||||
vi.mocked(mockConfig.getModel).mockReturnValue(PREVIEW_GEMINI_MODEL_AUTO);
|
||||
const mockApiResponse = {
|
||||
reasoning: 'Complex task',
|
||||
model_choice: 'pro',
|
||||
};
|
||||
vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(
|
||||
mockApiResponse,
|
||||
);
|
||||
|
||||
const decision = await strategy.route(
|
||||
mockContext,
|
||||
mockConfig,
|
||||
mockBaseLlmClient,
|
||||
);
|
||||
|
||||
expect(decision?.model).toBe(PREVIEW_GEMINI_3_1_MODEL);
|
||||
});
|
||||
|
||||
it('should route to PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL when Gemini 3.1 is launched and auth is USE_GEMINI', async () => {
|
||||
vi.mocked(mockConfig.getGemini31Launched).mockResolvedValue(true);
|
||||
vi.mocked(mockConfig.getModel).mockReturnValue(PREVIEW_GEMINI_MODEL_AUTO);
|
||||
vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({
|
||||
authType: AuthType.USE_GEMINI,
|
||||
});
|
||||
const mockApiResponse = {
|
||||
reasoning: 'Complex task',
|
||||
model_choice: 'pro',
|
||||
};
|
||||
vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(
|
||||
mockApiResponse,
|
||||
);
|
||||
|
||||
const decision = await strategy.route(
|
||||
mockContext,
|
||||
mockConfig,
|
||||
mockBaseLlmClient,
|
||||
);
|
||||
|
||||
expect(decision?.model).toBe(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -21,6 +21,7 @@ import {
|
||||
} from '../../utils/messageInspectors.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
import { LlmRole } from '../../telemetry/types.js';
|
||||
import { AuthType } from '../../core/contentGenerator.js';
|
||||
|
||||
// The number of recent history turns to provide to the router for context.
|
||||
const HISTORY_TURNS_FOR_CONTEXT = 4;
|
||||
@@ -169,9 +170,15 @@ export class ClassifierStrategy implements RoutingStrategy {
|
||||
|
||||
const reasoning = routerResponse.reasoning;
|
||||
const latencyMs = Date.now() - startTime;
|
||||
const useGemini3_1 = (await config.getGemini31Launched?.()) ?? false;
|
||||
const useCustomToolModel =
|
||||
useGemini3_1 &&
|
||||
config.getContentGeneratorConfig().authType === AuthType.USE_GEMINI;
|
||||
const selectedModel = resolveClassifierModel(
|
||||
model,
|
||||
routerResponse.model_choice,
|
||||
useGemini3_1,
|
||||
useCustomToolModel,
|
||||
);
|
||||
|
||||
return {
|
||||
|
||||
@@ -21,7 +21,10 @@ export class DefaultStrategy implements TerminalStrategy {
|
||||
config: Config,
|
||||
_baseLlmClient: BaseLlmClient,
|
||||
): Promise<RoutingDecision> {
|
||||
const defaultModel = resolveModel(config.getModel());
|
||||
const defaultModel = resolveModel(
|
||||
config.getModel(),
|
||||
config.getGemini31LaunchedSync?.() ?? false,
|
||||
);
|
||||
return {
|
||||
model: defaultModel,
|
||||
metadata: {
|
||||
|
||||
@@ -23,7 +23,10 @@ export class FallbackStrategy implements RoutingStrategy {
|
||||
_baseLlmClient: BaseLlmClient,
|
||||
): Promise<RoutingDecision | null> {
|
||||
const requestedModel = context.requestedModel ?? config.getModel();
|
||||
const resolvedModel = resolveModel(requestedModel);
|
||||
const resolvedModel = resolveModel(
|
||||
requestedModel,
|
||||
config.getGemini31LaunchedSync?.() ?? false,
|
||||
);
|
||||
const service = config.getModelAvailabilityService();
|
||||
const snapshot = service.snapshot(resolvedModel);
|
||||
|
||||
|
||||
@@ -12,6 +12,8 @@ import type { BaseLlmClient } from '../../core/baseLlmClient.js';
|
||||
import {
|
||||
PREVIEW_GEMINI_FLASH_MODEL,
|
||||
PREVIEW_GEMINI_MODEL,
|
||||
PREVIEW_GEMINI_3_1_MODEL,
|
||||
PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
|
||||
PREVIEW_GEMINI_MODEL_AUTO,
|
||||
DEFAULT_GEMINI_MODEL_AUTO,
|
||||
DEFAULT_GEMINI_MODEL,
|
||||
@@ -20,6 +22,7 @@ import { promptIdContext } from '../../utils/promptIdContext.js';
|
||||
import type { Content } from '@google/genai';
|
||||
import type { ResolvedModelConfig } from '../../services/modelConfigService.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
import { AuthType } from '../../core/contentGenerator.js';
|
||||
|
||||
vi.mock('../../core/baseLlmClient.js');
|
||||
|
||||
@@ -52,6 +55,10 @@ describe('NumericalClassifierStrategy', () => {
|
||||
getSessionId: vi.fn().mockReturnValue('control-group-id'), // Default to Control Group (Hash 71 >= 50)
|
||||
getNumericalRoutingEnabled: vi.fn().mockResolvedValue(true),
|
||||
getClassifierThreshold: vi.fn().mockResolvedValue(undefined),
|
||||
getGemini31Launched: vi.fn().mockResolvedValue(false),
|
||||
getContentGeneratorConfig: vi.fn().mockReturnValue({
|
||||
authType: AuthType.LOGIN_WITH_GOOGLE,
|
||||
}),
|
||||
} as unknown as Config;
|
||||
mockBaseLlmClient = {
|
||||
generateJson: vi.fn(),
|
||||
@@ -535,4 +542,68 @@ describe('NumericalClassifierStrategy', () => {
|
||||
),
|
||||
);
|
||||
});
|
||||
|
||||
describe('Gemini 3.1 and Custom Tools Routing', () => {
|
||||
it('should route to PREVIEW_GEMINI_3_1_MODEL when Gemini 3.1 is launched', async () => {
|
||||
vi.mocked(mockConfig.getGemini31Launched).mockResolvedValue(true);
|
||||
const mockApiResponse = {
|
||||
complexity_reasoning: 'Complex task',
|
||||
complexity_score: 80,
|
||||
};
|
||||
vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(
|
||||
mockApiResponse,
|
||||
);
|
||||
|
||||
const decision = await strategy.route(
|
||||
mockContext,
|
||||
mockConfig,
|
||||
mockBaseLlmClient,
|
||||
);
|
||||
|
||||
expect(decision?.model).toBe(PREVIEW_GEMINI_3_1_MODEL);
|
||||
});
|
||||
it('should route to PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL when Gemini 3.1 is launched and auth is USE_GEMINI', async () => {
|
||||
vi.mocked(mockConfig.getGemini31Launched).mockResolvedValue(true);
|
||||
vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({
|
||||
authType: AuthType.USE_GEMINI,
|
||||
});
|
||||
const mockApiResponse = {
|
||||
complexity_reasoning: 'Complex task',
|
||||
complexity_score: 80,
|
||||
};
|
||||
vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(
|
||||
mockApiResponse,
|
||||
);
|
||||
|
||||
const decision = await strategy.route(
|
||||
mockContext,
|
||||
mockConfig,
|
||||
mockBaseLlmClient,
|
||||
);
|
||||
|
||||
expect(decision?.model).toBe(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL);
|
||||
});
|
||||
|
||||
it('should NOT route to custom tools model when auth is USE_VERTEX_AI', async () => {
|
||||
vi.mocked(mockConfig.getGemini31Launched).mockResolvedValue(true);
|
||||
vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({
|
||||
authType: AuthType.USE_VERTEX_AI,
|
||||
});
|
||||
const mockApiResponse = {
|
||||
complexity_reasoning: 'Complex task',
|
||||
complexity_score: 80,
|
||||
};
|
||||
vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(
|
||||
mockApiResponse,
|
||||
);
|
||||
|
||||
const decision = await strategy.route(
|
||||
mockContext,
|
||||
mockConfig,
|
||||
mockBaseLlmClient,
|
||||
);
|
||||
|
||||
expect(decision?.model).toBe(PREVIEW_GEMINI_3_1_MODEL);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -17,6 +17,7 @@ import { createUserContent, Type } from '@google/genai';
|
||||
import type { Config } from '../../config/config.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
import { LlmRole } from '../../telemetry/types.js';
|
||||
import { AuthType } from '../../core/contentGenerator.js';
|
||||
|
||||
// The number of recent history turns to provide to the router for context.
|
||||
const HISTORY_TURNS_FOR_CONTEXT = 8;
|
||||
@@ -182,8 +183,16 @@ export class NumericalClassifierStrategy implements RoutingStrategy {
|
||||
config,
|
||||
config.getSessionId() || 'unknown-session',
|
||||
);
|
||||
|
||||
const selectedModel = resolveClassifierModel(model, modelAlias);
|
||||
const useGemini3_1 = (await config.getGemini31Launched?.()) ?? false;
|
||||
const useCustomToolModel =
|
||||
useGemini3_1 &&
|
||||
config.getContentGeneratorConfig().authType === AuthType.USE_GEMINI;
|
||||
const selectedModel = resolveClassifierModel(
|
||||
model,
|
||||
modelAlias,
|
||||
useGemini3_1,
|
||||
useCustomToolModel,
|
||||
);
|
||||
|
||||
const latencyMs = Date.now() - startTime;
|
||||
|
||||
|
||||
@@ -33,7 +33,10 @@ export class OverrideStrategy implements RoutingStrategy {
|
||||
|
||||
// Return the overridden model name.
|
||||
return {
|
||||
model: resolveModel(overrideModel),
|
||||
model: resolveModel(
|
||||
overrideModel,
|
||||
config.getGemini31LaunchedSync?.() ?? false,
|
||||
),
|
||||
metadata: {
|
||||
source: this.name,
|
||||
latencyMs: 0,
|
||||
|
||||
@@ -29,6 +29,7 @@ import {
|
||||
DEFAULT_GEMINI_MODEL,
|
||||
PREVIEW_GEMINI_MODEL,
|
||||
PREVIEW_GEMINI_FLASH_MODEL,
|
||||
PREVIEW_GEMINI_3_1_MODEL,
|
||||
} from '../config/models.js';
|
||||
import { PreCompressTrigger } from '../hooks/types.js';
|
||||
import { LlmRole } from '../telemetry/types.js';
|
||||
@@ -101,6 +102,7 @@ export function findCompressSplitPoint(
|
||||
export function modelStringToModelConfigAlias(model: string): string {
|
||||
switch (model) {
|
||||
case PREVIEW_GEMINI_MODEL:
|
||||
case PREVIEW_GEMINI_3_1_MODEL:
|
||||
return 'chat-compression-3-pro';
|
||||
case PREVIEW_GEMINI_FLASH_MODEL:
|
||||
return 'chat-compression-3-flash';
|
||||
|
||||
Reference in New Issue
Block a user