Skip to content

Commit e01574f

Browse files
authored
feat(amazonq): basic paginated suggestions using language server (aws#6803)
## Problem ## Solution Basic paginated inline suggestions implementation using language server (related: aws/language-servers#831) --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
1 parent d79c13c commit e01574f

File tree

9 files changed

+2992
-684
lines changed

9 files changed

+2992
-684
lines changed

package-lock.json

Lines changed: 2355 additions & 624 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

packages/amazonq/src/app/inline/completion.ts

Lines changed: 147 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -13,81 +13,175 @@ import {
1313
TextDocument,
1414
commands,
1515
languages,
16+
Disposable,
1617
} from 'vscode'
1718
import { LanguageClient } from 'vscode-languageclient'
18-
import {
19-
InlineCompletionListWithReferences,
20-
InlineCompletionWithReferencesParams,
21-
inlineCompletionWithReferencesRequestType,
22-
logInlineCompletionSessionResultsNotificationType,
23-
LogInlineCompletionSessionResultsParams,
24-
} from '@aws/language-server-runtimes/protocol'
19+
import { LogInlineCompletionSessionResultsParams } from '@aws/language-server-runtimes/protocol'
20+
import { SessionManager } from './sessionManager'
21+
import { RecommendationService } from './recommendationService'
2522
import { CodeWhispererConstants } from 'aws-core-vscode/codewhisperer'
2623

27-
export function registerInlineCompletion(languageClient: LanguageClient) {
28-
const inlineCompletionProvider = new AmazonQInlineCompletionItemProvider(languageClient)
29-
languages.registerInlineCompletionItemProvider(CodeWhispererConstants.platformLanguageIds, inlineCompletionProvider)
24+
export class InlineCompletionManager implements Disposable {
25+
private disposable: Disposable
26+
private inlineCompletionProvider: AmazonQInlineCompletionItemProvider
27+
private languageClient: LanguageClient
28+
private sessionManager: SessionManager
29+
private recommendationService: RecommendationService
30+
private readonly logSessionResultMessageName = 'aws/logInlineCompletionSessionResults'
31+
32+
constructor(languageClient: LanguageClient) {
33+
this.languageClient = languageClient
34+
this.sessionManager = new SessionManager()
35+
this.recommendationService = new RecommendationService(this.sessionManager)
36+
this.inlineCompletionProvider = new AmazonQInlineCompletionItemProvider(
37+
languageClient,
38+
this.recommendationService,
39+
this.sessionManager
40+
)
41+
this.disposable = languages.registerInlineCompletionItemProvider(
42+
CodeWhispererConstants.platformLanguageIds,
43+
this.inlineCompletionProvider
44+
)
45+
}
46+
47+
public dispose(): void {
48+
if (this.disposable) {
49+
this.disposable.dispose()
50+
}
51+
}
52+
53+
public registerInlineCompletion() {
54+
const onInlineAcceptance = async (
55+
sessionId: string,
56+
itemId: string,
57+
requestStartTime: number,
58+
firstCompletionDisplayLatency?: number
59+
) => {
60+
// TODO: also log the seen state for other suggestions in session
61+
const params: LogInlineCompletionSessionResultsParams = {
62+
sessionId: sessionId,
63+
completionSessionResult: {
64+
[itemId]: {
65+
seen: true,
66+
accepted: true,
67+
discarded: false,
68+
},
69+
},
70+
totalSessionDisplayTime: Date.now() - requestStartTime,
71+
firstCompletionDisplayLatency: firstCompletionDisplayLatency,
72+
}
73+
this.languageClient.sendNotification(this.logSessionResultMessageName, params)
74+
this.disposable.dispose()
75+
this.disposable = languages.registerInlineCompletionItemProvider(
76+
CodeWhispererConstants.platformLanguageIds,
77+
this.inlineCompletionProvider
78+
)
79+
}
80+
commands.registerCommand('aws.amazonq.acceptInline', onInlineAcceptance)
3081

31-
const onInlineAcceptance = async (
32-
sessionId: string,
33-
itemId: string,
34-
requestStartTime: number,
35-
firstCompletionDisplayLatency?: number
36-
) => {
37-
const params: LogInlineCompletionSessionResultsParams = {
38-
sessionId: sessionId,
39-
completionSessionResult: {
40-
[itemId]: {
41-
seen: true,
42-
accepted: true,
43-
discarded: false,
82+
const onInlineRejection = async () => {
83+
await commands.executeCommand('editor.action.inlineSuggest.hide')
84+
// TODO: also log the seen state for other suggestions in session
85+
this.disposable.dispose()
86+
this.disposable = languages.registerInlineCompletionItemProvider(
87+
CodeWhispererConstants.platformLanguageIds,
88+
this.inlineCompletionProvider
89+
)
90+
const sessionId = this.sessionManager.getActiveSession()?.sessionId
91+
const itemId = this.sessionManager.getActiveRecommendation()[0]?.itemId
92+
if (!sessionId || !itemId) {
93+
return
94+
}
95+
const params: LogInlineCompletionSessionResultsParams = {
96+
sessionId: sessionId,
97+
completionSessionResult: {
98+
[itemId]: {
99+
seen: true,
100+
accepted: false,
101+
discarded: false,
102+
},
44103
},
45-
},
46-
totalSessionDisplayTime: Date.now() - requestStartTime,
47-
firstCompletionDisplayLatency: firstCompletionDisplayLatency,
104+
}
105+
this.languageClient.sendNotification(this.logSessionResultMessageName, params)
48106
}
49-
languageClient.sendNotification(logInlineCompletionSessionResultsNotificationType as any, params)
107+
commands.registerCommand('aws.amazonq.rejectCodeSuggestion', onInlineRejection)
108+
109+
/*
110+
We have to overwrite the prev. and next. commands because the inlineCompletionProvider only contained the current item
111+
To show prev. and next. recommendation we need to re-register a new provider with the previous or next item
112+
*/
113+
114+
const swapProviderAndShow = async () => {
115+
await commands.executeCommand('editor.action.inlineSuggest.hide')
116+
this.disposable.dispose()
117+
this.disposable = languages.registerInlineCompletionItemProvider(
118+
CodeWhispererConstants.platformLanguageIds,
119+
new AmazonQInlineCompletionItemProvider(
120+
this.languageClient,
121+
this.recommendationService,
122+
this.sessionManager,
123+
false
124+
)
125+
)
126+
await commands.executeCommand('editor.action.inlineSuggest.trigger')
127+
}
128+
129+
const prevCommandHandler = async () => {
130+
this.sessionManager.decrementActiveIndex()
131+
await swapProviderAndShow()
132+
}
133+
commands.registerCommand('editor.action.inlineSuggest.showPrevious', prevCommandHandler)
134+
135+
const nextCommandHandler = async () => {
136+
this.sessionManager.incrementActiveIndex()
137+
await swapProviderAndShow()
138+
}
139+
commands.registerCommand('editor.action.inlineSuggest.showNext', nextCommandHandler)
50140
}
51-
commands.registerCommand('aws.sample-vscode-ext-amazonq.accept', onInlineAcceptance)
52141
}
53142

54143
export class AmazonQInlineCompletionItemProvider implements InlineCompletionItemProvider {
55-
constructor(private readonly languageClient: LanguageClient) {}
144+
constructor(
145+
private readonly languageClient: LanguageClient,
146+
private readonly recommendationService: RecommendationService,
147+
private readonly sessionManager: SessionManager,
148+
private readonly isNewSession: boolean = true
149+
) {}
56150

57151
async provideInlineCompletionItems(
58152
document: TextDocument,
59153
position: Position,
60154
context: InlineCompletionContext,
61155
token: CancellationToken
62156
): Promise<InlineCompletionItem[] | InlineCompletionList> {
63-
const requestStartTime = Date.now()
64-
const request: InlineCompletionWithReferencesParams = {
65-
textDocument: {
66-
uri: document.uri.toString(),
67-
},
68-
position,
69-
context,
157+
if (this.isNewSession) {
158+
// make service requests if it's a new session
159+
await this.recommendationService.getAllRecommendations(
160+
this.languageClient,
161+
document,
162+
position,
163+
context,
164+
token
165+
)
70166
}
71-
72-
const response = await this.languageClient.sendRequest(
73-
inlineCompletionWithReferencesRequestType as any,
74-
request,
75-
token
76-
)
77-
78-
const list: InlineCompletionListWithReferences = response as InlineCompletionListWithReferences
79-
this.languageClient.info(`Client: Received ${list.items.length} suggestions`)
80-
const firstCompletionDisplayLatency = Date.now() - requestStartTime
81-
82-
// Add completion session tracking and attach onAcceptance command to each item to record used decision
83-
for (const item of list.items) {
167+
// get active item from session for displaying
168+
const items = this.sessionManager.getActiveRecommendation()
169+
const session = this.sessionManager.getActiveSession()
170+
if (!session || !items.length) {
171+
return []
172+
}
173+
for (const item of items) {
84174
item.command = {
85-
command: 'aws.sample-vscode-ext-amazonq.accept',
175+
command: 'aws.amazonq.acceptInline',
86176
title: 'On acceptance',
87-
arguments: [list.sessionId, item.itemId, requestStartTime, firstCompletionDisplayLatency],
177+
arguments: [
178+
session.sessionId,
179+
item.itemId,
180+
session.requestStartTime,
181+
session.firstCompletionDisplayLatency,
182+
],
88183
}
89184
}
90-
91-
return list as InlineCompletionList
185+
return items as InlineCompletionItem[]
92186
}
93187
}
Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import {
7+
InlineCompletionListWithReferences,
8+
InlineCompletionWithReferencesParams,
9+
inlineCompletionWithReferencesRequestType,
10+
} from '@aws/language-server-runtimes/protocol'
11+
import { CancellationToken, InlineCompletionContext, Position, TextDocument } from 'vscode'
12+
import { LanguageClient } from 'vscode-languageclient'
13+
import { SessionManager } from './sessionManager'
14+
15+
export class RecommendationService {
16+
constructor(private readonly sessionManager: SessionManager) {}
17+
18+
async getAllRecommendations(
19+
languageClient: LanguageClient,
20+
document: TextDocument,
21+
position: Position,
22+
context: InlineCompletionContext,
23+
token: CancellationToken
24+
) {
25+
const request: InlineCompletionWithReferencesParams = {
26+
textDocument: {
27+
uri: document.uri.toString(),
28+
},
29+
position,
30+
context,
31+
}
32+
const requestStartTime = Date.now()
33+
34+
// Handle first request
35+
const firstResult: InlineCompletionListWithReferences = await languageClient.sendRequest(
36+
inlineCompletionWithReferencesRequestType as any,
37+
request,
38+
token
39+
)
40+
41+
const firstCompletionDisplayLatency = Date.now() - requestStartTime
42+
this.sessionManager.startSession(
43+
firstResult.sessionId,
44+
firstResult.items,
45+
requestStartTime,
46+
firstCompletionDisplayLatency
47+
)
48+
49+
if (firstResult.partialResultToken) {
50+
// If there are more results to fetch, handle them in the background
51+
this.processRemainingRequests(languageClient, request, firstResult, token).catch((error) => {
52+
languageClient.warn(`Error when getting suggestions: ${error}`)
53+
})
54+
} else {
55+
this.sessionManager.closeSession()
56+
}
57+
}
58+
59+
private async processRemainingRequests(
60+
languageClient: LanguageClient,
61+
initialRequest: InlineCompletionWithReferencesParams,
62+
firstResult: InlineCompletionListWithReferences,
63+
token: CancellationToken
64+
): Promise<void> {
65+
let nextToken = firstResult.partialResultToken
66+
while (nextToken) {
67+
const request = { ...initialRequest, partialResultToken: nextToken }
68+
const result: InlineCompletionListWithReferences = await languageClient.sendRequest(
69+
inlineCompletionWithReferencesRequestType as any,
70+
request,
71+
token
72+
)
73+
this.sessionManager.updateSessionSuggestions(result.items)
74+
nextToken = result.partialResultToken
75+
}
76+
this.sessionManager.closeSession()
77+
}
78+
}

0 commit comments

Comments
 (0)