From 49001e5237910cfdae20dde5219b3c6572bd99f5 Mon Sep 17 00:00:00 2001 From: Johannes Rieken Date: Thu, 8 Feb 2024 18:43:24 +0100 Subject: [PATCH] know for whom a LM request is made (#204744) --- .../workbench/api/browser/mainThreadChatProvider.ts | 6 +++--- src/vs/workbench/api/common/extHost.protocol.ts | 2 +- src/vs/workbench/api/common/extHostChatProvider.ts | 11 ++++++----- src/vs/workbench/contrib/chat/common/chatProvider.ts | 8 ++++---- src/vscode-dts/vscode.proposed.chatProvider.d.ts | 2 ++ 5 files changed, 16 insertions(+), 13 deletions(-) diff --git a/src/vs/workbench/api/browser/mainThreadChatProvider.ts b/src/vs/workbench/api/browser/mainThreadChatProvider.ts index 24ce8d39797..7c6534079e2 100644 --- a/src/vs/workbench/api/browser/mainThreadChatProvider.ts +++ b/src/vs/workbench/api/browser/mainThreadChatProvider.ts @@ -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 { 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); diff --git a/src/vs/workbench/api/common/extHost.protocol.ts b/src/vs/workbench/api/common/extHost.protocol.ts index e8b75cbceba..a34f7a83238 100644 --- a/src/vs/workbench/api/common/extHost.protocol.ts +++ b/src/vs/workbench/api/common/extHost.protocol.ts @@ -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; + $provideLanguageModelResponse(handle: number, requestId: number, from: ExtensionIdentifier, messages: IChatMessage[], options: { [name: string]: any }, token: CancellationToken): Promise; $handleResponseFragment(requestId: number, chunk: IChatResponseFragment): Promise; } diff --git a/src/vs/workbench/api/common/extHostChatProvider.ts b/src/vs/workbench/api/common/extHostChatProvider.ts index d9159a7350a..d47cf519006 100644 --- a/src/vs/workbench/api/common/extHostChatProvider.ts +++ b/src/vs/workbench/api/common/extHostChatProvider.ts @@ -121,7 +121,7 @@ export class ExtHostChatProvider implements ExtHostChatProviderShape { }); } - async $provideLanguageModelResponse(handle: number, requestId: number, messages: IChatMessage[], options: { [name: string]: any }, token: CancellationToken): Promise { + async $provideLanguageModelResponse(handle: number, requestId: number, from: ExtensionIdentifier, messages: IChatMessage[], options: { [name: string]: any }, token: CancellationToken): Promise { 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[] = []; diff --git a/src/vs/workbench/contrib/chat/common/chatProvider.ts b/src/vs/workbench/contrib/chat/common/chatProvider.ts index 88349952e4e..38be72eb51e 100644 --- a/src/vs/workbench/contrib/chat/common/chatProvider.ts +++ b/src/vs/workbench/contrib/chat/common/chatProvider.ts @@ -35,7 +35,7 @@ export interface IChatResponseProviderMetadata { export interface IChatResponseProvider { metadata: IChatResponseProviderMetadata; - provideChatResponse(messages: IChatMessage[], options: { [name: string]: any }, progress: IProgress, token: CancellationToken): Promise; + provideChatResponse(messages: IChatMessage[], from: ExtensionIdentifier, options: { [name: string]: any }, progress: IProgress, token: CancellationToken): Promise; } export const IChatProviderService = createDecorator('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, token: CancellationToken): Promise; + fetchChatResponse(identifier: string, from: ExtensionIdentifier, messages: IChatMessage[], options: { [name: string]: any }, progress: IProgress, token: CancellationToken): Promise; } 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, token: CancellationToken): Promise { + fetchChatResponse(identifier: string, from: ExtensionIdentifier, messages: IChatMessage[], options: { [name: string]: any }, progress: IProgress, token: CancellationToken): Promise { 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); } } diff --git a/src/vscode-dts/vscode.proposed.chatProvider.d.ts b/src/vscode-dts/vscode.proposed.chatProvider.d.ts index 3ff89cca683..1d218bc8e00 100644 --- a/src/vscode-dts/vscode.proposed.chatProvider.d.ts +++ b/src/vscode-dts/vscode.proposed.chatProvider.d.ts @@ -17,6 +17,8 @@ declare module 'vscode' { */ export interface ChatResponseProvider { provideChatResponse(messages: ChatMessage[], options: { [name: string]: any }, progress: Progress, token: CancellationToken): Thenable; + + provideLanguageModelResponse?(messages: ChatMessage[], options: { [name: string]: any }, extensionId: string, progress: Progress, token: CancellationToken): Thenable; } export interface ChatResponseProviderMetadata {