@codingame/monaco-vscode-extensions-service-override
Version:
VSCode public API plugged on the monaco editor - extensions service-override
167 lines (163 loc) • 8.9 kB
JavaScript
import { __decorate, __param } from '@codingame/monaco-vscode-api/external/tslib/tslib.es6';
import { CancellationToken } from '@codingame/monaco-vscode-api/vscode/vs/base/common/cancellation';
import { Disposable, DisposableMap, DisposableStore } from '@codingame/monaco-vscode-api/vscode/vs/base/common/lifecycle';
import { revive } from '@codingame/monaco-vscode-api/vscode/vs/base/common/marshalling';
import { ThemeIcon } from '@codingame/monaco-vscode-api/vscode/vs/base/common/themables';
import { isUriComponents, URI } from '@codingame/monaco-vscode-api/vscode/vs/base/common/uri';
import { ExtensionIdentifier } from '@codingame/monaco-vscode-api/vscode/vs/platform/extensions/common/extensions';
import { ILogService } from '@codingame/monaco-vscode-api/vscode/vs/platform/log/common/log.service';
import { IProductService } from '@codingame/monaco-vscode-api/vscode/vs/platform/product/common/productService.service';
import { toToolSetKey } from '@codingame/monaco-vscode-api/vscode/vs/workbench/contrib/chat/common/tools/languageModelToolsContribution';
import { toolResultHasBuffers, ToolDataSource, ToolSet } from '@codingame/monaco-vscode-api/vscode/vs/workbench/contrib/chat/common/tools/languageModelToolsService';
import { ILanguageModelToolsService } from '@codingame/monaco-vscode-api/vscode/vs/workbench/contrib/chat/common/tools/languageModelToolsService.service';
import { extHostNamedCustomer } from '../../services/extensions/common/extHostCustomers.js';
import { SerializableObjectWithBuffers } from '@codingame/monaco-vscode-api/vscode/vs/workbench/services/extensions/common/proxyIdentifier';
import { ExtHostContext, MainContext } from '@codingame/monaco-vscode-api/vscode/vs/workbench/api/common/extHost.protocol';
let MainThreadLanguageModelTools = class MainThreadLanguageModelTools extends Disposable {
constructor(extHostContext, _languageModelToolsService, _logService, _productService) {
super();
this._languageModelToolsService = _languageModelToolsService;
this._logService = _logService;
this._productService = _productService;
this._tools = this._register(( new DisposableMap()));
this._runningToolCalls = ( new Map());
this._proxy = ( extHostContext.getProxy(ExtHostContext.ExtHostLanguageModelTools));
this._register(
this._languageModelToolsService.onDidChangeTools(e => this._proxy.$onDidChangeTools(this.getToolDtos()))
);
}
getToolDtos() {
return ( Array.from(this._languageModelToolsService.getAllToolsIncludingDisabled()).map(tool => ({
id: tool.id,
displayName: tool.displayName,
toolReferenceName: tool.toolReferenceName,
legacyToolReferenceFullNames: tool.legacyToolReferenceFullNames,
fullReferenceName: tool.source.type === "mcp" ? this._languageModelToolsService.getFullReferenceName(tool) : undefined,
tags: tool.tags,
userDescription: tool.userDescription,
modelDescription: tool.modelDescription,
inputSchema: tool.inputSchema,
source: tool.source
})));
}
async $getTools() {
return this.getToolDtos();
}
async $invokeTool(dto, token) {
const result = await this._languageModelToolsService.invokeTool(
revive(dto),
(input, token) => this._proxy.$countTokensForInvocation(dto.callId, input, token),
token ?? CancellationToken.None
);
const out = {
content: result.content,
toolMetadata: result.toolMetadata,
toolResultError: result.toolResultError
};
return toolResultHasBuffers(result) ? ( new SerializableObjectWithBuffers(out)) : out;
}
$acceptToolProgress(callId, progress) {
this._runningToolCalls.get(callId)?.progress.report(progress);
}
$countTokensForInvocation(callId, input, token) {
const fn = this._runningToolCalls.get(callId);
if (!fn) {
throw ( new Error(`Tool invocation call ${callId} not found`));
}
return fn.countTokens(input, token);
}
$registerTool(id, hasHandleToolStream) {
const disposable = this._languageModelToolsService.registerToolImplementation(id, {
invoke: async (dto, countTokens, progress, token) => {
try {
this._runningToolCalls.set(dto.callId, {
countTokens,
progress
});
const resultSerialized = await this._proxy.$invokeTool(dto, token);
const resultDto = resultSerialized instanceof SerializableObjectWithBuffers ? resultSerialized.value : resultSerialized;
return revive(resultDto);
} finally {
this._runningToolCalls.delete(dto.callId);
}
},
prepareToolInvocation: (context, token) => this._proxy.$prepareToolInvocation(id, context, token),
handleToolStream: hasHandleToolStream ? (context, token) => this._proxy.$handleToolStream(id, context, token) : undefined
});
this._tools.set(id, disposable);
}
$registerToolWithDefinition(extensionId, definition, hasHandleToolStream) {
let icon;
if (definition.icon) {
if (ThemeIcon.isThemeIcon(definition.icon)) {
icon = definition.icon;
} else if (typeof definition.icon === "object" && definition.icon !== null && isUriComponents(definition.icon)) {
icon = {
dark: URI.revive(definition.icon)
};
} else {
const iconObj = definition.icon;
icon = {
dark: URI.revive(iconObj.dark),
light: iconObj.light ? URI.revive(iconObj.light) : undefined
};
}
}
const isBuiltinTool = this._productService.defaultChatAgent?.chatExtensionId ? ExtensionIdentifier.equals(extensionId, this._productService.defaultChatAgent.chatExtensionId) : false;
const source = isBuiltinTool ? ToolDataSource.Internal : revive(definition.source);
const toolData = {
id: definition.id,
displayName: definition.displayName,
toolReferenceName: definition.toolReferenceName,
legacyToolReferenceFullNames: definition.legacyToolReferenceFullNames,
tags: definition.tags,
userDescription: definition.userDescription,
modelDescription: definition.modelDescription,
inputSchema: definition.inputSchema,
source,
icon,
models: definition.models,
canBeReferencedInPrompt: !!definition.userDescription && !definition.toolSet
};
const id = definition.id;
const store = ( new DisposableStore());
store.add(this._languageModelToolsService.registerTool(toolData, {
invoke: async (dto, countTokens, progress, token) => {
try {
this._runningToolCalls.set(dto.callId, {
countTokens,
progress
});
const resultSerialized = await this._proxy.$invokeTool(dto, token);
const resultDto = resultSerialized instanceof SerializableObjectWithBuffers ? resultSerialized.value : resultSerialized;
return revive(resultDto);
} finally {
this._runningToolCalls.delete(dto.callId);
}
},
handleToolStream: hasHandleToolStream ? (context, token) => this._proxy.$handleToolStream(id, context, token) : undefined,
prepareToolInvocation: (context, token) => this._proxy.$prepareToolInvocation(id, context, token)
}));
if (definition.toolSet) {
const ts = this._languageModelToolsService.getToolSet(toToolSetKey(extensionId, definition.toolSet)) || this._languageModelToolsService.getToolSet(definition.toolSet);
if (!ts || !(ts instanceof ToolSet)) {
this._logService.warn(
`ToolSet ${definition.toolSet} not found for tool ${definition.id} from extension ${extensionId.value}`
);
} else {
store.add(ts.addTool(toolData));
}
}
this._tools.set(id, store);
}
$unregisterTool(name) {
this._tools.deleteAndDispose(name);
}
};
MainThreadLanguageModelTools = __decorate([
extHostNamedCustomer(MainContext.MainThreadLanguageModelTools),
( __param(1, ILanguageModelToolsService)),
( __param(2, ILogService)),
( __param(3, IProductService))
], MainThreadLanguageModelTools);
export { MainThreadLanguageModelTools };