Merge pull request aws#6931 from ctlai95/replace-read-messages

feat(chat): group consecutive read tool messages

Author: mr-lee

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

@@ -168,6 +168,35 @@ export class Connector extends BaseConnector {
         }
     }
 
+    private processToolMessage = async (messageData: any): Promise<void> => {
+        if (this.onChatAnswerUpdated === undefined) {
+            return
+        }
+        const answer: CWCChatItem = {
+            type: messageData.messageType,
+            messageId: messageData.messageID ?? messageData.triggerID,
+            body: messageData.message,
+            followUp: messageData.followUps,
+            canBeVoted: messageData.canBeVoted ?? false,
+            codeReference: messageData.codeReference,
+            userIntent: messageData.contextList,
+            codeBlockLanguage: messageData.codeBlockLanguage,
+            contextList: messageData.contextList,
+            title: messageData.title,
+            buttons: messageData.buttons,
+            fileList: messageData.fileList,
+            header: messageData.header ?? undefined,
+            padding: messageData.padding ?? undefined,
+            fullWidth: messageData.fullWidth ?? undefined,
+            codeBlockActions: messageData.codeBlockActions ?? undefined,
+        }
+        if (answer.messageId) {
+            this.storeChatItem(messageData.tabID, answer.messageId, answer)
+        }
+        this.onChatAnswerUpdated(messageData.tabID, answer)
+        return
+    }
+
     private storeChatItem(tabId: string, messageId: string, item: ChatItem): void {
         if (!this.chatItems.has(tabId)) {
             this.chatItems.set(tabId, new Map())
@@ -224,6 +253,11 @@ export class Connector extends BaseConnector {
             return
         }
 
+        if (messageData.type === 'toolMessage') {
+            await this.processToolMessage(messageData)
+            return
+        }
+
         if (messageData.type === 'editorContextCommandMessage') {
             await this.processEditorContextCommandMessage(messageData)
             return

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

@@ -331,7 +331,12 @@ export const createMynahUI = (
                     ...(item.followUp !== undefined ? { followUp: item.followUp } : {}),
                     ...(item.footer !== undefined ? { footer: item.footer } : {}),
                     ...(item.canBeVoted !== undefined ? { canBeVoted: item.canBeVoted } : {}),
+                    ...(item.fileList !== undefined ? { fileList: item.fileList } : {}),
                     ...(item.header !== undefined ? { header: item.header } : {}),
+                    ...(item.buttons !== undefined ? { buttons: item.buttons } : {}),
+                    ...(item.fullWidth !== undefined ? { fullWidth: item.fullWidth } : {}),
+                    ...(item.padding !== undefined ? { padding: item.padding } : {}),
+                    ...(item.codeBlockActions !== undefined ? { codeBlockActions: item.codeBlockActions } : {}),
                 })
             } else {
                 mynahUI.updateLastChatAnswer(tabID, {
@@ -341,6 +346,10 @@ export const createMynahUI = (
                     ...(item.footer !== undefined ? { footer: item.footer } : {}),
                     ...(item.canBeVoted !== undefined ? { canBeVoted: item.canBeVoted } : {}),
                     ...(item.header !== undefined ? { header: item.header } : {}),
+                    ...(item.buttons !== undefined ? { buttons: item.buttons } : {}),
+                    ...(item.fullWidth !== undefined ? { fullWidth: item.fullWidth } : {}),
+                    ...(item.padding !== undefined ? { padding: item.padding } : {}),
+                    ...(item.codeBlockActions !== undefined ? { codeBlockActions: item.codeBlockActions } : {}),
                 })
             }
         },

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

@@ -22,11 +22,13 @@ export class ChatSession {
      * _readFiles = list of files read from the project to gather context before generating response.
      * _showDiffOnFileWrite = Controls whether to show diff view (true) or file context view (false) to the user
      * _context = Additional context to be passed to the LLM for generating the response
+     * _messageIdToUpdate = messageId of a chat message to be updated, used for reducing consecutive tool messages
      */
     private _readFiles: string[] = []
     private _toolUse: ToolUse | undefined
     private _showDiffOnFileWrite: boolean = false
     private _context: PromptMessage['context']
+    private _messageIdToUpdate: string | undefined
 
     contexts: Map<string, { first: number; second: number }[]> = new Map()
     // TODO: doesn't handle the edge case when two files share the same relativePath string but from different root
@@ -52,6 +54,14 @@ export class ChatSession {
         this._context = context
     }
 
+    public get messageIdToUpdate(): string | undefined {
+        return this._messageIdToUpdate
+    }
+
+    public setMessageIdToUpdate(messageId: string | undefined) {
+        this._messageIdToUpdate = messageId
+    }
+
     public tokenSource!: vscode.CancellationTokenSource
 
     constructor() {

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

@@ -663,9 +663,16 @@ export class ChatController {
                     try {
                         await ToolUtils.validate(tool)
 
-                        const chatStream = new ChatStream(this.messenger, tabID, triggerID, toolUse, {
-                            requiresAcceptance: false,
-                        })
+                        const chatStream = new ChatStream(
+                            this.messenger,
+                            tabID,
+                            triggerID,
+                            // Pass in a different toolUseId so that the output does not overwrite
+                            // any previous messages
+                            { ...toolUse, toolUseId: `${toolUse.toolUseId}-output` },
+                            { requiresAcceptance: false },
+                            undefined
+                        )
                         const output = await ToolUtils.invoke(tool, chatStream)
 
                         toolResults.push({

packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts

@@ -14,6 +14,7 @@ import {
     OpenSettingsMessage,
     QuickActionMessage,
     ShowCustomFormMessage,
+    ToolMessage,
 } from '../../../view/connector/connector'
 import { EditorContextCommandType } from '../../../commands/registerCommands'
 import { ChatResponseStream as qdevChatResponseStream } from '@amzn/amazon-q-developer-streaming-client'
@@ -221,11 +222,35 @@ export class Messenger {
                             if (tool.type === ToolType.FsWrite) {
                                 session.setShowDiffOnFileWrite(true)
                             }
+                            if (
+                                tool.type === ToolType.FsWrite ||
+                                tool.type === ToolType.ExecuteBash ||
+                                eventCounts.has('assistantResponseEvent')
+                            ) {
+                                // FsWrite and ExecuteBash should never replace older messages
+                                // If the current stream also has assistantResponseEvent then reset this as well.
+                                session.setMessageIdToUpdate(undefined)
+                            }
                             const validation = ToolUtils.requiresAcceptance(tool)
 
-                            const chatStream = new ChatStream(this, tabID, triggerID, toolUse, validation)
+                            const chatStream = new ChatStream(
+                                this,
+                                tabID,
+                                triggerID,
+                                toolUse,
+                                validation,
+                                session.messageIdToUpdate
+                            )
                             await ToolUtils.queueDescription(tool, chatStream)
 
+                            if (
+                                session.messageIdToUpdate === undefined &&
+                                (tool.type === ToolType.FsRead || tool.type === ToolType.ListDirectory)
+                            ) {
+                                // Store the first messageId in a chain of tool uses
+                                session.setMessageIdToUpdate(toolUse.toolUseId)
+                            }
+
                             if (!validation.requiresAcceptance) {
                                 // Need separate id for read tool and safe bash command execution as 'confirm-tool-use' id is required to change button status from `Confirm` to `Confirmed` state in cwChatConnector.ts which will impact generic tool execution.
                                 this.dispatcher.sendCustomFormActionMessage(
@@ -429,12 +454,33 @@ export class Messenger {
         )
     }
 
+    public sendInitialToolMessage(tabID: string, triggerID: string, toolUseId: string | undefined) {
+        this.dispatcher.sendChatMessage(
+            new ChatMessage(
+                {
+                    message: '',
+                    messageType: 'answer',
+                    followUps: undefined,
+                    followUpsHeader: undefined,
+                    relatedSuggestions: undefined,
+                    triggerID,
+                    messageID: toolUseId ?? 'toolUse',
+                    userIntent: undefined,
+                    codeBlockLanguage: undefined,
+                    contextList: undefined,
+                },
+                tabID
+            )
+        )
+    }
+
     public sendPartialToolLog(
         message: string,
         tabID: string,
         triggerID: string,
         toolUse: ToolUse | undefined,
-        validation: CommandValidation
+        validation: CommandValidation,
+        messageIdToUpdate: string | undefined
     ) {
         const buttons: ChatItemButton[] = []
         let fileList: ChatItemContent['fileList'] = undefined
@@ -483,16 +529,16 @@ export class Messenger {
             })
         }
 
-        this.dispatcher.sendChatMessage(
-            new ChatMessage(
+        this.dispatcher.sendToolMessage(
+            new ToolMessage(
                 {
                     message: message,
                     messageType: 'answer-part',
                     followUps: undefined,
                     followUpsHeader: undefined,
                     relatedSuggestions: undefined,
                     triggerID,
-                    messageID: toolUse?.toolUseId ?? `tool-output`,
+                    messageID: messageIdToUpdate ?? toolUse?.toolUseId ?? '',
                     userIntent: undefined,
                     codeBlockLanguage: undefined,
                     contextList: undefined,

packages/core/src/codewhispererChat/tools/chatStream.ts

@@ -22,11 +22,16 @@ export class ChatStream extends Writable {
         private readonly triggerID: string,
         private readonly toolUse: ToolUse | undefined,
         private readonly validation: CommandValidation,
+        private readonly messageIdToUpdate: string | undefined,
         private readonly logger = getLogger('chatStream')
     ) {
         super()
         this.logger.debug(`ChatStream created for tabID: ${tabID}, triggerID: ${triggerID}`)
-        this.messenger.sendInitalStream(tabID, triggerID, undefined)
+        if (!messageIdToUpdate) {
+            // If messageIdToUpdate is undefined, we need to first create an empty message
+            // with messageId so it can be updated later
+            this.messenger.sendInitialToolMessage(tabID, triggerID, toolUse?.toolUseId)
+        }
     }
 
     override _write(chunk: Buffer, encoding: BufferEncoding, callback: (error?: Error | null) => void): void {
@@ -38,21 +43,13 @@ export class ChatStream extends Writable {
             this.tabID,
             this.triggerID,
             this.toolUse,
-            this.validation
+            this.validation,
+            this.messageIdToUpdate
         )
         callback()
     }
 
     override _final(callback: (error?: Error | null) => void): void {
-        if (this.accumulatedLogs.trim().length > 0) {
-            this.messenger.sendPartialToolLog(
-                this.accumulatedLogs,
-                this.tabID,
-                this.triggerID,
-                this.toolUse,
-                this.validation
-            )
-        }
         callback()
     }
 }

packages/core/src/codewhispererChat/view/connector/connector.ts

@@ -286,6 +286,10 @@ export class ChatMessage extends UiMessage {
     }
 }
 
+export class ToolMessage extends ChatMessage {
+    override type = 'toolMessage'
+}
+
 export interface FollowUp {
     readonly type: string
     readonly pillText: string
@@ -340,6 +344,10 @@ export class AppToWebViewMessageDispatcher {
         this.appsToWebViewMessagePublisher.publish(message)
     }
 
+    public sendToolMessage(message: ToolMessage) {
+        this.appsToWebViewMessagePublisher.publish(message)
+    }
+
     public sendEditorContextCommandMessage(message: EditorContextCommandMessage) {
         this.appsToWebViewMessagePublisher.publish(message)
     }