Tool call plumbing

This commit is contained in:
Rob Lourens
2024-10-10 15:36:04 -07:00
parent 716df78daf
commit b010088034
4 changed files with 72 additions and 47 deletions

View File

@ -1,16 +1,16 @@
{
"name": "chat-tool-sample",
"name": "chat-tools-sample",
"version": "0.1.0",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"name": "chat-tool-sample",
"name": "chat-tools-sample",
"version": "0.1.0",
"devDependencies": {
"@types/node": "^20.5.9",
"@types/vscode": "^1.94.0",
"@vscode/prompt-tsx": "^0.2.7-alpha",
"@vscode/prompt-tsx": "^0.3.0-alpha.2",
"eslint": "^7.22.0",
"run-script-os": "^1.1.6",
"tslint": "^6.1.3",
@ -219,9 +219,9 @@
"dev": true
},
"node_modules/@vscode/prompt-tsx": {
"version": "0.2.7-alpha",
"resolved": "https://registry.npmjs.org/@vscode/prompt-tsx/-/prompt-tsx-0.2.7-alpha.tgz",
"integrity": "sha512-E2jmCMoNbY1XYI2yiDs7WNZD6mbv5ILZxSc94DOLRY44Qw0ozNm6DLUthN+7pHVhYVjxkMLkvBBuCCF9W9isIg==",
"version": "0.3.0-alpha.2",
"resolved": "https://registry.npmjs.org/@vscode/prompt-tsx/-/prompt-tsx-0.3.0-alpha.2.tgz",
"integrity": "sha512-xRD82rlLx7250IPdWXn+MAPjfEM4H6KXCDodlV7hHHiSrSGVJWiY+WQrT0KGMjGfaX9A2UyF0MaJ9/FjXkH3aQ==",
"dev": true
},
"node_modules/acorn": {

View File

@ -41,7 +41,10 @@
"languageModelTools": [
{
"id": "chat-tools-sample_tabCount",
"tags": ["editors", "chat-tools-sample"],
"tags": [
"editors",
"chat-tools-sample"
],
"name": "tabCount",
"displayName": "Tab Count",
"modelDescription": "The number of active tabs in a tab group",
@ -78,7 +81,9 @@
"description": "Search for files that match this glob pattern"
}
},
"required": ["pattern"]
"required": [
"pattern"
]
},
"supportedContentTypes": [
"text/plain"
@ -120,7 +125,7 @@
"devDependencies": {
"@types/node": "^20.5.9",
"@types/vscode": "^1.94.0",
"@vscode/prompt-tsx": "^0.2.7-alpha",
"@vscode/prompt-tsx": "^0.3.0-alpha.2",
"eslint": "^7.22.0",
"run-script-os": "^1.1.6",
"tslint": "^6.1.3",

View File

@ -1,7 +1,7 @@
import * as vscode from 'vscode';
import { FindFilesTool, RunInTerminalTool, TabCountTool } from './tools';
import { renderPrompt } from '@vscode/prompt-tsx';
import { ToolUserPrompt } from './toolsPrompt';
import { ToolResultMetadata, ToolUserPrompt } from './toolsPrompt';
export function activate(context: vscode.ExtensionContext) {
registerChatTool(context);
@ -221,6 +221,7 @@ function registerChatParticipant2(context: vscode.ExtensionContext) {
model)
const toolReferences = [...request.toolReferences];
const accumulatedToolCalls = new Map<string, vscode.LanguageModelToolResult>();
const runWithFunctions = async (): Promise<void> => {
const requestedTool = toolReferences.shift();
if (requestedTool) {
@ -248,7 +249,7 @@ function registerChatParticipant2(context: vscode.ExtensionContext) {
}
if (toolCalls.length) {
messages = (await renderPrompt(
const result = (await renderPrompt(
ToolUserPrompt,
{
context: chatContext,
@ -256,7 +257,12 @@ function registerChatParticipant2(context: vscode.ExtensionContext) {
toolCalls,
},
{ modelMaxPromptTokens: model.maxInputTokens },
model)).messages;
model));
messages = result.messages;
const toolResultMetadata = result.metadatas.get(ToolResultMetadata)
if (toolResultMetadata) {
toolResultMetadata.resultMap.forEach((value, key) => accumulatedToolCalls.set(key, value));
}
// RE-enter
return runWithFunctions();

View File

@ -8,8 +8,9 @@ import {
PromptPiece,
PromptSizing,
UserMessage,
PromptMetadata,
} from '@vscode/prompt-tsx';
import { ToolMessage, ToolResult } from '@vscode/prompt-tsx/dist/base/promptElements';
import { Chunk, ToolMessage, ToolResult } from '@vscode/prompt-tsx/dist/base/promptElements';
import * as vscode from 'vscode';
export interface ToolUserProps extends BasePromptElementProps {
@ -56,43 +57,50 @@ interface ToolCallsProps extends BasePromptElementProps {
toolInvocationToken: vscode.ChatParticipantToolToken;
}
const agentSupportedContentTypes = [promptTsxContentType, 'text/plain'];
const dummyCancellationToken: vscode.CancellationToken = new vscode.CancellationTokenSource().token;
class ToolCalls extends PromptElement<ToolCallsProps, void> {
render(state: void, sizing: PromptSizing) {
async render(state: void, sizing: PromptSizing) {
if (!this.props.toolCalls.length) {
return undefined;
}
// 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 }));
// TODO@prompt-tsx- don't remove "empty" assistant messages!
return <>
const toolResultMap = new Map<string, vscode.LanguageModelToolResult>();
const budget = Math.floor(sizing.tokenBudget / this.props.toolCalls.length);
return <Chunk>
<meta value={new ToolResultMetadata(toolResultMap)}></meta>
<AssistantMessage toolCalls={assistantToolCalls}>todo</AssistantMessage>
{this.props.toolCalls.map(toolCall => {
const tool = vscode.lm.tools.find(t => t.id === toolCall.name);
if (!tool) {
console.error(`Tool not found: ${toolCall.name}`);
return undefined;
}
return <ToolCall tool={tool} toolCall={toolCall} toolInvocationToken={this.props.toolInvocationToken}></ToolCall>;
})}
{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,
};
const result = await this.renderOneToolCall(toolCall, toolCallSizing, this.props.toolInvocationToken);
toolResultMap.set(toolCall.toolCallId, result);
return result.message;
}))}
<UserMessage priority={100}>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.</UserMessage>
</>;
</Chunk>;
}
}
interface ToolCallProps extends BasePromptElementProps {
tool: vscode.LanguageModelToolDescription;
toolCall: vscode.LanguageModelToolCallPart;
toolInvocationToken: vscode.ChatParticipantToolToken;
}
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 <ToolMessage toolCallId={toolCall.toolCallId}>Tool not found</ToolMessage>;
}
const agentSupportedContentTypes = [promptTsxContentType, 'text/plain'];
const dummyCancellationToken: vscode.CancellationToken = new vscode.CancellationTokenSource().token;
class ToolCall extends PromptElement<ToolCallProps, void> {
async render(state: void, sizing: PromptSizing) {
const contentType = agentSupportedContentTypes.find(type => this.props.tool.supportedContentTypes.includes(type));
const contentType = agentSupportedContentTypes.find(type => tool.supportedContentTypes.includes(type));
if (!contentType) {
console.error(`Tool does not support any of the agent's content types: ${this.props.tool.id}`);
return <ToolMessage toolCallId={this.props.toolCall.toolCallId}>Tool unsupported</ToolMessage>;
console.error(`Tool does not support any of the agent's content types: ${tool.id}`);
return <ToolMessage toolCallId={toolCall.toolCallId}>Tool unsupported</ToolMessage>;
}
const tokenOptions: vscode.LanguageModelToolInvocationOptions<unknown>['tokenOptions'] = {
@ -100,14 +108,20 @@ class ToolCall extends PromptElement<ToolCallProps, void> {
countTokens: async (content: string) => sizing.countTokens(content),
};
const result = await vscode.lm.invokeTool(this.props.toolCall.name, { parameters: this.props.toolCall.parameters, requestedContentTypes: [contentType], toolInvocationToken: this.props.toolInvocationToken, tokenOptions }, dummyCancellationToken);
return <>
<ToolMessage toolCallId={this.props.toolCall.toolCallId}>
{contentType === 'text/plain' ?
result[contentType] :
<elementJSON data={result[contentType]}></elementJSON>}
</ToolMessage>
</>
const result = await vscode.lm.invokeTool(toolCall.name, { parameters: toolCall.parameters, requestedContentTypes: [contentType], toolInvocationToken: toolInvocationToken, tokenOptions }, dummyCancellationToken);
return <ToolMessage toolCallId={toolCall.toolCallId}>
{contentType === 'text/plain' ?
result[contentType] :
<elementJSON data={result[contentType]}></elementJSON>}
</ToolMessage>;
}
}
export class ToolResultMetadata extends PromptMetadata {
constructor(
public resultMap: Map<string, vscode.LanguageModelToolResult>,
) {
super();
}
}