From 9376bd7d8748b44c1eff6de218e2efcd38ecac16 Mon Sep 17 00:00:00 2001 From: Tai Lai Date: Thu, 3 Apr 2025 13:38:28 -0700 Subject: [PATCH] feat(chat): group consecutive read tool messages --- .../webview/ui/apps/cwChatConnector.ts | 27 +++++++++ packages/core/src/amazonq/webview/ui/main.ts | 1 + .../codewhispererChat/clients/chat/v0/chat.ts | 10 ++++ .../controllers/chat/controller.ts | 13 ++++- .../controllers/chat/messenger/messenger.ts | 56 +++++++++++++++++-- .../src/codewhispererChat/tools/chatStream.ts | 19 +++---- .../view/connector/connector.ts | 8 +++ 7 files changed, 115 insertions(+), 19 deletions(-) diff --git a/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts b/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts index 4565e6cbffc..4a2f1a7e2ef 100644 --- a/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts +++ b/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts @@ -154,6 +154,28 @@ export class Connector extends BaseConnector { } } + private processToolMessage = async (messageData: any): Promise => { + 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, + } + this.onChatAnswerUpdated(messageData.tabID, answer) + return + } + processContextCommandData(messageData: any) { if (messageData.data) { this.onContextCommandDataReceived(messageData.data) @@ -199,6 +221,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 diff --git a/packages/core/src/amazonq/webview/ui/main.ts b/packages/core/src/amazonq/webview/ui/main.ts index 05ae716a798..19c936852d1 100644 --- a/packages/core/src/amazonq/webview/ui/main.ts +++ b/packages/core/src/amazonq/webview/ui/main.ts @@ -331,6 +331,7 @@ 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 } : {}), }) } else { mynahUI.updateLastChatAnswer(tabID, { diff --git a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts b/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts index f8f08f71148..35948e8eb79 100644 --- a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts +++ b/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 = 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() { diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts index da6e11bec5f..3b6b1ee06b1 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/controller.ts +++ b/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({ diff --git a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts index 942a2944542..7bfc93fa3ea 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts +++ b/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 + ) 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 @@ -474,8 +520,8 @@ export class Messenger { }) } - this.dispatcher.sendChatMessage( - new ChatMessage( + this.dispatcher.sendToolMessage( + new ToolMessage( { message: message, messageType: 'answer-part', @@ -483,7 +529,7 @@ export class Messenger { followUpsHeader: undefined, relatedSuggestions: undefined, triggerID, - messageID: toolUse?.toolUseId ?? `tool-output`, + messageID: messageIdToUpdate ?? toolUse?.toolUseId ?? '', userIntent: undefined, codeBlockLanguage: undefined, contextList: undefined, diff --git a/packages/core/src/codewhispererChat/tools/chatStream.ts b/packages/core/src/codewhispererChat/tools/chatStream.ts index 07af534c06e..99fae93e6fa 100644 --- a/packages/core/src/codewhispererChat/tools/chatStream.ts +++ b/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() } } diff --git a/packages/core/src/codewhispererChat/view/connector/connector.ts b/packages/core/src/codewhispererChat/view/connector/connector.ts index 62176ccb1a1..945b087aa5b 100644 --- a/packages/core/src/codewhispererChat/view/connector/connector.ts +++ b/packages/core/src/codewhispererChat/view/connector/connector.ts @@ -253,6 +253,10 @@ export class ChatMessage extends UiMessage { } } +export class ToolMessage extends ChatMessage { + override type = 'toolMessage' +} + export interface FollowUp { readonly type: string readonly pillText: string @@ -307,6 +311,10 @@ export class AppToWebViewMessageDispatcher { this.appsToWebViewMessagePublisher.publish(message) } + public sendToolMessage(message: ToolMessage) { + this.appsToWebViewMessagePublisher.publish(message) + } + public sendEditorContextCommandMessage(message: EditorContextCommandMessage) { this.appsToWebViewMessagePublisher.publish(message) }