know for whom a LM request is made (#204744)

This commit is contained in:
Johannes Rieken 2024-02-08 18:43:24 +01:00 committed by GitHub
parent 884acabd70
commit 49001e5237
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 16 additions and 13 deletions

View File

@ -48,11 +48,11 @@ export class MainThreadChatProvider implements MainThreadChatProviderShape {
$registerProvider(handle: number, identifier: string, metadata: IChatResponseProviderMetadata): void {
const registration = this._chatProviderService.registerChatResponseProvider(identifier, {
metadata,
provideChatResponse: async (messages, options, progress, token) => {
provideChatResponse: async (messages, from, options, progress, token) => {
const requestId = (Math.random() * 1e6) | 0;
this._pendingProgress.set(requestId, progress);
try {
await this._proxy.$provideLanguageModelResponse(handle, requestId, messages, options, token);
await this._proxy.$provideLanguageModelResponse(handle, requestId, from, messages, options, token);
} finally {
this._pendingProgress.delete(requestId);
}
@ -80,7 +80,7 @@ export class MainThreadChatProvider implements MainThreadChatProviderShape {
async $fetchResponse(extension: ExtensionIdentifier, providerId: string, requestId: number, messages: IChatMessage[], options: {}, token: CancellationToken): Promise<any> {
this._logService.debug('[CHAT] extension request STARTED', extension.value, requestId);
const task = this._chatProviderService.fetchChatResponse(providerId, messages, options, new Progress(value => {
const task = this._chatProviderService.fetchChatResponse(providerId, extension, messages, options, new Progress(value => {
this._proxy.$handleResponseFragment(requestId, value);
}), token);

View File

@ -1183,7 +1183,7 @@ export interface MainThreadChatProviderShape extends IDisposable {
export interface ExtHostChatProviderShape {
$updateLanguageModels(data: { added?: string[]; removed?: string[] }): void;
$updateAccesslist(data: { extension: ExtensionIdentifier; enabled: boolean }[]): void;
$provideLanguageModelResponse(handle: number, requestId: number, messages: IChatMessage[], options: { [name: string]: any }, token: CancellationToken): Promise<any>;
$provideLanguageModelResponse(handle: number, requestId: number, from: ExtensionIdentifier, messages: IChatMessage[], options: { [name: string]: any }, token: CancellationToken): Promise<any>;
$handleResponseFragment(requestId: number, chunk: IChatResponseFragment): Promise<void>;
}

View File

@ -121,7 +121,7 @@ export class ExtHostChatProvider implements ExtHostChatProviderShape {
});
}
async $provideLanguageModelResponse(handle: number, requestId: number, messages: IChatMessage[], options: { [name: string]: any }, token: CancellationToken): Promise<any> {
async $provideLanguageModelResponse(handle: number, requestId: number, from: ExtensionIdentifier, messages: IChatMessage[], options: { [name: string]: any }, token: CancellationToken): Promise<any> {
const data = this._languageModels.get(handle);
if (!data) {
return;
@ -134,14 +134,15 @@ export class ExtHostChatProvider implements ExtHostChatProviderShape {
this._proxy.$handleProgressChunk(requestId, { index: fragment.index, part: fragment.part });
});
return data.provider.provideChatResponse(messages.map(typeConvert.ChatMessage.to), options, progress, token);
if (data.provider.provideLanguageModelResponse) {
return data.provider.provideLanguageModelResponse(messages.map(typeConvert.ChatMessage.to), options, ExtensionIdentifier.toKey(from), progress, token);
} else {
return data.provider.provideChatResponse(messages.map(typeConvert.ChatMessage.to), options, progress, token);
}
}
//#region --- making request
$updateLanguageModels(data: { added?: string[] | undefined; removed?: string[] | undefined }): void {
const added: string[] = [];
const removed: string[] = [];

View File

@ -35,7 +35,7 @@ export interface IChatResponseProviderMetadata {
export interface IChatResponseProvider {
metadata: IChatResponseProviderMetadata;
provideChatResponse(messages: IChatMessage[], options: { [name: string]: any }, progress: IProgress<IChatResponseFragment>, token: CancellationToken): Promise<any>;
provideChatResponse(messages: IChatMessage[], from: ExtensionIdentifier, options: { [name: string]: any }, progress: IProgress<IChatResponseFragment>, token: CancellationToken): Promise<any>;
}
export const IChatProviderService = createDecorator<IChatProviderService>('chatProviderService');
@ -52,7 +52,7 @@ export interface IChatProviderService {
registerChatResponseProvider(identifier: string, provider: IChatResponseProvider): IDisposable;
fetchChatResponse(identifier: string, messages: IChatMessage[], options: { [name: string]: any }, progress: IProgress<IChatResponseFragment>, token: CancellationToken): Promise<any>;
fetchChatResponse(identifier: string, from: ExtensionIdentifier, messages: IChatMessage[], options: { [name: string]: any }, progress: IProgress<IChatResponseFragment>, token: CancellationToken): Promise<any>;
}
export class ChatProviderService implements IChatProviderService {
@ -89,11 +89,11 @@ export class ChatProviderService implements IChatProviderService {
});
}
fetchChatResponse(identifier: string, messages: IChatMessage[], options: { [name: string]: any }, progress: IProgress<IChatResponseFragment>, token: CancellationToken): Promise<any> {
fetchChatResponse(identifier: string, from: ExtensionIdentifier, messages: IChatMessage[], options: { [name: string]: any }, progress: IProgress<IChatResponseFragment>, token: CancellationToken): Promise<any> {
const provider = this._providers.get(identifier);
if (!provider) {
throw new Error(`Chat response provider with identifier ${identifier} is not registered.`);
}
return provider.provideChatResponse(messages, options, progress, token);
return provider.provideChatResponse(messages, from, options, progress, token);
}
}

View File

@ -17,6 +17,8 @@ declare module 'vscode' {
*/
export interface ChatResponseProvider {
provideChatResponse(messages: ChatMessage[], options: { [name: string]: any }, progress: Progress<ChatResponseFragment>, token: CancellationToken): Thenable<any>;
provideLanguageModelResponse?(messages: ChatMessage[], options: { [name: string]: any }, extensionId: string, progress: Progress<ChatResponseFragment>, token: CancellationToken): Thenable<any>;
}
export interface ChatResponseProviderMetadata {