Skip to content

Commit 3c8e737

Browse files
committed
Initial Agentic loop Setup
1 parent 8e906b2 commit 3c8e737

File tree

12 files changed

+736
-28
lines changed

12 files changed

+736
-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/controllers/chat/chatRequest/converter.ts

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,18 @@ import {
77
ConversationState,
88
CursorState,
99
DocumentSymbol,
10+
EnvState,
1011
RelevantTextDocument,
12+
ShellState,
1113
SymbolType,
1214
TextDocument,
15+
Tool,
1316
} from '@amzn/codewhisperer-streaming'
1417
import { ChatTriggerType, TriggerPayload } from '../model'
1518
import { undefinedIfEmpty } from '../../../../shared/utilities/textUtilities'
19+
import { tryGetCurrentWorkingDirectory } from '../../../../shared/utilities/workspaceUtils'
20+
import toolsJson from '../../../tools/tool_index.json'
21+
import { getOperatingSystem } from '../../../../shared/telemetry/util'
1622

1723
const fqnNameSizeDownLimit = 1
1824
const fqnNameSizeUpLimit = 256
@@ -37,6 +43,14 @@ export const supportedLanguagesList = [
3743
const filePathSizeLimit = 4_000
3844
const customerMessageSizeLimit = 4_000
3945

46+
interface ToolSpec {
47+
name: string
48+
description: string
49+
// eslint-disable-next-line @typescript-eslint/naming-convention
50+
input_schema: Record<string, any>
51+
[key: string]: any
52+
}
53+
4054
export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { conversationState: ConversationState } {
4155
let document: TextDocument | undefined = undefined
4256
let cursorState: CursorState | undefined = undefined
@@ -102,6 +116,15 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c
102116
const customizationArn: string | undefined = undefinedIfEmpty(triggerPayload.customization.arn)
103117
const chatTriggerType = triggerPayload.trigger === ChatTriggerType.InlineChatMessage ? 'INLINE_CHAT' : 'MANUAL'
104118

119+
const tools: Tool[] = Object.entries(toolsJson as Record<string, ToolSpec>).map(([toolName, toolSpec]) => ({
120+
toolSpecification: {
121+
...toolSpec,
122+
// Use the key as name if not already defined in the spec
123+
name: toolSpec.name || toolName,
124+
inputSchema: { json: toolSpec.input_schema },
125+
},
126+
}))
127+
105128
return {
106129
conversationState: {
107130
currentMessage: {
@@ -116,9 +139,16 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c
116139
relevantDocuments,
117140
useRelevantDocuments,
118141
},
142+
envState: buildEnvState(),
143+
shellState: buildShellState(),
119144
additionalContext: triggerPayload.additionalContents,
145+
tools,
146+
...(triggerPayload.toolResults !== undefined &&
147+
triggerPayload.toolResults !== null && { toolResults: triggerPayload.toolResults }),
120148
},
121149
userIntent: triggerPayload.userIntent,
150+
...(triggerPayload.origin !== undefined &&
151+
triggerPayload.origin !== null && { origin: triggerPayload.origin }),
122152
},
123153
},
124154
chatTriggerType,
@@ -127,3 +157,26 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c
127157
},
128158
}
129159
}
160+
161+
/**
162+
* Helper function to build environment state
163+
*/
164+
export function buildEnvState(): EnvState {
165+
return {
166+
operatingSystem: getOperatingSystem(),
167+
currentWorkingDirectory: tryGetCurrentWorkingDirectory(),
168+
}
169+
}
170+
171+
/**
172+
* Helper function to build shell state
173+
*/
174+
export function buildShellState(): ShellState {
175+
// In a real implementation, you would detect the shell
176+
// This is a simplified version
177+
const shellName = process.env.SHELL || 'bash'
178+
return {
179+
shellName: shellName.split('/').pop() || 'bash',
180+
shellHistory: undefined,
181+
}
182+
}

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
)

0 commit comments

Comments
 (0)