Skip to content

Commit 8822a39

Browse files
laileni-awslaileni-aws
authored
Merge branch 'feature/agentic-chat' into part1
2 parents f2d1467 + 2dbe485 commit 8822a39

File tree

11 files changed

+590
-28
lines changed

11 files changed

+590
-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'
@@ -21,6 +25,7 @@ export class ChatSession {
2125
private _listOfReadFiles: string[] = []
2226
private _filePath: string | undefined
2327
private _tempFilePath: string | undefined
28+
private _toolUse: ToolUse | undefined
2429

2530
contexts: Map<string, { first: number; second: number }[]> = new Map()
2631
// TODO: doesn't handle the edge case when two files share the same relativePath string but from different root
@@ -30,6 +35,14 @@ export class ChatSession {
3035
return this.sessionId
3136
}
3237

38+
public get toolUse(): ToolUse | undefined {
39+
return this._toolUse
40+
}
41+
42+
public setToolUse(toolUse: ToolUse | undefined) {
43+
this._toolUse = toolUse
44+
}
45+
3346
public tokenSource!: vscode.CancellationTokenSource
3447

3548
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
@@ -46,7 +46,7 @@ import { EditorContextCommand } from '../../commands/registerCommands'
4646
import { PromptsGenerator } from './prompts/promptsGenerator'
4747
import { TriggerEventsStorage } from '../../storages/triggerEvents'
4848
import { SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
49-
import { CodeWhispererStreamingServiceException } from '@amzn/codewhisperer-streaming'
49+
import { CodeWhispererStreamingServiceException, Origin, ToolResult } from '@amzn/codewhisperer-streaming'
5050
import { UserIntentRecognizer } from './userIntent/userIntentRecognizer'
5151
import { CWCTelemetryHelper, recordTelemetryChatRunCommand } from './telemetryHelper'
5252
import { CodeWhispererTracker } from '../../../codewhisperer/tracker/codewhispererTracker'
@@ -83,6 +83,7 @@ import {
8383
import { ChatSession } from '../../clients/chat/v0/chat'
8484
import { ChatHistoryManager } from '../../storages/chatHistory'
8585
import { amazonQTabSuffix } from '../../../shared/constants'
86+
import { FsRead, FsReadParams } from '../../tools/fsRead'
8687

8788
export interface ChatControllerMessagePublishers {
8889
readonly processPromptChatMessage: MessagePublisher<PromptMessage>
@@ -630,6 +631,8 @@ export class ChatController {
630631
// Reset the filePaths to undefined
631632
this.sessionStorage.getSession(message.tabID ?? '').setFilePath(undefined)
632633
this.sessionStorage.getSession(message.tabID ?? '').setTempFilePath(undefined)
634+
} else if (message.action.id === 'confirm-tool-use') {
635+
await this.processToolUseMessage(message)
633636
}
634637
}
635638

@@ -887,12 +890,110 @@ export class ChatController {
887890
}
888891
}
889892

893+
private async processToolUseMessage(message: CustomFormActionMessage) {
894+
const tabID = message.tabID
895+
if (!tabID) {
896+
return
897+
}
898+
this.editorContextExtractor
899+
.extractContextForTrigger('ChatMessage')
900+
.then(async (context) => {
901+
const triggerID = randomUUID()
902+
this.triggerEventsStorage.addTriggerEvent({
903+
id: triggerID,
904+
tabID: message.tabID,
905+
message: undefined,
906+
type: 'chat_message',
907+
context,
908+
})
909+
const session = this.sessionStorage.getSession(tabID)
910+
const toolUse = session.toolUse
911+
if (!toolUse || !toolUse.input) {
912+
return
913+
}
914+
session.setToolUse(undefined)
915+
916+
let result: any
917+
const toolResults: ToolResult[] = []
918+
try {
919+
switch (toolUse.name) {
920+
// case 'execute_bash': {
921+
// const executeBash = new ExecuteBash(toolUse.input as unknown as ExecuteBashParams)
922+
// await executeBash.validate()
923+
// result = await executeBash.invoke(process.stdout)
924+
// break
925+
// }
926+
case 'fs_read': {
927+
const fsRead = new FsRead(toolUse.input as unknown as FsReadParams)
928+
await fsRead.validate()
929+
result = await fsRead.invoke()
930+
break
931+
}
932+
// case 'fs_write': {
933+
// const fsWrite = new FsWrite(toolUse.input as unknown as FsWriteParams)
934+
// const ctx = new DefaultContext()
935+
// result = await fsWrite.invoke(ctx, process.stdout)
936+
// break
937+
// }
938+
// case 'open_file': {
939+
// result = await openFile(toolUse.input as unknown as OpenFileParams)
940+
// break
941+
// }
942+
default:
943+
break
944+
}
945+
toolResults.push({
946+
content: [
947+
result.output.kind === 'text'
948+
? { text: result.output.content }
949+
: { json: result.output.content },
950+
],
951+
toolUseId: toolUse.toolUseId,
952+
status: 'success',
953+
})
954+
} catch (e: any) {
955+
toolResults.push({ content: [{ text: e.message }], toolUseId: toolUse.toolUseId, status: 'error' })
956+
}
957+
958+
this.chatHistoryManager.appendUserMessage({
959+
userInputMessage: {
960+
content: 'Tool Results',
961+
userIntent: undefined,
962+
origin: Origin.IDE,
963+
},
964+
})
965+
966+
await this.generateResponse(
967+
{
968+
message: 'Tool Results',
969+
trigger: ChatTriggerType.ChatMessage,
970+
query: undefined,
971+
codeSelection: context?.focusAreaContext?.selectionInsideExtendedCodeBlock,
972+
fileText: context?.focusAreaContext?.extendedCodeBlock,
973+
fileLanguage: context?.activeFileContext?.fileLanguage,
974+
filePath: context?.activeFileContext?.filePath,
975+
matchPolicy: context?.activeFileContext?.matchPolicy,
976+
codeQuery: context?.focusAreaContext?.names,
977+
userIntent: undefined,
978+
customization: getSelectedCustomization(),
979+
context: undefined,
980+
toolResults: toolResults,
981+
origin: Origin.IDE,
982+
},
983+
triggerID
984+
)
985+
})
986+
.catch((e) => {
987+
this.processException(e, tabID)
988+
})
989+
}
990+
890991
private async processPromptMessageAsNewThread(message: PromptMessage) {
891992
const session = this.sessionStorage.getSession(message.tabID)
892993
session.clearListOfReadFiles()
893994
this.editorContextExtractor
894995
.extractContextForTrigger('ChatMessage')
895-
.then((context) => {
996+
.then(async (context) => {
896997
const triggerID = randomUUID()
897998
this.triggerEventsStorage.addTriggerEvent({
898999
id: triggerID,
@@ -905,9 +1006,10 @@ export class ChatController {
9051006
userInputMessage: {
9061007
content: message.message,
9071008
userIntent: message.userIntent,
1009+
origin: Origin.IDE,
9081010
},
9091011
})
910-
return this.generateResponse(
1012+
await this.generateResponse(
9111013
{
9121014
message: message.message,
9131015
trigger: ChatTriggerType.ChatMessage,
@@ -922,6 +1024,7 @@ export class ChatController {
9221024
customization: getSelectedCustomization(),
9231025
context: message.context,
9241026
chatHistory: this.chatHistoryManager.getHistory(),
1027+
origin: Origin.IDE,
9251028
},
9261029
triggerID
9271030
)

‎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'
@@ -135,6 +136,8 @@ export class Messenger {
135136
let followUps: FollowUp[] = []
136137
let relatedSuggestions: Suggestion[] = []
137138
let codeBlockLanguage: string = 'plaintext'
139+
let toolUseInput = ''
140+
const toolUse: ToolUse = { toolUseId: undefined, name: undefined, input: undefined }
138141

139142
if (response.message === undefined) {
140143
throw new ToolkitError(
@@ -162,7 +165,7 @@ export class Messenger {
162165
})
163166

164167
const eventCounts = new Map<string, number>()
165-
waitUntil(
168+
await waitUntil(
166169
async () => {
167170
for await (const chatEvent of response.message!) {
168171
for (const key of keys(chatEvent)) {
@@ -192,6 +195,53 @@ export class Messenger {
192195
]
193196
}
194197

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

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

0 commit comments

Comments
 (0)