Skip to content

feat(chat): Initial Agentic Chat loop Setup #6844

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jpinkney-aws merged 4 commits into from
Mar 25, 2025
Merged

feat(chat): Initial Agentic Chat loop Setup #6844

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3c8e737
Initial Agentic loop Setup
ashishrp-aws Mar 25, 2025
f0d6050
Merge branch 'feature/agentic-chat' into feature/agentic-chat
ashishrp-aws Mar 25, 2025
b4b7a3d
Removing Environment state objects from converter and addressing othe…
ashishrp-aws Mar 25, 2025
545870e
Removed merge conflicts
ashishrp-aws Mar 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 14 additions & 1 deletion packages/core/src/codewhispererChat/clients/chat/v0/chat.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,11 @@
*/

import { SendMessageCommandOutput, SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
import { GenerateAssistantResponseCommandOutput, GenerateAssistantResponseRequest } from '@amzn/codewhisperer-streaming'
import {
GenerateAssistantResponseCommandOutput,
GenerateAssistantResponseRequest,
ToolUse,
} from '@amzn/codewhisperer-streaming'
import * as vscode from 'vscode'
import { ToolkitError } from '../../../../shared/errors'
import { createCodeWhispererChatStreamingClient } from '../../../../shared/clients/codewhispererChatClient'
Expand All @@ -13,6 +17,7 @@ import { UserWrittenCodeTracker } from '../../../../codewhisperer/tracker/userWr

export class ChatSession {
private sessionId?: string
private _toolUse: ToolUse | 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
Expand All @@ -22,6 +27,14 @@ export class ChatSession {
return this.sessionId
}

public get toolUse(): ToolUse | undefined {
return this._toolUse
}

public setToolUse(toolUse: ToolUse | undefined) {
this._toolUse = toolUse
}

public tokenSource!: vscode.CancellationTokenSource

constructor() {
Expand Down
9 changes: 9 additions & 0 deletions packages/core/src/codewhispererChat/constants.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
*/
import * as path from 'path'
import fs from '../shared/fs/fs'
import { Tool } from '@amzn/codewhisperer-streaming'
import toolsJson from '../codewhispererChat/tools/tool_index.json'

export const promptFileExtension = '.md'

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

export const createSavedPromptCommandId = 'create-saved-prompt'

export const tools: Tool[] = Object.entries(toolsJson).map(([, toolSpec]) => ({
toolSpecification: {
...toolSpec,
inputSchema: { json: toolSpec.inputSchema },
Copy link
Contributor

@ctlai95 ctlai95 Mar 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: I think we don't need to do this anymore because we changed the tool index from index_schema to indexSchema

},
}))
7 changes: 7 additions & 0 deletions packages/core/src/codewhispererChat/controllers/chat/chatRequest/converter.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import {
} from '@amzn/codewhisperer-streaming'
import { ChatTriggerType, TriggerPayload } from '../model'
import { undefinedIfEmpty } from '../../../../shared/utilities/textUtilities'
import { tools } from '../../../constants'

const fqnNameSizeDownLimit = 1
const fqnNameSizeUpLimit = 256
Expand Down Expand Up @@ -115,10 +116,16 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c
cursorState,
relevantDocuments,
useRelevantDocuments,
// TODO: Need workspace folders here after model update.
},
additionalContext: triggerPayload.additionalContents,
tools,
...(triggerPayload.toolResults !== undefined &&
triggerPayload.toolResults !== null && { toolResults: triggerPayload.toolResults }),
},
userIntent: triggerPayload.userIntent,
...(triggerPayload.origin !== undefined &&
triggerPayload.origin !== null && { origin: triggerPayload.origin }),
},
},
chatTriggerType,
Expand Down
109 changes: 106 additions & 3 deletions packages/core/src/codewhispererChat/controllers/chat/controller.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ import { EditorContextCommand } from '../../commands/registerCommands'
import { PromptsGenerator } from './prompts/promptsGenerator'
import { TriggerEventsStorage } from '../../storages/triggerEvents'
import { SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
import { CodeWhispererStreamingServiceException } from '@amzn/codewhisperer-streaming'
import { CodeWhispererStreamingServiceException, Origin, ToolResult } from '@amzn/codewhisperer-streaming'
import { UserIntentRecognizer } from './userIntent/userIntentRecognizer'
import { CWCTelemetryHelper, recordTelemetryChatRunCommand } from './telemetryHelper'
import { CodeWhispererTracker } from '../../../codewhisperer/tracker/codewhispererTracker'
Expand Down Expand Up @@ -81,6 +81,7 @@ import {
} from '../../constants'
import { ChatSession } from '../../clients/chat/v0/chat'
import { ChatHistoryManager } from '../../storages/chatHistory'
import { FsRead, FsReadParams } from '../../tools/fsRead'

export interface ChatControllerMessagePublishers {
readonly processPromptChatMessage: MessagePublisher<PromptMessage>
Expand Down Expand Up @@ -577,6 +578,8 @@ export class ChatController {
const newFileDoc = await vscode.workspace.openTextDocument(newFilePath)
await vscode.window.showTextDocument(newFileDoc)
telemetry.ui_click.emit({ elementId: 'amazonq_createSavedPrompt' })
} else if (message.action.id === 'confirm-tool-use') {
await this.processToolUseMessage(message)
}
}

Expand Down Expand Up @@ -834,10 +837,108 @@ export class ChatController {
}
}

private async processToolUseMessage(message: CustomFormActionMessage) {
const tabID = message.tabID
if (!tabID) {
return
}
this.editorContextExtractor
.extractContextForTrigger('ChatMessage')
.then(async (context) => {
const triggerID = randomUUID()
this.triggerEventsStorage.addTriggerEvent({
id: triggerID,
tabID: message.tabID,
message: undefined,
type: 'chat_message',
context,
})
const session = this.sessionStorage.getSession(tabID)
const toolUse = session.toolUse
if (!toolUse || !toolUse.input) {
return
}
session.setToolUse(undefined)

let result: any
const toolResults: ToolResult[] = []
try {
switch (toolUse.name) {
Copy link

@mr-lee mr-lee Mar 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

did you comment things out here intentionally?

Copy link
Contributor Author

@ashishrp-aws ashishrp-aws Mar 25, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes....write tools aren't merged yet.

// case 'execute_bash': {
// const executeBash = new ExecuteBash(toolUse.input as unknown as ExecuteBashParams)
// await executeBash.validate()
// result = await executeBash.invoke(process.stdout)
// break
// }
case 'fs_read': {
const fsRead = new FsRead(toolUse.input as unknown as FsReadParams)
await fsRead.validate()
result = await fsRead.invoke()
break
}
// case 'fs_write': {
// const fsWrite = new FsWrite(toolUse.input as unknown as FsWriteParams)
// const ctx = new DefaultContext()
// result = await fsWrite.invoke(ctx, process.stdout)
// break
// }
// case 'open_file': {
// result = await openFile(toolUse.input as unknown as OpenFileParams)
// break
// }
default:
break
}
toolResults.push({
content: [
result.output.kind === 'text'
? { text: result.output.content }
: { json: result.output.content },
],
toolUseId: toolUse.toolUseId,
status: 'success',
})
} catch (e: any) {
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',
trigger: ChatTriggerType.ChatMessage,
query: undefined,
codeSelection: context?.focusAreaContext?.selectionInsideExtendedCodeBlock,
fileText: context?.focusAreaContext?.extendedCodeBlock,
fileLanguage: context?.activeFileContext?.fileLanguage,
filePath: context?.activeFileContext?.filePath,
matchPolicy: context?.activeFileContext?.matchPolicy,
codeQuery: context?.focusAreaContext?.names,
userIntent: undefined,
customization: getSelectedCustomization(),
context: undefined,
toolResults: toolResults,
origin: Origin.IDE,
},
triggerID
)
})
.catch((e) => {
this.processException(e, tabID)
})
}

private async processPromptMessageAsNewThread(message: PromptMessage) {
this.editorContextExtractor
.extractContextForTrigger('ChatMessage')
.then((context) => {
.then(async (context) => {
const triggerID = randomUUID()
this.triggerEventsStorage.addTriggerEvent({
id: triggerID,
Expand All @@ -850,9 +951,10 @@ export class ChatController {
userInputMessage: {
content: message.message,
userIntent: message.userIntent,
origin: Origin.IDE,
},
})
return this.generateResponse(
await this.generateResponse(
{
message: message.message,
trigger: ChatTriggerType.ChatMessage,
Expand All @@ -867,6 +969,7 @@ export class ChatController {
customization: getSelectedCustomization(),
context: message.context,
chatHistory: this.chatHistoryManager.getHistory(),
origin: Origin.IDE,
},
triggerID
)
Expand Down
117 changes: 115 additions & 2 deletions packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ import {
ChatResponseStream as cwChatResponseStream,
CodeWhispererStreamingServiceException,
SupplementaryWebLink,
ToolUse,
} from '@amzn/codewhisperer-streaming'
import { ChatMessage, ErrorMessage, FollowUp, Suggestion } from '../../../view/connector/connector'
import { ChatSession } from '../../../clients/chat/v0/chat'
Expand Down Expand Up @@ -131,6 +132,8 @@ export class Messenger {
let followUps: FollowUp[] = []
let relatedSuggestions: Suggestion[] = []
let codeBlockLanguage: string = 'plaintext'
let toolUseInput = ''
const toolUse: ToolUse = { toolUseId: undefined, name: undefined, input: undefined }

if (response.message === undefined) {
throw new ToolkitError(
Expand Down Expand Up @@ -158,7 +161,7 @@ export class Messenger {
})

const eventCounts = new Map<string, number>()
waitUntil(
await waitUntil(
async () => {
for await (const chatEvent of response.message!) {
for (const key of keys(chatEvent)) {
Expand Down Expand Up @@ -188,6 +191,53 @@ export class Messenger {
]
}

const cwChatEvent: cwChatResponseStream = chatEvent
if (
cwChatEvent.toolUseEvent?.input !== undefined &&
cwChatEvent.toolUseEvent.input.length > 0 &&
!cwChatEvent.toolUseEvent.stop
) {
toolUseInput += cwChatEvent.toolUseEvent.input
}

if (cwChatEvent.toolUseEvent?.stop) {
toolUse.input = JSON.parse(toolUseInput)
toolUse.toolUseId = cwChatEvent.toolUseEvent.toolUseId ?? ''
toolUse.name = cwChatEvent.toolUseEvent.name ?? ''
session.setToolUse(toolUse)

const message = this.getToolUseMessage(toolUse)
// const isConfirmationRequired = this.getIsConfirmationRequired(toolUse)

this.dispatcher.sendChatMessage(
new ChatMessage(
{
message,
messageType: 'answer',
followUps: undefined,
followUpsHeader: undefined,
relatedSuggestions: undefined,
codeReference,
triggerID,
messageID: toolUse.toolUseId,
userIntent: triggerPayload.userIntent,
codeBlockLanguage: codeBlockLanguage,
contextList: undefined,
// TODO: confirmation buttons
},
tabID
)
)
// TODO: setup permission action
// if (!isConfirmationRequired) {
// this.dispatcher.sendCustomFormActionMessage(
// new CustomFormActionMessage(tabID, {
// id: 'confirm-tool-use',
// })
// )
// }
}

if (
chatEvent.assistantResponseEvent?.content !== undefined &&
chatEvent.assistantResponseEvent.content.length > 0
Expand Down Expand Up @@ -338,7 +388,7 @@ export class Messenger {
messageId: messageID,
content: message,
references: codeReference,
// TODO: Add tools data and follow up prompt details
toolUses: [{ ...toolUse }],
},
})

Expand Down Expand Up @@ -533,4 +583,67 @@ export class Messenger {
new ShowCustomFormMessage(tabID, formItems, buttons, title, description)
)
}

// TODO: Make this cleaner
// private getIsConfirmationRequired(toolUse: ToolUse) {
// if (toolUse.name === 'execute_bash') {
// const executeBash = new ExecuteBash(toolUse.input as unknown as ExecuteBashParams)
// return executeBash.requiresAcceptance()
// }
// return toolUse.name === 'fs_write'
// }
private getToolUseMessage(toolUse: ToolUse) {
if (toolUse.name === 'fs_read') {
return `Reading the file at \`${(toolUse.input as any)?.path}\` using the \`fs_read\` tool.`
}
// if (toolUse.name === 'execute_bash') {
// const input = toolUse.input as unknown as ExecuteBashParams
// return `Executing the bash command
// \`\`\`bash
// ${input.command}
// \`\`\`
// using the \`execute_bash\` tool.`
// }
// if (toolUse.name === 'fs_write') {
// const input = toolUse.input as unknown as FsWriteParams
// switch (input.command) {
// case 'create': {
// return `Writing
// \`\`\`
// ${input.file_text}
// \`\`\`
// into the file at \`${input.path}\` using the \`fs_write\` tool.`
// }
// case 'str_replace': {
// return `Replacing
// \`\`\`
// ${input.old_str}
// \`\`\`
// with
// \`\`\`
// ${input.new_str}
// \`\`\`
// at \`${input.path}\` using the \`fs_write\` tool.`
// }
// case 'insert': {
// return `Inserting
// \`\`\`
// ${input.new_str}
// \`\`\`
// at line
// \`\`\`
// ${input.insert_line}
// \`\`\`
// at \`${input.path}\` using the \`fs_write\` tool.`
// }
// case 'append': {
// return `Appending
// \`\`\`
// ${input.new_str}
// \`\`\`
// at \`${input.path}\` using the \`fs_write\` tool.`
// }
// }
// }
}
}
Loading
Loading