Skip to content

Commit d1b3cc1

Browse files
zixlin7zixlin7
authored
Add project context for Q chat (#4661)
Add option for users to use local workspace indexing to provide more context to chat
1 parent f97e9be commit d1b3cc1

File tree

37 files changed

+1567
-146
lines changed

37 files changed

+1567
-146
lines changed

‎.changes/next-release/feature-528029bf-b5c5-4c8a-9ea5-f27112432af3.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type" : "feature",
3+
"description" : "Add support for [Amazon Q Chat Workspace Context](https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/workspace-context.html). Customers can use @workspace to ask questions regarding local workspace."
4+
}

‎gradle/libs.versions.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ sshd = "2.12.1"
3131
wiremock = "3.3.1"
3232
undercouch-download = "5.2.1"
3333
zjsonpatch = "0.4.16"
34+
nimbus-jose-jwt = "9.24.4"
3435

3536
[libraries]
3637
assertj = { module = "org.assertj:assertj-core", version.ref = "assertJ" }
@@ -106,6 +107,7 @@ sshd-scp = { module = "org.apache.sshd:sshd-scp", version.ref = "sshd" }
106107
sshd-sftp = { module = "org.apache.sshd:sshd-sftp", version.ref = "sshd" }
107108
wiremock = { module = "org.wiremock:wiremock", version.ref = "wiremock" }
108109
zjsonpatch = { module = "com.flipkart.zjsonpatch:zjsonpatch", version.ref = "zjsonpatch" }
110+
nimbus-jose-jwt = {module = "com.nimbusds:nimbus-jose-jwt", version.ref = "nimbus-jose-jwt"}
109111

110112
[bundles]
111113
jackson = ["jackson-datetime", "jackson-kotlin", "jackson-yaml", "jackson-xml"]

‎plugins/amazonq/chat/jetbrains-community/build.gradle.kts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ dependencies {
1919
implementation(project(":plugin-amazonq:shared:jetbrains-community"))
2020
// everything references codewhisperer, which is not ideal
2121
implementation(project(":plugin-amazonq:codewhisperer:jetbrains-community"))
22+
implementation(libs.nimbus.jose.jwt)
2223

2324
compileOnly(project(":plugin-core:jetbrains-community"))
2425

‎plugins/amazonq/chat/jetbrains-community/resources/META-INF/plugin-chat.xml‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,9 @@
77

88
<projectListeners>
99
<listener class="software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindowListener"
10-
topic="com.intellij.openapi.wm.ex.ToolWindowManagerListener"/>
10+
topic="com.intellij.openapi.wm.ex.ToolWindowManagerListener"/>
11+
<listener class="software.aws.toolkits.jetbrains.services.cwc.editor.context.project.ProjectContextEditorListener"
12+
topic="com.intellij.openapi.fileEditor.FileEditorManagerListener"/>
1113
</projectListeners>
1214

1315
<extensions defaultExtensionNs="com.intellij">

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/App.kt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ class App : AmazonQApp {
4141
"chat-item-voted" to IncomingCwcMessage.ChatItemVoted::class,
4242
"chat-item-feedback" to IncomingCwcMessage.ChatItemFeedback::class,
4343
"ui-focus" to IncomingCwcMessage.UIFocus::class,
44+
"open-settings" to IncomingCwcMessage.OpenSettings::class,
4445
"auth-follow-up-was-clicked" to IncomingCwcMessage.AuthFollowUpWasClicked::class,
4546

4647
// JB specific (not in vscode)
@@ -72,6 +73,7 @@ class App : AmazonQApp {
7273
is IncomingCwcMessage.ChatItemFeedback -> inboundAppMessagesHandler.processChatItemFeedback(message)
7374
is IncomingCwcMessage.UIFocus -> inboundAppMessagesHandler.processUIFocus(message)
7475
is IncomingCwcMessage.AuthFollowUpWasClicked -> inboundAppMessagesHandler.processAuthFollowUpClick(message)
76+
is IncomingCwcMessage.OpenSettings -> inboundAppMessagesHandler.processOpenSettings(message)
7577
is OnboardingPageInteraction -> inboundAppMessagesHandler.processOnboardingPageInteraction(message)
7678

7779
// JB specific (not in vscode)

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/InboundAppMessagesHandler.kt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ interface InboundAppMessagesHandler {
2525
suspend fun processUIFocus(message: IncomingCwcMessage.UIFocus)
2626
suspend fun processAuthFollowUpClick(message: IncomingCwcMessage.AuthFollowUpWasClicked)
2727
suspend fun processOnboardingPageInteraction(message: OnboardingPageInteraction)
28+
suspend fun processOpenSettings(message: IncomingCwcMessage.OpenSettings)
2829

2930
// JB specific (not in vscode)
3031
suspend fun processClearQuickAction(message: IncomingCwcMessage.ClearChat)

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/clients/chat/model/Requests.kt‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import com.fasterxml.jackson.annotation.JsonProperty
77
import software.amazon.awssdk.services.codewhispererstreaming.model.UserIntent
88
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
99
import software.aws.toolkits.jetbrains.services.cwc.editor.context.ActiveFileContext
10+
import software.aws.toolkits.jetbrains.services.cwc.editor.context.project.RelevantDocument
1011

1112
enum class TriggerType {
1213
Click,
@@ -21,7 +22,8 @@ data class ChatRequestData(
2122
val activeFileContext: ActiveFileContext,
2223
val userIntent: UserIntent?,
2324
val triggerType: TriggerType,
24-
val customization: CodeWhispererCustomization?
25+
val customization: CodeWhispererCustomization?,
26+
val relevantTextDocuments: List<RelevantDocument>
2527
)
2628

2729
interface CodeNames {

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/clients/chat/v1/ChatSessionV1.kt‎

Lines changed: 61 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import software.amazon.awssdk.services.codewhispererstreaming.model.GenerateAssi
2727
import software.amazon.awssdk.services.codewhispererstreaming.model.Position
2828
import software.amazon.awssdk.services.codewhispererstreaming.model.ProgrammingLanguage
2929
import software.amazon.awssdk.services.codewhispererstreaming.model.Range
30+
import software.amazon.awssdk.services.codewhispererstreaming.model.RelevantTextDocument
3031
import software.amazon.awssdk.services.codewhispererstreaming.model.SupplementaryWebLinksEvent
3132
import software.amazon.awssdk.services.codewhispererstreaming.model.SymbolType
3233
import software.amazon.awssdk.services.codewhispererstreaming.model.TextDocument
@@ -50,6 +51,7 @@ import software.aws.toolkits.jetbrains.services.cwc.clients.chat.model.Reference
5051
import software.aws.toolkits.jetbrains.services.cwc.clients.chat.model.SuggestedFollowUp
5152
import software.aws.toolkits.jetbrains.services.cwc.clients.chat.model.Suggestion
5253
import software.aws.toolkits.jetbrains.services.cwc.editor.context.ActiveFileContext
54+
import software.aws.toolkits.jetbrains.services.cwc.editor.context.project.RelevantDocument
5355

5456
class ChatSessionV1(
5557
private val project: Project,
@@ -60,7 +62,6 @@ class ChatSessionV1(
6062
override fun chat(data: ChatRequestData): Flow<ChatResponseEvent> = callbackFlow {
6163
var requestId: String = ""
6264
var statusCode: Int = 0
63-
6465
val responseHandler = GenerateAssistantResponseResponseHandler.builder()
6566
.onResponse {
6667
requestId = it.responseMetadata().requestId()
@@ -197,9 +198,7 @@ class ChatSessionV1(
197198
*/
198199
private fun ChatRequestData.toChatRequest(): GenerateAssistantResponseRequest {
199200
val userInputMessageContextBuilder = UserInputMessageContext.builder()
200-
if (activeFileContext.fileContext != null) {
201-
userInputMessageContextBuilder.editorState(activeFileContext.toEditorState())
202-
}
201+
userInputMessageContextBuilder.editorState(activeFileContext.toEditorState(relevantTextDocuments))
203202
val userInputMessageContext = userInputMessageContextBuilder.build()
204203

205204
val userInput = UserInputMessage.builder()
@@ -218,67 +217,74 @@ class ChatSessionV1(
218217
.build()
219218
}
220219

221-
private fun ActiveFileContext.toEditorState(): EditorState {
222-
// Cursor State
223-
val start = focusAreaContext?.codeSelectionRange?.start
224-
val end = focusAreaContext?.codeSelectionRange?.end
220+
private fun ActiveFileContext.toEditorState(relevantDocuments: List<RelevantDocument>): EditorState {
221+
val editorStateBuilder = EditorState.builder()
222+
if (fileContext != null) {
223+
val cursorStateBuilder = CursorState.builder()
224+
// Cursor State
225+
val start = focusAreaContext?.codeSelectionRange?.start
226+
val end = focusAreaContext?.codeSelectionRange?.end
225227

226-
val cursorStateBuilder = CursorState.builder()
228+
if (start != null && end != null) {
229+
cursorStateBuilder.range(
230+
Range.builder()
231+
.start(
232+
Position.builder()
233+
.line(start.row)
234+
.character(start.column)
235+
.build(),
236+
)
237+
.end(
238+
Position.builder()
239+
.line(end.row)
240+
.character(end.column)
241+
.build(),
242+
).build(),
243+
)
244+
}
245+
editorStateBuilder.cursorState(cursorStateBuilder.build())
227246

228-
if (start != null && end != null) {
229-
cursorStateBuilder.range(
230-
Range.builder()
231-
.start(
232-
Position.builder()
233-
.line(start.row)
234-
.character(start.column)
235-
.build(),
236-
)
237-
.end(
238-
Position.builder()
239-
.line(end.row)
240-
.character(end.column)
241-
.build(),
242-
).build(),
243-
)
244-
}
247+
// Code Names -> DocumentSymbols
248+
val documentBuilder = TextDocument.builder()
249+
val codeNames = focusAreaContext?.codeNames
245250

246-
// Code Names -> DocumentSymbols
247-
val codeNames = focusAreaContext?.codeNames
248-
val documentBuilder = TextDocument.builder()
251+
val documentSymbolList = codeNames?.fullyQualifiedNames?.used?.map {
252+
DocumentSymbol.builder()
253+
.name(it.symbol?.joinToString(separator = "."))
254+
.type(SymbolType.USAGE)
255+
.source(it.source?.joinToString(separator = "."))
256+
.build()
257+
}?.filter { it.name().length in ChatConstants.FQN_SIZE_MIN until ChatConstants.FQN_SIZE_LIMIT }.orEmpty()
258+
documentBuilder.documentSymbols(documentSymbolList)
249259

250-
val documentSymbolList = codeNames?.fullyQualifiedNames?.used?.map {
251-
DocumentSymbol.builder()
252-
.name(it.symbol?.joinToString(separator = "."))
253-
.type(SymbolType.USAGE)
254-
.source(it.source?.joinToString(separator = "."))
255-
.build()
256-
}?.filter { it.name().length in ChatConstants.FQN_SIZE_MIN until ChatConstants.FQN_SIZE_LIMIT }.orEmpty()
257-
documentBuilder.documentSymbols(documentSymbolList)
260+
// File Text
261+
val trimmedFileText = focusAreaContext?.trimmedSurroundingFileText
262+
documentBuilder.text(trimmedFileText)
258263

259-
// File Text
260-
val trimmedFileText = focusAreaContext?.trimmedSurroundingFileText
261-
documentBuilder.text(trimmedFileText)
264+
// Programming Language
265+
val programmingLanguage = fileContext.fileLanguage
266+
if (programmingLanguage != null && validLanguages.contains(programmingLanguage)) {
267+
documentBuilder.programmingLanguage(
268+
ProgrammingLanguage.builder()
269+
.languageName(programmingLanguage).build(),
270+
)
271+
}
262272

263-
// Programming Language
264-
val programmingLanguage = fileContext?.fileLanguage
265-
if (programmingLanguage != null && validLanguages.contains(programmingLanguage)) {
266-
documentBuilder.programmingLanguage(
267-
ProgrammingLanguage.builder()
268-
.languageName(programmingLanguage).build(),
269-
)
273+
// Relative File Path
274+
val filePath = fileContext.filePath
275+
if (filePath != null) {
276+
documentBuilder.relativeFilePath(filePath.take(ChatConstants.FILE_PATH_SIZE_LIMIT))
277+
}
278+
editorStateBuilder.document(documentBuilder.build())
270279
}
271280

272-
// Relative File Path
273-
val filePath = fileContext?.filePath
274-
if (filePath != null) {
275-
documentBuilder.relativeFilePath(filePath.take(ChatConstants.FILE_PATH_SIZE_LIMIT))
281+
// Relevant Documents
282+
val documents: List<RelevantTextDocument> = relevantDocuments.map { doc ->
283+
RelevantTextDocument.builder().text(doc.text).relativeFilePath(doc.relativeFilePath.take(ChatConstants.FILE_PATH_SIZE_LIMIT)).build()
276284
}
277285

278-
return EditorState.builder()
279-
.cursorState(cursorStateBuilder.build())
280-
.document(documentBuilder.build())
281-
.build()
286+
editorStateBuilder.relevantDocuments(documents)
287+
return editorStateBuilder.build()
282288
}
283289

284290
private fun UserIntent?.toFollowUpType() = when (this) {

0 commit comments

Comments
 (0)