UNPKG

@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
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 };