Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
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
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
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 } {
Copy link
Contributor

Choose a reason for hiding this comment

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

Can we move the changes to a new function triggerPayloadToAgenticChatRequest

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Why do we need a new one? we can use the same function and add the API request blocks based on triggerPayload right?

Copy link
Contributor

Choose a reason for hiding this comment

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

I see, then should tools also be part of triggerPayload? Because it's not needed for mynah

Copy link
Contributor Author

Choose a reason for hiding this comment

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

not needed, but i can make sure that Origin is present to exclude tools from the request. I'll do it in followup

let document: TextDocument | undefined = undefined
let cursorState: CursorState | undefined = undefined
Expand Down Expand Up @@ -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]) => ({
Copy link

Choose a reason for hiding this comment

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

I feel like this won't be easily scalable to other use cases like MCP where arbitrary tools can be added. Is there a better way to do this?

Copy link
Contributor Author

@ashishrp-aws ashishrp-aws Mar 25, 2025

Choose a reason for hiding this comment

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

We can have a number of static tools. I have moved this to constants. Agreed will check for a better way to handle tools.

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: {
Expand All @@ -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,
Expand All @@ -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,
}
}
109 changes: 106 additions & 3 deletions packages/core/src/codewhispererChat/controllers/chat/controller.ts
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

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

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
Loading
Loading