Update models menu dialog to support manual model selection and group… (#80)

* Update models menu dialog to support manual model selection and group by model family

* fix tests

* update codebase investigator model setting

* fix test

* fix test

* update generated golden file

* regenerate scheme and doc for settings

* use Preview Auto if previewFeatures is set to true
This commit is contained in:
Sehoon Shon
2025-12-11 09:57:27 -05:00
committed by Tommaso Sciortino
parent 5eb817c426
commit ad994cfe8b
27 changed files with 456 additions and 538 deletions
@@ -12,7 +12,7 @@ import {
LS_TOOL_NAME,
READ_FILE_TOOL_NAME,
} from '../tools/tool-names.js';
import { GEMINI_MODEL_ALIAS_PRO } from '../config/models.js';
import { DEFAULT_GEMINI_MODEL } from '../config/models.js';
describe('CodebaseInvestigatorAgent', () => {
it('should have the correct agent definition', () => {
@@ -26,7 +26,7 @@ describe('CodebaseInvestigatorAgent', () => {
).toBe(true);
expect(CodebaseInvestigatorAgent.outputConfig?.outputName).toBe('report');
expect(CodebaseInvestigatorAgent.modelConfig?.model).toBe(
GEMINI_MODEL_ALIAS_PRO,
DEFAULT_GEMINI_MODEL,
);
expect(CodebaseInvestigatorAgent.toolConfig?.tools).toEqual([
LS_TOOL_NAME,
@@ -11,7 +11,7 @@ import {
LS_TOOL_NAME,
READ_FILE_TOOL_NAME,
} from '../tools/tool-names.js';
import { GEMINI_MODEL_ALIAS_PRO } from '../config/models.js';
import { DEFAULT_GEMINI_MODEL } from '../config/models.js';
import { z } from 'zod';
// Define a type that matches the outputConfig schema for type safety.
@@ -69,7 +69,7 @@ export const CodebaseInvestigatorAgent: AgentDefinition<
processOutput: (output) => JSON.stringify(output, null, 2),
modelConfig: {
model: GEMINI_MODEL_ALIAS_PRO,
model: DEFAULT_GEMINI_MODEL,
temp: 0.1,
top_p: 0.95,
thinkingBudget: -1,
+2 -1
View File
@@ -10,6 +10,7 @@ import { makeFakeConfig } from '../test-utils/config.js';
import type { AgentDefinition } from './types.js';
import type { Config } from '../config/config.js';
import { debugLogger } from '../utils/debugLogger.js';
import { DEFAULT_GEMINI_MODEL } from '../config/models.js';
// A test-only subclass to expose the protected `registerAgent` method.
class TestableAgentRegistry extends AgentRegistry {
@@ -78,7 +79,7 @@ describe('AgentRegistry', () => {
model: 'gemini-3-pro-preview',
codebaseInvestigatorSettings: {
enabled: true,
model: 'pro',
model: DEFAULT_GEMINI_MODEL,
},
});
const previewRegistry = new TestableAgentRegistry(previewConfig);
+9 -4
View File
@@ -10,9 +10,10 @@ import { CodebaseInvestigatorAgent } from './codebase-investigator.js';
import { type z } from 'zod';
import { debugLogger } from '../utils/debugLogger.js';
import {
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_MODEL_AUTO,
GEMINI_MODEL_ALIAS_PRO,
PREVIEW_GEMINI_MODEL,
PREVIEW_GEMINI_MODEL_AUTO,
} from '../config/models.js';
import type { ModelConfigAlias } from '../services/modelConfigService.js';
@@ -59,10 +60,14 @@ export class AgentRegistry {
// If the user is using the preview model for the main agent, force the sub-agent to use it too
// if it's configured to use 'pro' or 'auto'.
if (this.config.getModel() === PREVIEW_GEMINI_MODEL) {
if (
this.config.getModel() === PREVIEW_GEMINI_MODEL ||
this.config.getModel() === PREVIEW_GEMINI_MODEL_AUTO
) {
if (
model === GEMINI_MODEL_ALIAS_PRO ||
model === DEFAULT_GEMINI_MODEL_AUTO
model === PREVIEW_GEMINI_MODEL_AUTO ||
model === DEFAULT_GEMINI_MODEL_AUTO ||
model === DEFAULT_GEMINI_MODEL
) {
model = PREVIEW_GEMINI_MODEL;
}
@@ -33,11 +33,7 @@ export function resolvePolicyChain(
// Switch to getActiveModel()
const activeModel =
preferredModel ??
getEffectiveModel(
config.isInFallbackMode(),
config.getModel(),
config.getPreviewFeatures(),
);
getEffectiveModel(config.getModel(), config.isInFallbackMode());
if (activeModel === 'auto') {
return [...chain];
+80 -184
View File
@@ -7,14 +7,15 @@
import { describe, it, expect } from 'vitest';
import {
getEffectiveModel,
isGemini2Model,
DEFAULT_GEMINI_MODEL,
PREVIEW_GEMINI_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
PREVIEW_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_FLASH_LITE_MODEL,
GEMINI_MODEL_ALIAS_PRO,
GEMINI_MODEL_ALIAS_FLASH,
GEMINI_MODEL_ALIAS_FLASH_LITE,
supportsMultimodalFunctionResponse,
PREVIEW_GEMINI_MODEL_AUTO,
DEFAULT_GEMINI_MODEL_AUTO,
} from './models.js';
describe('supportsMultimodalFunctionResponse', () => {
@@ -35,209 +36,104 @@ describe('supportsMultimodalFunctionResponse', () => {
describe('getEffectiveModel', () => {
describe('When NOT in fallback mode', () => {
const isInFallbackMode = false;
const useFallbackModel = false;
it('should return the Pro model when Pro is requested', () => {
it('should return the Preview Pro model when auto-preview is requested', () => {
const model = getEffectiveModel(
isInFallbackMode,
DEFAULT_GEMINI_MODEL,
false,
PREVIEW_GEMINI_MODEL_AUTO,
useFallbackModel,
);
expect(model).toBe(PREVIEW_GEMINI_MODEL);
});
it('should return the Default Pro model when auto-default is requested', () => {
const model = getEffectiveModel(
DEFAULT_GEMINI_MODEL_AUTO,
useFallbackModel,
);
expect(model).toBe(DEFAULT_GEMINI_MODEL);
});
it('should return the Flash model when Flash is requested', () => {
const model = getEffectiveModel(
isInFallbackMode,
DEFAULT_GEMINI_FLASH_MODEL,
false,
it('should return the requested model as-is for explicit specific models', () => {
expect(getEffectiveModel(DEFAULT_GEMINI_MODEL, useFallbackModel)).toBe(
DEFAULT_GEMINI_MODEL,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
});
it('should return the Lite model when Lite is requested', () => {
const model = getEffectiveModel(
isInFallbackMode,
DEFAULT_GEMINI_FLASH_LITE_MODEL,
false,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_LITE_MODEL);
expect(
getEffectiveModel(DEFAULT_GEMINI_FLASH_MODEL, useFallbackModel),
).toBe(DEFAULT_GEMINI_FLASH_MODEL);
expect(
getEffectiveModel(DEFAULT_GEMINI_FLASH_LITE_MODEL, useFallbackModel),
).toBe(DEFAULT_GEMINI_FLASH_LITE_MODEL);
});
it('should return a custom model name when requested', () => {
const customModel = 'custom-model-v1';
const model = getEffectiveModel(isInFallbackMode, customModel, false);
const model = getEffectiveModel(customModel, useFallbackModel);
expect(model).toBe(customModel);
});
describe('with preview features', () => {
it('should return the preview model when pro alias is requested', () => {
const model = getEffectiveModel(
isInFallbackMode,
GEMINI_MODEL_ALIAS_PRO,
true,
);
expect(model).toBe(PREVIEW_GEMINI_MODEL);
});
it('should return the default pro model when pro alias is requested and preview is off', () => {
const model = getEffectiveModel(
isInFallbackMode,
GEMINI_MODEL_ALIAS_PRO,
false,
);
expect(model).toBe(DEFAULT_GEMINI_MODEL);
});
it('should return the flash model when flash is requested and preview is on', () => {
const model = getEffectiveModel(
isInFallbackMode,
GEMINI_MODEL_ALIAS_FLASH,
true,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
});
it('should return the flash model when lite is requested and preview is on', () => {
const model = getEffectiveModel(
isInFallbackMode,
GEMINI_MODEL_ALIAS_FLASH_LITE,
true,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_LITE_MODEL);
});
it('should return the flash model when the flash model name is explicitly requested and preview is on', () => {
const model = getEffectiveModel(
isInFallbackMode,
DEFAULT_GEMINI_FLASH_MODEL,
true,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
});
it('should return the lite model when the lite model name is requested and preview is on', () => {
const model = getEffectiveModel(
isInFallbackMode,
DEFAULT_GEMINI_FLASH_LITE_MODEL,
true,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_LITE_MODEL);
});
it('should return the default gemini model when the model is explicitly set and preview is on', () => {
const model = getEffectiveModel(
isInFallbackMode,
DEFAULT_GEMINI_MODEL,
true,
);
expect(model).toBe(DEFAULT_GEMINI_MODEL);
});
});
});
describe('When IN fallback mode', () => {
const isInFallbackMode = true;
const useFallbackModel = true;
it('should downgrade the Pro model to the Flash model', () => {
it('should return the Preview Flash model when auto-preview is requested', () => {
const model = getEffectiveModel(
isInFallbackMode,
PREVIEW_GEMINI_MODEL_AUTO,
useFallbackModel,
);
expect(model).toBe(PREVIEW_GEMINI_FLASH_MODEL);
});
it('should return the Default Flash model when auto-default is requested', () => {
const model = getEffectiveModel(
DEFAULT_GEMINI_MODEL_AUTO,
useFallbackModel,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
});
it('should return the requested model as-is for explicit specific models', () => {
expect(getEffectiveModel(DEFAULT_GEMINI_MODEL, useFallbackModel)).toBe(
DEFAULT_GEMINI_MODEL,
false,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
expect(
getEffectiveModel(DEFAULT_GEMINI_FLASH_MODEL, useFallbackModel),
).toBe(DEFAULT_GEMINI_FLASH_MODEL);
expect(
getEffectiveModel(DEFAULT_GEMINI_FLASH_LITE_MODEL, useFallbackModel),
).toBe(DEFAULT_GEMINI_FLASH_LITE_MODEL);
});
it('should return the Flash model when Flash is requested', () => {
const model = getEffectiveModel(
isInFallbackMode,
DEFAULT_GEMINI_FLASH_MODEL,
false,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
});
it('should HONOR the Lite model when Lite is requested', () => {
const model = getEffectiveModel(
isInFallbackMode,
DEFAULT_GEMINI_FLASH_LITE_MODEL,
false,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_LITE_MODEL);
});
it('should HONOR any model with "lite" in its name', () => {
const customLiteModel = 'gemini-2.5-custom-lite-vNext';
const model = getEffectiveModel(isInFallbackMode, customLiteModel, false);
expect(model).toBe(customLiteModel);
});
it('should downgrade any other custom model to the Flash model', () => {
const customModel = 'custom-model-v1-unlisted';
const model = getEffectiveModel(isInFallbackMode, customModel, false);
expect(model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
});
describe('with preview features', () => {
it('should downgrade the Pro alias to the Flash model', () => {
const model = getEffectiveModel(
isInFallbackMode,
GEMINI_MODEL_ALIAS_PRO,
true,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
});
it('should return the Flash alias when requested', () => {
const model = getEffectiveModel(
isInFallbackMode,
GEMINI_MODEL_ALIAS_FLASH,
true,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
});
it('should return the Lite alias when requested', () => {
const model = getEffectiveModel(
isInFallbackMode,
GEMINI_MODEL_ALIAS_FLASH_LITE,
true,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_LITE_MODEL);
});
it('should downgrade the default Gemini model to the Flash model', () => {
const model = getEffectiveModel(
isInFallbackMode,
DEFAULT_GEMINI_MODEL,
true,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
});
it('should return the default Flash model when requested', () => {
const model = getEffectiveModel(
isInFallbackMode,
DEFAULT_GEMINI_FLASH_MODEL,
true,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
});
it('should return the default Lite model when requested', () => {
const model = getEffectiveModel(
isInFallbackMode,
DEFAULT_GEMINI_FLASH_LITE_MODEL,
true,
);
expect(model).toBe(DEFAULT_GEMINI_FLASH_LITE_MODEL);
});
it('should downgrade any other custom model to the Flash model', () => {
const customModel = 'custom-model-v1-unlisted';
const model = getEffectiveModel(isInFallbackMode, customModel, true);
expect(model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
});
it('should return custom model name as-is', () => {
const customModel = 'custom-model-v1';
const model = getEffectiveModel(customModel, useFallbackModel);
expect(model).toBe(customModel);
});
});
});
describe('isGemini2Model', () => {
it('should return true for gemini-2.5-pro', () => {
expect(isGemini2Model('gemini-2.5-pro')).toBe(true);
});
it('should return true for gemini-2.5-flash', () => {
expect(isGemini2Model('gemini-2.5-flash')).toBe(true);
});
it('should return true for gemini-2.0-flash', () => {
expect(isGemini2Model('gemini-2.0-flash')).toBe(true);
});
it('should return false for gemini-1.5-pro', () => {
expect(isGemini2Model('gemini-1.5-pro')).toBe(false);
});
it('should return false for gemini-3-pro', () => {
expect(isGemini2Model('gemini-3-pro')).toBe(false);
});
it('should return false for arbitrary strings', () => {
expect(isGemini2Model('gpt-4')).toBe(false);
});
});
+33 -55
View File
@@ -17,83 +17,61 @@ export const VALID_GEMINI_MODELS = new Set([
DEFAULT_GEMINI_FLASH_LITE_MODEL,
]);
export const DEFAULT_GEMINI_MODEL_AUTO = 'auto';
// Model aliases for user convenience.
export const GEMINI_MODEL_ALIAS_PRO = 'pro';
export const GEMINI_MODEL_ALIAS_FLASH = 'flash';
export const GEMINI_MODEL_ALIAS_FLASH_LITE = 'flash-lite';
export const PREVIEW_GEMINI_MODEL_AUTO = 'auto-gemini-3';
export const DEFAULT_GEMINI_MODEL_AUTO = 'auto-gemini-2.5';
export const DEFAULT_GEMINI_EMBEDDING_MODEL = 'gemini-embedding-001';
// Cap the thinking at 8192 to prevent run-away thinking loops.
export const DEFAULT_THINKING_MODE = 8192;
/**
* Resolves the requested model alias (e.g., 'auto', 'pro', 'flash', 'flash-lite')
* to a concrete model name, considering preview features.
*
* @param requestedModel The model alias or concrete model name requested by the user.
* @param previewFeaturesEnabled A boolean indicating if preview features are enabled.
* @returns The resolved concrete model name.
*/
export function resolveModel(
requestedModel: string,
previewFeaturesEnabled: boolean | undefined,
): string {
switch (requestedModel) {
case DEFAULT_GEMINI_MODEL_AUTO:
case GEMINI_MODEL_ALIAS_PRO: {
return previewFeaturesEnabled
? PREVIEW_GEMINI_MODEL
: DEFAULT_GEMINI_MODEL;
}
case GEMINI_MODEL_ALIAS_FLASH: {
return DEFAULT_GEMINI_FLASH_MODEL;
}
case GEMINI_MODEL_ALIAS_FLASH_LITE: {
return DEFAULT_GEMINI_FLASH_LITE_MODEL;
}
default: {
return requestedModel;
}
}
}
/**
* Determines the effective model to use, applying fallback logic if necessary.
*
* When fallback mode is active, this function enforces the use of the standard
* fallback model. However, it makes an exception for "lite" models (any model
* with "lite" in its name), allowing them to be used to preserve cost savings.
* This ensures that "pro" models are always downgraded, while "lite" model
* requests are honored.
* fallback model.
*
* @param isInFallbackMode Whether the application is in fallback mode.
* @param requestedModel The model that was originally requested.
* @param previewFeaturesEnabled A boolean indicating if preview features are enabled.
* @param isInFallbackMode Whether the application is in fallback mode.
* @returns The effective model name.
*/
export function getEffectiveModel(
isInFallbackMode: boolean,
requestedModel: string,
previewFeaturesEnabled: boolean | undefined,
useFallbackModel: boolean,
): string {
const resolvedModel = resolveModel(requestedModel, previewFeaturesEnabled);
// If we are not in fallback mode, simply use the resolved model.
if (!isInFallbackMode) {
return resolvedModel;
if (!useFallbackModel) {
switch (requestedModel) {
case PREVIEW_GEMINI_MODEL_AUTO:
return PREVIEW_GEMINI_MODEL;
case DEFAULT_GEMINI_MODEL_AUTO:
return DEFAULT_GEMINI_MODEL;
default:
return requestedModel;
}
}
// If a "lite" model is requested, honor it. This allows for variations of
// lite models without needing to list them all as constants.
if (resolvedModel.includes('lite')) {
return resolvedModel;
// Fallback model for corresponding model family. We are doing fallback only
// for Auto modes
switch (requestedModel) {
case PREVIEW_GEMINI_MODEL_AUTO:
return PREVIEW_GEMINI_FLASH_MODEL;
case DEFAULT_GEMINI_MODEL_AUTO:
return DEFAULT_GEMINI_FLASH_MODEL;
default:
return requestedModel;
}
}
// Default fallback for Gemini CLI.
return DEFAULT_GEMINI_FLASH_MODEL;
export function getDisplayString(model: string) {
switch (model) {
case PREVIEW_GEMINI_MODEL_AUTO:
return 'Auto (Gemini 3)';
case DEFAULT_GEMINI_MODEL_AUTO:
return 'Auto (Gemini 2.5)';
default:
return model;
}
}
/**
+1 -5
View File
@@ -396,11 +396,7 @@ export class GeminiClient {
}
const configModel = this.config.getModel();
return getEffectiveModel(
this.config.isInFallbackMode(),
configModel,
this.config.getPreviewFeatures(),
);
return getEffectiveModel(configModel, this.config.isInFallbackMode());
}
async *sendMessageStream(
+1 -2
View File
@@ -118,9 +118,8 @@ export async function createContentGenerator(
}
const version = await getVersion();
const model = getEffectiveModel(
gcConfig.isInFallbackMode(),
gcConfig.getModel(),
gcConfig.getPreviewFeatures(),
gcConfig.isInFallbackMode(),
);
const customHeadersEnv =
process.env['GEMINI_CLI_CUSTOM_HEADERS'] || undefined;
+14 -11
View File
@@ -21,7 +21,9 @@ import {
DEFAULT_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_MODEL,
DEFAULT_THINKING_MODE,
PREVIEW_GEMINI_FLASH_MODEL,
PREVIEW_GEMINI_MODEL,
PREVIEW_GEMINI_MODEL_AUTO,
} from '../config/models.js';
import { AuthType } from './contentGenerator.js';
import { TerminalQuotaError } from '../utils/googleQuotaErrors.js';
@@ -155,7 +157,8 @@ describe('GeminiChat', () => {
getUserTier: vi.fn().mockReturnValue(undefined),
modelConfigService: {
getResolvedConfig: vi.fn().mockImplementation((modelConfigKey) => {
const thinkingConfig = modelConfigKey.model.startsWith('gemini-3')
const model = modelConfigKey.model ?? mockConfig.getModel();
const thinkingConfig = model.startsWith('gemini-3')
? {
thinkingLevel: ThinkingLevel.HIGH,
}
@@ -163,7 +166,7 @@ describe('GeminiChat', () => {
thinkingBudget: DEFAULT_THINKING_MODE,
};
return {
model: modelConfigKey.model,
model,
generateContentConfig: {
temperature: 0,
thinkingConfig,
@@ -1859,7 +1862,7 @@ describe('GeminiChat', () => {
} as unknown as GenerateContentResponse;
it('should use the FLASH model when in fallback mode (sendMessageStream)', async () => {
vi.mocked(mockConfig.getModel).mockReturnValue('gemini-pro');
vi.mocked(mockConfig.getModel).mockReturnValue('auto-gemini-2.5');
vi.mocked(mockConfig.isInFallbackMode).mockReturnValue(true);
vi.mocked(mockContentGenerator.generateContentStream).mockImplementation(
async () =>
@@ -1869,7 +1872,7 @@ describe('GeminiChat', () => {
);
const stream = await chat.sendMessageStream(
{ model: 'test-model' },
{ model: 'auto-gemini-2.5' },
'test message',
'prompt-id-res3',
new AbortController().signal,
@@ -1983,7 +1986,7 @@ describe('GeminiChat', () => {
expect(modelTurn.parts![0].text).toBe('Success on retry');
});
it('should switch to DEFAULT_GEMINI_FLASH_MODEL and use thinkingBudget when falling back from a gemini-3 model', async () => {
it('should switch to PREVIEW_GEMINI_FLASH_MODEL and use thinkingLevel when falling back from a gemini-3 model', async () => {
// ARRANGE
const authType = AuthType.LOGIN_WITH_GOOGLE;
vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({
@@ -2020,7 +2023,7 @@ describe('GeminiChat', () => {
// ACT
const stream = await chat.sendMessageStream(
{ model: 'gemini-3-test-model' }, // Start with a gemini-3 model
{ model: PREVIEW_GEMINI_MODEL_AUTO }, // Start with a gemini-3 model
'test fallback thinking',
'prompt-id-fb3',
new AbortController().signal,
@@ -2040,7 +2043,7 @@ describe('GeminiChat', () => {
).toHaveBeenNthCalledWith(
1,
expect.objectContaining({
model: 'gemini-3-test-model',
model: PREVIEW_GEMINI_MODEL,
config: expect.objectContaining({
thinkingConfig: {
thinkingBudget: undefined,
@@ -2051,17 +2054,17 @@ describe('GeminiChat', () => {
'prompt-id-fb3',
);
// Second call: DEFAULT_GEMINI_FLASH_MODEL (due to fallback), thinkingBudget set (due to fix)
// Second call: PREVIEW_GEMINI_FLASH_MODEL (due to fallback), thinkingLevel set
expect(
mockContentGenerator.generateContentStream,
).toHaveBeenNthCalledWith(
2,
expect.objectContaining({
model: DEFAULT_GEMINI_FLASH_MODEL,
model: PREVIEW_GEMINI_FLASH_MODEL,
config: expect.objectContaining({
thinkingConfig: {
thinkingBudget: DEFAULT_THINKING_MODE,
thinkingLevel: undefined,
thinkingBudget: undefined,
thinkingLevel: ThinkingLevel.HIGH,
},
}),
}),
+1 -5
View File
@@ -454,11 +454,7 @@ export class GeminiChat {
}
}
} else {
modelToUse = getEffectiveModel(
this.config.isInFallbackMode(),
model,
this.config.getPreviewFeatures(),
);
modelToUse = getEffectiveModel(model, this.config.isInFallbackMode());
// Preview Model Bypass Logic:
// If we are in "Preview Model Bypass Mode" (transient failure), we force downgrade to 2.5 Pro
+1 -2
View File
@@ -107,9 +107,8 @@ export function getCoreSystemPrompt(
// TODO(joshualitt): Replace with system instructions on model configs.
const desiredModel = getEffectiveModel(
config.isInFallbackMode(),
config.getModel(),
config.getPreviewFeatures(),
config.isInFallbackMode(),
);
const isGemini3 = desiredModel === PREVIEW_GEMINI_MODEL;
@@ -16,6 +16,7 @@ import {
import {
DEFAULT_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_MODEL_AUTO,
} from '../../config/models.js';
import { promptIdContext } from '../../utils/promptIdContext.js';
import type { Content } from '@google/genai';
@@ -50,6 +51,7 @@ describe('ClassifierStrategy', () => {
modelConfigService: {
getResolvedConfig: vi.fn().mockReturnValue(mockResolvedConfig),
},
getModel: () => DEFAULT_GEMINI_MODEL_AUTO,
getPreviewFeatures: () => false,
} as unknown as Config;
mockBaseLlmClient = {
@@ -12,11 +12,7 @@ import type {
RoutingDecision,
RoutingStrategy,
} from '../routingStrategy.js';
import {
GEMINI_MODEL_ALIAS_FLASH,
GEMINI_MODEL_ALIAS_PRO,
resolveModel,
} from '../../config/models.js';
import { getEffectiveModel } from '../../config/models.js';
import { createUserContent, Type } from '@google/genai';
import type { Config } from '../../config/config.js';
import {
@@ -174,10 +170,7 @@ export class ClassifierStrategy implements RoutingStrategy {
if (routerResponse.model_choice === FLASH_MODEL) {
return {
model: resolveModel(
GEMINI_MODEL_ALIAS_FLASH,
config.getPreviewFeatures(),
),
model: getEffectiveModel(config.getModel(), true),
metadata: {
source: 'Classifier',
latencyMs,
@@ -186,10 +179,7 @@ export class ClassifierStrategy implements RoutingStrategy {
};
} else {
return {
model: resolveModel(
GEMINI_MODEL_ALIAS_PRO,
config.getPreviewFeatures(),
),
model: getEffectiveModel(config.getModel(), false),
metadata: {
source: 'Classifier',
reasoning,
@@ -13,6 +13,9 @@ import {
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_FLASH_LITE_MODEL,
DEFAULT_GEMINI_MODEL_AUTO,
PREVIEW_GEMINI_MODEL_AUTO,
PREVIEW_GEMINI_FLASH_MODEL,
} from '../../config/models.js';
describe('FallbackStrategy', () => {
@@ -32,11 +35,10 @@ describe('FallbackStrategy', () => {
});
describe('when in fallback mode', () => {
it('should downgrade a pro model to the flash model', async () => {
it('should downgrade a default auto model to the flash model', async () => {
const mockConfig = {
isInFallbackMode: () => true,
getModel: () => DEFAULT_GEMINI_MODEL,
getPreviewFeatures: () => false,
getModel: () => DEFAULT_GEMINI_MODEL_AUTO,
} as Config;
const decision = await strategy.route(
@@ -51,6 +53,42 @@ describe('FallbackStrategy', () => {
expect(decision?.metadata.reasoning).toContain('In fallback mode');
});
it('should downgrade a preview auto model to the preview flash model', async () => {
const mockConfig = {
isInFallbackMode: () => true,
getModel: () => PREVIEW_GEMINI_MODEL_AUTO,
} as Config;
const decision = await strategy.route(
mockContext,
mockConfig,
mockClient,
);
expect(decision).not.toBeNull();
expect(decision?.model).toBe(PREVIEW_GEMINI_FLASH_MODEL);
expect(decision?.metadata.source).toBe('fallback');
expect(decision?.metadata.reasoning).toContain('In fallback mode');
});
it('should not downgrade a pro model to the flash model', async () => {
const mockConfig = {
isInFallbackMode: () => true,
getModel: () => DEFAULT_GEMINI_MODEL,
} as Config;
const decision = await strategy.route(
mockContext,
mockConfig,
mockClient,
);
expect(decision).not.toBeNull();
expect(decision?.model).toBe(DEFAULT_GEMINI_MODEL);
expect(decision?.metadata.source).toBe('fallback');
expect(decision?.metadata.reasoning).toContain('In fallback mode');
});
it('should honor a lite model request', async () => {
const mockConfig = {
isInFallbackMode: () => true,
@@ -28,9 +28,8 @@ export class FallbackStrategy implements RoutingStrategy {
}
const effectiveModel = getEffectiveModel(
isInFallbackMode,
config.getModel(),
config.getPreviewFeatures(),
isInFallbackMode,
);
return {
model: effectiveModel,
@@ -7,7 +7,8 @@
import type { Config } from '../../config/config.js';
import {
DEFAULT_GEMINI_MODEL_AUTO,
resolveModel,
getEffectiveModel,
PREVIEW_GEMINI_MODEL_AUTO,
} from '../../config/models.js';
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
import type {
@@ -30,11 +31,15 @@ export class OverrideStrategy implements RoutingStrategy {
const overrideModel = config.getModel();
// If the model is 'auto' we should pass to the next strategy.
if (overrideModel === DEFAULT_GEMINI_MODEL_AUTO) return null;
if (
overrideModel === DEFAULT_GEMINI_MODEL_AUTO ||
overrideModel === PREVIEW_GEMINI_MODEL_AUTO
)
return null;
// Return the overridden model name.
return {
model: resolveModel(overrideModel, config.getPreviewFeatures()),
model: getEffectiveModel(overrideModel, false),
metadata: {
source: this.name,
latencyMs: 0,