Merging Tai Agentic Tool exe Changes (#14)

* wip

* wip

* wip

---------

Co-authored-by: Tai Lai <[email protected]>

Author: laileni-aws

packages/core/src/codewhisperer/client/user-service-2.json

@@ -4,7 +4,12 @@
  */
 
 import { SendMessageCommandOutput, SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
-import { GenerateAssistantResponseCommandOutput, GenerateAssistantResponseRequest } from '@amzn/codewhisperer-streaming'
+import {
+    ChatMessage,
+    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 +18,8 @@ import { UserWrittenCodeTracker } from '../../../../codewhisperer/tracker/userWr
 
 export class ChatSession {
     private sessionId?: string
+    private _toolUse: ToolUse | undefined
+    private _chatHistory: ChatMessage[] = []
 
     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
@@ -21,6 +28,12 @@ export class ChatSession {
     public get sessionIdentifier(): string | undefined {
         return this.sessionId
     }
+    public get toolUse(): ToolUse | undefined {
+        return this._toolUse
+    }
+    public get chatHistory(): ChatMessage[] {
+        return this._chatHistory
+    }
 
     public tokenSource!: vscode.CancellationTokenSource
 
@@ -35,6 +48,31 @@ export class ChatSession {
     public setSessionID(id?: string) {
         this.sessionId = id
     }
+    public setToolUse(toolUse: ToolUse | undefined) {
+        this._toolUse = toolUse
+    }
+    public pushToChatHistory(message: ChatMessage | undefined) {
+        if (message === undefined) {
+            return
+        }
+        this._chatHistory.push(this.formatChatHistoryMessage(message))
+    }
+
+    private formatChatHistoryMessage(message: ChatMessage): ChatMessage {
+        if (message.userInputMessage !== undefined) {
+            return {
+                userInputMessage: {
+                    ...message.userInputMessage,
+                    userInputMessageContext: {
+                        ...message.userInputMessage.userInputMessageContext,
+                        tools: undefined,
+                    },
+                },
+            }
+        }
+        return message
+    }
+
     async chatIam(chatRequest: SendMessageRequest): Promise<SendMessageCommandOutput> {
         const client = await createQDeveloperStreamingClient()
 
@@ -60,11 +98,13 @@ export class ChatSession {
     async chatSso(chatRequest: GenerateAssistantResponseRequest): Promise<GenerateAssistantResponseCommandOutput> {
         const client = await createCodeWhispererChatStreamingClient()
 
-        if (this.sessionId !== undefined && chatRequest.conversationState !== undefined) {
+        if (this.sessionId !== '' && this.sessionId !== undefined && chatRequest.conversationState !== undefined) {
             chatRequest.conversationState.conversationId = this.sessionId
         }
 
         const response = await client.generateAssistantResponse(chatRequest)
+        // eslint-disable-next-line aws-toolkits/no-console-log
+        console.log(response.$metadata.requestId)
         if (!response.generateAssistantResponseResponse) {
             throw new ToolkitError(
                 `Empty chat response. Session id: ${this.sessionId} Request ID: ${response.$metadata.requestId}`

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

@@ -4,15 +4,19 @@
  */
 
 import {
+    ChatMessage,
     ConversationState,
     CursorState,
     DocumentSymbol,
     RelevantTextDocument,
     SymbolType,
     TextDocument,
+    Tool,
 } from '@amzn/codewhisperer-streaming'
 import { ChatTriggerType, TriggerPayload } from '../model'
 import { undefinedIfEmpty } from '../../../../shared/utilities/textUtilities'
+import toolsJson from '../../../tools/tool_index.json'
+import { tryGetCurrentWorkingDirectory } from '../../../../shared/utilities/workspaceUtils'
 
 const fqnNameSizeDownLimit = 1
 const fqnNameSizeUpLimit = 256
@@ -126,3 +130,114 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c
         },
     }
 }
+
+export function triggerPayloadToAgenticChatRequest(
+    triggerPayload: TriggerPayload,
+    chatHistory?: ChatMessage[]
+): { conversationState: ConversationState } {
+    let document: TextDocument | undefined = undefined
+    let cursorState: CursorState | undefined = undefined
+
+    if (triggerPayload.filePath !== undefined && triggerPayload.filePath !== '') {
+        const documentSymbolFqns: DocumentSymbol[] = []
+        if (triggerPayload.codeQuery?.fullyQualifiedNames?.used) {
+            for (const fqn of triggerPayload.codeQuery.fullyQualifiedNames.used) {
+                const elem = {
+                    name: fqn.symbol?.join('.') ?? '',
+                    type: SymbolType.USAGE,
+                    source: fqn.source?.join('.'),
+                }
+
+                if (
+                    elem.name.length >= fqnNameSizeDownLimit &&
+                    elem.name.length < fqnNameSizeUpLimit &&
+                    (elem.source === undefined ||
+                        (elem.source.length >= fqnNameSizeDownLimit && elem.source.length < fqnNameSizeUpLimit))
+                ) {
+                    documentSymbolFqns.push(elem)
+                }
+            }
+        }
+
+        let programmingLanguage
+        if (
+            triggerPayload.fileLanguage !== undefined &&
+            triggerPayload.fileLanguage !== '' &&
+            supportedLanguagesList.includes(triggerPayload.fileLanguage)
+        ) {
+            programmingLanguage = { languageName: triggerPayload.fileLanguage }
+        }
+
+        document = {
+            relativeFilePath: triggerPayload.filePath ? triggerPayload.filePath.substring(0, filePathSizeLimit) : '',
+            text: triggerPayload.fileText,
+            programmingLanguage: programmingLanguage,
+            documentSymbols: documentSymbolFqns,
+        }
+
+        if (triggerPayload.codeSelection?.start) {
+            cursorState = {
+                range: {
+                    start: {
+                        line: triggerPayload.codeSelection.start.line,
+                        character: triggerPayload.codeSelection.start.character,
+                    },
+                    end: {
+                        line: triggerPayload.codeSelection.end.line,
+                        character: triggerPayload.codeSelection.end.character,
+                    },
+                },
+            }
+        }
+    }
+
+    const relevantDocuments: RelevantTextDocument[] = triggerPayload.relevantTextDocuments
+        ? triggerPayload.relevantTextDocuments
+        : []
+    const useRelevantDocuments = triggerPayload.useRelevantDocuments
+    // service will throw validation exception if string is empty
+    const customizationArn: string | undefined = undefinedIfEmpty(triggerPayload.customization.arn)
+    const chatTriggerType = triggerPayload.trigger === ChatTriggerType.InlineChatMessage ? 'INLINE_CHAT' : 'MANUAL'
+
+    const operatingSystem = 'macos'
+    const currentWorkingDirectory = tryGetCurrentWorkingDirectory()
+
+    const tools: Tool[] = Object.entries(toolsJson).map(([, toolSpec]) => ({
+        toolSpecification: {
+            ...toolSpec,
+            inputSchema: { json: toolSpec.input_schema },
+        },
+    }))
+    return {
+        conversationState: {
+            currentMessage: {
+                userInputMessage: {
+                    content: triggerPayload.message
+                        ? triggerPayload.message.substring(0, customerMessageSizeLimit)
+                        : '',
+                    userInputMessageContext: {
+                        editorState: {
+                            document,
+                            cursorState,
+                            relevantDocuments,
+                            useRelevantDocuments,
+                        },
+                        envState: {
+                            operatingSystem,
+                            currentWorkingDirectory,
+                            environmentVariables: [],
+                        },
+                        additionalContext: triggerPayload.additionalContents,
+                        tools,
+                        toolResults: triggerPayload.toolResults,
+                    },
+                    userIntent: triggerPayload.userIntent,
+                    origin: 'IDE',
+                },
+            },
+            chatTriggerType,
+            customizationArn: customizationArn,
+            history: chatHistory,
+        },
+    }
+}

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

@@ -50,7 +50,7 @@ import { UserIntentRecognizer } from './userIntent/userIntentRecognizer'
 import { CWCTelemetryHelper, recordTelemetryChatRunCommand } from './telemetryHelper'
 import { CodeWhispererTracker } from '../../../codewhisperer/tracker/codewhispererTracker'
 import { getLogger } from '../../../shared/logger/logger'
-import { triggerPayloadToChatRequest } from './chatRequest/converter'
+import { triggerPayloadToAgenticChatRequest, triggerPayloadToChatRequest } from './chatRequest/converter'
 import { AuthUtil } from '../../../codewhisperer/util/authUtil'
 import { openUrl } from '../../../shared/utilities/vsCodeUtils'
 import { randomUUID } from '../../../shared/crypto'
@@ -80,6 +80,8 @@ import {
     contextMaxLength,
 } from '../../constants'
 import { ChatSession } from '../../clients/chat/v0/chat'
+import { FsRead, FsReadParams } from '../../tools/fsRead'
+import { tryGetCurrentWorkingDirectory } from '../../../shared/utilities/workspaceUtils'
 
 export interface ChatControllerMessagePublishers {
     readonly processPromptChatMessage: MessagePublisher<PromptMessage>
@@ -815,7 +817,7 @@ export class ChatController {
     private async processPromptMessageAsNewThread(message: PromptMessage) {
         this.editorContextExtractor
             .extractContextForTrigger('ChatMessage')
-            .then((context) => {
+            .then(async (context) => {
                 const triggerID = randomUUID()
                 this.triggerEventsStorage.addTriggerEvent({
                     id: triggerID,
@@ -824,7 +826,7 @@ export class ChatController {
                     type: 'chat_message',
                     context,
                 })
-                return this.generateResponse(
+                await this.generateResponse(
                     {
                         message: message.message,
                         trigger: ChatTriggerType.ChatMessage,
@@ -839,8 +841,57 @@ export class ChatController {
                         customization: getSelectedCustomization(),
                         context: message.context,
                     },
-                    triggerID
+                    triggerID,
+                    true
                 )
+
+                const session = this.sessionStorage.getSession(message.tabID)
+                while (true) {
+                    const toolUse = session.toolUse
+                    if (!toolUse || !toolUse.input) {
+                        break
+                    }
+                    session.setToolUse(undefined)
+
+                    let result: any
+                    switch (toolUse.name) {
+                        case 'fs_read': {
+                            const fsRead = new FsRead(toolUse.input as unknown as FsReadParams)
+                            result = await fsRead.invoke({
+                                env: { currentDir: () => tryGetCurrentWorkingDirectory() },
+                            })
+                            break
+                        }
+                        default:
+                            break
+                    }
+
+                    await this.generateResponse(
+                        {
+                            message: '',
+                            trigger: ChatTriggerType.ChatMessage,
+                            query: message.message,
+                            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: this.userIntentRecognizer.getFromPromptChatMessage(message),
+                            customization: getSelectedCustomization(),
+                            context: message.context,
+                            toolResults: [
+                                {
+                                    content: [{ text: result.output.content }],
+                                    toolUseId: toolUse.toolUseId,
+                                    status: 'success',
+                                },
+                            ],
+                        },
+                        triggerID,
+                        true
+                    )
+                }
             })
             .catch((e) => {
                 this.processException(e, message.tabID)
@@ -1025,7 +1076,8 @@ export class ChatController {
 
     private async generateResponse(
         triggerPayload: TriggerPayload & { projectContextQueryLatencyMs?: number },
-        triggerID: string
+        triggerID: string,
+        agentic?: boolean
     ) {
         const triggerEvent = this.triggerEventsStorage.getTriggerEvent(triggerID)
         if (triggerEvent === undefined) {
@@ -1038,7 +1090,7 @@ export class ChatController {
 
         if (triggerEvent.tabID === undefined) {
             setTimeout(() => {
-                this.generateResponse(triggerPayload, triggerID).catch((e) => {
+                this.generateResponse(triggerPayload, triggerID, agentic).catch((e) => {
                     getLogger().error('generateResponse failed: %s', (e as Error).message)
                 })
             }, 20)
@@ -1103,7 +1155,9 @@ export class ChatController {
 
         triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments || [])
 
-        const request = triggerPayloadToChatRequest(triggerPayload)
+        const request = agentic
+            ? triggerPayloadToAgenticChatRequest(triggerPayload, session.chatHistory)
+            : triggerPayloadToChatRequest(triggerPayload)
 
         if (triggerPayload.documentReferences !== undefined) {
             const relativePathsOfMergedRelevantDocuments = triggerPayload.documentReferences.map(
@@ -1157,6 +1211,7 @@ export class ChatController {
                     response.$metadata.requestId
                 } metadata: ${inspect(response.$metadata, { depth: 12 })}`
             )
+            session.pushToChatHistory(request.conversationState.currentMessage)
             await this.messenger.sendAIResponse(response, session, tabID, triggerID, triggerPayload)
         } catch (e: any) {
             this.telemetryHelper.recordMessageResponseError(triggerPayload, tabID, getHttpStatusCode(e) ?? 0)