Add ModelChain support to ModelConfigService and make ModelDialog dynamic (#22914)

This commit is contained in:
kevinjwang1
2026-03-19 15:22:26 -07:00
committed by GitHub
parent 0e66f545ca
commit 06a7873c51
11 changed files with 1014 additions and 18 deletions
@@ -19,6 +19,8 @@ import {
PREVIEW_GEMINI_3_1_MODEL,
} from '../config/models.js';
import { AuthType } from '../core/contentGenerator.js';
import { ModelConfigService } from '../services/modelConfigService.js';
import { DEFAULT_MODEL_CONFIGS } from '../config/defaultModelConfigs.js';
const createMockConfig = (overrides: Partial<Config> = {}): Config => {
const config = {
@@ -163,6 +165,66 @@ describe('policyHelpers', () => {
});
});
describe('resolvePolicyChain behavior is identical between dynamic and legacy implementations', () => {
const testCases = [
{ name: 'Default Auto', model: DEFAULT_GEMINI_MODEL_AUTO },
{ name: 'Gemini 3 Auto', model: 'auto-gemini-3' },
{ name: 'Flash Lite', model: DEFAULT_GEMINI_FLASH_LITE_MODEL },
{
name: 'Gemini 3 Auto (3.1 Enabled)',
model: 'auto-gemini-3',
useGemini31: true,
},
{
name: 'Gemini 3 Auto (3.1 + Custom Tools)',
model: 'auto-gemini-3',
useGemini31: true,
authType: AuthType.USE_GEMINI,
},
{
name: 'Gemini 3 Auto (No Access)',
model: 'auto-gemini-3',
hasAccess: false,
},
{ name: 'Concrete Model (2.5 Pro)', model: 'gemini-2.5-pro' },
{ name: 'Custom Model', model: 'my-custom-model' },
{
name: 'Wrap Around',
model: DEFAULT_GEMINI_MODEL_AUTO,
wrapsAround: true,
},
];
testCases.forEach(
({ name, model, useGemini31, hasAccess, authType, wrapsAround }) => {
it(`achieves parity for: ${name}`, () => {
const createBaseConfig = (dynamic: boolean) =>
createMockConfig({
getExperimentalDynamicModelConfiguration: () => dynamic,
getModel: () => model,
getGemini31LaunchedSync: () => useGemini31 ?? false,
getHasAccessToPreviewModel: () => hasAccess ?? true,
getContentGeneratorConfig: () => ({ authType }),
modelConfigService: new ModelConfigService(DEFAULT_MODEL_CONFIGS),
});
const legacyChain = resolvePolicyChain(
createBaseConfig(false),
model,
wrapsAround,
);
const dynamicChain = resolvePolicyChain(
createBaseConfig(true),
model,
wrapsAround,
);
expect(dynamicChain).toEqual(legacyChain);
});
},
);
});
describe('buildFallbackPolicyContext', () => {
it('returns remaining candidates after the failed model', () => {
const chain = [
@@ -53,12 +53,57 @@ export function resolvePolicyChain(
useGemini31,
useCustomToolModel,
hasAccessToPreview,
config,
);
const isAutoPreferred = preferredModel
? isAutoModel(preferredModel, config)
: false;
const isAutoConfigured = isAutoModel(configuredModel, config);
// --- DYNAMIC PATH ---
if (config.getExperimentalDynamicModelConfiguration?.() === true) {
const context = {
useGemini3_1: useGemini31,
useCustomTools: useCustomToolModel,
};
if (resolvedModel === DEFAULT_GEMINI_FLASH_LITE_MODEL) {
chain = config.modelConfigService.resolveChain('lite', context);
} else if (
isGemini3Model(resolvedModel, config) ||
isAutoModel(preferredModel ?? '', config) ||
isAutoModel(configuredModel, config)
) {
// 1. Try to find a chain specifically for the current configured alias
if (
isAutoModel(configuredModel, config) &&
config.modelConfigService.getModelChain(configuredModel)
) {
chain = config.modelConfigService.resolveChain(
configuredModel,
context,
);
}
// 2. Fallback to family-based auto-routing
if (!chain) {
const previewEnabled =
hasAccessToPreview &&
(isGemini3Model(resolvedModel, config) ||
preferredModel === PREVIEW_GEMINI_MODEL_AUTO ||
configuredModel === PREVIEW_GEMINI_MODEL_AUTO);
const chainKey = previewEnabled ? 'preview' : 'default';
chain = config.modelConfigService.resolveChain(chainKey, context);
}
}
if (!chain) {
// No matching modelChains found, default to single model chain
chain = createSingleModelChain(modelFromConfig);
}
return applyDynamicSlicing(chain, resolvedModel, wrapsAround);
}
// --- LEGACY PATH ---
if (resolvedModel === DEFAULT_GEMINI_FLASH_LITE_MODEL) {
chain = getFlashLitePolicyChain();
} else if (
@@ -90,7 +135,17 @@ export function resolvePolicyChain(
} else {
chain = createSingleModelChain(modelFromConfig);
}
return applyDynamicSlicing(chain, resolvedModel, wrapsAround);
}
/**
* Applies active-index slicing and wrap-around logic to a chain template.
*/
function applyDynamicSlicing(
chain: ModelPolicy[],
resolvedModel: string,
wrapsAround: boolean,
): ModelPolicyChain {
const activeIndex = chain.findIndex(
(policy) => policy.model === resolvedModel,
);
+5
View File
@@ -994,6 +994,10 @@ export class Config implements McpContext, AgentLoopContext {
...DEFAULT_MODEL_CONFIGS.classifierIdResolutions,
...modelConfigServiceConfig.classifierIdResolutions,
};
const mergedModelChains = {
...DEFAULT_MODEL_CONFIGS.modelChains,
...modelConfigServiceConfig.modelChains,
};
modelConfigServiceConfig = {
// Preserve other user settings like customAliases
@@ -1007,6 +1011,7 @@ export class Config implements McpContext, AgentLoopContext {
modelDefinitions: mergedModelDefinitions,
modelIdResolutions: mergedModelIdResolutions,
classifierIdResolutions: mergedClassifierIdResolutions,
modelChains: mergedModelChains,
};
}
+145 -1
View File
@@ -251,6 +251,13 @@ export const DEFAULT_MODEL_CONFIGS: ModelConfigServiceConfig = {
],
modelDefinitions: {
// Concrete Models
'gemini-3.1-flash-lite-preview': {
tier: 'flash-lite',
family: 'gemini-3',
isPreview: true,
isVisible: true,
features: { thinking: false, multimodalToolUse: true },
},
'gemini-3.1-pro-preview': {
tier: 'pro',
family: 'gemini-3',
@@ -331,7 +338,7 @@ export const DEFAULT_MODEL_CONFIGS: ModelConfigServiceConfig = {
isPreview: true,
isVisible: true,
dialogDescription:
'Let Gemini CLI decide the best model for the task: gemini-3.1-pro, gemini-3-flash',
'Let Gemini CLI decide the best model for the task: gemini-3-pro, gemini-3-flash',
features: { thinking: true, multimodalToolUse: false },
},
'auto-gemini-2.5': {
@@ -345,6 +352,27 @@ export const DEFAULT_MODEL_CONFIGS: ModelConfigServiceConfig = {
},
},
modelIdResolutions: {
'gemini-3.1-pro-preview': {
default: 'gemini-3.1-pro-preview',
contexts: [
{ condition: { hasAccessToPreview: false }, target: 'gemini-2.5-pro' },
],
},
'gemini-3.1-pro-preview-customtools': {
default: 'gemini-3.1-pro-preview-customtools',
contexts: [
{ condition: { hasAccessToPreview: false }, target: 'gemini-2.5-pro' },
],
},
'gemini-3-flash-preview': {
default: 'gemini-3-flash-preview',
contexts: [
{
condition: { hasAccessToPreview: false },
target: 'gemini-2.5-flash',
},
],
},
'gemini-3-pro-preview': {
default: 'gemini-3-pro-preview',
contexts: [
@@ -451,4 +479,120 @@ export const DEFAULT_MODEL_CONFIGS: ModelConfigServiceConfig = {
],
},
},
modelChains: {
preview: [
{
model: 'gemini-3-pro-preview',
actions: {
terminal: 'prompt',
transient: 'prompt',
not_found: 'prompt',
unknown: 'prompt',
},
stateTransitions: {
terminal: 'terminal',
transient: 'terminal',
not_found: 'terminal',
unknown: 'terminal',
},
},
{
model: 'gemini-3-flash-preview',
isLastResort: true,
actions: {
terminal: 'prompt',
transient: 'prompt',
not_found: 'prompt',
unknown: 'prompt',
},
stateTransitions: {
terminal: 'terminal',
transient: 'terminal',
not_found: 'terminal',
unknown: 'terminal',
},
},
],
default: [
{
model: 'gemini-2.5-pro',
actions: {
terminal: 'prompt',
transient: 'prompt',
not_found: 'prompt',
unknown: 'prompt',
},
stateTransitions: {
terminal: 'terminal',
transient: 'terminal',
not_found: 'terminal',
unknown: 'terminal',
},
},
{
model: 'gemini-2.5-flash',
isLastResort: true,
actions: {
terminal: 'prompt',
transient: 'prompt',
not_found: 'prompt',
unknown: 'prompt',
},
stateTransitions: {
terminal: 'terminal',
transient: 'terminal',
not_found: 'terminal',
unknown: 'terminal',
},
},
],
lite: [
{
model: 'gemini-2.5-flash-lite',
actions: {
terminal: 'silent',
transient: 'silent',
not_found: 'silent',
unknown: 'silent',
},
stateTransitions: {
terminal: 'terminal',
transient: 'terminal',
not_found: 'terminal',
unknown: 'terminal',
},
},
{
model: 'gemini-2.5-flash',
actions: {
terminal: 'silent',
transient: 'silent',
not_found: 'silent',
unknown: 'silent',
},
stateTransitions: {
terminal: 'terminal',
transient: 'terminal',
not_found: 'terminal',
unknown: 'terminal',
},
},
{
model: 'gemini-2.5-pro',
isLastResort: true,
actions: {
terminal: 'silent',
transient: 'silent',
not_found: 'silent',
unknown: 'silent',
},
stateTransitions: {
terminal: 'terminal',
transient: 'terminal',
not_found: 'terminal',
unknown: 'terminal',
},
},
],
},
};
-8
View File
@@ -190,14 +190,6 @@ describe('Dynamic Configuration Parity', () => {
}
});
it('supportsModernFeatures should match legacy behavior', () => {
for (const model of modelsToTest) {
const legacy = supportsModernFeatures(model);
const dynamic = supportsModernFeatures(model);
expect(dynamic).toBe(legacy);
}
});
it('supportsMultimodalFunctionResponse should match legacy behavior', () => {
for (const model of modelsToTest) {
const legacy = supportsMultimodalFunctionResponse(model, legacyConfig);
+14 -1
View File
@@ -102,11 +102,24 @@ export function resolveModel(
config?: ModelCapabilityContext,
): string {
if (config?.getExperimentalDynamicModelConfiguration?.() === true) {
return config.modelConfigService.resolveModelId(requestedModel, {
const resolved = config.modelConfigService.resolveModelId(requestedModel, {
useGemini3_1,
useCustomTools: useCustomToolModel,
hasAccessToPreview,
});
if (!hasAccessToPreview && isPreviewModel(resolved, config)) {
// Fallback for unknown preview models.
if (resolved.includes('flash-lite')) {
return DEFAULT_GEMINI_FLASH_LITE_MODEL;
}
if (resolved.includes('flash')) {
return DEFAULT_GEMINI_FLASH_MODEL;
}
return DEFAULT_GEMINI_MODEL;
}
return resolved;
}
let resolved: string;
@@ -5,6 +5,7 @@
*/
import type { GenerateContentConfig } from '@google/genai';
import type { ModelPolicy } from '../availability/modelPolicy.js';
// The primary key for the ModelConfig is the model string. However, we also
// support a secondary key to limit the override scope, typically an agent name.
@@ -111,6 +112,7 @@ export interface ModelConfigServiceConfig {
modelDefinitions?: Record<string, ModelDefinition>;
modelIdResolutions?: Record<string, ModelResolution>;
classifierIdResolutions?: Record<string, ModelResolution>;
modelChains?: Record<string, ModelPolicy[]>;
}
const MAX_ALIAS_CHAIN_DEPTH = 100;
@@ -221,6 +223,29 @@ export class ModelConfigService {
return resolution.default;
}
getModelChain(chainName: string): ModelPolicy[] | undefined {
return this.config.modelChains?.[chainName];
}
/**
* Fetches a chain template and resolves all model IDs within it
* based on the provided context.
*/
resolveChain(
chainName: string,
context: ResolutionContext = {},
): ModelPolicy[] | undefined {
const template = this.config.modelChains?.[chainName];
if (!template) {
return undefined;
}
// Map through the template and resolve each model ID
return template.map((policy) => ({
...policy,
model: this.resolveModelId(policy.model, context),
}));
}
registerRuntimeModelConfig(aliasName: string, alias: ModelConfigAlias): void {
this.runtimeAliases[aliasName] = alias;
}