From ddc6407f74510aaeb5e3453ee37cf94a06654302 Mon Sep 17 00:00:00 2001 From: Rob Lourens Date: Thu, 10 Oct 2024 16:20:10 -0700 Subject: [PATCH] Tool calling over multiple rounds --- chat-tools-sample/src/extension.ts | 20 +++++--- chat-tools-sample/src/toolsPrompt.tsx | 69 ++++++++++++++++++--------- 2 files changed, 61 insertions(+), 28 deletions(-) diff --git a/chat-tools-sample/src/extension.ts b/chat-tools-sample/src/extension.ts index 038a0657..41acf62d 100644 --- a/chat-tools-sample/src/extension.ts +++ b/chat-tools-sample/src/extension.ts @@ -1,7 +1,7 @@ import * as vscode from 'vscode'; import { FindFilesTool, RunInTerminalTool, TabCountTool } from './tools'; import { renderPrompt } from '@vscode/prompt-tsx'; -import { ToolResultMetadata, ToolUserPrompt } from './toolsPrompt'; +import { ToolCallRound, ToolResultMetadata, ToolUserPrompt } from './toolsPrompt'; export function activate(context: vscode.ExtensionContext) { registerChatTool(context); @@ -215,13 +215,15 @@ function registerChatParticipant2(context: vscode.ExtensionContext) { { context: chatContext, request, - toolCalls: [], + toolCallRounds: [], + toolCallResults: new Map() }, { modelMaxPromptTokens: model.maxInputTokens }, model) const toolReferences = [...request.toolReferences]; const accumulatedToolCalls = new Map(); + const toolCallRounds: ToolCallRound[] = []; const runWithFunctions = async (): Promise => { const requestedTool = toolReferences.shift(); if (requestedTool) { @@ -249,19 +251,25 @@ function registerChatParticipant2(context: vscode.ExtensionContext) { } if (toolCalls.length) { + toolCallRounds.push({ + response: responseStr, + toolCalls + }); const result = (await renderPrompt( ToolUserPrompt, { context: chatContext, request, - toolCalls, + toolCallRounds, + toolCallResults: accumulatedToolCalls }, { modelMaxPromptTokens: model.maxInputTokens }, model)); messages = result.messages; - const toolResultMetadata = result.metadatas.get(ToolResultMetadata) - if (toolResultMetadata) { - toolResultMetadata.resultMap.forEach((value, key) => accumulatedToolCalls.set(key, value)); + const toolResultMetadata = result.metadatas.getAll(ToolResultMetadata) + if (toolResultMetadata?.length) { + // TODO flatten + toolResultMetadata.forEach(meta => meta.resultMap.forEach((value, key) => accumulatedToolCalls.set(key, value))); } // RE-enter diff --git a/chat-tools-sample/src/toolsPrompt.tsx b/chat-tools-sample/src/toolsPrompt.tsx index 3fe68a8e..f576701d 100644 --- a/chat-tools-sample/src/toolsPrompt.tsx +++ b/chat-tools-sample/src/toolsPrompt.tsx @@ -13,10 +13,16 @@ import { import { Chunk, ToolMessage, ToolResult } from '@vscode/prompt-tsx/dist/base/promptElements'; import * as vscode from 'vscode'; +export interface ToolCallRound { + response: string; + toolCalls: vscode.LanguageModelToolCallPart[]; +} + export interface ToolUserProps extends BasePromptElementProps { request: vscode.ChatRequest; context: vscode.ChatContext; - toolCalls: vscode.LanguageModelToolCallPart[]; + toolCallRounds: ToolCallRound[]; + toolCallResults: Map; } export class ToolUserPrompt extends PromptElement { @@ -46,14 +52,19 @@ export class ToolUserPrompt extends PromptElement { priority={30} /> {this.props.request.prompt} - + + ); } } interface ToolCallsProps extends BasePromptElementProps { - toolCalls: vscode.LanguageModelToolCallPart[]; + toolCallRounds: ToolCallRound[]; + toolCallResults: Map; toolInvocationToken: vscode.ChatParticipantToolToken; } @@ -62,45 +73,54 @@ const dummyCancellationToken: vscode.CancellationToken = new vscode.Cancellation class ToolCalls extends PromptElement { async render(state: void, sizing: PromptSizing) { - if (!this.props.toolCalls.length) { + if (!this.props.toolCallRounds.length) { return undefined; } - + + return <> + {await Promise.all(this.props.toolCallRounds.map(round => this.renderOneToolCallRound(round, sizing)))} + + } + + private async renderOneToolCallRound(round: ToolCallRound, sizing: PromptSizing) { // TODO- prompt-tsx export this type? // TODO- at what level do the parameters get stringified? - const assistantToolCalls: any[] = this.props.toolCalls.map(tc => ({ type: 'function', function: { name: tc.name, arguments: JSON.stringify(tc.parameters) }, id: tc.toolCallId })); + const assistantToolCalls: any[] = round.toolCalls.map(tc => ({ type: 'function', function: { name: tc.name, arguments: JSON.stringify(tc.parameters) }, id: tc.toolCallId })); // TODO@prompt-tsx- don't remove "empty" assistant messages! const toolResultMap = new Map(); - const budget = Math.floor(sizing.tokenBudget / this.props.toolCalls.length); + const budget = Math.floor(sizing.tokenBudget / round.toolCalls.length); + // TODO@prompt-tsx- this would be a bit easier with a ToolCall element, but we can only return one instance of a metadata right now + const toolCallSizing: PromptSizing = { + ...sizing, + tokenBudget: budget, + }; + return - todo - {await Promise.all(this.props.toolCalls.map(async toolCall => { - // TODO@prompt-tsx- this would be a bit easier with a ToolCall element, but we can only return one instance of a metadata right now - const toolCallSizing: PromptSizing = { - ...sizing, - tokenBudget: budget, - }; + {round.response || 'TODO'} + {await Promise.all(round.toolCalls.map(async toolCall => { const result = await this.renderOneToolCall(toolCall, toolCallSizing, this.props.toolInvocationToken); - toolResultMap.set(toolCall.toolCallId, result); + if (result.toolResult) { + toolResultMap.set(toolCall.toolCallId, result.toolResult); + } return result.message; }))} Above is the result of calling one or more tools. The user cannot see the results, so you should explain them to the user if referencing them in your answer. ; } - private async renderOneToolCall(toolCall: vscode.LanguageModelToolCallPart, sizing: PromptSizing, toolInvocationToken: vscode.ChatParticipantToolToken): Promise<{ message: ToolMessage; toolResult: vscode.LanguageModelToolResult; }> { + private async renderOneToolCall(toolCall: vscode.LanguageModelToolCallPart, sizing: PromptSizing, toolInvocationToken: vscode.ChatParticipantToolToken): Promise<{ message: ToolMessage; toolResult?: vscode.LanguageModelToolResult; }> { const tool = vscode.lm.tools.find(t => t.id === toolCall.name); if (!tool) { console.error(`Tool not found: ${toolCall.name}`); - return Tool not found; + return { message: Tool not found }; } const contentType = agentSupportedContentTypes.find(type => tool.supportedContentTypes.includes(type)); if (!contentType) { console.error(`Tool does not support any of the agent's content types: ${tool.id}`); - return Tool unsupported; + return { message: Tool unsupported }; } const tokenOptions: vscode.LanguageModelToolInvocationOptions['tokenOptions'] = { @@ -108,12 +128,17 @@ class ToolCalls extends PromptElement { countTokens: async (content: string) => sizing.countTokens(content), }; - const result = await vscode.lm.invokeTool(toolCall.name, { parameters: toolCall.parameters, requestedContentTypes: [contentType], toolInvocationToken: toolInvocationToken, tokenOptions }, dummyCancellationToken); - return + const toolResult = this.props.toolCallResults.get(toolCall.toolCallId) ?? + await vscode.lm.invokeTool(toolCall.name, { parameters: toolCall.parameters, requestedContentTypes: [contentType], toolInvocationToken: toolInvocationToken, tokenOptions }, dummyCancellationToken); + const message = {contentType === 'text/plain' ? - result[contentType] : - } + toolResult[contentType] : + } ; + return { + message, + toolResult: toolResult + }; } }