refactor(cw): Rename the mergedTextDocument field. (#6666)

leigaol · web-flow · commit ae20598f35d7 · 2025-02-24T18:50:41.000-05:00
## Problem The mergedRelevantTextDocument should have a better name. It is actually the documents that are referenced. ## Solution --- - 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.
diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts
@@ -29,7 +29,7 @@ import {
     ViewDiff,
     AcceptDiff,
     QuickCommandGroupActionClick,
-    MergedRelevantDocument,
+    DocumentReference,
     FileClick,
     RelevantTextDocumentAddition,
 } from './model'
@@ -974,53 +974,53 @@ export class ChatController {
         }
 
         const relativePaths = await this.resolveContextCommandPayload(triggerPayload)
-        // TODO: resolve the context into real context up to 90k
         triggerPayload.useRelevantDocuments = false
-        triggerPayload.mergedRelevantDocuments = []
-        if (triggerPayload.message) {
-            triggerPayload.useRelevantDocuments = triggerPayload.context?.some(
-                (context) => typeof context !== 'string' && context.command === '@workspace'
-            )
-            if (triggerPayload.useRelevantDocuments) {
-                triggerPayload.message = triggerPayload.message.replace(/workspace/, '')
-                if (CodeWhispererSettings.instance.isLocalIndexEnabled()) {
-                    const start = performance.now()
-                    triggerPayload.relevantTextDocuments = await LspController.instance.query(triggerPayload.message)
-                    triggerPayload.mergedRelevantDocuments = this.mergeRelevantTextDocuments(
-                        triggerPayload.relevantTextDocuments
+        triggerPayload.documentReferences = []
+        triggerPayload.useRelevantDocuments = triggerPayload.context?.some(
+            (context) => typeof context !== 'string' && context.command === '@workspace'
+        )
+        if (triggerPayload.useRelevantDocuments && triggerPayload.message) {
+            triggerPayload.message = triggerPayload.message.replace(/workspace/, '')
+            if (CodeWhispererSettings.instance.isLocalIndexEnabled()) {
+                const start = performance.now()
+                triggerPayload.relevantTextDocuments = await LspController.instance.query(triggerPayload.message)
+
+                for (const doc of triggerPayload.relevantTextDocuments) {
+                    getLogger().info(
+                        `amazonq: Using workspace files ${doc.relativeFilePath}, content(partial): ${doc.text?.substring(0, 200)}, start line: ${doc.startLine}, end line: ${doc.endLine}`
                     )
-                    for (const doc of triggerPayload.relevantTextDocuments) {
-                        getLogger().info(
-                            `amazonq: Using workspace files ${doc.relativeFilePath}, content(partial): ${doc.text?.substring(0, 200)}, start line: ${doc.startLine}, end line: ${doc.endLine}`
-                        )
-                    }
-                    triggerPayload.projectContextQueryLatencyMs = performance.now() - start
-                } else {
-                    this.messenger.sendOpenSettingsMessage(triggerID, tabID)
-                    return
                 }
+                triggerPayload.projectContextQueryLatencyMs = performance.now() - start
+            } else {
+                this.messenger.sendOpenSettingsMessage(triggerID, tabID)
+                return
             }
         }
+        triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments || [])
+
+        // TODO: make sure the user input + current focused document + relevantDocument + additionalContext
+        // combined does not exceed 100k characters before generating the request payload.
+        // Do truncation and make sure triggerPayload.documentReferences is up-to-date after truncation
 
         const request = triggerPayloadToChatRequest(triggerPayload)
         const session = this.sessionStorage.getSession(tabID)
 
         session.currentContextId++
         session.contexts.set(session.currentContextId, new Map())
-        if (triggerPayload.mergedRelevantDocuments !== undefined) {
-            const relativePathsOfMergedRelevantDocuments = triggerPayload.mergedRelevantDocuments.map(
+        if (triggerPayload.documentReferences !== undefined) {
+            const relativePathsOfMergedRelevantDocuments = triggerPayload.documentReferences.map(
                 (doc) => doc.relativeFilePath
             )
             for (const relativePath of relativePaths) {
                 if (!relativePathsOfMergedRelevantDocuments.includes(relativePath)) {
-                    triggerPayload.mergedRelevantDocuments.push({
+                    triggerPayload.documentReferences.push({
                         relativeFilePath: relativePath,
                         lineRanges: [{ first: -1, second: -1 }],
                     })
                 }
             }
-            if (triggerPayload.mergedRelevantDocuments) {
-                for (const doc of triggerPayload.mergedRelevantDocuments) {
+            if (triggerPayload.documentReferences) {
+                for (const doc of triggerPayload.documentReferences) {
                     const currentContext = session.contexts.get(session.currentContextId)
                     if (currentContext) {
                         currentContext.set(doc.relativeFilePath, doc.lineRanges)
@@ -1037,7 +1037,7 @@ export class ChatController {
         let response: MessengerResponseType | undefined = undefined
         session.createNewTokenSource()
         try {
-            this.messenger.sendInitalStream(tabID, triggerID, triggerPayload.mergedRelevantDocuments)
+            this.messenger.sendInitalStream(tabID, triggerID, triggerPayload.documentReferences)
             this.telemetryHelper.setConversationStreamStartTime(tabID)
             if (isSsoConnection(AuthUtil.instance.conn)) {
                 const { $metadata, generateAssistantResponseResponse } = await session.chatSso(request)
@@ -1068,11 +1068,9 @@ export class ChatController {
         }
     }
 
-    private mergeRelevantTextDocuments(
-        documents: RelevantTextDocumentAddition[] | undefined
-    ): MergedRelevantDocument[] | undefined {
-        if (documents === undefined) {
-            return undefined
+    private mergeRelevantTextDocuments(documents: RelevantTextDocumentAddition[]): DocumentReference[] {
+        if (documents.length === 0) {
+            return []
         }
         return Object.entries(
             documents.reduce<Record<string, { first: number; second: number }[]>>((acc, doc) => {
diff --git a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts
@@ -25,7 +25,7 @@ import { ChatMessage, ErrorMessage, FollowUp, Suggestion } from '../../../view/c
 import { ChatSession } from '../../../clients/chat/v0/chat'
 import { ChatException } from './model'
 import { CWCTelemetryHelper } from '../telemetryHelper'
-import { ChatPromptCommandType, MergedRelevantDocument, TriggerPayload } from '../model'
+import { ChatPromptCommandType, DocumentReference, TriggerPayload } from '../model'
 import { getHttpStatusCode, getRequestId, ToolkitError } from '../../../../shared/errors'
 import { keys } from '../../../../shared/utilities/tsUtils'
 import { getLogger } from '../../../../shared/logger/logger'
@@ -69,7 +69,7 @@ export class Messenger {
     public sendInitalStream(
         tabID: string,
         triggerID: string,
-        mergedRelevantDocuments: MergedRelevantDocument[] | undefined
+        mergedRelevantDocuments: DocumentReference[] | undefined
     ) {
         this.dispatcher.sendChatMessage(
             new ChatMessage(
diff --git a/packages/core/src/codewhispererChat/controllers/chat/model.ts b/packages/core/src/codewhispererChat/controllers/chat/model.ts
@@ -189,15 +189,17 @@ export interface TriggerPayload {
     readonly context?: string[] | QuickActionCommand[]
     relevantTextDocuments?: RelevantTextDocumentAddition[]
     additionalContents?: AdditionalContentEntry[]
-    mergedRelevantDocuments?: MergedRelevantDocument[]
+    // a reference to all documents used in chat payload
+    // for providing better context transparency
+    documentReferences?: DocumentReference[]
     useRelevantDocuments?: boolean
     traceId?: string
 }
 
 // TODO move this to API definition (or just use this across the codebase)
 export type RelevantTextDocumentAddition = RelevantTextDocument & { startLine: number; endLine: number }
 
-export interface MergedRelevantDocument {
+export interface DocumentReference {
     readonly relativeFilePath: string
     readonly lineRanges: Array<{ first: number; second: number }>
 }
diff --git a/packages/core/src/codewhispererChat/view/connector/connector.ts b/packages/core/src/codewhispererChat/view/connector/connector.ts
@@ -8,7 +8,7 @@ import { MessagePublisher } from '../../../amazonq/messages/messagePublisher'
 import { EditorContextCommandType } from '../../commands/registerCommands'
 import { AuthFollowUpType } from '../../../amazonq/auth/model'
 import { ChatItemButton, ChatItemFormItem, MynahUIDataModel, QuickActionCommand } from '@aws/mynah-ui'
-import { MergedRelevantDocument } from '../../controllers/chat/model'
+import { DocumentReference } from '../../controllers/chat/model'
 
 class UiMessage {
     readonly time: number = Date.now()
@@ -207,7 +207,7 @@ export interface ChatMessageProps {
     readonly messageID: string
     readonly userIntent: string | undefined
     readonly codeBlockLanguage: string | undefined
-    readonly contextList: MergedRelevantDocument[] | undefined
+    readonly contextList: DocumentReference[] | undefined
 }
 
 export class ChatMessage extends UiMessage {
@@ -222,7 +222,7 @@ export class ChatMessage extends UiMessage {
     readonly messageID: string | undefined
     readonly userIntent: string | undefined
     readonly codeBlockLanguage: string | undefined
-    readonly contextList: MergedRelevantDocument[] | undefined
+    readonly contextList: DocumentReference[] | undefined
     override type = 'chatMessage'
 
     constructor(props: ChatMessageProps, tabID: string) {
