Skip to content

Commit 2dbe485

Browse files
ashishrp-awsashishrp-aws
authored
feat(chat): Initial Agentic Chat loop Setup (aws#6844)
## Problem Setting up new Q Agentic Chat loop with tool use. ## Solution Initial Agentic Chat loop Setup with tool use --- - 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.
1 parent 567a954 commit 2dbe485

File tree

9 files changed

+390
-28
lines changed

9 files changed

+390
-28
lines changed

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

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,11 @@
44
*/
55

66
import { SendMessageCommandOutput, SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
7-
import { GenerateAssistantResponseCommandOutput, GenerateAssistantResponseRequest } from '@amzn/codewhisperer-streaming'
7+
import {
8+
GenerateAssistantResponseCommandOutput,
9+
GenerateAssistantResponseRequest,
10+
ToolUse,
11+
} from '@amzn/codewhisperer-streaming'
812
import * as vscode from 'vscode'
913
import { ToolkitError } from '../../../../shared/errors'
1014
import { createCodeWhispererChatStreamingClient } from '../../../../shared/clients/codewhispererChatClient'
@@ -13,6 +17,7 @@ import { UserWrittenCodeTracker } from '../../../../codewhisperer/tracker/userWr
1317

1418
export class ChatSession {
1519
private sessionId?: string
20+
private _toolUse: ToolUse | undefined
1621

1722
contexts: Map<string, { first: number; second: number }[]> = new Map()
1823
// TODO: doesn't handle the edge case when two files share the same relativePath string but from different root
@@ -22,6 +27,14 @@ export class ChatSession {
2227
return this.sessionId
2328
}
2429

30+
public get toolUse(): ToolUse | undefined {
31+
return this._toolUse
32+
}
33+
34+
public setToolUse(toolUse: ToolUse | undefined) {
35+
this._toolUse = toolUse
36+
}
37+
2538
public tokenSource!: vscode.CancellationTokenSource
2639

2740
constructor() {

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
*/
55
import * as path from 'path'
66
import fs from '../shared/fs/fs'
7+
import { Tool } from '@amzn/codewhisperer-streaming'
8+
import toolsJson from '../codewhispererChat/tools/tool_index.json'
79

810
export const promptFileExtension = '.md'
911

@@ -19,3 +21,10 @@ export const getUserPromptsDirectory = () => {
1921
}
2022

2123
export const createSavedPromptCommandId = 'create-saved-prompt'
24+
25+
export const tools: Tool[] = Object.entries(toolsJson).map(([, toolSpec]) => ({
26+
toolSpecification: {
27+
...toolSpec,
28+
inputSchema: { json: toolSpec.inputSchema },
29+
},
30+
}))

‎packages/core/src/codewhispererChat/controllers/chat/chatRequest/converter.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import {
1313
} from '@amzn/codewhisperer-streaming'
1414
import { ChatTriggerType, TriggerPayload } from '../model'
1515
import { undefinedIfEmpty } from '../../../../shared/utilities/textUtilities'
16+
import { tools } from '../../../constants'
1617

1718
const fqnNameSizeDownLimit = 1
1819
const fqnNameSizeUpLimit = 256
@@ -115,10 +116,16 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c
115116
cursorState,
116117
relevantDocuments,
117118
useRelevantDocuments,
119+
// TODO: Need workspace folders here after model update.
118120
},
119121
additionalContext: triggerPayload.additionalContents,
122+
tools,
123+
...(triggerPayload.toolResults !== undefined &&
124+
triggerPayload.toolResults !== null && { toolResults: triggerPayload.toolResults }),
120125
},
121126
userIntent: triggerPayload.userIntent,
127+
...(triggerPayload.origin !== undefined &&
128+
triggerPayload.origin !== null && { origin: triggerPayload.origin }),
122129
},
123130
},
124131
chatTriggerType,

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

Lines changed: 106 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ import { EditorContextCommand } from '../../commands/registerCommands'
4545
import { PromptsGenerator } from './prompts/promptsGenerator'
4646
import { TriggerEventsStorage } from '../../storages/triggerEvents'
4747
import { SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
48-
import { CodeWhispererStreamingServiceException } from '@amzn/codewhisperer-streaming'
48+
import { CodeWhispererStreamingServiceException, Origin, ToolResult } from '@amzn/codewhisperer-streaming'
4949
import { UserIntentRecognizer } from './userIntent/userIntentRecognizer'
5050
import { CWCTelemetryHelper, recordTelemetryChatRunCommand } from './telemetryHelper'
5151
import { CodeWhispererTracker } from '../../../codewhisperer/tracker/codewhispererTracker'
@@ -81,6 +81,7 @@ import {
8181
} from '../../constants'
8282
import { ChatSession } from '../../clients/chat/v0/chat'
8383
import { ChatHistoryManager } from '../../storages/chatHistory'
84+
import { FsRead, FsReadParams } from '../../tools/fsRead'
8485

8586
export interface ChatControllerMessagePublishers {
8687
readonly processPromptChatMessage: MessagePublisher<PromptMessage>
@@ -577,6 +578,8 @@ export class ChatController {
577578
const newFileDoc = await vscode.workspace.openTextDocument(newFilePath)
578579
await vscode.window.showTextDocument(newFileDoc)
579580
telemetry.ui_click.emit({ elementId: 'amazonq_createSavedPrompt' })
581+
} else if (message.action.id === 'confirm-tool-use') {
582+
await this.processToolUseMessage(message)
580583
}
581584
}
582585

@@ -834,10 +837,108 @@ export class ChatController {
834837
}
835838
}
836839

840+
private async processToolUseMessage(message: CustomFormActionMessage) {
841+
const tabID = message.tabID
842+
if (!tabID) {
843+
return
844+
}
845+
this.editorContextExtractor
846+
.extractContextForTrigger('ChatMessage')
847+
.then(async (context) => {
848+
const triggerID = randomUUID()
849+
this.triggerEventsStorage.addTriggerEvent({
850+
id: triggerID,
851+
tabID: message.tabID,
852+
message: undefined,
853+
type: 'chat_message',
854+
context,
855+
})
856+
const session = this.sessionStorage.getSession(tabID)
857+
const toolUse = session.toolUse
858+
if (!toolUse || !toolUse.input) {
859+
return
860+
}
861+
session.setToolUse(undefined)
862+
863+
let result: any
864+
const toolResults: ToolResult[] = []
865+
try {
866+
switch (toolUse.name) {
867+
// case 'execute_bash': {
868+
// const executeBash = new ExecuteBash(toolUse.input as unknown as ExecuteBashParams)
869+
// await executeBash.validate()
870+
// result = await executeBash.invoke(process.stdout)
871+
// break
872+
// }
873+
case 'fs_read': {
874+
const fsRead = new FsRead(toolUse.input as unknown as FsReadParams)
875+
await fsRead.validate()
876+
result = await fsRead.invoke()
877+
break
878+
}
879+
// case 'fs_write': {
880+
// const fsWrite = new FsWrite(toolUse.input as unknown as FsWriteParams)
881+
// const ctx = new DefaultContext()
882+
// result = await fsWrite.invoke(ctx, process.stdout)
883+
// break
884+
// }
885+
// case 'open_file': {
886+
// result = await openFile(toolUse.input as unknown as OpenFileParams)
887+
// break
888+
// }
889+
default:
890+
break
891+
}
892+
toolResults.push({
893+
content: [
894+
result.output.kind === 'text'
895+
? { text: result.output.content }
896+
: { json: result.output.content },
897+
],
898+
toolUseId: toolUse.toolUseId,
899+
status: 'success',
900+
})
901+
} catch (e: any) {
902+
toolResults.push({ content: [{ text: e.message }], toolUseId: toolUse.toolUseId, status: 'error' })
903+
}
904+
905+
this.chatHistoryManager.appendUserMessage({
906+
userInputMessage: {
907+
content: 'Tool Results',
908+
userIntent: undefined,
909+
origin: Origin.IDE,
910+
},
911+
})
912+
913+
await this.generateResponse(
914+
{
915+
message: 'Tool Results',
916+
trigger: ChatTriggerType.ChatMessage,
917+
query: undefined,
918+
codeSelection: context?.focusAreaContext?.selectionInsideExtendedCodeBlock,
919+
fileText: context?.focusAreaContext?.extendedCodeBlock,
920+
fileLanguage: context?.activeFileContext?.fileLanguage,
921+
filePath: context?.activeFileContext?.filePath,
922+
matchPolicy: context?.activeFileContext?.matchPolicy,
923+
codeQuery: context?.focusAreaContext?.names,
924+
userIntent: undefined,
925+
customization: getSelectedCustomization(),
926+
context: undefined,
927+
toolResults: toolResults,
928+
origin: Origin.IDE,
929+
},
930+
triggerID
931+
)
932+
})
933+
.catch((e) => {
934+
this.processException(e, tabID)
935+
})
936+
}
937+
837938
private async processPromptMessageAsNewThread(message: PromptMessage) {
838939
this.editorContextExtractor
839940
.extractContextForTrigger('ChatMessage')
840-
.then((context) => {
941+
.then(async (context) => {
841942
const triggerID = randomUUID()
842943
this.triggerEventsStorage.addTriggerEvent({
843944
id: triggerID,
@@ -850,9 +951,10 @@ export class ChatController {
850951
userInputMessage: {
851952
content: message.message,
852953
userIntent: message.userIntent,
954+
origin: Origin.IDE,
853955
},
854956
})
855-
return this.generateResponse(
957+
await this.generateResponse(
856958
{
857959
message: message.message,
858960
trigger: ChatTriggerType.ChatMessage,
@@ -867,6 +969,7 @@ export class ChatController {
867969
customization: getSelectedCustomization(),
868970
context: message.context,
869971
chatHistory: this.chatHistoryManager.getHistory(),
972+
origin: Origin.IDE,
870973
},
871974
triggerID
872975
)

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

Lines changed: 115 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import {
2020
ChatResponseStream as cwChatResponseStream,
2121
CodeWhispererStreamingServiceException,
2222
SupplementaryWebLink,
23+
ToolUse,
2324
} from '@amzn/codewhisperer-streaming'
2425
import { ChatMessage, ErrorMessage, FollowUp, Suggestion } from '../../../view/connector/connector'
2526
import { ChatSession } from '../../../clients/chat/v0/chat'
@@ -131,6 +132,8 @@ export class Messenger {
131132
let followUps: FollowUp[] = []
132133
let relatedSuggestions: Suggestion[] = []
133134
let codeBlockLanguage: string = 'plaintext'
135+
let toolUseInput = ''
136+
const toolUse: ToolUse = { toolUseId: undefined, name: undefined, input: undefined }
134137

135138
if (response.message === undefined) {
136139
throw new ToolkitError(
@@ -158,7 +161,7 @@ export class Messenger {
158161
})
159162

160163
const eventCounts = new Map<string, number>()
161-
waitUntil(
164+
await waitUntil(
162165
async () => {
163166
for await (const chatEvent of response.message!) {
164167
for (const key of keys(chatEvent)) {
@@ -188,6 +191,53 @@ export class Messenger {
188191
]
189192
}
190193

194+
const cwChatEvent: cwChatResponseStream = chatEvent
195+
if (
196+
cwChatEvent.toolUseEvent?.input !== undefined &&
197+
cwChatEvent.toolUseEvent.input.length > 0 &&
198+
!cwChatEvent.toolUseEvent.stop
199+
) {
200+
toolUseInput += cwChatEvent.toolUseEvent.input
201+
}
202+
203+
if (cwChatEvent.toolUseEvent?.stop) {
204+
toolUse.input = JSON.parse(toolUseInput)
205+
toolUse.toolUseId = cwChatEvent.toolUseEvent.toolUseId ?? ''
206+
toolUse.name = cwChatEvent.toolUseEvent.name ?? ''
207+
session.setToolUse(toolUse)
208+
209+
const message = this.getToolUseMessage(toolUse)
210+
// const isConfirmationRequired = this.getIsConfirmationRequired(toolUse)
211+
212+
this.dispatcher.sendChatMessage(
213+
new ChatMessage(
214+
{
215+
message,
216+
messageType: 'answer',
217+
followUps: undefined,
218+
followUpsHeader: undefined,
219+
relatedSuggestions: undefined,
220+
codeReference,
221+
triggerID,
222+
messageID: toolUse.toolUseId,
223+
userIntent: triggerPayload.userIntent,
224+
codeBlockLanguage: codeBlockLanguage,
225+
contextList: undefined,
226+
// TODO: confirmation buttons
227+
},
228+
tabID
229+
)
230+
)
231+
// TODO: setup permission action
232+
// if (!isConfirmationRequired) {
233+
// this.dispatcher.sendCustomFormActionMessage(
234+
// new CustomFormActionMessage(tabID, {
235+
// id: 'confirm-tool-use',
236+
// })
237+
// )
238+
// }
239+
}
240+
191241
if (
192242
chatEvent.assistantResponseEvent?.content !== undefined &&
193243
chatEvent.assistantResponseEvent.content.length > 0
@@ -338,7 +388,7 @@ export class Messenger {
338388
messageId: messageID,
339389
content: message,
340390
references: codeReference,
341-
// TODO: Add tools data and follow up prompt details
391+
toolUses: [{ ...toolUse }],
342392
},
343393
})
344394

@@ -533,4 +583,67 @@ export class Messenger {
533583
new ShowCustomFormMessage(tabID, formItems, buttons, title, description)
534584
)
535585
}
586+
587+
// TODO: Make this cleaner
588+
// private getIsConfirmationRequired(toolUse: ToolUse) {
589+
// if (toolUse.name === 'execute_bash') {
590+
// const executeBash = new ExecuteBash(toolUse.input as unknown as ExecuteBashParams)
591+
// return executeBash.requiresAcceptance()
592+
// }
593+
// return toolUse.name === 'fs_write'
594+
// }
595+
private getToolUseMessage(toolUse: ToolUse) {
596+
if (toolUse.name === 'fs_read') {
597+
return `Reading the file at \`${(toolUse.input as any)?.path}\` using the \`fs_read\` tool.`
598+
}
599+
// if (toolUse.name === 'execute_bash') {
600+
// const input = toolUse.input as unknown as ExecuteBashParams
601+
// return `Executing the bash command
602+
// \`\`\`bash
603+
// ${input.command}
604+
// \`\`\`
605+
// using the \`execute_bash\` tool.`
606+
// }
607+
// if (toolUse.name === 'fs_write') {
608+
// const input = toolUse.input as unknown as FsWriteParams
609+
// switch (input.command) {
610+
// case 'create': {
611+
// return `Writing
612+
// \`\`\`
613+
// ${input.file_text}
614+
// \`\`\`
615+
// into the file at \`${input.path}\` using the \`fs_write\` tool.`
616+
// }
617+
// case 'str_replace': {
618+
// return `Replacing
619+
// \`\`\`
620+
// ${input.old_str}
621+
// \`\`\`
622+
// with
623+
// \`\`\`
624+
// ${input.new_str}
625+
// \`\`\`
626+
// at \`${input.path}\` using the \`fs_write\` tool.`
627+
// }
628+
// case 'insert': {
629+
// return `Inserting
630+
// \`\`\`
631+
// ${input.new_str}
632+
// \`\`\`
633+
// at line
634+
// \`\`\`
635+
// ${input.insert_line}
636+
// \`\`\`
637+
// at \`${input.path}\` using the \`fs_write\` tool.`
638+
// }
639+
// case 'append': {
640+
// return `Appending
641+
// \`\`\`
642+
// ${input.new_str}
643+
// \`\`\`
644+
// at \`${input.path}\` using the \`fs_write\` tool.`
645+
// }
646+
// }
647+
// }
648+
}
536649
}

0 commit comments

Comments
 (0)