Skip to content

Commit

Permalink
Joh/languageModelInformation (microsoft#210490)
Browse files Browse the repository at this point in the history 
* add `LanguageModelInformation` interface

* tweaks to the language models metadata API

* update docs

* 💄

* add removal comments
  • Loading branch information
@jrieken
jrieken authored Apr 16, 2024
1 parent 26cf440 commit e700ece
Show file tree
Hide file tree
Showing 7 changed files with 117 additions and 54 deletions.
2 changes: 1 addition & 1 deletion src/vs/workbench/api/browser/mainThreadLanguageModels.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ export class MainThreadLanguageModels implements MainThreadLanguageModelsShape {
}


$countTokens(provider: string, value: string, token: CancellationToken): Promise<number> {
$countTokens(provider: string, value: string | IChatMessage, token: CancellationToken): Promise<number> {
return this._chatProviderService.computeTokenLength(provider, value, token);
}

Expand Down
17 changes: 13 additions & 4 deletions src/vs/workbench/api/common/extHost.api.impl.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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);
}
};

Expand Down Expand Up @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions src/vs/workbench/api/common/extHost.protocol.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1200,15 +1200,15 @@ export interface MainThreadLanguageModelsShape extends IDisposable {
$prepareChatAccess(extension: ExtensionIdentifier, providerId: string, justification?: string): Promise<ILanguageModelChatMetadata | undefined>;
$fetchResponse(extension: ExtensionIdentifier, provider: string, requestId: number, messages: IChatMessage[], options: {}, token: CancellationToken): Promise<any>;

$countTokens(provider: string, value: string, token: CancellationToken): Promise<number>;
$countTokens(provider: string, value: string | IChatMessage, token: CancellationToken): Promise<number>;
}

export interface ExtHostLanguageModelsShape {
$updateLanguageModels(data: { added?: ILanguageModelChatMetadata[]; removed?: string[] }): void;
$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<any>;
$handleResponseFragment(requestId: number, chunk: IChatResponseFragment): Promise<void>;
$provideTokenLength(handle: number, value: string, token: CancellationToken): Promise<number>;
$provideTokenLength(handle: number, value: string | IChatMessage, token: CancellationToken): Promise<number>;
}

export interface IExtensionChatAgentMetadata extends Dto<IChatAgentMetadata> {
Expand Down
52 changes: 33 additions & 19 deletions src/vs/workbench/api/common/extHostLanguageModels.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
});

Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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<number> {

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<Readonly<IExtensionDescription>>();
Expand Down Expand Up @@ -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<number> {

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);
}
};
}
Expand Down
12 changes: 7 additions & 5 deletions src/vs/workbench/contrib/chat/common/languageModels.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -40,7 +42,7 @@ export interface ILanguageModelChatMetadata {
export interface ILanguageModelChat {
metadata: ILanguageModelChatMetadata;
provideChatResponse(messages: IChatMessage[], from: ExtensionIdentifier, options: { [name: string]: any }, progress: IProgress<IChatResponseFragment>, token: CancellationToken): Promise<any>;
provideTokenCount(str: string, token: CancellationToken): Promise<number>;
provideTokenCount(message: string | IChatMessage, token: CancellationToken): Promise<number>;
}

export const ILanguageModelsService = createDecorator<ILanguageModelsService>('ILanguageModelsService');
Expand All @@ -59,7 +61,7 @@ export interface ILanguageModelsService {

makeLanguageModelChatRequest(identifier: string, from: ExtensionIdentifier, messages: IChatMessage[], options: { [name: string]: any }, progress: IProgress<IChatResponseFragment>, token: CancellationToken): Promise<any>;

computeTokenLength(identifier: string, message: string, token: CancellationToken): Promise<number>;
computeTokenLength(identifier: string, message: string | IChatMessage, token: CancellationToken): Promise<number>;
}

export class LanguageModelsService implements ILanguageModelsService {
Expand Down Expand Up @@ -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<number> {
computeTokenLength(identifier: string, message: string | IChatMessage, token: CancellationToken): Promise<number> {
const provider = this._providers.get(identifier);
if (!provider) {
throw new Error(`Chat response provider with identifier ${identifier} is not registered.`);
Expand Down
6 changes: 5 additions & 1 deletion src/vscode-dts/vscode.proposed.chatProvider.d.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ declare module 'vscode' {
export interface ChatResponseProvider {
provideLanguageModelResponse2(messages: LanguageModelChatMessage[], options: { [name: string]: any }, extensionId: string, progress: Progress<ChatResponseFragment>, token: CancellationToken): Thenable<any>;

provideTokenCount(text: string, token: CancellationToken): Thenable<number>;
provideTokenCount(text: string | LanguageModelChatMessage, token: CancellationToken): Thenable<number>;
}

export interface ChatResponseProviderMetadata {
Expand All @@ -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.
Expand Down
78 changes: 56 additions & 22 deletions src/vscode-dts/vscode.proposed.languageModels.d.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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<LanguageModelChangeEvent>;

/**
* 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.
*
Expand All @@ -198,14 +245,16 @@ declare module 'vscode' {
export function sendChatRequest(languageModel: string, messages: LanguageModelChatMessage[], options: LanguageModelChatRequestOptions, token: CancellationToken): Thenable<LanguageModelChatResponse>;

/**
* 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.
* 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 onDidChangeLanguageModels: Event<LanguageModelChangeEvent>;
export function computeTokenLength(languageModel: string, text: string | LanguageModelChatMessage, token?: CancellationToken): Thenable<number>;
}

/**
Expand All @@ -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<number>;
}

export interface ExtensionContext {
Expand Down

0 comments on commit e700ece

Please sign in to comment.