diff --git a/docs/cli/model-routing.md b/docs/cli/model-routing.md index a29dd98a9b9..7f75ebabe81 100644 --- a/docs/cli/model-routing.md +++ b/docs/cli/model-routing.md @@ -26,6 +26,37 @@ policies. the CLI will use an available fallback model for the current turn or the remainder of the session. +### Numerical Routing Configuration + +When numerical routing is active, Gemini CLI scores request complexity from 1 +to 100. By default, the CLI uses a threshold to choose between Flash and Pro +models. You can customize this routing logic using `routing.numericalRules` in +your `settings.json` file. + +This allows you to define score bands and map them to specific model IDs or +aliases: + +```json +{ + "routing": { + "numericalRules": [ + { "maxScore": 10, "model": "local-gemma" }, + { "maxScore": 30, "model": "gemini-2.5-flash" }, + { "maxScore": 80, "model": "gemini-3-flash-preview" }, + { "maxScore": 100, "model": "gemini-3-pro-preview" } + ] + } +} +``` + +- **`maxScore`**: The upper complexity score boundary (inclusive, 1–100) for + which this rule applies. +- **`model`**: The model ID or alias to route to. + +Rules are evaluated in ascending order of their `maxScore`. If a request's score +exceeds all configured `maxScore` bands, it falls back to the model configured +in the final score band. + ### Local Model Routing (Experimental) Gemini CLI supports using a local model for routing decisions. When configured, diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index 2153b4d07ba..8206352f6d6 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -94,6 +94,15 @@ their corresponding top-level category object in your `settings.json` file. +#### `routing` + +- **`routing.numericalRules`** (array): + - **Description:** Score bands matched in ascending maxScore order. Each + maxScore is inclusive. Scores above all configured bands use the model in + the final band. + - **Default:** `[]` + - **Requires restart:** Yes + #### `policyPaths` - **`policyPaths`** (array): diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 3c6672e755b..5b4fb672d92 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -1104,6 +1104,7 @@ export async function loadCliConfig( format: (argv.outputFormat ?? settings.output?.format) as OutputFormat, }, gemmaModelRouter: settings.experimental?.gemmaModelRouter, + routing: settings.routing, adk: settings.experimental?.adk, fakeResponses: argv.fakeResponses, fakeResponsesNonStrict: argv.fakeResponsesNonStrict, diff --git a/packages/cli/src/config/settings-validation.test.ts b/packages/cli/src/config/settings-validation.test.ts index 01cd545cdea..41e86ec611f 100644 --- a/packages/cli/src/config/settings-validation.test.ts +++ b/packages/cli/src/config/settings-validation.test.ts @@ -538,6 +538,92 @@ describe('settings-validation', () => { expect(formatted).toContain('...and 1 more errors.'); } }); + + it('should accept valid routing.numericalRules configuration', () => { + const validSettings = { + routing: { + numericalRules: [ + { maxScore: 10, model: 'local-gemma' }, + { maxScore: 30, model: 'gemini-2.5-flash' }, + { maxScore: 80, model: 'gemini-3-flash-preview' }, + { maxScore: 100, model: 'gemini-3-pro-preview' }, + ], + }, + }; + + const result = validateSettings(validSettings); + expect(result.success).toBe(true); + }); + + it('should accept empty routing.numericalRules array', () => { + const validSettings = { + routing: { + numericalRules: [], + }, + }; + + const result = validateSettings(validSettings); + expect(result.success).toBe(true); + }); + + it('should reject routing when it is not an object', () => { + const invalidSettings = { + routing: 'invalid-string', + }; + + const result = validateSettings(invalidSettings); + expect(result.success).toBe(false); + expect(result.error?.issues[0]?.path).toEqual(['routing']); + }); + + it('should reject routing.numericalRules when it is not an array', () => { + const invalidSettings = { + routing: { + numericalRules: 'not-an-array', + }, + }; + + const result = validateSettings(invalidSettings); + expect(result.success).toBe(false); + expect(result.error?.issues[0]?.path).toEqual([ + 'routing', + 'numericalRules', + ]); + }); + + it('should reject routing.numericalRules when maxScore is not a number', () => { + const invalidSettings = { + routing: { + numericalRules: [{ maxScore: 'not-a-number', model: 'flash' }], + }, + }; + + const result = validateSettings(invalidSettings); + expect(result.success).toBe(false); + expect(result.error?.issues[0]?.path).toEqual([ + 'routing', + 'numericalRules', + 0, + 'maxScore', + ]); + }); + + it('should reject routing.numericalRules when model is not a string', () => { + const invalidSettings = { + routing: { + numericalRules: [{ maxScore: 50, model: 12345 }], + }, + }; + + const result = validateSettings(invalidSettings); + expect(result.success).toBe(false); + expect(result.error?.issues[0]?.path).toEqual([ + 'routing', + 'numericalRules', + 0, + 'model', + ]); + }); }); describe('settingsZodSchema', () => { diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index 442b07b106f..c85a189961f 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -17,6 +17,7 @@ import { type MCPServerConfig, type RequiredMcpServerConfig, type BugCommandSettings, + type NumericalRoutingRule, type TelemetrySettings, type AuthType, type AgentOverride, @@ -173,6 +174,49 @@ const SETTINGS_SCHEMA = { }, }, + routing: { + type: 'object', + label: 'Routing', + category: 'Model', + requiresRestart: true, + default: {}, + description: 'Configure numerical complexity score routing.', + showInDialog: false, + properties: { + numericalRules: { + type: 'array', + label: 'Numerical Routing Rules', + category: 'Model', + requiresRestart: true, + default: [] as NumericalRoutingRule[], + description: + 'Score bands matched in ascending maxScore order. Each maxScore is inclusive. Scores above all configured bands use the model in the final band.', + showInDialog: false, + items: { + type: 'object', + properties: { + maxScore: { + type: 'number', + label: 'Maximum Score', + category: 'Model', + requiresRestart: true, + default: 100, + description: 'Maximum score for this routing rule.', + }, + model: { + type: 'string', + label: 'Model ID', + category: 'Model', + requiresRestart: true, + default: '', + description: 'Model alias or ID to route to.', + }, + }, + }, + }, + }, + }, + policyPaths: pathArraySetting( 'Policy Paths', 'Additional policy files or directories to load.', diff --git a/packages/cli/src/test-utils/mockConfig.ts b/packages/cli/src/test-utils/mockConfig.ts index 94f15d8ac4e..282bb874e8c 100644 --- a/packages/cli/src/test-utils/mockConfig.ts +++ b/packages/cli/src/test-utils/mockConfig.ts @@ -123,6 +123,8 @@ export const createMockConfig = (overrides: Partial = {}): Config => getCompressionThreshold: vi.fn().mockResolvedValue(undefined), getUserCaching: vi.fn().mockResolvedValue(false), getNumericalRoutingEnabled: vi.fn().mockResolvedValue(false), + hasCustomNumericalRoutingRules: vi.fn().mockReturnValue(false), + getNumericalRoutingRules: vi.fn().mockReturnValue([]), getClassifierThreshold: vi.fn().mockResolvedValue(undefined), getBannerTextNoCapacityIssues: vi.fn().mockResolvedValue(''), getBannerTextCapacityIssues: vi.fn().mockResolvedValue(''), diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 747ba4abf13..480b6e710b8 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -697,6 +697,40 @@ describe('Server Config (config.ts)', () => { }); }); + describe('numericalRoutingRules', () => { + it('should return empty rules and hasCustomNumericalRoutingRules false when not configured', () => { + const config = new Config(baseParams); + expect(config.hasCustomNumericalRoutingRules()).toBe(false); + expect(config.getNumericalRoutingRules()).toEqual([]); + }); + + it('should return configured rules and hasCustomNumericalRoutingRules true', () => { + const rules = [ + { maxScore: 20, model: 'gemini-2.5-flash' }, + { maxScore: 80, model: 'gemini-3-pro' }, + ]; + const config = new Config({ + ...baseParams, + routing: { + numericalRules: rules, + }, + }); + expect(config.hasCustomNumericalRoutingRules()).toBe(true); + expect(config.getNumericalRoutingRules()).toEqual(rules); + }); + + it('should return false for hasCustomNumericalRoutingRules when numericalRules is empty', () => { + const config = new Config({ + ...baseParams, + routing: { + numericalRules: [], + }, + }); + expect(config.hasCustomNumericalRoutingRules()).toBe(false); + expect(config.getNumericalRoutingRules()).toEqual([]); + }); + }); + describe('getGemini31LaunchedSync', () => { it.each([AuthType.USE_GEMINI, AuthType.USE_VERTEX_AI, AuthType.GATEWAY])( 'should return true for %s', diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 7f4b3b02dc3..612350194dc 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -237,6 +237,19 @@ export interface OutputSettings { format?: OutputFormat; } +export interface NumericalRoutingRule { + /** Inclusive upper bound for the complexity score (1-100). */ + maxScore: number; + /** Model ID or configured model alias to use for this score band. */ + model: string; +} + +export interface RoutingSettings { + numericalRules?: NumericalRoutingRule[]; +} + +export type NumericalRoutingSettings = RoutingSettings; + export interface GemmaModelRouterSettings { enabled?: boolean; autoStartServer?: boolean; @@ -695,6 +708,7 @@ export interface ConfigParameters { directWebFetch?: boolean; policyUpdateConfirmationRequest?: PolicyUpdateConfirmationRequest; output?: OutputSettings; + routing?: RoutingSettings; gemmaModelRouter?: GemmaModelRouterSettings; adk?: ADKSettings; disableModelRouterForAuth?: AuthType[]; @@ -930,6 +944,7 @@ export class Config implements McpContext, AgentLoopContext { | undefined; private readonly outputSettings: OutputSettings; + private readonly routingSettings?: RoutingSettings; private readonly gemmaModelRouter: GemmaModelRouterSettings; private readonly agentSessionNoninteractiveEnabled: boolean; private readonly agentSessionInteractiveEnabled: boolean; @@ -1370,6 +1385,7 @@ export class Config implements McpContext, AgentLoopContext { this.outputSettings = { format: params.output?.format ?? OutputFormat.TEXT, }; + this.routingSettings = params.routing; this.gemmaModelRouter = { enabled: params.gemmaModelRouter?.enabled ?? false, autoStartServer: params.gemmaModelRouter?.autoStartServer ?? true, @@ -3920,6 +3936,14 @@ export class Config implements McpContext, AgentLoopContext { return this.gemmaModelRouter; } + getNumericalRoutingRules(): NumericalRoutingRule[] { + return this.routingSettings?.numericalRules ?? []; + } + + hasCustomNumericalRoutingRules(): boolean { + return this.getNumericalRoutingRules().length > 0; + } + getAgentSessionNoninteractiveEnabled(): boolean { return ( process.env['GEMINI_CLI_EXP_AGENT'] === 'true' || diff --git a/packages/core/src/routing/modelRouterService.test.ts b/packages/core/src/routing/modelRouterService.test.ts index 4e0c32c62ff..8c7873a63fc 100644 --- a/packages/core/src/routing/modelRouterService.test.ts +++ b/packages/core/src/routing/modelRouterService.test.ts @@ -55,6 +55,9 @@ describe('ModelRouterService', () => { mockLocalLiteRtLmClient, ); vi.spyOn(mockConfig, 'getNumericalRoutingEnabled').mockResolvedValue(true); + vi.spyOn(mockConfig, 'hasCustomNumericalRoutingRules').mockReturnValue( + false, + ); vi.spyOn(mockConfig, 'getResolvedClassifierThreshold').mockResolvedValue( 90, ); @@ -144,6 +147,25 @@ describe('ModelRouterService', () => { expect(compositeStrategyArgs[1]).toBe('agent-router'); }); + it('should prioritize numerical routing when custom rules are configured', () => { + vi.spyOn(mockConfig, 'hasCustomNumericalRoutingRules').mockReturnValue( + true, + ); + vi.mocked(CompositeStrategy).mockClear(); + + service = new ModelRouterService(mockConfig); + + const compositeStrategyArgs = vi.mocked(CompositeStrategy).mock.calls[0]; + const childStrategies = compositeStrategyArgs[0]; + + expect(childStrategies[0]).toBeInstanceOf(FallbackStrategy); + expect(childStrategies[1]).toBeInstanceOf(OverrideStrategy); + expect(childStrategies[2]).toBeInstanceOf(ApprovalModeStrategy); + expect(childStrategies[3]).toBeInstanceOf(NumericalClassifierStrategy); + expect(childStrategies[4]).toBeInstanceOf(ClassifierStrategy); + expect(childStrategies[5]).toBeInstanceOf(DefaultStrategy); + }); + describe('route()', () => { const strategyDecision: RoutingDecision = { model: 'strategy-chosen-model', diff --git a/packages/core/src/routing/modelRouterService.ts b/packages/core/src/routing/modelRouterService.ts index 30e2bb9f8d5..b42de057743 100644 --- a/packages/core/src/routing/modelRouterService.ts +++ b/packages/core/src/routing/modelRouterService.ts @@ -46,6 +46,13 @@ export class ModelRouterService { // Approval mode is next. strategies.push(new ApprovalModeStrategy()); + const hasCustomNumericalRules = + this.config.hasCustomNumericalRoutingRules?.() ?? false; + if (hasCustomNumericalRules) { + // Custom numerical rules take precedence so they can target any model. + strategies.push(new NumericalClassifierStrategy()); + } + // Then, if enabled, the Gemma classifier is used. if (this.config.getGemmaModelRouterSettings()?.enabled) { strategies.push(new GemmaClassifierStrategy()); @@ -54,8 +61,10 @@ export class ModelRouterService { // The generic classifier is next. strategies.push(new ClassifierStrategy()); - // The numerical classifier is next. - strategies.push(new NumericalClassifierStrategy()); + // Preserve the existing default numerical routing order. + if (!hasCustomNumericalRules) { + strategies.push(new NumericalClassifierStrategy()); + } // The default strategy is the terminal strategy. const terminalStrategy = new DefaultStrategy(); diff --git a/packages/core/src/routing/strategies/numericalClassifierStrategy.test.ts b/packages/core/src/routing/strategies/numericalClassifierStrategy.test.ts index c4a9f595f0d..720a1bc3b0d 100644 --- a/packages/core/src/routing/strategies/numericalClassifierStrategy.test.ts +++ b/packages/core/src/routing/strategies/numericalClassifierStrategy.test.ts @@ -61,6 +61,8 @@ describe('NumericalClassifierStrategy', () => { getModel: vi.fn().mockReturnValue(PREVIEW_GEMINI_MODEL_AUTO), getSessionId: vi.fn().mockReturnValue('control-group-id'), // Default to Control Group (Hash 71 >= 50) getNumericalRoutingEnabled: vi.fn().mockResolvedValue(true), + hasCustomNumericalRoutingRules: vi.fn().mockReturnValue(false), + getNumericalRoutingRules: vi.fn().mockReturnValue([]), getResolvedClassifierThreshold: vi.fn().mockResolvedValue(90), getClassifierThreshold: vi.fn().mockResolvedValue(undefined), getGemini31Launched: vi.fn().mockResolvedValue(false), @@ -215,6 +217,431 @@ describe('NumericalClassifierStrategy', () => { }); }); + describe('Custom Numerical Routing Rules', () => { + describe('Range Boundaries', () => { + beforeEach(() => { + vi.mocked(mockConfig.hasCustomNumericalRoutingRules).mockReturnValue( + true, + ); + vi.mocked(mockConfig.getNumericalRoutingRules).mockReturnValue([ + { maxScore: 10, model: 'local-gemma' }, + { maxScore: 30, model: 'gemini-2.5-flash' }, + { maxScore: 80, model: 'gemini-3-flash-preview' }, + { maxScore: 100, model: 'gemini-3-pro-preview' }, + ]); + }); + + it('routes lower bound score (1) to the first band', async () => { + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Minimal complexity', + complexity_score: 1, + }); + + const decision = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decision).toEqual({ + model: 'local-gemma', + metadata: { + source: 'NumericalClassifier (Custom)', + latencyMs: expect.any(Number), + reasoning: + '[Score: 1 / Max score: 10 / Model: local-gemma] Minimal complexity', + }, + }); + }); + + it('routes exact upper boundary of first band (10) to the first band', async () => { + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Simple task', + complexity_score: 10, + }); + + const decision = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decision?.model).toBe('local-gemma'); + expect(decision?.metadata.source).toBe('NumericalClassifier (Custom)'); + expect(decision?.metadata.reasoning).toContain( + '[Score: 10 / Max score: 10 / Model: local-gemma]', + ); + }); + + it('routes just above boundary (11) to the second band', async () => { + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Slightly more than trivial', + complexity_score: 11, + }); + + const decision = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decision?.model).toBe('gemini-2.5-flash'); + expect(decision?.metadata.reasoning).toContain( + '[Score: 11 / Max score: 30 / Model: gemini-2.5-flash]', + ); + }); + + it('routes exact boundary (30) to the second band', async () => { + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Moderate task', + complexity_score: 30, + }); + + const decision = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decision?.model).toBe('gemini-2.5-flash'); + expect(decision?.metadata.reasoning).toContain( + '[Score: 30 / Max score: 30 / Model: gemini-2.5-flash]', + ); + }); + + it('routes exact boundary (80) to the third band', async () => { + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'High complexity task', + complexity_score: 80, + }); + + const decision = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decision?.model).toBe('gemini-3-flash-preview'); + expect(decision?.metadata.reasoning).toContain( + '[Score: 80 / Max score: 80 / Model: gemini-3-flash-preview]', + ); + }); + + it('routes just above boundary (81) to the fourth band', async () => { + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Very complex task', + complexity_score: 81, + }); + + const decision = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decision?.model).toBe('gemini-3-pro-preview'); + expect(decision?.metadata.reasoning).toContain( + '[Score: 81 / Max score: 100 / Model: gemini-3-pro-preview]', + ); + }); + + it('routes maximum score (100) to the fourth band', async () => { + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Maximum complexity', + complexity_score: 100, + }); + + const decision = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decision?.model).toBe('gemini-3-pro-preview'); + expect(decision?.metadata.reasoning).toContain( + '[Score: 100 / Max score: 100 / Model: gemini-3-pro-preview]', + ); + }); + }); + + describe('Unordered Rules', () => { + it('correctly sorts rules by maxScore ascending before matching', async () => { + vi.mocked(mockConfig.hasCustomNumericalRoutingRules).mockReturnValue( + true, + ); + vi.mocked(mockConfig.getNumericalRoutingRules).mockReturnValue([ + { maxScore: 80, model: 'gemini-3-flash-preview' }, + { maxScore: 10, model: 'local-gemma' }, + { maxScore: 30, model: 'gemini-2.5-flash' }, + { maxScore: 100, model: 'gemini-3-pro-preview' }, + ]); + + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Task with score 25', + complexity_score: 25, + }); + + const decision = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decision?.model).toBe('gemini-2.5-flash'); + expect(decision?.metadata.reasoning).toContain( + '[Score: 25 / Max score: 30 / Model: gemini-2.5-flash]', + ); + }); + }); + + describe('Fallback Behavior', () => { + it('falls back to the final band model when score exceeds all configured max scores', async () => { + vi.mocked(mockConfig.hasCustomNumericalRoutingRules).mockReturnValue( + true, + ); + vi.mocked(mockConfig.getNumericalRoutingRules).mockReturnValue([ + { maxScore: 10, model: 'gemini-2.5-flash-lite' }, + { maxScore: 50, model: 'gemini-2.5-flash' }, + ]); + + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Score above all rules', + complexity_score: 75, + }); + + const decision = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decision?.model).toBe('gemini-2.5-flash'); + expect(decision?.metadata.source).toBe('NumericalClassifier (Custom)'); + expect(decision?.metadata.reasoning).toContain( + '[Score: 75 / Above final max score: 50 / Model: gemini-2.5-flash]', + ); + }); + + it('falls back to default Flash/Pro threshold routing when custom rules array is empty', async () => { + vi.mocked(mockConfig.hasCustomNumericalRoutingRules).mockReturnValue( + true, + ); + vi.mocked(mockConfig.getNumericalRoutingRules).mockReturnValue([]); + vi.mocked(mockConfig.getResolvedClassifierThreshold).mockResolvedValue( + 90, + ); + + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Standard task', + complexity_score: 50, + }); + + const decision = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decision?.model).toBe(PREVIEW_GEMINI_FLASH_MODEL); + expect(decision?.metadata.source).toBe('NumericalClassifier (Default)'); + expect(decision?.metadata.reasoning).toContain( + '[Score: 50 / Threshold: 90]', + ); + }); + + it('filters out invalid rules and uses remaining valid rules', async () => { + vi.mocked(mockConfig.hasCustomNumericalRoutingRules).mockReturnValue( + true, + ); + vi.mocked(mockConfig.getNumericalRoutingRules).mockReturnValue([ + { maxScore: NaN, model: 'invalid-model' }, + { maxScore: -10, model: 'negative-score-model' }, + { maxScore: 150, model: 'out-of-range-model' }, + { maxScore: 20, model: '' }, + { maxScore: 40, model: 'gemini-2.5-flash' }, + ]); + + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Valid rule match', + complexity_score: 35, + }); + + const decision = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decision?.model).toBe('gemini-2.5-flash'); + expect(decision?.metadata.source).toBe('NumericalClassifier (Custom)'); + expect(decision?.metadata.reasoning).toContain( + '[Score: 35 / Max score: 40 / Model: gemini-2.5-flash]', + ); + }); + + it('falls back to default Flash/Pro routing when all configured rules are invalid', async () => { + vi.mocked(mockConfig.hasCustomNumericalRoutingRules).mockReturnValue( + true, + ); + vi.mocked(mockConfig.getNumericalRoutingRules).mockReturnValue([ + { maxScore: NaN, model: 'invalid' }, + { maxScore: 0, model: 'zero-max-score' }, + { maxScore: 50, model: ' ' }, + ]); + vi.mocked(mockConfig.getResolvedClassifierThreshold).mockResolvedValue( + 90, + ); + + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Fall back to default', + complexity_score: 95, + }); + + const decision = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decision?.model).toBe(PREVIEW_GEMINI_MODEL); + expect(decision?.metadata.source).toBe('NumericalClassifier (Default)'); + expect(decision?.metadata.reasoning).toContain( + '[Score: 95 / Threshold: 90]', + ); + }); + + it('bypasses (returns null) when selected custom model is unavailable in ModelAvailabilityService', async () => { + vi.mocked(mockConfig.hasCustomNumericalRoutingRules).mockReturnValue( + true, + ); + vi.mocked(mockConfig.getNumericalRoutingRules).mockReturnValue([ + { maxScore: 50, model: 'unavailable-model' }, + ]); + + const availabilityService = new ModelAvailabilityService(); + vi.spyOn(availabilityService, 'snapshot').mockReturnValue({ + available: false, + reason: 'quota', + }); + vi.mocked(mockConfig.getModelAvailabilityService).mockReturnValue( + availabilityService, + ); + + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Task with unavailable model', + complexity_score: 30, + }); + + const decision = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decision).toBeNull(); + }); + }); + + describe('Model Resolution and Aliases', () => { + beforeEach(() => { + vi.mocked(mockConfig.hasCustomNumericalRoutingRules).mockReturnValue( + true, + ); + }); + + it('resolves model aliases like "flash" and "pro"', async () => { + vi.mocked(mockConfig.getNumericalRoutingRules).mockReturnValue([ + { maxScore: 40, model: 'flash' }, + { maxScore: 100, model: 'pro' }, + ]); + + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Alias test flash', + complexity_score: 20, + }); + + const decisionFlash = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decisionFlash?.model).toBe(PREVIEW_GEMINI_FLASH_MODEL); + + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Alias test pro', + complexity_score: 80, + }); + + const decisionPro = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decisionPro?.model).toBe(PREVIEW_GEMINI_MODEL); + }); + + it('normalizes model identifiers with models/ prefix', async () => { + vi.mocked(mockConfig.getNumericalRoutingRules).mockReturnValue([ + { maxScore: 100, model: 'models/gemini-2.5-flash' }, + ]); + + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Models prefix test', + complexity_score: 50, + }); + + const decision = await strategy.route( + mockContext, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decision?.model).toBe('gemini-2.5-flash'); + }); + + it('applies custom rules even when requestedModel is not Gemini 3', async () => { + vi.mocked(mockConfig.getModel).mockReturnValue('gemini-2.5-flash'); + vi.mocked(mockConfig.getNumericalRoutingRules).mockReturnValue([ + { maxScore: 50, model: 'gemini-2.5-flash-lite' }, + { maxScore: 100, model: 'gemini-2.5-pro' }, + ]); + + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue({ + complexity_reasoning: 'Non-gemini-3 base model test', + complexity_score: 80, + }); + + const decision = await strategy.route( + { ...mockContext, requestedModel: 'gemini-2.5-flash' }, + mockConfig, + mockBaseLlmClient, + mockLocalLiteRtLmClient, + ); + + expect(decision?.model).toBe('gemini-2.5-pro'); + expect(decision?.metadata.source).toBe('NumericalClassifier (Custom)'); + }); + }); + }); + describe('Remote Threshold Logic', () => { it('should use the remote CLASSIFIER_THRESHOLD if provided (int value)', async () => { vi.mocked(mockConfig.getClassifierThreshold).mockResolvedValue(70); diff --git a/packages/core/src/routing/strategies/numericalClassifierStrategy.ts b/packages/core/src/routing/strategies/numericalClassifierStrategy.ts index fc648d0b6d6..d85714336dd 100644 --- a/packages/core/src/routing/strategies/numericalClassifierStrategy.ts +++ b/packages/core/src/routing/strategies/numericalClassifierStrategy.ts @@ -12,7 +12,11 @@ import type { RoutingDecision, RoutingStrategy, } from '../routingStrategy.js'; -import { resolveClassifierModel, isGemini3Model } from '../../config/models.js'; +import { + resolveClassifierModel, + resolveModel, + isGemini3Model, +} from '../../config/models.js'; import { createUserContent, Type } from '@google/genai'; import type { Config } from '../../config/config.js'; import { @@ -114,7 +118,25 @@ export class NumericalClassifierStrategy implements RoutingStrategy { return null; } - if (!isGemini3Model(model, config)) { + const hasCustomRules = config.hasCustomNumericalRoutingRules?.() ?? false; + const routingRules = config.getNumericalRoutingRules?.() ?? []; + + const validRules = hasCustomRules + ? [...routingRules] + .filter( + (rule) => + Number.isFinite(rule.maxScore) && + rule.maxScore >= 1 && + rule.maxScore <= 100 && + typeof rule.model === 'string' && + rule.model.trim().length > 0, + ) + .sort((a, b) => a.maxScore - b.maxScore) + : []; + + const hasActiveCustomRules = validRules.length > 0; + + if (!hasActiveCustomRules && !isGemini3Model(model, config)) { return null; } @@ -178,26 +200,62 @@ export class NumericalClassifierStrategy implements RoutingStrategy { const routerResponse = ClassifierResponseSchema.parse(jsonResponse); const score = routerResponse.complexity_score; - const { threshold, groupLabel, modelAlias } = - await this.getRoutingDecision(score, config); const [useGemini3_1, useCustomToolModel] = await Promise.all([ config.getGemini31Launched(), config.getUseCustomToolModel(), ]); const useLatestFlash = config.hasLatestFlashGAAccess?.() ?? false; const useLatestFlashLite = config.hasLatestFlashLiteGAAccess?.() ?? false; - const selectedModel = normalizeModelId( - resolveClassifierModel( - normalizeModelId(model), + + let selectedModel: string; + let groupLabel: string; + let reasonPrefix: string; + + if (hasActiveCustomRules) { + const configuredRule = validRules.find( + (rule) => score <= rule.maxScore, + ); + const selectedRule = configuredRule ?? validRules.at(-1)!; + + selectedModel = normalizeModelId( + resolveModel( + selectedRule.model.trim(), + useGemini3_1, + useCustomToolModel, + config.getHasAccessToPreviewModel?.() ?? true, + config, + useLatestFlash, + useLatestFlashLite, + ), + ); + + groupLabel = 'Custom'; + reasonPrefix = configuredRule + ? `[Score: ${score} / Max score: ${selectedRule.maxScore} / Model: ${selectedRule.model}]` + : `[Score: ${score} / Above final max score: ${selectedRule.maxScore} / Model: ${selectedRule.model}]`; + } else { + const { + threshold, + groupLabel: defaultGroupLabel, modelAlias, - useGemini3_1, - useCustomToolModel, - config.getHasAccessToPreviewModel?.() ?? true, - config, - useLatestFlash, - useLatestFlashLite, - ), - ); + } = await this.getRoutingDecision(score, config); + groupLabel = defaultGroupLabel; + + selectedModel = normalizeModelId( + resolveClassifierModel( + normalizeModelId(model), + modelAlias, + useGemini3_1, + useCustomToolModel, + config.getHasAccessToPreviewModel?.() ?? true, + config, + useLatestFlash, + useLatestFlashLite, + ), + ); + + reasonPrefix = `[Score: ${score} / Threshold: ${threshold}]`; + } const service = config.getModelAvailabilityService(); const snapshot = service.snapshot(selectedModel); @@ -216,7 +274,7 @@ export class NumericalClassifierStrategy implements RoutingStrategy { metadata: { source: `NumericalClassifier (${groupLabel})`, latencyMs, - reasoning: `[Score: ${score} / Threshold: ${threshold}] ${routerResponse.complexity_reasoning}`, + reasoning: `${reasonPrefix} ${routerResponse.complexity_reasoning}`, }, }; } catch (error) { diff --git a/schemas/settings.schema.json b/schemas/settings.schema.json index 821b8ac499d..cf396d702cd 100644 --- a/schemas/settings.schema.json +++ b/schemas/settings.schema.json @@ -22,6 +22,43 @@ "$ref": "#/$defs/MCPServerConfig" } }, + "routing": { + "title": "Routing", + "description": "Configure numerical complexity score routing.", + "markdownDescription": "Configure numerical complexity score routing.\n\n- Category: `Model`\n- Requires restart: `yes`\n- Default: `{}`", + "default": {}, + "type": "object", + "properties": { + "numericalRules": { + "title": "Numerical Routing Rules", + "description": "Score bands matched in ascending maxScore order. Each maxScore is inclusive. Scores above all configured bands use the model in the final band.", + "markdownDescription": "Score bands matched in ascending maxScore order. Each maxScore is inclusive. Scores above all configured bands use the model in the final band.\n\n- Category: `Model`\n- Requires restart: `yes`\n- Default: `[]`", + "default": [], + "type": "array", + "items": { + "type": "object", + "properties": { + "maxScore": { + "title": "Maximum Score", + "description": "Maximum score for this routing rule.", + "markdownDescription": "Maximum score for this routing rule.\n\n- Category: `Model`\n- Requires restart: `yes`\n- Default: `100`", + "default": 100, + "type": "number" + }, + "model": { + "title": "Model ID", + "description": "Model alias or ID to route to.", + "markdownDescription": "Model alias or ID to route to.\n\n- Category: `Model`\n- Requires restart: `yes`\n- Default: ``", + "default": "", + "type": "string" + } + }, + "additionalProperties": false + } + } + }, + "additionalProperties": false + }, "policyPaths": { "title": "Policy Paths", "description": "Additional policy files or directories to load.",