Skip to content

Commit 2292bd7

Browse files
authored
feat: nep auto trigger (#2424)
## Problem The existing edit trigger criteria is too loose (1) non empty right context (2) edit history in the past 20s. ## Solution Build a logistic regression classifier and use users editor state as the features of the formula. <!--- REMINDER: - Read CONTRIBUTING.md first. - Add test coverage for your changes. - Link to related issues/commits. - Testing: how did you test your changes? - Screenshots if applicable --> ## License By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
1 parent 2c33b38 commit 2292bd7

File tree

9 files changed

+1192
-136
lines changed

9 files changed

+1192
-136
lines changed

server/aws-lsp-codewhisperer/src/language-server/inline-completion/auto-trigger/editPredictionAutoTrigger.test.ts

Lines changed: 510 additions & 3 deletions
Large diffs are not rendered by default.

server/aws-lsp-codewhisperer/src/language-server/inline-completion/auto-trigger/editPredictionAutoTrigger.ts

Lines changed: 414 additions & 3 deletions
Large diffs are not rendered by default.

server/aws-lsp-codewhisperer/src/language-server/inline-completion/handler/editCompletionHandler.test.ts

Lines changed: 20 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import * as sinon from 'sinon'
66
import { CodeWhispererSession, SessionData, SessionManager } from '../session/sessionManager'
77
import { HELLO_WORLD_IN_CSHARP } from '../../../shared/testUtils'
88
import { CodeWhispererServiceToken } from '../../../shared/codeWhispererService'
9+
import * as EditAutotrigger from '../auto-trigger/editPredictionAutoTrigger'
910

1011
describe('EditCompletionHandler', () => {
1112
let handler: EditCompletionHandler
@@ -54,7 +55,7 @@ describe('EditCompletionHandler', () => {
5455
}
5556
amazonQServiceManager = { getCodewhispererService: sinon.stub().returns(codeWhispererService) }
5657
cursorTracker = { trackPosition: sinon.stub() }
57-
recentEditsTracker = {}
58+
recentEditsTracker = { generateEditBasedContext: sinon.stub() }
5859
rejectedEditTracker = { isSimilarToRejected: sinon.stub().returns(false) }
5960
telemetry = { emitMetric: sinon.stub() }
6061
telemetryService = { emitUserTriggerDecision: sinon.stub() }
@@ -225,7 +226,7 @@ describe('EditCompletionHandler', () => {
225226
})
226227
})
227228

228-
describe('documentChanged', () => {
229+
describe.skip('documentChanged', () => {
229230
it('should set hasDocumentChangedSinceInvocation when waiting', () => {
230231
handler['debounceTimeout'] = setTimeout(() => {}, 1000) as any
231232
handler['isWaiting'] = true
@@ -337,12 +338,23 @@ describe('EditCompletionHandler', () => {
337338
position: { line: 0, character: 0 },
338339
context: { triggerKind: InlineCompletionTriggerKind.Automatic },
339340
}
341+
342+
afterEach('teardown', function () {
343+
sinon.restore()
344+
})
345+
346+
function aTriggerStub(flag: boolean): EditAutotrigger.EditClassifier {
347+
return {
348+
shouldTriggerNep: sinon
349+
.stub()
350+
.returns({ score: 0, threshold: EditAutotrigger.EditClassifier.THRESHOLD, shouldTrigger: flag }),
351+
} as any as EditAutotrigger.EditClassifier
352+
}
353+
340354
it('should return empty result when shouldTriggerEdits returns false', async () => {
341355
workspace.getWorkspaceFolder.returns(undefined)
342356

343-
const shouldTriggerEditsStub = sinon
344-
.stub(require('../utils/triggerUtils'), 'shouldTriggerEdits')
345-
.returns(false)
357+
sinon.stub(EditAutotrigger, 'EditClassifier').returns(aTriggerStub(false))
346358

347359
const result = await handler._invoke(
348360
params as any,
@@ -354,15 +366,12 @@ describe('EditCompletionHandler', () => {
354366
)
355367

356368
assert.deepEqual(result, EMPTY_RESULT)
357-
shouldTriggerEditsStub.restore()
358369
})
359370

360371
it('should create session and call generateSuggestions when trigger is valid', async () => {
361372
workspace.getWorkspaceFolder.returns(undefined)
362373

363-
const shouldTriggerEditsStub = sinon
364-
.stub(require('../utils/triggerUtils'), 'shouldTriggerEdits')
365-
.returns(true)
374+
sinon.stub(EditAutotrigger, 'EditClassifier').returns(aTriggerStub(true))
366375
codeWhispererService.constructSupplementalContext.resolves(null)
367376
codeWhispererService.generateSuggestions.resolves({
368377
suggestions: [{ itemId: 'item-1', content: 'test content' }],
@@ -380,7 +389,6 @@ describe('EditCompletionHandler', () => {
380389

381390
assert.strictEqual(result.items.length, 1)
382391
sinon.assert.called(codeWhispererService.generateSuggestions)
383-
shouldTriggerEditsStub.restore()
384392
})
385393

386394
it('should handle active session and emit telemetry', async () => {
@@ -391,9 +399,7 @@ describe('EditCompletionHandler', () => {
391399
if (currentSession) {
392400
sessionManager.activateSession(currentSession)
393401
}
394-
const shouldTriggerEditsStub = sinon
395-
.stub(require('../utils/triggerUtils'), 'shouldTriggerEdits')
396-
.returns(true)
402+
sinon.stub(EditAutotrigger, 'EditClassifier').returns(aTriggerStub(true))
397403
codeWhispererService.constructSupplementalContext.resolves(null)
398404
codeWhispererService.generateSuggestions.resolves({
399405
suggestions: [{ itemId: 'item-1', content: 'test content' }],
@@ -410,15 +416,12 @@ describe('EditCompletionHandler', () => {
410416
)
411417

412418
assert.strictEqual(currentSession?.state, 'DISCARD')
413-
shouldTriggerEditsStub.restore()
414419
})
415420

416421
it('should handle supplemental context when available', async () => {
417422
workspace.getWorkspaceFolder.returns(undefined)
418423

419-
const shouldTriggerEditsStub = sinon
420-
.stub(require('../utils/triggerUtils'), 'shouldTriggerEdits')
421-
.returns(true)
424+
sinon.stub(EditAutotrigger, 'EditClassifier').returns(aTriggerStub(true))
422425
codeWhispererService.constructSupplementalContext.resolves({
423426
items: [{ content: 'context', filePath: 'file.ts' }],
424427
supContextData: { isUtg: false },
@@ -438,7 +441,6 @@ describe('EditCompletionHandler', () => {
438441
)
439442

440443
sinon.assert.calledWith(codeWhispererService.generateSuggestions, sinon.match.has('supplementalContexts'))
441-
shouldTriggerEditsStub.restore()
442444
})
443445
})
444446

server/aws-lsp-codewhisperer/src/language-server/inline-completion/handler/editCompletionHandler.ts

Lines changed: 53 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ import { CodeWhispererSession, SessionManager } from '../session/sessionManager'
2222
import { CursorTracker } from '../tracker/cursorTracker'
2323
import { CodewhispererLanguage, getSupportedLanguageId } from '../../../shared/languageDetection'
2424
import { WorkspaceFolderManager } from '../../workspaceContext/workspaceFolderManager'
25-
import { shouldTriggerEdits } from '../utils/triggerUtils'
25+
import { inferTriggerChar } from '../utils/triggerUtils'
2626
import {
2727
emitEmptyUserTriggerDecisionTelemetry,
2828
emitServiceInvocationFailure,
@@ -36,9 +36,10 @@ import { RejectedEditTracker } from '../tracker/rejectedEditTracker'
3636
import { getErrorMessage, hasConnectionExpired } from '../../../shared/utils'
3737
import { AmazonQError, AmazonQServiceConnectionExpiredError } from '../../../shared/amazonQServiceManager/errors'
3838
import { DocumentChangedListener } from '../documentChangedListener'
39-
import { EMPTY_RESULT, EDIT_DEBOUNCE_INTERVAL_MS } from '../contants/constants'
39+
import { EMPTY_RESULT } from '../contants/constants'
4040
import { StreakTracker } from '../tracker/streakTracker'
4141
import { processEditSuggestion } from '../utils/diffUtils'
42+
import { EditClassifier } from '../auto-trigger/editPredictionAutoTrigger'
4243

4344
export class EditCompletionHandler {
4445
private readonly editsEnabled: boolean
@@ -79,14 +80,15 @@ export class EditCompletionHandler {
7980
* Also as a followup, ideally it should be a message/event publish/subscribe pattern instead of manual invocation like this
8081
*/
8182
documentChanged() {
82-
if (this.debounceTimeout) {
83-
if (this.isWaiting) {
84-
this.hasDocumentChangedSinceInvocation = true
85-
} else {
86-
this.logging.info(`refresh and debounce edits suggestion for another ${EDIT_DEBOUNCE_INTERVAL_MS}`)
87-
this.debounceTimeout.refresh()
88-
}
89-
}
83+
// TODO: Remove this entirely once we are sure we dont need debounce
84+
// if (this.debounceTimeout) {
85+
// if (this.isWaiting) {
86+
// this.hasDocumentChangedSinceInvocation = true
87+
// } else {
88+
// this.logging.info(`refresh and debounce edits suggestion for another ${EDIT_DEBOUNCE_INTERVAL_MS}`)
89+
// this.debounceTimeout.refresh()
90+
// }
91+
// }
9092
}
9193

9294
async onEditCompletion(
@@ -171,40 +173,18 @@ export class EditCompletionHandler {
171173
}
172174
}
173175

174-
return new Promise<InlineCompletionListWithReferences>(async resolve => {
175-
this.debounceTimeout = setTimeout(async () => {
176-
try {
177-
this.isWaiting = true
178-
const result = await this._invoke(
179-
params,
180-
startPreprocessTimestamp,
181-
token,
182-
textDocument,
183-
inferredLanguageId,
184-
currentSession
185-
).finally(() => {
186-
this.isWaiting = false
187-
})
188-
if (this.hasDocumentChangedSinceInvocation) {
189-
this.logging.info(
190-
'EditCompletionHandler - Document changed during execution, resolving empty result'
191-
)
192-
resolve({
193-
sessionId: SessionManager.getInstance('EDITS').getActiveSession()?.id ?? '',
194-
items: [],
195-
})
196-
} else {
197-
this.logging.info('EditCompletionHandler - No document changes, resolving result')
198-
resolve(result)
199-
}
200-
} finally {
201-
this.debounceTimeout = undefined
202-
this.hasDocumentChangedSinceInvocation = false
203-
}
204-
}, EDIT_DEBOUNCE_INTERVAL_MS)
205-
}).finally(() => {
176+
try {
177+
return await this._invoke(
178+
params,
179+
startPreprocessTimestamp,
180+
token,
181+
textDocument,
182+
inferredLanguageId,
183+
currentSession
184+
)
185+
} finally {
206186
this.isInProgress = false
207-
})
187+
}
208188
}
209189

210190
async _invoke(
@@ -218,53 +198,57 @@ export class EditCompletionHandler {
218198
// Build request context
219199
const isAutomaticLspTriggerKind = params.context.triggerKind == InlineCompletionTriggerKind.Automatic
220200
const maxResults = isAutomaticLspTriggerKind ? 1 : 5
221-
const fileContext = getFileContext({
201+
const fileContextClss = getFileContext({
222202
textDocument,
223203
inferredLanguageId,
224204
position: params.position,
225205
workspaceFolder: this.workspace.getWorkspaceFolder(textDocument.uri),
226206
})
227207

208+
// TODO: Parametrize these to a util function, duplicate code as inineCompletionHandler
209+
const triggerCharacters = inferTriggerChar(fileContextClss, params.documentChangeParams)
210+
228211
const workspaceState = WorkspaceFolderManager.getInstance()?.getWorkspaceState()
229212
const workspaceId = workspaceState?.webSocketClient?.isConnected() ? workspaceState.workspaceId : undefined
230213

231-
const qEditsTrigger = shouldTriggerEdits(
232-
this.codeWhispererService,
233-
fileContext,
234-
params,
235-
this.cursorTracker,
236-
this.recentEditsTracker,
237-
this.sessionManager,
238-
true
214+
const recentEdits = await this.recentEditsTracker.generateEditBasedContext(textDocument)
215+
const classifier = new EditClassifier(
216+
{
217+
fileContext: fileContextClss,
218+
triggerChar: triggerCharacters,
219+
recentEdits: recentEdits,
220+
recentDecisions: this.sessionManager.userDecisionLog.map(it => it.decision),
221+
},
222+
this.logging
239223
)
240224

241-
if (!qEditsTrigger) {
225+
const qEditsTrigger = classifier.shouldTriggerNep()
226+
227+
if (!qEditsTrigger.shouldTrigger) {
242228
return EMPTY_RESULT
243229
}
244230

245231
const generateCompletionReq: GenerateSuggestionsRequest = {
246-
fileContext: fileContext,
232+
fileContext: fileContextClss.toServiceModel(),
247233
maxResults: maxResults,
248234
predictionTypes: ['EDITS'],
249235
workspaceId: workspaceId,
250236
}
251237

252-
if (qEditsTrigger) {
253-
generateCompletionReq.editorState = {
254-
document: {
255-
relativeFilePath: textDocument.uri,
256-
programmingLanguage: {
257-
languageName: generateCompletionReq.fileContext?.programmingLanguage?.languageName,
258-
},
259-
text: textDocument.getText(),
238+
generateCompletionReq.editorState = {
239+
document: {
240+
relativeFilePath: textDocument.uri,
241+
programmingLanguage: {
242+
languageName: generateCompletionReq.fileContext?.programmingLanguage?.languageName,
260243
},
261-
cursorState: {
262-
position: {
263-
line: params.position.line,
264-
character: params.position.character,
265-
},
244+
text: textDocument.getText(),
245+
},
246+
cursorState: {
247+
position: {
248+
line: params.position.line,
249+
character: params.position.character,
266250
},
267-
}
251+
},
268252
}
269253

270254
const supplementalContext = await this.codeWhispererService.constructSupplementalContext(
@@ -306,7 +290,7 @@ export class EditCompletionHandler {
306290
startPreprocessTimestamp: startPreprocessTimestamp,
307291
startPosition: params.position,
308292
triggerType: isAutomaticLspTriggerKind ? 'AutoTrigger' : 'OnDemand',
309-
language: fileContext.programmingLanguage.languageName,
293+
language: fileContextClss.programmingLanguage.languageName,
310294
requestContext: generateCompletionReq,
311295
autoTriggerType: undefined,
312296
triggerCharacter: '',

server/aws-lsp-codewhisperer/src/language-server/inline-completion/handler/inlineCompletionHandler.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,7 @@ export class InlineCompletionHandler {
191191
inferredLanguageId,
192192
position: params.position,
193193
workspaceFolder: this.workspace.getWorkspaceFolder(textDocument.uri),
194-
})
194+
}).toServiceModel()
195195
}
196196

197197
const workspaceState = WorkspaceFolderManager.getInstance()?.getWorkspaceState()

server/aws-lsp-codewhisperer/src/language-server/inline-completion/session/sessionManager.ts

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -309,6 +309,10 @@ export class SessionManager {
309309
private sessionsLog: CodeWhispererSession[] = []
310310
private maxHistorySize = 5
311311
// TODO, for user decision telemetry: accepted suggestions (not necessarily the full corresponding session) should be stored for 5 minutes
312+
private _userDecisionLog: { sessionId: string; decision: UserTriggerDecision }[] = []
313+
get userDecisionLog() {
314+
return [...this._userDecisionLog]
315+
}
312316

313317
private constructor() {}
314318

@@ -352,6 +356,19 @@ export class SessionManager {
352356

353357
closeSession(session: CodeWhispererSession) {
354358
session.close()
359+
360+
// Note: it has to be called after session.close() as getAggregatedUserTriggerDecision() will return undefined if getAggregatedUserTriggerDecision() is called before session is closed
361+
const decision = session.getAggregatedUserTriggerDecision()
362+
// As we only care about AR here, pushing Accept/Reject only
363+
if (decision === 'Accept' || decision === 'Reject') {
364+
if (this._userDecisionLog.length === 5) {
365+
this._userDecisionLog.shift()
366+
}
367+
this._userDecisionLog.push({
368+
sessionId: session.codewhispererSessionId ?? 'undefined',
369+
decision: decision,
370+
})
371+
}
355372
}
356373

357374
discardSession(session: CodeWhispererSession) {

0 commit comments

Comments
 (0)