From e700ece8e329253b9ebaf003380c73c3280ae033 Mon Sep 17 00:00:00 2001 From: Johannes Rieken Date: Tue, 16 Apr 2024 10:22:06 -0700 Subject: [PATCH] Joh/languageModelInformation (#210490) * add `LanguageModelInformation` interface * tweaks to the language models metadata API * update docs * :lipstick: * add removal comments --- .../api/browser/mainThreadLanguageModels.ts | 2 +- .../workbench/api/common/extHost.api.impl.ts | 17 +++- .../workbench/api/common/extHost.protocol.ts | 4 +- .../api/common/extHostLanguageModels.ts | 52 ++++++++----- .../contrib/chat/common/languageModels.ts | 12 +-- .../vscode.proposed.chatProvider.d.ts | 6 +- .../vscode.proposed.languageModels.d.ts | 78 +++++++++++++------ 7 files changed, 117 insertions(+), 54 deletions(-) diff --git a/src/vs/workbench/api/browser/mainThreadLanguageModels.ts b/src/vs/workbench/api/browser/mainThreadLanguageModels.ts index 028ef5db32c..593c943ff14 100644 --- a/src/vs/workbench/api/browser/mainThreadLanguageModels.ts +++ b/src/vs/workbench/api/browser/mainThreadLanguageModels.ts @@ -123,7 +123,7 @@ export class MainThreadLanguageModels implements MainThreadLanguageModelsShape { } - $countTokens(provider: string, value: string, token: CancellationToken): Promise { + $countTokens(provider: string, value: string | IChatMessage, token: CancellationToken): Promise { return this._chatProviderService.computeTokenLength(provider, value, token); } diff --git a/src/vs/workbench/api/common/extHost.api.impl.ts b/src/vs/workbench/api/common/extHost.api.impl.ts index 4e8dab3e724..ed36aa572ca 100644 --- a/src/vs/workbench/api/common/extHost.api.impl.ts +++ b/src/vs/workbench/api/common/extHost.api.impl.ts @@ -3,7 +3,7 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import { CancellationTokenSource } from 'vs/base/common/cancellation'; +import { CancellationToken, CancellationTokenSource } from 'vs/base/common/cancellation'; import * as errors from 'vs/base/common/errors'; import { Emitter, Event } from 'vs/base/common/event'; import { combinedDisposable } from 'vs/base/common/lifecycle'; @@ -1448,6 +1448,15 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I sendChatRequest(languageModel: string, messages: vscode.LanguageModelChatMessage[], options: vscode.LanguageModelChatRequestOptions, token: vscode.CancellationToken) { checkProposedApiEnabled(extension, 'languageModels'); return extHostLanguageModels.sendChatRequest(extension, languageModel, messages, options, token); + }, + computeTokenLength(languageModel: string, text: string | vscode.LanguageModelChatMessage, token?: vscode.CancellationToken) { + checkProposedApiEnabled(extension, 'languageModels'); + token ??= CancellationToken.None; + return extHostLanguageModels.computeTokenLength(languageModel, text, token); + }, + getLanguageModelInformation(languageModel: string) { + checkProposedApiEnabled(extension, 'languageModels'); + return extHostLanguageModels.getLanguageModelInfo(languageModel); } }; @@ -1703,9 +1712,9 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I LanguageModelChatSystemMessage: extHostTypes.LanguageModelChatSystemMessage, LanguageModelChatUserMessage: extHostTypes.LanguageModelChatUserMessage, LanguageModelChatAssistantMessage: extHostTypes.LanguageModelChatAssistantMessage, - LanguageModelSystemMessage: extHostTypes.LanguageModelChatSystemMessage, - LanguageModelUserMessage: extHostTypes.LanguageModelChatUserMessage, - LanguageModelAssistantMessage: extHostTypes.LanguageModelChatAssistantMessage, + LanguageModelSystemMessage: extHostTypes.LanguageModelChatSystemMessage, // TODO@jrieken REMOVE + LanguageModelUserMessage: extHostTypes.LanguageModelChatUserMessage, // TODO@jrieken REMOVE + LanguageModelAssistantMessage: extHostTypes.LanguageModelChatAssistantMessage, // TODO@jrieken REMOVE LanguageModelError: extHostTypes.LanguageModelError, NewSymbolName: extHostTypes.NewSymbolName, NewSymbolNameTag: extHostTypes.NewSymbolNameTag, diff --git a/src/vs/workbench/api/common/extHost.protocol.ts b/src/vs/workbench/api/common/extHost.protocol.ts index 4e875689459..39e85e70c23 100644 --- a/src/vs/workbench/api/common/extHost.protocol.ts +++ b/src/vs/workbench/api/common/extHost.protocol.ts @@ -1200,7 +1200,7 @@ export interface MainThreadLanguageModelsShape extends IDisposable { $prepareChatAccess(extension: ExtensionIdentifier, providerId: string, justification?: string): Promise; $fetchResponse(extension: ExtensionIdentifier, provider: string, requestId: number, messages: IChatMessage[], options: {}, token: CancellationToken): Promise; - $countTokens(provider: string, value: string, token: CancellationToken): Promise; + $countTokens(provider: string, value: string | IChatMessage, token: CancellationToken): Promise; } export interface ExtHostLanguageModelsShape { @@ -1208,7 +1208,7 @@ export interface ExtHostLanguageModelsShape { $updateModelAccesslist(data: { from: ExtensionIdentifier; to: ExtensionIdentifier; enabled: boolean }[]): void; $provideLanguageModelResponse(handle: number, requestId: number, from: ExtensionIdentifier, messages: IChatMessage[], options: { [name: string]: any }, token: CancellationToken): Promise; $handleResponseFragment(requestId: number, chunk: IChatResponseFragment): Promise; - $provideTokenLength(handle: number, value: string, token: CancellationToken): Promise; + $provideTokenLength(handle: number, value: string | IChatMessage, token: CancellationToken): Promise; } export interface IExtensionChatAgentMetadata extends Dto { diff --git a/src/vs/workbench/api/common/extHostLanguageModels.ts b/src/vs/workbench/api/common/extHostLanguageModels.ts index 623ad2b6e3b..263c1578c86 100644 --- a/src/vs/workbench/api/common/extHostLanguageModels.ts +++ b/src/vs/workbench/api/common/extHostLanguageModels.ts @@ -155,7 +155,9 @@ export class ExtHostLanguageModels implements ExtHostLanguageModelsShape { this._proxy.$registerLanguageModelProvider(handle, identifier, { extension: extension.identifier, identifier: identifier, - model: metadata.name ?? '', + name: metadata.name ?? '', + version: metadata.version, + tokens: metadata.tokens, auth }); @@ -201,7 +203,7 @@ export class ExtHostLanguageModels implements ExtHostLanguageModelsShape { if (data.added) { for (const metadata of data.added) { this._allLanguageModelData.set(metadata.identifier, metadata); - added.push(metadata.model); + added.push(metadata.identifier); } } if (data.removed) { @@ -359,7 +361,36 @@ export class ExtHostLanguageModels implements ExtHostLanguageModelsShape { this._logService.error(err); } } + } + async computeTokenLength(languageModelId: string, value: string | vscode.LanguageModelChatMessage, token: vscode.CancellationToken): Promise { + + const data = this._allLanguageModelData.get(languageModelId); + if (!data) { + throw LanguageModelError.NotFound(`Language model '${languageModelId}' is unknown.`); + } + + const local = Iterable.find(this._languageModels.values(), candidate => candidate.languageModelId === languageModelId); + if (local) { + // stay inside the EH + return local.provider.provideTokenCount(value, token); + } + + return this._proxy.$countTokens(data.identifier, (typeof value === 'string' ? value : typeConvert.LanguageModelMessage.from(value)), token); + } + + getLanguageModelInfo(languageModelId: string): vscode.LanguageModelInformation | undefined { + const data = this._allLanguageModelData.get(languageModelId); + if (!data) { + return undefined; + } + + return { + id: data.identifier, + name: data.name, + version: data.version, + tokens: data.tokens, + }; } private readonly _languageAccessInformationExtensions = new Set>(); @@ -391,23 +422,6 @@ export class ExtHostLanguageModels implements ExtHostLanguageModelsShape { return undefined; } return list.has(data.extension); - }, - async computeTokenLength(languageModelId: string, value: string, token?: vscode.CancellationToken): Promise { - - token ??= CancellationToken.None; - - const data = that._allLanguageModelData.get(languageModelId); - if (!data) { - throw LanguageModelError.NotFound(`Language model '${languageModelId}' is unknown.`); - } - - const local = Iterable.find(that._languageModels.values(), candidate => candidate.languageModelId === languageModelId); - if (local) { - // stay inside the EH - return local.provider.provideTokenCount(value, token); - } - - return that._proxy.$countTokens(data.identifier, value, token); } }; } diff --git a/src/vs/workbench/contrib/chat/common/languageModels.ts b/src/vs/workbench/contrib/chat/common/languageModels.ts index 5c95913f016..a538c2aa93a 100644 --- a/src/vs/workbench/contrib/chat/common/languageModels.ts +++ b/src/vs/workbench/contrib/chat/common/languageModels.ts @@ -29,8 +29,10 @@ export interface IChatResponseFragment { export interface ILanguageModelChatMetadata { readonly extension: ExtensionIdentifier; readonly identifier: string; - readonly model: string; - readonly description?: string; + readonly name: string; + readonly version: string; + readonly tokens: number; + readonly auth?: { readonly providerLabel: string; readonly accountLabel?: string; @@ -40,7 +42,7 @@ export interface ILanguageModelChatMetadata { export interface ILanguageModelChat { metadata: ILanguageModelChatMetadata; provideChatResponse(messages: IChatMessage[], from: ExtensionIdentifier, options: { [name: string]: any }, progress: IProgress, token: CancellationToken): Promise; - provideTokenCount(str: string, token: CancellationToken): Promise; + provideTokenCount(message: string | IChatMessage, token: CancellationToken): Promise; } export const ILanguageModelsService = createDecorator('ILanguageModelsService'); @@ -59,7 +61,7 @@ export interface ILanguageModelsService { makeLanguageModelChatRequest(identifier: string, from: ExtensionIdentifier, messages: IChatMessage[], options: { [name: string]: any }, progress: IProgress, token: CancellationToken): Promise; - computeTokenLength(identifier: string, message: string, token: CancellationToken): Promise; + computeTokenLength(identifier: string, message: string | IChatMessage, token: CancellationToken): Promise; } export class LanguageModelsService implements ILanguageModelsService { @@ -104,7 +106,7 @@ export class LanguageModelsService implements ILanguageModelsService { return provider.provideChatResponse(messages, from, options, progress, token); } - computeTokenLength(identifier: string, message: string, token: CancellationToken): Promise { + computeTokenLength(identifier: string, message: string | IChatMessage, token: CancellationToken): Promise { const provider = this._providers.get(identifier); if (!provider) { throw new Error(`Chat response provider with identifier ${identifier} is not registered.`); diff --git a/src/vscode-dts/vscode.proposed.chatProvider.d.ts b/src/vscode-dts/vscode.proposed.chatProvider.d.ts index 3a4d53afbbf..da2c522419a 100644 --- a/src/vscode-dts/vscode.proposed.chatProvider.d.ts +++ b/src/vscode-dts/vscode.proposed.chatProvider.d.ts @@ -18,7 +18,7 @@ declare module 'vscode' { export interface ChatResponseProvider { provideLanguageModelResponse2(messages: LanguageModelChatMessage[], options: { [name: string]: any }, extensionId: string, progress: Progress, token: CancellationToken): Thenable; - provideTokenCount(text: string, token: CancellationToken): Thenable; + provideTokenCount(text: string | LanguageModelChatMessage, token: CancellationToken): Thenable; } export interface ChatResponseProviderMetadata { @@ -29,6 +29,10 @@ declare module 'vscode' { // TODO@API rename to model name: string; + version: string; + + tokens: number; + /** * When present, this gates the use of `requestLanguageModelAccess` behind an authorization flow where * the user must approve of another extension accessing the models contributed by this extension. diff --git a/src/vscode-dts/vscode.proposed.languageModels.d.ts b/src/vscode-dts/vscode.proposed.languageModels.d.ts index 4544241b130..73f4d95414c 100644 --- a/src/vscode-dts/vscode.proposed.languageModels.d.ts +++ b/src/vscode-dts/vscode.proposed.languageModels.d.ts @@ -98,9 +98,38 @@ declare module 'vscode' { */ export type LanguageModelChatMessage = LanguageModelChatSystemMessage | LanguageModelChatUserMessage | LanguageModelChatAssistantMessage; + /** + * Represents information about a registered language model. + */ + export interface LanguageModelInformation { + /** + * The identifier of the language model. + */ + readonly id: string; + + /** + * The human-readable name of the language model. + */ + readonly name: string; + + /** + * The version of the language model. + */ + readonly version: string; + + /** + * The number of available tokens that can be used when sending requests + * to the language model. + * + * @see {@link lm.sendChatRequest} + */ + readonly tokens: number; + } + /** * An event describing the change in the set of available language models. */ + // TODO@API use LanguageModelInformation instead of string? export interface LanguageModelChangeEvent { /** * Added language models. @@ -173,6 +202,24 @@ declare module 'vscode' { */ export namespace lm { + /** + * The identifiers of all language models that are currently available. + */ + export const languageModels: readonly string[]; + + /** + * An event that is fired when the set of available language models changes. + */ + export const onDidChangeLanguageModels: Event; + + /** + * Retrieve information about a language model. + * + * @param languageModel A language model identifier. + * @returns A {@link LanguageModelInformation} instance or `undefined` if the language model does not exist. + */ + export function getLanguageModelInformation(languageModel: string): LanguageModelInformation | undefined; + /** * Make a chat request using a language model. * @@ -198,14 +245,16 @@ declare module 'vscode' { export function sendChatRequest(languageModel: string, messages: LanguageModelChatMessage[], options: LanguageModelChatRequestOptions, token: CancellationToken): Thenable; /** - * The identifiers of all language models that are currently available. + * Uses the language model specific tokenzier and computes the length in token of a given message. + * + * *Note* that this function will throw when the language model does not exist. + * + * @param languageModel A language model identifier. + * @param text A string or a message instance. + * @param token Optional cancellation token. + * @returns A thenable that resolves to the length of the message in tokens. */ - export const languageModels: readonly string[]; - - /** - * An event that is fired when the set of available language models changes. - */ - export const onDidChangeLanguageModels: Event; + export function computeTokenLength(languageModel: string, text: string | LanguageModelChatMessage, token?: CancellationToken): Thenable; } /** @@ -228,21 +277,6 @@ declare module 'vscode' { * model does not exist or consent hasn't been asked for. */ canSendRequest(languageModelId: string): boolean | undefined; - - // TODO@API SYNC or ASYNC? - // TODO@API future - // retrieveQuota(languageModelId: string): { remaining: number; resets: Date }; - - // TODO@API SHOULD THIS BE in vscode.lm? - // TODO@API should this check for access/permissions? - /** - * - * Compute the token length for the given text - * @param languageModelId - * @param text - * @param token - */ - computeTokenLength(languageModelId: string, text: string, token?: CancellationToken): Thenable; } export interface ExtensionContext {