Merge pull request aws#6851 from ashishrp-aws/feature/agentic-chat

feat(chat): fix for q agentic chat loop

Author: mr-lee

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

@@ -79,6 +79,7 @@ import {
     aditionalContentNameLimit,
     additionalContentInnerContextLimit,
     contextMaxLength,
+    tools,
 } from '../../constants'
 import { ChatSession } from '../../clients/chat/v0/chat'
 import { ChatHistoryManager } from '../../storages/chatHistory'
@@ -775,13 +776,6 @@ export class ChatController {
                     command,
                 })
 
-                this.chatHistoryManager.appendUserMessage({
-                    userInputMessage: {
-                        content: prompt,
-                        userIntent: this.userIntentRecognizer.getFromContextMenuCommand(command),
-                    },
-                })
-
                 return this.generateResponse(
                     {
                         message: prompt,
@@ -861,13 +855,6 @@ export class ChatController {
                 context: lastTriggerEvent.context,
             })
 
-            this.chatHistoryManager.appendUserMessage({
-                userInputMessage: {
-                    content: message.message,
-                    userIntent: message.userIntent,
-                },
-            })
-
             return this.generateResponse(
                 {
                     message: message.message,
@@ -923,7 +910,7 @@ export class ChatController {
                         //     result = await executeBash.invoke(process.stdout)
                         //     break
                         // }
-                        case 'fs_read': {
+                        case 'fsRead': {
                             const fsRead = new FsRead(toolUse.input as unknown as FsReadParams)
                             await fsRead.validate()
                             result = await fsRead.invoke()
@@ -955,17 +942,9 @@ export class ChatController {
                     toolResults.push({ content: [{ text: e.message }], toolUseId: toolUse.toolUseId, status: 'error' })
                 }
 
-                this.chatHistoryManager.appendUserMessage({
-                    userInputMessage: {
-                        content: 'Tool Results',
-                        userIntent: undefined,
-                        origin: Origin.IDE,
-                    },
-                })
-
                 await this.generateResponse(
                     {
-                        message: 'Tool Results',
+                        message: '',
                         trigger: ChatTriggerType.ChatMessage,
                         query: undefined,
                         codeSelection: context?.focusAreaContext?.selectionInsideExtendedCodeBlock,
@@ -979,6 +958,7 @@ export class ChatController {
                         context: undefined,
                         toolResults: toolResults,
                         origin: Origin.IDE,
+                        chatHistory: this.chatHistoryManager.getHistory(),
                     },
                     triggerID
                 )
@@ -1002,13 +982,6 @@ export class ChatController {
                     type: 'chat_message',
                     context,
                 })
-                this.chatHistoryManager.appendUserMessage({
-                    userInputMessage: {
-                        content: message.message,
-                        userIntent: message.userIntent,
-                        origin: Origin.IDE,
-                    },
-                })
                 await this.generateResponse(
                     {
                         message: message.message,
@@ -1348,6 +1321,18 @@ export class ChatController {
             this.telemetryHelper.recordEnterFocusConversation(triggerEvent.tabID)
             this.telemetryHelper.recordStartConversation(triggerEvent, triggerPayload)
 
+            this.chatHistoryManager.appendUserMessage({
+                userInputMessage: {
+                    content: triggerPayload.message,
+                    userIntent: triggerPayload.userIntent,
+                    ...(triggerPayload.origin && { origin: triggerPayload.origin }),
+                    userInputMessageContext: {
+                        tools: tools,
+                        ...(triggerPayload.toolResults && { toolResults: triggerPayload.toolResults }),
+                    },
+                },
+            })
+
             getLogger().info(
                 `response to tab: ${tabID} conversationID: ${session.sessionIdentifier} requestID: ${
                     response.$metadata.requestId

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

@@ -9,6 +9,7 @@ import {
     AuthNeededException,
     CodeReference,
     ContextCommandData,
+    CustomFormActionMessage,
     EditorContextCommandMessage,
     OpenSettingsMessage,
     QuickActionMessage,
@@ -232,13 +233,14 @@ export class Messenger {
                                 tabID
                             )
                         )
+
+                        this.dispatcher.sendCustomFormActionMessage(
+                            new CustomFormActionMessage(tabID, {
+                                id: 'confirm-tool-use',
+                            })
+                        )
                         // TODO: setup permission action
                         // if (!isConfirmationRequired) {
-                        //     this.dispatcher.sendCustomFormActionMessage(
-                        //         new CustomFormActionMessage(tabID, {
-                        //             id: 'confirm-tool-use',
-                        //         })
-                        //     )
                         // }
                     }
 
@@ -392,7 +394,9 @@ export class Messenger {
                         messageId: messageID,
                         content: message,
                         references: codeReference,
-                        toolUses: [{ ...toolUse }],
+                        ...(toolUse &&
+                            toolUse.input !== undefined &&
+                            toolUse.input !== '' && { toolUses: [{ ...toolUse }] }),
                     },
                 })
 

packages/core/src/codewhispererChat/storages/chatHistory.ts

@@ -69,7 +69,7 @@ export class ChatHistoryManager {
         if (!newMessage.userInputMessage?.content || newMessage.userInputMessage?.content.trim() === '') {
             this.logger.warn('input must not be empty when adding new messages')
         }
-        this.history.push(this.lastUserMessage)
+        this.history.push(this.formatChatHistoryMessage(this.lastUserMessage))
     }
 
     /**
@@ -196,4 +196,19 @@ export class ChatHistoryManager {
             this.lastUserMessage.userInputMessage = msg
         }
     }
+
+    private formatChatHistoryMessage(message: ChatMessage): ChatMessage {
+        if (message.userInputMessage !== undefined) {
+            return {
+                userInputMessage: {
+                    ...message.userInputMessage,
+                    userInputMessageContext: {
+                        ...message.userInputMessage.userInputMessageContext,
+                        tools: undefined,
+                    },
+                },
+            }
+        }
+        return message
+    }
 }