|
| 1 | +/*--------------------------------------------------------------------------------------------- |
| 2 | + * Copyright (c) Microsoft Corporation. All rights reserved. |
| 3 | + * Licensed under the MIT License. See License.txt in the project root for license information. |
| 4 | + *--------------------------------------------------------------------------------------------*/ |
| 5 | + |
| 6 | +import type { LanguageModelChat } from 'vscode'; |
| 7 | +import { IEndpointProvider } from '../../../platform/endpoint/common/endpointProvider'; |
| 8 | +import { IVSCodeExtensionContext } from '../../../platform/extContext/common/extensionContext'; |
| 9 | +import { IChatEndpoint } from '../../../platform/networking/common/networking'; |
| 10 | +import { createServiceIdentifier } from '../../../util/common/services'; |
| 11 | +import { LRUCache } from '../../../util/vs/base/common/map'; |
| 12 | +import { mapValues } from '../../../util/vs/base/common/objects'; |
| 13 | +import { EditTools as _EditTools, EDIT_TOOL_LEARNING_STATES, IEditToolLearningData, LearningConfig, State } from './editToolLearningStates'; |
| 14 | +import { ToolName } from './toolNames'; |
| 15 | + |
| 16 | +export type EditTools = _EditTools; |
| 17 | + |
| 18 | +const CACHE_STORAGE_KEY = 'editToolLearning_cache'; |
| 19 | + |
| 20 | +function mapToolsRecord<I, O>(record: { [K in EditTools]?: I }, fn: (input: I, tool: EditTools) => O) { |
| 21 | + return mapValues(record, (value, key) => fn(value!, key as EditTools)) as { [K in EditTools]?: O }; |
| 22 | +} |
| 23 | + |
| 24 | +interface IStoredToolData { |
| 25 | + state: State; |
| 26 | + tools: { [K in EditTools]?: { successBitset: string; attempts: number } }; |
| 27 | +} |
| 28 | + |
| 29 | +export const IEditToolLearningService = createServiceIdentifier<IEditToolLearningService>('IEditToolLearningService'); |
| 30 | + |
| 31 | +export interface IEditToolLearningService { |
| 32 | + readonly _serviceBrand: undefined; |
| 33 | + getPreferredEditTool(model: LanguageModelChat): Promise<EditTools[] | undefined>; |
| 34 | + getPreferredEndpointEditTool(model: IChatEndpoint): EditTools[] | undefined; |
| 35 | + didMakeEdit(model: LanguageModelChat, tool: EditTools, success: boolean): void; |
| 36 | +} |
| 37 | + |
| 38 | +function addToWindow(window: bigint, bit: bigint): bigint { |
| 39 | + // Shift left to make room for new bit, add the bit, then mask to WINDOW_SIZE |
| 40 | + const mask = (1n << BigInt(LearningConfig.WINDOW_SIZE)) - 1n; |
| 41 | + return ((window << 1n) | bit) & mask; |
| 42 | +} |
| 43 | + |
| 44 | +export class EditToolLearningService implements IEditToolLearningService { |
| 45 | + readonly _serviceBrand: undefined; |
| 46 | + |
| 47 | + private _cache?: LRUCache<string, IEditToolLearningData>; |
| 48 | + |
| 49 | + constructor( |
| 50 | + @IVSCodeExtensionContext private readonly _context: IVSCodeExtensionContext, |
| 51 | + @IEndpointProvider private readonly _endpointProvider: IEndpointProvider, |
| 52 | + ) { } |
| 53 | + |
| 54 | + async getPreferredEditTool(model: LanguageModelChat): Promise<EditTools[] | undefined> { |
| 55 | + const endpoint = await this._endpointProvider.getChatEndpoint(model); |
| 56 | + return this.getPreferredEndpointEditTool(endpoint); |
| 57 | + } |
| 58 | + |
| 59 | + getPreferredEndpointEditTool(endpoint: IChatEndpoint): EditTools[] | undefined { |
| 60 | + if (!endpoint.isExtensionContributed) { |
| 61 | + return undefined; |
| 62 | + } |
| 63 | + |
| 64 | + const hardcoded = this._getHardcodedPreferences(endpoint.model); |
| 65 | + if (hardcoded) { |
| 66 | + return hardcoded; |
| 67 | + } |
| 68 | + |
| 69 | + const learningData = this._getModelLearningData(endpoint.model); |
| 70 | + return this._computePreferences(learningData); |
| 71 | + } |
| 72 | + |
| 73 | + async didMakeEdit(model: LanguageModelChat, tool: EditTools, success: boolean): Promise<void> { |
| 74 | + const endpoint = await this._endpointProvider.getChatEndpoint(model); |
| 75 | + |
| 76 | + if (!endpoint.isExtensionContributed || this._getHardcodedPreferences(endpoint.family)) { |
| 77 | + return; |
| 78 | + } |
| 79 | + |
| 80 | + const learningData = this._getModelLearningData(model.id); |
| 81 | + this._recordEdit(learningData, tool, success); |
| 82 | + await this._saveModelLearningData(model.id, learningData); |
| 83 | + } |
| 84 | + |
| 85 | + private _getHardcodedPreferences(family: string): EditTools[] | undefined { |
| 86 | + const lowerFamily = family.toLowerCase(); |
| 87 | + |
| 88 | + if (lowerFamily.includes('gpt') || lowerFamily.includes('openai')) { |
| 89 | + return [ToolName.ApplyPatch]; |
| 90 | + } |
| 91 | + |
| 92 | + if (lowerFamily.includes('sonnet')) { |
| 93 | + return [ToolName.ReplaceString, ToolName.MultiReplaceString]; |
| 94 | + } |
| 95 | + |
| 96 | + return undefined; |
| 97 | + } |
| 98 | + |
| 99 | + private _computePreferences(data: IEditToolLearningData): EditTools[] | undefined { |
| 100 | + return EDIT_TOOL_LEARNING_STATES[data.state].allowedTools; |
| 101 | + } |
| 102 | + |
| 103 | + private _checkStateTransitions(data: IEditToolLearningData): State { |
| 104 | + const currentConfig = EDIT_TOOL_LEARNING_STATES[data.state]; |
| 105 | + |
| 106 | + for (const [targetState, condition] of Object.entries(currentConfig.transitions)) { |
| 107 | + if (condition(data)) { |
| 108 | + return Number(targetState) as State; |
| 109 | + } |
| 110 | + } |
| 111 | + |
| 112 | + return data.state; // No transition |
| 113 | + } |
| 114 | + |
| 115 | + private _recordEdit(data: IEditToolLearningData, tool: EditTools, success: boolean): void { |
| 116 | + const successBit = success ? 1n : 0n; |
| 117 | + const toolData = (data.tools[tool] ??= { successBitset: 0n, attempts: 0 }); |
| 118 | + toolData.successBitset = addToWindow(toolData.successBitset, successBit); |
| 119 | + toolData.attempts++; |
| 120 | + |
| 121 | + const newState = this._checkStateTransitions(data); |
| 122 | + if (newState !== data.state) { |
| 123 | + data.state = newState; |
| 124 | + data.tools = {}; |
| 125 | + } |
| 126 | + } |
| 127 | + |
| 128 | + private _getCache(): LRUCache<string, IEditToolLearningData> { |
| 129 | + if (!this._cache) { |
| 130 | + this._cache = this._loadCacheFromStorage(); |
| 131 | + } |
| 132 | + return this._cache; |
| 133 | + } |
| 134 | + |
| 135 | + private _loadCacheFromStorage(): LRUCache<string, IEditToolLearningData> { |
| 136 | + const cache = new LRUCache<string, IEditToolLearningData>(LearningConfig.CACHE_SIZE); |
| 137 | + const storedCacheData = this._context.globalState.get<{ entries: [string, IStoredToolData][] }>(CACHE_STORAGE_KEY); |
| 138 | + |
| 139 | + if (!storedCacheData?.entries) { |
| 140 | + return cache; |
| 141 | + } |
| 142 | + |
| 143 | + for (const [modelId, storedData] of storedCacheData.entries) { |
| 144 | + const data: IEditToolLearningData = { |
| 145 | + state: storedData.state, |
| 146 | + tools: mapToolsRecord(storedData.tools, r => ({ |
| 147 | + successBitset: BigInt(r.successBitset), |
| 148 | + attempts: r.attempts, |
| 149 | + })), |
| 150 | + }; |
| 151 | + cache.set(modelId, data); |
| 152 | + } |
| 153 | + |
| 154 | + return cache; |
| 155 | + } |
| 156 | + |
| 157 | + private async _saveCacheToStorage(): Promise<void> { |
| 158 | + if (!this._cache) { |
| 159 | + return; |
| 160 | + } |
| 161 | + |
| 162 | + const entries: [string, IStoredToolData][] = Array.from(this._cache.entries(), ([modelId, data]) => { |
| 163 | + const storedData = { |
| 164 | + state: data.state, |
| 165 | + tools: mapToolsRecord(data.tools, r => ({ |
| 166 | + successBitset: '0x' + r.successBitset.toString(16), |
| 167 | + attempts: r.attempts |
| 168 | + })), |
| 169 | + }; |
| 170 | + return [modelId, storedData]; |
| 171 | + }); |
| 172 | + |
| 173 | + await this._context.globalState.update(CACHE_STORAGE_KEY, { entries }); |
| 174 | + } |
| 175 | + |
| 176 | + private async _saveModelLearningData(modelId: string, data: IEditToolLearningData): Promise<void> { |
| 177 | + const cache = this._getCache(); |
| 178 | + cache.set(modelId, data); |
| 179 | + await this._saveCacheToStorage(); |
| 180 | + } |
| 181 | + |
| 182 | + private _getModelLearningData(modelId: string): IEditToolLearningData { |
| 183 | + const cache = this._getCache(); |
| 184 | + |
| 185 | + let data = cache.get(modelId); |
| 186 | + if (!data) { |
| 187 | + data = { state: State.Initial, tools: {} }; |
| 188 | + cache.set(modelId, data); |
| 189 | + } |
| 190 | + return data; |
| 191 | + } |
| 192 | +} |
0 commit comments