feat(core): Plumbing for late resolution of model configs. (#14597)

This commit is contained in:
joshualitt
2025-12-10 09:36:27 -08:00
committed by GitHub
parent 68ebf5d655
commit c8b688655c
10 changed files with 493 additions and 9 deletions

View File

@@ -524,6 +524,11 @@ their corresponding top-level category object in your `settings.json` file.
with (and override) the built-in aliases.
- **Default:** `{}`
- **`modelConfigs.customOverrides`** (array):
- **Description:** Custom model config overrides. These are merged with (and
added to) the built-in overrides.
- **Default:** `[]`
- **`modelConfigs.overrides`** (array):
- **Description:** Apply specific configuration overrides based on matches,
with a primary key of model (or alias). The most specific match will be

View File

@@ -742,6 +742,16 @@ const SETTINGS_SCHEMA = {
'Custom named presets for model configs. These are merged with (and override) the built-in aliases.',
showInDialog: false,
},
customOverrides: {
type: 'array',
label: 'Custom Model Config Overrides',
category: 'Model',
requiresRestart: false,
default: [],
description:
'Custom model config overrides. These are merged with (and added to) the built-in overrides.',
showInDialog: false,
},
overrides: {
type: 'array',
label: 'Model Config Overrides',

View File

@@ -1344,6 +1344,30 @@ describe('Generation Config Merging (HACK)', () => {
expect(serviceConfig.overrides).toEqual(userOverrides);
});
it('should merge default overrides when user provides only aliases', () => {
const userAliases = {
'my-alias': {
modelConfig: { model: 'my-model' },
},
};
const params: ConfigParameters = {
...baseParams,
modelConfigServiceConfig: {
aliases: userAliases,
},
};
const config = new Config(params);
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const serviceConfig = (config.modelConfigService as any).config;
// Assert that the user's aliases are present
expect(serviceConfig.aliases).toEqual(userAliases);
// Assert that the default overrides are present
expect(serviceConfig.overrides).toEqual(DEFAULT_MODEL_CONFIGS.overrides);
});
it('should use user-provided aliases if they exist', () => {
const userAliases = {
'my-alias': {

View File

@@ -629,11 +629,19 @@ export class Config {
// TODO(12593): Fix the settings loading logic to properly merge defaults and
// remove this hack.
let modelConfigServiceConfig = params.modelConfigServiceConfig;
if (modelConfigServiceConfig && !modelConfigServiceConfig.aliases) {
modelConfigServiceConfig = {
...modelConfigServiceConfig,
aliases: DEFAULT_MODEL_CONFIGS.aliases,
};
if (modelConfigServiceConfig) {
if (!modelConfigServiceConfig.aliases) {
modelConfigServiceConfig = {
...modelConfigServiceConfig,
aliases: DEFAULT_MODEL_CONFIGS.aliases,
};
}
if (!modelConfigServiceConfig.overrides) {
modelConfigServiceConfig = {
...modelConfigServiceConfig,
overrides: DEFAULT_MODEL_CONFIGS.overrides,
};
}
}
this.modelConfigService = new ModelConfigService(

View File

@@ -209,4 +209,14 @@ export const DEFAULT_MODEL_CONFIGS: ModelConfigServiceConfig = {
},
},
},
overrides: [
{
match: { model: 'chat-base', isRetry: true },
modelConfig: {
generateContentConfig: {
temperature: 1,
},
},
},
],
};

View File

@@ -18,6 +18,14 @@ const GOLDEN_FILE_PATH = path.resolve(
'resolved-aliases.golden.json',
);
const RETRY_GOLDEN_FILE_PATH = path.resolve(
process.cwd(),
'src',
'services',
'test-data',
'resolved-aliases-retry.golden.json',
);
describe('ModelConfigService Golden Test', () => {
it('should match the golden file for resolved default aliases', async () => {
const service = new ModelConfigService(DEFAULT_MODEL_CONFIGS);
@@ -60,4 +68,49 @@ describe('ModelConfigService Golden Test', () => {
'Golden file mismatch. If the new resolved aliases are correct, run the test with `UPDATE_GOLDENS=true` to regenerate the golden file.',
).toEqual(goldenData);
});
it('should match the golden file for resolved default aliases with isRetry=true', async () => {
const service = new ModelConfigService(DEFAULT_MODEL_CONFIGS);
const aliases = Object.keys(DEFAULT_MODEL_CONFIGS.aliases ?? {});
const resolvedAliases: Record<string, unknown> = {};
for (const alias of aliases) {
resolvedAliases[alias] =
// eslint-disable-next-line @typescript-eslint/no-explicit-any
(service as any).internalGetResolvedConfig({
model: alias,
isRetry: true,
});
}
if (process.env['UPDATE_GOLDENS']) {
await fs.mkdir(path.dirname(RETRY_GOLDEN_FILE_PATH), { recursive: true });
await fs.writeFile(
RETRY_GOLDEN_FILE_PATH,
JSON.stringify(resolvedAliases, null, 2),
'utf-8',
);
// In update mode, we pass the test after writing the file.
return;
}
let goldenContent: string;
try {
goldenContent = await fs.readFile(RETRY_GOLDEN_FILE_PATH, 'utf-8');
} catch (e) {
if ((e as NodeJS.ErrnoException).code === 'ENOENT') {
throw new Error(
'Golden file not found. Run with `UPDATE_GOLDENS=true` to create it.',
);
}
throw e;
}
const goldenData = JSON.parse(goldenContent);
expect(
resolvedAliases,
'Golden file mismatch. If the new resolved aliases are correct, run the test with `UPDATE_GOLDENS=true` to regenerate the golden file.',
).toEqual(goldenData);
});
});

View File

@@ -697,4 +697,122 @@ describe('ModelConfigService', () => {
});
});
});
describe('custom overrides', () => {
it('should apply custom overrides on top of defaults', () => {
const config: ModelConfigServiceConfig = {
aliases: {
'test-alias': {
modelConfig: {
model: 'gemini-test',
generateContentConfig: { temperature: 0.5 },
},
},
},
overrides: [
{
match: { model: 'test-alias' },
modelConfig: { generateContentConfig: { temperature: 0.6 } },
},
],
customOverrides: [
{
match: { model: 'test-alias' },
modelConfig: { generateContentConfig: { temperature: 0.7 } },
},
],
};
const service = new ModelConfigService(config);
const resolved = service.getResolvedConfig({ model: 'test-alias' });
// Custom overrides should be appended to overrides, so they win
expect(resolved.generateContentConfig.temperature).toBe(0.7);
});
});
describe('retry behavior', () => {
it('should apply retry-specific overrides when isRetry is true', () => {
const config: ModelConfigServiceConfig = {
aliases: {
'test-model': {
modelConfig: {
model: 'gemini-test',
generateContentConfig: {
temperature: 0.5,
},
},
},
},
overrides: [
{
match: { model: 'test-model', isRetry: true },
modelConfig: {
generateContentConfig: {
temperature: 1.0,
},
},
},
],
};
const service = new ModelConfigService(config);
// Normal request
const normal = service.getResolvedConfig({ model: 'test-model' });
expect(normal.generateContentConfig.temperature).toBe(0.5);
// Retry request
const retry = service.getResolvedConfig({
model: 'test-model',
isRetry: true,
});
expect(retry.generateContentConfig.temperature).toBe(1.0);
});
it('should prioritize retry overrides over generic overrides', () => {
const config: ModelConfigServiceConfig = {
aliases: {
'test-model': {
modelConfig: {
model: 'gemini-test',
generateContentConfig: {
temperature: 0.5,
},
},
},
},
overrides: [
// Generic override for this model
{
match: { model: 'test-model' },
modelConfig: {
generateContentConfig: {
temperature: 0.7,
},
},
},
// Retry-specific override
{
match: { model: 'test-model', isRetry: true },
modelConfig: {
generateContentConfig: {
temperature: 1.0,
},
},
},
],
};
const service = new ModelConfigService(config);
// Normal request - hits generic override
const normal = service.getResolvedConfig({ model: 'test-model' });
expect(normal.generateContentConfig.temperature).toBe(0.7);
// Retry request - hits retry override (more specific)
const retry = service.getResolvedConfig({
model: 'test-model',
isRetry: true,
});
expect(retry.generateContentConfig.temperature).toBe(1.0);
});
});
});

View File

@@ -21,6 +21,11 @@ export interface ModelConfigKey {
// model calls made by this specific subagent, and no others, while still
// ensuring model configs are fully orthogonal to the agents who use them.
overrideScope?: string;
// Indicates whether this configuration request is happening during a retry attempt.
// This allows overrides to specify different settings (e.g., higher temperature)
// specifically for retry scenarios.
isRetry?: boolean;
}
export interface ModelConfig {
@@ -32,6 +37,7 @@ export interface ModelConfigOverride {
match: {
model?: string; // Can be a model name or an alias
overrideScope?: string;
isRetry?: boolean;
};
modelConfig: ModelConfig;
}
@@ -45,6 +51,7 @@ export interface ModelConfigServiceConfig {
aliases?: Record<string, ModelConfigAlias>;
customAliases?: Record<string, ModelConfigAlias>;
overrides?: ModelConfigOverride[];
customOverrides?: ModelConfigOverride[];
}
export type ResolvedModelConfig = _ResolvedModelConfig & {
@@ -105,12 +112,18 @@ export class ModelConfigService {
generateContentConfig: GenerateContentConfig;
} {
const config = this.config || {};
const { aliases = {}, customAliases = {}, overrides = [] } = config;
const {
aliases = {},
customAliases = {},
overrides = [],
customOverrides = [],
} = config;
const allAliases = {
...aliases,
...customAliases,
...this.runtimeAliases,
};
const allOverrides = [...overrides, ...customOverrides];
let baseModel: string | undefined = context.model;
let resolvedConfig: GenerateContentConfig = {};
@@ -135,7 +148,7 @@ export class ModelConfigService {
};
// Step 2: Override Application
const matches = overrides
const matches = allOverrides
.map((override, index) => {
const matchEntries = Object.entries(override.match);
if (matchEntries.length === 0) {

View File

@@ -0,0 +1,222 @@
{
"base": {
"generateContentConfig": {
"temperature": 0,
"topP": 1
}
},
"chat-base": {
"generateContentConfig": {
"temperature": 1,
"topP": 0.95,
"thinkingConfig": {
"includeThoughts": true
},
"topK": 64
}
},
"chat-base-2.5": {
"generateContentConfig": {
"temperature": 1,
"topP": 0.95,
"thinkingConfig": {
"includeThoughts": true,
"thinkingBudget": 8192
},
"topK": 64
}
},
"chat-base-3": {
"generateContentConfig": {
"temperature": 1,
"topP": 0.95,
"thinkingConfig": {
"includeThoughts": true,
"thinkingLevel": "HIGH"
},
"topK": 64
}
},
"gemini-3-pro-preview": {
"model": "gemini-3-pro-preview",
"generateContentConfig": {
"temperature": 1,
"topP": 0.95,
"thinkingConfig": {
"includeThoughts": true,
"thinkingLevel": "HIGH"
},
"topK": 64
}
},
"gemini-2.5-pro": {
"model": "gemini-2.5-pro",
"generateContentConfig": {
"temperature": 1,
"topP": 0.95,
"thinkingConfig": {
"includeThoughts": true,
"thinkingBudget": 8192
},
"topK": 64
}
},
"gemini-2.5-flash": {
"model": "gemini-2.5-flash",
"generateContentConfig": {
"temperature": 1,
"topP": 0.95,
"thinkingConfig": {
"includeThoughts": true,
"thinkingBudget": 8192
},
"topK": 64
}
},
"gemini-2.5-flash-lite": {
"model": "gemini-2.5-flash-lite",
"generateContentConfig": {
"temperature": 1,
"topP": 0.95,
"thinkingConfig": {
"includeThoughts": true,
"thinkingBudget": 8192
},
"topK": 64
}
},
"gemini-2.5-flash-base": {
"model": "gemini-2.5-flash",
"generateContentConfig": {
"temperature": 0,
"topP": 1
}
},
"classifier": {
"model": "gemini-2.5-flash-lite",
"generateContentConfig": {
"temperature": 0,
"topP": 1,
"maxOutputTokens": 1024,
"thinkingConfig": {
"thinkingBudget": 512
}
}
},
"prompt-completion": {
"model": "gemini-2.5-flash-lite",
"generateContentConfig": {
"temperature": 0.3,
"topP": 1,
"maxOutputTokens": 16000,
"thinkingConfig": {
"thinkingBudget": 0
}
}
},
"edit-corrector": {
"model": "gemini-2.5-flash-lite",
"generateContentConfig": {
"temperature": 0,
"topP": 1,
"thinkingConfig": {
"thinkingBudget": 0
}
}
},
"summarizer-default": {
"model": "gemini-2.5-flash-lite",
"generateContentConfig": {
"temperature": 0,
"topP": 1,
"maxOutputTokens": 2000
}
},
"summarizer-shell": {
"model": "gemini-2.5-flash-lite",
"generateContentConfig": {
"temperature": 0,
"topP": 1,
"maxOutputTokens": 2000
}
},
"web-search": {
"model": "gemini-2.5-flash",
"generateContentConfig": {
"temperature": 0,
"topP": 1,
"tools": [
{
"googleSearch": {}
}
]
}
},
"web-fetch": {
"model": "gemini-2.5-flash",
"generateContentConfig": {
"temperature": 0,
"topP": 1,
"tools": [
{
"urlContext": {}
}
]
}
},
"web-fetch-fallback": {
"model": "gemini-2.5-flash",
"generateContentConfig": {
"temperature": 0,
"topP": 1
}
},
"loop-detection": {
"model": "gemini-2.5-flash",
"generateContentConfig": {
"temperature": 0,
"topP": 1
}
},
"loop-detection-double-check": {
"model": "gemini-2.5-pro",
"generateContentConfig": {
"temperature": 0,
"topP": 1
}
},
"llm-edit-fixer": {
"model": "gemini-2.5-flash",
"generateContentConfig": {
"temperature": 0,
"topP": 1
}
},
"next-speaker-checker": {
"model": "gemini-2.5-flash",
"generateContentConfig": {
"temperature": 0,
"topP": 1
}
},
"chat-compression-3-pro": {
"model": "gemini-3-pro-preview",
"generateContentConfig": {}
},
"chat-compression-2.5-pro": {
"model": "gemini-2.5-pro",
"generateContentConfig": {}
},
"chat-compression-2.5-flash": {
"model": "gemini-2.5-flash",
"generateContentConfig": {}
},
"chat-compression-2.5-flash-lite": {
"model": "gemini-2.5-flash-lite",
"generateContentConfig": {}
},
"chat-compression-default": {
"model": "gemini-2.5-pro",
"generateContentConfig": {}
}
}

File diff suppressed because one or more lines are too long