From b269464df27044c0bf14ef72bb8ece6d1f00bc49 Mon Sep 17 00:00:00 2001 From: Adam Weidman Date: Tue, 10 Feb 2026 14:43:18 -0500 Subject: [PATCH] feat: refine adaptive thinking budget (stickiness, context-awareness, recursion prevention) --- packages/core/src/core/client.ts | 43 ++++++++++++------ .../services/adaptiveBudgetService.test.ts | 21 +++------ .../src/services/adaptiveBudgetService.ts | 45 ++++++++++++++----- .../core/src/services/modelConfigService.ts | 30 ++++++++----- 4 files changed, 90 insertions(+), 49 deletions(-) diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index a2a1cd8141..9f7b73b107 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -65,6 +65,7 @@ import { resolveModel } from '../config/models.js'; import type { RetryAvailabilityContext } from '../utils/retry.js'; import { partToString } from '../utils/partUtils.js'; import { coreEvents, CoreEvent } from '../utils/events.js'; +import type { AdaptiveBudgetResult } from '../services/adaptiveBudgetService.js'; const MAX_TURNS = 100; @@ -89,6 +90,7 @@ export class GeminiClient { private readonly toolOutputMaskingService: ToolOutputMaskingService; private lastPromptId: string; private currentSequenceModel: string | null = null; + private currentAdaptiveConfig: AdaptiveBudgetResult | null = null; private lastSentIdeContext: IdeContext | undefined; private forceFullIdeContext = true; @@ -647,22 +649,34 @@ export class GeminiClient { // Adaptive Thinking Budget Integration if ( !isInvalidStreamRetry && - this.config.getAdaptiveThinkingConfig().enabled + this.config.getAdaptiveThinkingConfig().enabled && + prompt_id !== 'adaptive-budget-classifier' ) { - const userMessage = partListUnionToString(request); - if (userMessage) { - const adaptiveConfig = await this.config - .getAdaptiveBudgetService() - .determineAdaptiveConfig(userMessage, modelToUse); + if (this.currentAdaptiveConfig) { + modelConfigKey.thinkingBudget = + this.currentAdaptiveConfig.thinkingBudget; + modelConfigKey.thinkingLevel = this.currentAdaptiveConfig.thinkingLevel; + } else { + const userMessage = partListUnionToString(request); + if (userMessage) { + const adaptiveConfig = await this.config + .getAdaptiveBudgetService() + .determineAdaptiveConfig( + userMessage, + modelToUse, + this.getHistory(), + ); - if (adaptiveConfig) { - modelConfigKey.thinkingBudget = adaptiveConfig.thinkingBudget; - modelConfigKey.thinkingLevel = adaptiveConfig.thinkingLevel; - this.getChat().recordAdaptiveThinking({ - complexity: adaptiveConfig.complexity, - thinkingBudget: adaptiveConfig.thinkingBudget, - thinkingLevel: adaptiveConfig.thinkingLevel, - }); + if (adaptiveConfig) { + this.currentAdaptiveConfig = adaptiveConfig; + modelConfigKey.thinkingBudget = adaptiveConfig.thinkingBudget; + modelConfigKey.thinkingLevel = adaptiveConfig.thinkingLevel; + this.getChat().recordAdaptiveThinking({ + complexity: adaptiveConfig.complexity, + thinkingBudget: adaptiveConfig.thinkingBudget, + thinkingLevel: adaptiveConfig.thinkingLevel, + }); + } } } } @@ -806,6 +820,7 @@ export class GeminiClient { this.hookStateMap.delete(this.lastPromptId); this.lastPromptId = prompt_id; this.currentSequenceModel = null; + this.currentAdaptiveConfig = null; } if (hooksEnabled && messageBus) { diff --git a/packages/core/src/services/adaptiveBudgetService.test.ts b/packages/core/src/services/adaptiveBudgetService.test.ts index f8c549e6de..a4e97017e1 100644 --- a/packages/core/src/services/adaptiveBudgetService.test.ts +++ b/packages/core/src/services/adaptiveBudgetService.test.ts @@ -9,7 +9,6 @@ import { ComplexityLevel, } from './adaptiveBudgetService.js'; import type { Config } from '../config/config.js'; -import { ThinkingLevel } from '@google/genai'; describe('AdaptiveBudgetService', () => { it('should map complexity levels to correct V2 budgets', () => { @@ -22,18 +21,10 @@ describe('AdaptiveBudgetService', () => { it('should map complexity levels to correct V3 levels', () => { const service = new AdaptiveBudgetService({} as Config); - expect(service.getThinkingLevelV3(ComplexityLevel.SIMPLE)).toBe( - ThinkingLevel.LOW, - ); - expect(service.getThinkingLevelV3(ComplexityLevel.MODERATE)).toBe( - ThinkingLevel.LOW, - ); - expect(service.getThinkingLevelV3(ComplexityLevel.HIGH)).toBe( - ThinkingLevel.HIGH, - ); - expect(service.getThinkingLevelV3(ComplexityLevel.EXTREME)).toBe( - ThinkingLevel.HIGH, - ); + expect(service.getThinkingLevelV3(ComplexityLevel.SIMPLE)).toBe('LOW'); + expect(service.getThinkingLevelV3(ComplexityLevel.MODERATE)).toBe('LOW'); + expect(service.getThinkingLevelV3(ComplexityLevel.HIGH)).toBe('HIGH'); + expect(service.getThinkingLevelV3(ComplexityLevel.EXTREME)).toBe('HIGH'); }); it('should determine adaptive config based on LLM response', async () => { @@ -55,6 +46,7 @@ describe('AdaptiveBudgetService', () => { const result = await service.determineAdaptiveConfig( 'Complex task', 'gemini-2.5-pro', + [], ); expect(result?.complexity).toBe(ComplexityLevel.HIGH); @@ -79,10 +71,11 @@ describe('AdaptiveBudgetService', () => { const result = await service.determineAdaptiveConfig( 'Hi', 'gemini-3-pro-preview', + [], ); expect(result?.complexity).toBe(ComplexityLevel.SIMPLE); - expect(result?.thinkingLevel).toBe(ThinkingLevel.LOW); + expect(result?.thinkingLevel).toBe('LOW'); expect(result?.thinkingBudget).toBeUndefined(); }); }); diff --git a/packages/core/src/services/adaptiveBudgetService.ts b/packages/core/src/services/adaptiveBudgetService.ts index f92697e9a3..53c9c447b2 100644 --- a/packages/core/src/services/adaptiveBudgetService.ts +++ b/packages/core/src/services/adaptiveBudgetService.ts @@ -6,7 +6,11 @@ import type { Config } from '../config/config.js'; import { debugLogger } from '../utils/debugLogger.js'; import { isGemini2Model, isPreviewModel } from '../config/models.js'; -import { ThinkingLevel } from '@google/genai'; +import type { Content } from '@google/genai'; +import { + isFunctionCall, + isFunctionResponse, +} from '../utils/messageInspectors.js'; export enum ComplexityLevel { SIMPLE = 1, @@ -22,17 +26,20 @@ export const BUDGET_MAPPING_V2: Record = { [ComplexityLevel.EXTREME]: 32768, }; -export const LEVEL_MAPPING_V3: Record = { - [ComplexityLevel.SIMPLE]: ThinkingLevel.LOW, - [ComplexityLevel.MODERATE]: ThinkingLevel.LOW, - [ComplexityLevel.HIGH]: ThinkingLevel.HIGH, - [ComplexityLevel.EXTREME]: ThinkingLevel.HIGH, +export const LEVEL_MAPPING_V3: Record = { + [ComplexityLevel.SIMPLE]: 'LOW', + [ComplexityLevel.MODERATE]: 'LOW', + [ComplexityLevel.HIGH]: 'HIGH', + [ComplexityLevel.EXTREME]: 'HIGH', }; +const HISTORY_TURNS_FOR_CONTEXT = 4; +const HISTORY_SEARCH_WINDOW = 20; + export interface AdaptiveBudgetResult { complexity: ComplexityLevel; thinkingBudget?: number; - thinkingLevel?: ThinkingLevel; + thinkingLevel?: string; strategyNote?: string; } @@ -51,6 +58,7 @@ export class AdaptiveBudgetService { async determineAdaptiveConfig( userPrompt: string, model: string, + history: Content[], ): Promise { const { classifierModel } = this.config.getAdaptiveThinkingConfig(); @@ -59,6 +67,18 @@ export class AdaptiveBudgetService { debugLogger.debug( `AdaptiveBudgetService: Classifying prompt complexity using ${classifierModel}...`, ); + + // situational context: provide the last N turns of the history + const historySlice = history.slice(-HISTORY_SEARCH_WINDOW); + + // Filter out tool-related turns to keep context focused on conversation. + const cleanHistory = historySlice.filter( + (content) => !isFunctionCall(content) && !isFunctionResponse(content), + ); + + // Take the last N turns from the *cleaned* history. + const finalHistory = cleanHistory.slice(-HISTORY_TURNS_FOR_CONTEXT); + const systemPrompt = `You are a complexity classifier for a coding assistant. Analyze the user's request and determine the complexity of the task. Output ONLY a single integer from 1 to 4 based on the following scale: @@ -73,7 +93,10 @@ Complexity Level:`; const response = await llm.generateContent({ modelConfigKey: { model: classifierModel }, - contents: [{ role: 'user', parts: [{ text: systemPrompt }] }], + contents: [ + ...finalHistory, + { role: 'user', parts: [{ text: systemPrompt }] }, + ], promptId: 'adaptive-budget-classifier', abortSignal: new AbortController().signal, }); @@ -99,7 +122,7 @@ Complexity Level:`; // Determine mapping based on model version // Gemini 3 uses ThinkingLevel, Gemini 2.x uses thinkingBudget if (isPreviewModel(model)) { - result.thinkingLevel = LEVEL_MAPPING_V3[level] ?? ThinkingLevel.HIGH; + result.thinkingLevel = LEVEL_MAPPING_V3[level] ?? 'HIGH'; } else if (isGemini2Model(model)) { result.thinkingBudget = BUDGET_MAPPING_V2[level]; } @@ -126,7 +149,7 @@ Complexity Level:`; return BUDGET_MAPPING_V2[level]; } - getThinkingLevelV3(level: ComplexityLevel): ThinkingLevel { - return LEVEL_MAPPING_V3[level] ?? ThinkingLevel.HIGH; + getThinkingLevelV3(level: ComplexityLevel): string { + return LEVEL_MAPPING_V3[level] ?? 'HIGH'; } } diff --git a/packages/core/src/services/modelConfigService.ts b/packages/core/src/services/modelConfigService.ts index 387bd90975..fe4f15ece7 100644 --- a/packages/core/src/services/modelConfigService.ts +++ b/packages/core/src/services/modelConfigService.ts @@ -29,7 +29,7 @@ export interface ModelConfigKey { // Dynamic thinking configuration determined at runtime (e.g. via complexity classification) thinkingBudget?: number; - thinkingLevel?: ThinkingLevel; + thinkingLevel?: ThinkingLevel | string; } export interface ModelConfig { @@ -159,17 +159,27 @@ export class ModelConfigService { context.thinkingBudget !== undefined || context.thinkingLevel !== undefined ) { + const thinkingConfig = { + ...(currentConfig.generateContentConfig?.thinkingConfig as object), + ...(context.thinkingBudget !== undefined + ? { thinkingBudget: context.thinkingBudget } + : {}), + ...(context.thinkingLevel !== undefined + ? { thinkingLevel: context.thinkingLevel } + : {}), + }; + + // Ensure we don't have BOTH if one was a default from an alias + if (context.thinkingLevel) { + delete (thinkingConfig as Record).thinkingBudget; + } + if (context.thinkingBudget) { + delete (thinkingConfig as Record).thinkingLevel; + } + currentConfig.generateContentConfig = { ...currentConfig.generateContentConfig, - thinkingConfig: { - ...(currentConfig.generateContentConfig?.thinkingConfig as object), - ...(context.thinkingBudget !== undefined - ? { thinkingBudget: context.thinkingBudget } - : {}), - ...(context.thinkingLevel !== undefined - ? { thinkingLevel: context.thinkingLevel } - : {}), - }, + thinkingConfig: thinkingConfig as Record, }; }