File size: 4,829 Bytes
6202252 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | import * as vscode from 'vscode';
export function registerChatTools(context: vscode.ExtensionContext) {
context.subscriptions.push(vscode.lm.registerTool('chat-tools-sample_tabCount', new TabCountTool()));
context.subscriptions.push(vscode.lm.registerTool('chat-tools-sample_findFiles', new FindFilesTool()));
context.subscriptions.push(vscode.lm.registerTool('chat-tools-sample_runInTerminal', new RunInTerminalTool()));
}
interface ITabCountParameters {
tabGroup?: number;
}
export class TabCountTool implements vscode.LanguageModelTool<ITabCountParameters> {
async invoke(
options: vscode.LanguageModelToolInvocationOptions<ITabCountParameters>,
_token: vscode.CancellationToken
) {
const params = options.input;
if (typeof params.tabGroup === 'number') {
const group = vscode.window.tabGroups.all[Math.max(params.tabGroup - 1, 0)];
const nth =
params.tabGroup === 1
? '1st'
: params.tabGroup === 2
? '2nd'
: params.tabGroup === 3
? '3rd'
: `${params.tabGroup}th`;
return new vscode.LanguageModelToolResult([new vscode.LanguageModelTextPart(`There are ${group.tabs.length} tabs open in the ${nth} tab group.`)]);
} else {
const group = vscode.window.tabGroups.activeTabGroup;
return new vscode.LanguageModelToolResult([new vscode.LanguageModelTextPart(`There are ${group.tabs.length} tabs open.`)]);
}
}
async prepareInvocation(
options: vscode.LanguageModelToolInvocationPrepareOptions<ITabCountParameters>,
_token: vscode.CancellationToken
) {
const confirmationMessages = {
title: 'Count the number of open tabs',
message: new vscode.MarkdownString(
`Count the number of open tabs?` +
(options.input.tabGroup !== undefined
? ` in tab group ${options.input.tabGroup}`
: '')
),
};
return {
invocationMessage: 'Counting the number of tabs',
confirmationMessages,
};
}
}
interface IFindFilesParameters {
pattern: string;
}
export class FindFilesTool implements vscode.LanguageModelTool<IFindFilesParameters> {
async invoke(
options: vscode.LanguageModelToolInvocationOptions<IFindFilesParameters>,
token: vscode.CancellationToken
) {
const params = options.input as IFindFilesParameters;
const files = await vscode.workspace.findFiles(
params.pattern,
'**/node_modules/**',
undefined,
token
);
const strFiles = files.map((f) => f.fsPath).join('\n');
return new vscode.LanguageModelToolResult([new vscode.LanguageModelTextPart(`Found ${files.length} files matching "${params.pattern}":\n${strFiles}`)]);
}
async prepareInvocation(
options: vscode.LanguageModelToolInvocationPrepareOptions<IFindFilesParameters>,
_token: vscode.CancellationToken
) {
return {
invocationMessage: `Searching workspace for "${options.input.pattern}"`,
};
}
}
interface IRunInTerminalParameters {
command: string;
}
async function waitForShellIntegration(
terminal: vscode.Terminal,
timeout: number
): Promise<void> {
let resolve: () => void;
let reject: (e: Error) => void;
const p = new Promise<void>((_resolve, _reject) => {
resolve = _resolve;
reject = _reject;
});
const timer = setTimeout(() => reject(new Error('Could not run terminal command: shell integration is not enabled')), timeout);
const listener = vscode.window.onDidChangeTerminalShellIntegration((e) => {
if (e.terminal === terminal) {
clearTimeout(timer);
listener.dispose();
resolve();
}
});
await p;
}
export class RunInTerminalTool
implements vscode.LanguageModelTool<IRunInTerminalParameters> {
async invoke(
options: vscode.LanguageModelToolInvocationOptions<IRunInTerminalParameters>,
_token: vscode.CancellationToken
) {
const params = options.input as IRunInTerminalParameters;
const terminal = vscode.window.createTerminal('Language Model Tool User');
terminal.show();
try {
await waitForShellIntegration(terminal, 5000);
} catch (e) {
return new vscode.LanguageModelToolResult([new vscode.LanguageModelTextPart((e as Error).message)]);
}
const execution = terminal.shellIntegration!.executeCommand(params.command);
const terminalStream = execution.read();
let terminalResult = '';
for await (const chunk of terminalStream) {
terminalResult += chunk;
}
return new vscode.LanguageModelToolResult([new vscode.LanguageModelTextPart(terminalResult)]);
}
async prepareInvocation(
options: vscode.LanguageModelToolInvocationPrepareOptions<IRunInTerminalParameters>,
_token: vscode.CancellationToken
) {
const confirmationMessages = {
title: 'Run command in terminal',
message: new vscode.MarkdownString(
`Run this command in a terminal?` +
`\n\n\`\`\`\n${options.input.command}\n\`\`\`\n`
),
};
return {
invocationMessage: `Running command in terminal`,
confirmationMessages,
};
}
}
|