feat(models): support Gemini 3.1 Pro Preview and fixes (#19676)

This commit is contained in:
Sehoon Shon
2026-02-20 14:19:21 -05:00
committed by GitHub
parent 788a40c445
commit f97b04cc9a
25 changed files with 670 additions and 50 deletions
@@ -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,