Initial Agentic loop Setup

Author: ashishrp-aws

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

@@ -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'
@@ -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
@@ -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() {

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

@@ -7,12 +7,18 @@ import {
     ConversationState,
     CursorState,
     DocumentSymbol,
+    EnvState,
     RelevantTextDocument,
+    ShellState,
     SymbolType,
     TextDocument,
+    Tool,
 } from '@amzn/codewhisperer-streaming'
 import { ChatTriggerType, TriggerPayload } from '../model'
 import { undefinedIfEmpty } from '../../../../shared/utilities/textUtilities'
+import { tryGetCurrentWorkingDirectory } from '../../../../shared/utilities/workspaceUtils'
+import toolsJson from '../../../tools/tool_index.json'
+import { getOperatingSystem } from '../../../../shared/telemetry/util'
 
 const fqnNameSizeDownLimit = 1
 const fqnNameSizeUpLimit = 256
@@ -37,6 +43,14 @@ export const supportedLanguagesList = [
 const filePathSizeLimit = 4_000
 const customerMessageSizeLimit = 4_000
 
+interface ToolSpec {
+    name: string
+    description: string
+    // eslint-disable-next-line @typescript-eslint/naming-convention
+    input_schema: Record<string, any>
+    [key: string]: any
+}
+
 export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { conversationState: ConversationState } {
     let document: TextDocument | undefined = undefined
     let cursorState: CursorState | undefined = undefined
@@ -102,6 +116,15 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c
     const customizationArn: string | undefined = undefinedIfEmpty(triggerPayload.customization.arn)
     const chatTriggerType = triggerPayload.trigger === ChatTriggerType.InlineChatMessage ? 'INLINE_CHAT' : 'MANUAL'
 
+    const tools: Tool[] = Object.entries(toolsJson as Record<string, ToolSpec>).map(([toolName, toolSpec]) => ({
+        toolSpecification: {
+            ...toolSpec,
+            // Use the key as name if not already defined in the spec
+            name: toolSpec.name || toolName,
+            inputSchema: { json: toolSpec.input_schema },
+        },
+    }))
+
     return {
         conversationState: {
             currentMessage: {
@@ -116,9 +139,16 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c
                             relevantDocuments,
                             useRelevantDocuments,
                         },
+                        envState: buildEnvState(),
+                        shellState: buildShellState(),
                         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,
@@ -127,3 +157,26 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c
         },
     }
 }
+
+/**
+ * Helper function to build environment state
+ */
+export function buildEnvState(): EnvState {
+    return {
+        operatingSystem: getOperatingSystem(),
+        currentWorkingDirectory: tryGetCurrentWorkingDirectory(),
+    }
+}
+
+/**
+ * Helper function to build shell state
+ */
+export function buildShellState(): ShellState {
+    // In a real implementation, you would detect the shell
+    // This is a simplified version
+    const shellName = process.env.SHELL || 'bash'
+    return {
+        shellName: shellName.split('/').pop() || 'bash',
+        shellHistory: undefined,
+    }
+}

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

@@ -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'
@@ -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>
@@ -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)
         }
     }
 
@@ -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) {
+                        // 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,
@@ -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,
@@ -867,6 +969,7 @@ export class ChatController {
                         customization: getSelectedCustomization(),
                         context: message.context,
                         chatHistory: this.chatHistoryManager.getHistory(),
+                        origin: Origin.IDE,
                     },
                     triggerID
                 )