Skip to content

Commit 280f4e5

Browse files
andrewyuqandrewyuq
authored
Add context transparency feature to Q chat @workspace (#9)
* Add context transparency feature to Q chat @workspace * remove console log * remove console log * remove unused FileClickMessage
1 parent e20f552 commit 280f4e5

File tree

14 files changed

+214
-6
lines changed

14 files changed

+214
-6
lines changed

‎packages/core/src/amazonq/lsp/lspController.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@ export interface Chunk {
3131
readonly context?: string
3232
readonly relativePath?: string
3333
readonly programmingLanguage?: string
34+
readonly startLine?: number
35+
readonly endLine?: number
3436
}
3537

3638
export interface Content {
@@ -292,11 +294,15 @@ export class LspController {
292294
programmingLanguage: {
293295
languageName: chunk.programmingLanguage,
294296
},
297+
startLine: chunk.startLine ?? -1,
298+
endLine: chunk.endLine ?? -1,
295299
})
296300
} else {
297301
resp.push({
298302
text: text,
299303
relativeFilePath: chunk.relativePath ? chunk.relativePath : path.basename(chunk.filePath),
304+
startLine: chunk.startLine ?? -1,
305+
endLine: chunk.endLine ?? -1,
300306
})
301307
}
302308
}

‎packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,7 @@ export class Connector extends BaseConnector {
9696
codeReference: messageData.codeReference,
9797
userIntent: messageData.userIntent,
9898
codeBlockLanguage: messageData.codeBlockLanguage,
99+
contextList: messageData.contextList,
99100
}
100101

101102
// If it is not there we will not set it
@@ -230,4 +231,14 @@ export class Connector extends BaseConnector {
230231
tabID: tabId,
231232
})
232233
}
234+
235+
onFileClick = (tabID: string, filePath: string, messageId?: string) => {
236+
this.sendMessageToExtension({
237+
command: 'file-click',
238+
tabID,
239+
messageId,
240+
filePath,
241+
tabType: 'cwc',
242+
})
243+
}
233244
}

‎packages/core/src/amazonq/webview/ui/connector.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,12 @@ export interface CWCChatItem extends ChatItem {
6262
traceId?: string
6363
userIntent?: UserIntent
6464
codeBlockLanguage?: string
65+
contextList?: Context[]
66+
}
67+
68+
export interface Context {
69+
relativeFilePath: string
70+
lineRanges: Array<{ first: number; second: number }> // List of [startLine, endLine] tuples
6571
}
6672

6773
export interface ConnectorProps {
@@ -604,6 +610,9 @@ export class Connector {
604610
case 'doc':
605611
this.docChatConnector.onOpenDiff(tabID, filePath, deleted)
606612
break
613+
case 'cwc':
614+
this.cwChatConnector.onFileClick(tabID, filePath, messageId)
615+
break
607616
}
608617
}
609618

‎packages/core/src/amazonq/webview/ui/main.ts‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -368,6 +368,30 @@ export const createMynahUI = (
368368
return
369369
}
370370

371+
if (item.contextList !== undefined && item.contextList.length > 0) {
372+
item.header = {
373+
fileList: {
374+
fileTreeTitle: '',
375+
filePaths: item.contextList.map((file) => file.relativeFilePath),
376+
rootFolderTitle: 'Context',
377+
collapsedByDefault: true,
378+
hideFileCount: true,
379+
details: Object.fromEntries(
380+
item.contextList.map((file) => [
381+
file.relativeFilePath,
382+
{
383+
label: file.lineRanges
384+
.map((range) => `line ${range.first} - ${range.second}`)
385+
.join(', '),
386+
description: file.relativeFilePath,
387+
clickable: true,
388+
},
389+
])
390+
),
391+
},
392+
}
393+
}
394+
371395
if (
372396
item.body !== undefined ||
373397
item.relatedContent !== undefined ||

‎packages/core/src/amazonqFeatureDev/client/codewhispererruntime-2022-11-11.json‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2587,6 +2587,12 @@
25872587
},
25882588
"documentSymbols": {
25892589
"shape": "DocumentSymbols"
2590+
},
2591+
"startLine": {
2592+
"shape": "Integer"
2593+
},
2594+
"endLine": {
2595+
"shape": "Integer"
25902596
}
25912597
}
25922598
},

‎packages/core/src/codewhisperer/client/user-service-2.json‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1907,6 +1907,12 @@
19071907
"documentSymbols": {
19081908
"shape": "DocumentSymbols",
19091909
"documentation": "<p>DocumentSymbols parsed from a text document</p>"
1910+
},
1911+
"startLine": {
1912+
"shape": "Integer"
1913+
},
1914+
"endLine": {
1915+
"shape": "Integer"
19101916
}
19111917
},
19121918
"documentation": "<p>Represents an IDE retrieved relevant Text Document / File</p>"

‎packages/core/src/codewhispererChat/app.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import {
2727
UIFocusMessage,
2828
AcceptDiff,
2929
QuickCommandGroupActionClick,
30+
FileClick,
3031
} from './controllers/chat/model'
3132
import { EditorContextCommand, registerCommands } from './commands/registerCommands'
3233
import { ContextSelectedMessage, CustomFormActionMessage } from './view/connector/connector'
@@ -54,6 +55,7 @@ export function init(appContext: AmazonQAppInitContext) {
5455
processQuickCommandGroupActionClicked: new EventEmitter<QuickCommandGroupActionClick>(),
5556
processCustomFormAction: new EventEmitter<CustomFormActionMessage>(),
5657
processContextSelected: new EventEmitter<ContextSelectedMessage>(),
58+
processFileClick: new EventEmitter<FileClick>(),
5759
}
5860

5961
const cwChatControllerMessageListeners = {
@@ -114,6 +116,9 @@ export function init(appContext: AmazonQAppInitContext) {
114116
processContextSelected: new MessageListener<ContextSelectedMessage>(
115117
cwChatControllerEventEmitters.processContextSelected
116118
),
119+
processFileClick: new MessageListener<FileClick>(
120+
cwChatControllerEventEmitters.processFileClick
121+
),
117122
}
118123

119124
const cwChatControllerMessagePublishers = {
@@ -176,6 +181,9 @@ export function init(appContext: AmazonQAppInitContext) {
176181
processContextSelected: new MessagePublisher<ContextSelectedMessage>(
177182
cwChatControllerEventEmitters.processContextSelected
178183
),
184+
processFileClick: new MessagePublisher<FileClick>(
185+
cwChatControllerEventEmitters.processFileClick
186+
),
179187
}
180188

181189
new CwChatController(

‎packages/core/src/codewhispererChat/clients/chat/v0/chat.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@ import { UserWrittenCodeTracker } from '../../../../codewhisperer/tracker/userWr
1414
export class ChatSession {
1515
private sessionId?: string
1616

17+
contexts: Map<number, Map<string, { first: number; second: number }[]>> = new Map()
18+
currentContextId: number = 0
1719
public get sessionIdentifier(): string | undefined {
1820
return this.sessionId
1921
}

‎packages/core/src/codewhispererChat/controllers/chat/controller.ts‎

Lines changed: 98 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55
import * as path from 'path'
6-
import { Event as VSCodeEvent, Uri } from 'vscode'
6+
import { Event as VSCodeEvent, Uri, workspace, window, ViewColumn, Position, Selection } from 'vscode'
77
import { EditorContextExtractor } from '../../editor/context/extractor'
88
import { ChatSessionStorage } from '../../storages/chatSession'
99
import { Messenger, MessengerResponseType, StaticTextResponseType } from './messenger/messenger'
@@ -28,6 +28,8 @@ import {
2828
ViewDiff,
2929
AcceptDiff,
3030
QuickCommandGroupActionClick,
31+
MergedRelevantDocument,
32+
FileClick,
3133
} from './model'
3234
import {
3335
AppToWebViewMessageDispatcher,
@@ -41,7 +43,7 @@ import { EditorContextCommand } from '../../commands/registerCommands'
4143
import { PromptsGenerator } from './prompts/promptsGenerator'
4244
import { TriggerEventsStorage } from '../../storages/triggerEvents'
4345
import { SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
44-
import { CodeWhispererStreamingServiceException } from '@amzn/codewhisperer-streaming'
46+
import { CodeWhispererStreamingServiceException, RelevantTextDocument } from '@amzn/codewhisperer-streaming'
4547
import { UserIntentRecognizer } from './userIntent/userIntentRecognizer'
4648
import { CWCTelemetryHelper, recordTelemetryChatRunCommand } from './telemetryHelper'
4749
import { CodeWhispererTracker } from '../../../codewhisperer/tracker/codewhispererTracker'
@@ -90,6 +92,7 @@ export interface ChatControllerMessagePublishers {
9092
readonly processQuickCommandGroupActionClicked: MessagePublisher<QuickCommandGroupActionClick>
9193
readonly processCustomFormAction: MessagePublisher<CustomFormActionMessage>
9294
readonly processContextSelected: MessagePublisher<ContextSelectedMessage>
95+
readonly processFileClick: MessagePublisher<FileClick>
9396
}
9497

9598
export interface ChatControllerMessageListeners {
@@ -114,6 +117,7 @@ export interface ChatControllerMessageListeners {
114117
readonly processQuickCommandGroupActionClicked: MessageListener<QuickCommandGroupActionClick>
115118
readonly processCustomFormAction: MessageListener<CustomFormActionMessage>
116119
readonly processContextSelected: MessageListener<ContextSelectedMessage>
120+
readonly processFileClick: MessageListener<FileClick>
117121
}
118122

119123
export class ChatController {
@@ -243,6 +247,9 @@ export class ChatController {
243247
this.chatControllerMessageListeners.processContextSelected.onMessage((data) => {
244248
return this.processContextSelected(data)
245249
})
250+
this.chatControllerMessageListeners.processFileClick.onMessage((data) => {
251+
return this.processFileClickMessage(data)
252+
})
246253
}
247254

248255
private processFooterInfoLinkClick(click: FooterInfoLinkClick) {
@@ -524,6 +531,7 @@ export class ChatController {
524531
`Create a saved prompt`
525532
)
526533
}
534+
527535
private processQuickCommandGroupActionClicked(message: QuickCommandGroupActionClick) {
528536
if (message.actionId === 'create-prompt') {
529537
this.handlePromptCreate(message.tabID)
@@ -558,6 +566,36 @@ export class ChatController {
558566
this.handlePromptCreate(message.tabID)
559567
}
560568
}
569+
private async processFileClickMessage(message: FileClick) {
570+
let session = this.sessionStorage.getSession(message.tabID)
571+
// TODO remove currentContextId but use messageID to track context for each answer message
572+
const lineRanges = session.contexts.get(session.currentContextId)?.get(message.filePath)
573+
574+
if (!lineRanges) {
575+
return
576+
}
577+
const projectRoot = workspace.workspaceFolders?.[0]?.uri.fsPath
578+
if (!projectRoot) {
579+
return
580+
}
581+
582+
const absoluteFilePath = path.join(projectRoot, message.filePath)
583+
584+
// Open the file in VSCode
585+
const document = await workspace.openTextDocument(absoluteFilePath)
586+
const editor = await window.showTextDocument(document, ViewColumn.Active)
587+
588+
// Create multiple selections based on line ranges
589+
const selections: Selection[] = lineRanges.map(({ first, second }) => {
590+
const startPosition = new Position(first - 1, 0) // Convert 1-based to 0-based
591+
const endPosition = new Position(second - 1, document.lineAt(second - 1).range.end.character)
592+
return new Selection(startPosition.line, startPosition.character, endPosition.line, endPosition.character)
593+
})
594+
595+
// Apply multiple selections to the editor using the new API
596+
editor.selection = selections[0] // Set the first selection as active
597+
editor.selections = selections // Apply multiple selections
598+
}
561599

562600
private processException(e: any, tabID: string) {
563601
let errorMessage = ''
@@ -880,9 +918,12 @@ export class ChatController {
880918
if (CodeWhispererSettings.instance.isLocalIndexEnabled()) {
881919
const start = performance.now()
882920
triggerPayload.relevantTextDocuments = await LspController.instance.query(triggerPayload.message)
921+
triggerPayload.mergedRelevantDocuments = this.mergeRelevantTextDocuments(
922+
triggerPayload.relevantTextDocuments
923+
)
883924
for (const doc of triggerPayload.relevantTextDocuments) {
884925
getLogger().info(
885-
`amazonq: Using workspace files ${doc.relativeFilePath}, content(partial): ${doc.text?.substring(0, 200)}`
926+
`amazonq: Using workspace files ${doc.relativeFilePath}, content(partial): ${doc.text?.substring(0, 200)}, start line: ${doc.startLine}, end line: ${doc.endLine}`
886927
)
887928
}
888929
triggerPayload.projectContextQueryLatencyMs = performance.now() - start
@@ -904,12 +945,25 @@ export class ChatController {
904945
},
905946
{ timeout: 500, interval: 200, truthy: false }
906947
)
948+
triggerPayload.mergedRelevantDocuments = this.mergeRelevantTextDocuments(
949+
triggerPayload.relevantTextDocuments
950+
)
907951
triggerPayload.projectContextQueryLatencyMs = performance.now() - start
908952
}
909953
}
910954

911955
const request = triggerPayloadToChatRequest(triggerPayload)
912956
const session = this.sessionStorage.getSession(tabID)
957+
958+
session.currentContextId++
959+
session.contexts.set(session.currentContextId, new Map())
960+
triggerPayload.mergedRelevantDocuments?.forEach((doc) => {
961+
const currentContext = session.contexts.get(session.currentContextId)
962+
if (currentContext) {
963+
currentContext.set(doc.relativeFilePath, doc.lineRanges)
964+
}
965+
})
966+
913967
getLogger().info(
914968
`request from tab: ${tabID} conversationID: ${session.sessionIdentifier} request: ${inspect(request, {
915969
depth: 12,
@@ -918,7 +972,7 @@ export class ChatController {
918972
let response: MessengerResponseType | undefined = undefined
919973
session.createNewTokenSource()
920974
try {
921-
this.messenger.sendInitalStream(tabID, triggerID)
975+
this.messenger.sendInitalStream(tabID, triggerID, triggerPayload.mergedRelevantDocuments)
922976
this.telemetryHelper.setConversationStreamStartTime(tabID)
923977
if (isSsoConnection(AuthUtil.instance.conn)) {
924978
const { $metadata, generateAssistantResponseResponse } = await session.chatSso(request)
@@ -948,4 +1002,44 @@ export class ChatController {
9481002
this.processException(e, tabID)
9491003
}
9501004
}
1005+
1006+
private mergeRelevantTextDocuments(
1007+
documents: RelevantTextDocument[] | undefined
1008+
): MergedRelevantDocument[] | undefined {
1009+
if (documents === undefined) {
1010+
return undefined
1011+
}
1012+
return Object.entries(
1013+
documents.reduce<Record<string, { first: number; second: number }[]>>((acc, doc) => {
1014+
if (!doc.relativeFilePath || doc.startLine === undefined || doc.endLine === undefined) {
1015+
return acc // Skip invalid documents
1016+
}
1017+
1018+
if (!acc[doc.relativeFilePath]) {
1019+
acc[doc.relativeFilePath] = []
1020+
}
1021+
acc[doc.relativeFilePath].push({ first: doc.startLine, second: doc.endLine })
1022+
return acc
1023+
}, {})
1024+
).map(([filePath, ranges]) => {
1025+
// Sort by startLine
1026+
const sortedRanges = ranges.sort((a, b) => a.first - b.first)
1027+
1028+
const mergedRanges: { first: number; second: number }[] = []
1029+
for (const { first, second } of sortedRanges) {
1030+
if (mergedRanges.length === 0 || mergedRanges[mergedRanges.length - 1].second < first - 1) {
1031+
// If no overlap, add new range
1032+
mergedRanges.push({ first, second })
1033+
} else {
1034+
// Merge overlapping or consecutive ranges
1035+
mergedRanges[mergedRanges.length - 1].second = Math.max(
1036+
mergedRanges[mergedRanges.length - 1].second,
1037+
second
1038+
)
1039+
}
1040+
}
1041+
1042+
return { relativeFilePath: filePath, lineRanges: mergedRanges }
1043+
})
1044+
}
9511045
}

0 commit comments

Comments
 (0)