mirror of
https://github.com/microsoft/vscode.git
synced 2026-09-13 05:01:08 +01:00
363 lines
14 KiB
TypeScript
363 lines
14 KiB
TypeScript
/*---------------------------------------------------------------------------------------------
|
|
* Copyright (c) Microsoft Corporation. All rights reserved.
|
|
* Licensed under the MIT License. See License.txt in the project root for license information.
|
|
*--------------------------------------------------------------------------------------------*/
|
|
|
|
import type * as vscode from 'vscode';
|
|
import { raceCancellation } from '../../../base/common/async.js';
|
|
import { CancellationToken } from '../../../base/common/cancellation.js';
|
|
import { CancellationError } from '../../../base/common/errors.js';
|
|
import { IDisposable, toDisposable } from '../../../base/common/lifecycle.js';
|
|
import { revive } from '../../../base/common/marshalling.js';
|
|
import { generateUuid } from '../../../base/common/uuid.js';
|
|
import { IExtensionDescription } from '../../../platform/extensions/common/extensions.js';
|
|
import { IPreparedToolInvocation, IStreamedToolInvocation, isToolInvocationContext, IToolInvocation, IToolInvocationContext, IToolInvocationPreparationContext, IToolInvocationStreamContext, IToolResult, ToolInvocationPresentation } from '../../contrib/chat/common/tools/languageModelToolsService.js';
|
|
import { ExtensionEditToolId, InternalEditToolId } from '../../contrib/chat/common/tools/builtinTools/editFileTool.js';
|
|
import { InternalFetchWebPageToolId } from '../../contrib/chat/common/tools/builtinTools/tools.js';
|
|
import { SearchExtensionsToolId } from '../../contrib/extensions/common/searchExtensionsTool.js';
|
|
import { checkProposedApiEnabled, isProposedApiEnabled } from '../../services/extensions/common/extensions.js';
|
|
import { Dto, SerializableObjectWithBuffers } from '../../services/extensions/common/proxyIdentifier.js';
|
|
import { ExtHostLanguageModelToolsShape, IMainContext, IToolDataDto, IToolDefinitionDto, MainContext, MainThreadLanguageModelToolsShape } from './extHost.protocol.js';
|
|
import { ExtHostLanguageModels } from './extHostLanguageModels.js';
|
|
import * as typeConvert from './extHostTypeConverters.js';
|
|
import { URI } from '../../../base/common/uri.js';
|
|
|
|
class Tool {
|
|
|
|
private _data: IToolDataDto;
|
|
private _apiObject: vscode.LanguageModelToolInformation | undefined;
|
|
private _apiObjectWithChatParticipantAdditions: vscode.LanguageModelToolInformation | undefined;
|
|
|
|
constructor(data: IToolDataDto) {
|
|
this._data = data;
|
|
}
|
|
|
|
update(newData: IToolDataDto): void {
|
|
this._data = newData;
|
|
this._apiObject = undefined;
|
|
this._apiObjectWithChatParticipantAdditions = undefined;
|
|
}
|
|
|
|
get data(): IToolDataDto {
|
|
return this._data;
|
|
}
|
|
|
|
get apiObject(): vscode.LanguageModelToolInformation {
|
|
if (!this._apiObject) {
|
|
this._apiObject = Object.freeze({
|
|
name: this._data.id,
|
|
description: this._data.modelDescription,
|
|
inputSchema: this._data.inputSchema,
|
|
tags: this._data.tags ?? [],
|
|
source: undefined
|
|
});
|
|
}
|
|
return this._apiObject;
|
|
}
|
|
|
|
get apiObjectWithChatParticipantAdditions() {
|
|
if (!this._apiObjectWithChatParticipantAdditions) {
|
|
this._apiObjectWithChatParticipantAdditions = Object.freeze({
|
|
name: this._data.id,
|
|
description: this._data.modelDescription,
|
|
inputSchema: this._data.inputSchema,
|
|
tags: this._data.tags ?? [],
|
|
source: typeConvert.LanguageModelToolSource.to(this._data.source)
|
|
});
|
|
}
|
|
return this._apiObjectWithChatParticipantAdditions;
|
|
}
|
|
}
|
|
|
|
export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape {
|
|
/** A map of tools that were registered in this EH */
|
|
private readonly _registeredTools = new Map<string, { extension: IExtensionDescription; tool: vscode.LanguageModelTool<Object> }>();
|
|
private readonly _proxy: MainThreadLanguageModelToolsShape;
|
|
private readonly _tokenCountFuncs = new Map</* call ID */string, (text: string, token?: vscode.CancellationToken) => Thenable<number>>();
|
|
|
|
/** A map of all known tools, from other EHs or registered in vscode core */
|
|
private readonly _allTools = new Map<string, Tool>();
|
|
|
|
constructor(
|
|
mainContext: IMainContext,
|
|
private readonly _languageModels: ExtHostLanguageModels,
|
|
) {
|
|
this._proxy = mainContext.getProxy(MainContext.MainThreadLanguageModelTools);
|
|
|
|
this._proxy.$getTools().then(tools => {
|
|
for (const tool of tools) {
|
|
this._allTools.set(tool.id, new Tool(revive(tool)));
|
|
}
|
|
});
|
|
}
|
|
|
|
async $countTokensForInvocation(callId: string, input: string, token: CancellationToken): Promise<number> {
|
|
const fn = this._tokenCountFuncs.get(callId);
|
|
if (!fn) {
|
|
throw new Error(`Tool invocation call ${callId} not found`);
|
|
}
|
|
|
|
return await fn(input, token);
|
|
}
|
|
|
|
async invokeTool(extension: IExtensionDescription, toolIdOrInfo: string | vscode.LanguageModelToolInformation, options: vscode.LanguageModelToolInvocationOptions<any>, token?: CancellationToken): Promise<vscode.LanguageModelToolResult> {
|
|
const toolId = typeof toolIdOrInfo === 'string' ? toolIdOrInfo : toolIdOrInfo.name;
|
|
const callId = generateUuid();
|
|
if (options.tokenizationOptions) {
|
|
this._tokenCountFuncs.set(callId, options.tokenizationOptions.countTokens);
|
|
}
|
|
|
|
try {
|
|
if (options.toolInvocationToken && !isToolInvocationContext(options.toolInvocationToken)) {
|
|
throw new Error(`Invalid tool invocation token`);
|
|
}
|
|
|
|
if ((toolId === InternalEditToolId || toolId === ExtensionEditToolId) && !isProposedApiEnabled(extension, 'chatParticipantPrivate')) {
|
|
throw new Error(`Invalid tool: ${toolId}`);
|
|
}
|
|
|
|
// Making the round trip here because not all tools were necessarily registered in this EH
|
|
const result = await this._proxy.$invokeTool({
|
|
toolId,
|
|
callId,
|
|
parameters: options.input,
|
|
tokenBudget: options.tokenizationOptions?.tokenBudget,
|
|
context: options.toolInvocationToken as IToolInvocationContext | undefined,
|
|
chatRequestId: isProposedApiEnabled(extension, 'chatParticipantPrivate') ? options.chatRequestId : undefined,
|
|
chatInteractionId: isProposedApiEnabled(extension, 'chatParticipantPrivate') ? options.chatInteractionId : undefined,
|
|
subAgentInvocationId: isProposedApiEnabled(extension, 'chatParticipantPrivate') ? options.subAgentInvocationId : undefined,
|
|
chatStreamToolCallId: isProposedApiEnabled(extension, 'chatParticipantAdditions') ? options.chatStreamToolCallId : undefined,
|
|
preToolUseResult: isProposedApiEnabled(extension, 'chatParticipantPrivate') ? options.preToolUseResult : undefined,
|
|
}, token);
|
|
|
|
const dto: Dto<IToolResult> = result instanceof SerializableObjectWithBuffers ? result.value : result;
|
|
return typeConvert.LanguageModelToolResult.to(revive(dto));
|
|
} finally {
|
|
this._tokenCountFuncs.delete(callId);
|
|
}
|
|
}
|
|
|
|
$onDidChangeTools(tools: IToolDataDto[]): void {
|
|
|
|
const oldTools = new Set(this._allTools.keys());
|
|
|
|
for (const tool of tools) {
|
|
oldTools.delete(tool.id);
|
|
const existing = this._allTools.get(tool.id);
|
|
if (existing) {
|
|
existing.update(tool);
|
|
} else {
|
|
this._allTools.set(tool.id, new Tool(revive(tool)));
|
|
}
|
|
}
|
|
|
|
for (const id of oldTools) {
|
|
this._allTools.delete(id);
|
|
}
|
|
}
|
|
|
|
getTools(extension: IExtensionDescription): vscode.LanguageModelToolInformation[] {
|
|
const hasParticipantAdditions = isProposedApiEnabled(extension, 'chatParticipantPrivate');
|
|
return Array.from(this._allTools.values())
|
|
.map(tool => hasParticipantAdditions ? tool.apiObjectWithChatParticipantAdditions : tool.apiObject)
|
|
.filter(tool => {
|
|
switch (tool.name) {
|
|
case InternalEditToolId:
|
|
case ExtensionEditToolId:
|
|
case InternalFetchWebPageToolId:
|
|
case SearchExtensionsToolId:
|
|
return isProposedApiEnabled(extension, 'chatParticipantPrivate');
|
|
default:
|
|
return true;
|
|
}
|
|
});
|
|
}
|
|
|
|
async $invokeTool(dto: Dto<IToolInvocation>, token: CancellationToken): Promise<Dto<IToolResult> | SerializableObjectWithBuffers<Dto<IToolResult>>> {
|
|
const item = this._registeredTools.get(dto.toolId);
|
|
if (!item) {
|
|
throw new Error(`Unknown tool ${dto.toolId}`);
|
|
}
|
|
|
|
const options: vscode.LanguageModelToolInvocationOptions<Object> = {
|
|
input: dto.parameters,
|
|
toolInvocationToken: revive(dto.context) as unknown as vscode.ChatParticipantToolToken | undefined,
|
|
};
|
|
if (isProposedApiEnabled(item.extension, 'chatParticipantPrivate')) {
|
|
options.chatRequestId = dto.chatRequestId;
|
|
options.chatInteractionId = dto.chatInteractionId;
|
|
options.chatSessionResource = URI.revive(dto.context?.sessionResource);
|
|
options.subAgentInvocationId = dto.subAgentInvocationId;
|
|
}
|
|
|
|
if (isProposedApiEnabled(item.extension, 'chatParticipantAdditions') && dto.modelId) {
|
|
options.model = await this.getModel(dto.modelId, item.extension);
|
|
}
|
|
if (isProposedApiEnabled(item.extension, 'chatParticipantAdditions') && dto.chatStreamToolCallId) {
|
|
options.chatStreamToolCallId = dto.chatStreamToolCallId;
|
|
}
|
|
|
|
if (dto.tokenBudget !== undefined) {
|
|
options.tokenizationOptions = {
|
|
tokenBudget: dto.tokenBudget,
|
|
countTokens: this._tokenCountFuncs.get(dto.callId) || ((value, token = CancellationToken.None) =>
|
|
this._proxy.$countTokensForInvocation(dto.callId, value, token))
|
|
};
|
|
}
|
|
|
|
let progress: vscode.Progress<{ message?: string | vscode.MarkdownString; increment?: number }> | undefined;
|
|
if (isProposedApiEnabled(item.extension, 'toolProgress')) {
|
|
let lastProgress: number | undefined;
|
|
progress = {
|
|
report: value => {
|
|
if (value.increment !== undefined) {
|
|
lastProgress = (lastProgress ?? 0) + value.increment;
|
|
}
|
|
|
|
this._proxy.$acceptToolProgress(dto.callId, {
|
|
message: typeConvert.MarkdownString.fromStrict(value.message),
|
|
progress: lastProgress === undefined ? undefined : lastProgress / 100,
|
|
});
|
|
}
|
|
};
|
|
}
|
|
|
|
// todo: 'any' cast because TS can't handle the overloads
|
|
// eslint-disable-next-line local/code-no-any-casts
|
|
const extensionResult = await raceCancellation(Promise.resolve((item.tool.invoke as any)(options, token, progress!)), token);
|
|
if (!extensionResult) {
|
|
throw new CancellationError();
|
|
}
|
|
|
|
return typeConvert.LanguageModelToolResult.from(extensionResult, item.extension);
|
|
}
|
|
|
|
private async getModel(modelId: string, extension: IExtensionDescription): Promise<vscode.LanguageModelChat> {
|
|
let model: vscode.LanguageModelChat | undefined;
|
|
if (modelId) {
|
|
model = await this._languageModels.getLanguageModelByIdentifier(extension, modelId);
|
|
}
|
|
if (!model) {
|
|
model = await this._languageModels.getDefaultLanguageModel(extension);
|
|
if (!model) {
|
|
throw new Error('Language model unavailable');
|
|
}
|
|
}
|
|
|
|
return model;
|
|
}
|
|
|
|
async $handleToolStream(toolId: string, context: IToolInvocationStreamContext, token: CancellationToken): Promise<IStreamedToolInvocation | undefined> {
|
|
const item = this._registeredTools.get(toolId);
|
|
if (!item) {
|
|
throw new Error(`Unknown tool ${toolId}`);
|
|
}
|
|
|
|
// Only call handleToolStream if it's defined on the tool
|
|
if (!item.tool.handleToolStream) {
|
|
return undefined;
|
|
}
|
|
|
|
// Ensure the chatParticipantAdditions API is enabled
|
|
checkProposedApiEnabled(item.extension, 'chatParticipantAdditions');
|
|
|
|
const options: vscode.LanguageModelToolInvocationStreamOptions<any> = {
|
|
rawInput: context.rawInput,
|
|
chatRequestId: context.chatRequestId,
|
|
chatSessionResource: context.chatSessionResource,
|
|
chatInteractionId: context.chatInteractionId
|
|
};
|
|
|
|
const result = await item.tool.handleToolStream(options, token);
|
|
if (!result) {
|
|
return undefined;
|
|
}
|
|
|
|
return {
|
|
invocationMessage: typeConvert.MarkdownString.fromStrict(result.invocationMessage)
|
|
};
|
|
}
|
|
|
|
async $prepareToolInvocation(toolId: string, context: IToolInvocationPreparationContext, token: CancellationToken): Promise<IPreparedToolInvocation | undefined> {
|
|
const item = this._registeredTools.get(toolId);
|
|
if (!item) {
|
|
throw new Error(`Unknown tool ${toolId}`);
|
|
}
|
|
|
|
const options: vscode.LanguageModelToolInvocationPrepareOptions<any> = {
|
|
input: context.parameters,
|
|
chatRequestId: context.chatRequestId,
|
|
chatSessionResource: context.chatSessionResource,
|
|
chatInteractionId: context.chatInteractionId,
|
|
forceConfirmationReason: context.forceConfirmationReason
|
|
};
|
|
if (context.forceConfirmationReason) {
|
|
checkProposedApiEnabled(item.extension, 'chatParticipantPrivate');
|
|
}
|
|
if (item.tool.prepareInvocation) {
|
|
const result = await item.tool.prepareInvocation(options, token);
|
|
if (!result) {
|
|
return undefined;
|
|
}
|
|
|
|
if (result.pastTenseMessage || result.presentation) {
|
|
checkProposedApiEnabled(item.extension, 'chatParticipantPrivate');
|
|
}
|
|
|
|
return {
|
|
confirmationMessages: result.confirmationMessages ? {
|
|
title: typeof result.confirmationMessages.title === 'string' ? result.confirmationMessages.title : typeConvert.MarkdownString.from(result.confirmationMessages.title),
|
|
message: typeof result.confirmationMessages.message === 'string' ? result.confirmationMessages.message : typeConvert.MarkdownString.from(result.confirmationMessages.message),
|
|
} : undefined,
|
|
invocationMessage: typeConvert.MarkdownString.fromStrict(result.invocationMessage),
|
|
pastTenseMessage: typeConvert.MarkdownString.fromStrict(result.pastTenseMessage),
|
|
presentation: result.presentation as ToolInvocationPresentation | undefined,
|
|
};
|
|
}
|
|
|
|
return undefined;
|
|
}
|
|
|
|
registerTool(extension: IExtensionDescription, id: string, tool: vscode.LanguageModelTool<any>): IDisposable {
|
|
this._registeredTools.set(id, { extension, tool });
|
|
this._proxy.$registerTool(id, typeof tool.handleToolStream === 'function');
|
|
|
|
return toDisposable(() => {
|
|
this._registeredTools.delete(id);
|
|
this._proxy.$unregisterTool(id);
|
|
});
|
|
}
|
|
|
|
registerToolDefinition(extension: IExtensionDescription, definition: vscode.LanguageModelToolDefinition, tool: vscode.LanguageModelTool<any>): IDisposable {
|
|
checkProposedApiEnabled(extension, 'languageModelToolSupportsModel');
|
|
|
|
const id = definition.name;
|
|
|
|
// Convert the definition to a DTO
|
|
const dto: IToolDefinitionDto = {
|
|
id,
|
|
displayName: definition.displayName,
|
|
toolReferenceName: definition.toolReferenceName,
|
|
userDescription: definition.userDescription,
|
|
modelDescription: definition.description,
|
|
inputSchema: definition.inputSchema as object,
|
|
source: {
|
|
type: 'extension',
|
|
label: extension.displayName ?? extension.name,
|
|
extensionId: extension.identifier,
|
|
},
|
|
icon: typeConvert.IconPath.from(definition.icon),
|
|
models: definition.models,
|
|
toolSet: definition.toolSet,
|
|
};
|
|
|
|
this._registeredTools.set(id, { extension, tool });
|
|
this._proxy.$registerToolWithDefinition(extension.identifier, dto, typeof tool.handleToolStream === 'function');
|
|
|
|
return toDisposable(() => {
|
|
this._registeredTools.delete(id);
|
|
this._proxy.$unregisterTool(id);
|
|
});
|
|
}
|
|
}
|