Skip to content

Commit

Permalink
Basic 'chat agent' API (#193152)
Browse files Browse the repository at this point in the history
* Very basic chat agent API/UX

* Add custom name/avatar, and restore them in persisted sessions

* Show agent subcommands on the top level

* Show editor decorations for subcommands

* Fix unit tests

* Implement unregister

* Revert slash command content widget change, still used by inline editor

* Remove content widget reference

* Fix leaked disposable
  • Loading branch information
roblourens authored Sep 15, 2023
1 parent 6d8c845 commit bc6b75e
Show file tree
Hide file tree
Showing 15 changed files with 716 additions and 52 deletions.
1 change: 1 addition & 0 deletions src/vs/workbench/api/browser/extensionHost.contribution.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import './mainThreadLocalization';
import './mainThreadBulkEdits';
import './mainThreadChatProvider';
import './mainThreadChatSlashCommands';
import './mainThreadChatAgents';
import './mainThreadChatVariables';
import './mainThreadCodeInsets';
import './mainThreadCLICommands';
Expand Down
65 changes: 65 additions & 0 deletions src/vs/workbench/api/browser/mainThreadChatAgents.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
/*---------------------------------------------------------------------------------------------
* Copyright (c) Microsoft Corporation. All rights reserved.
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/

import { DisposableMap } from 'vs/base/common/lifecycle';
import { revive } from 'vs/base/common/marshalling';
import { IProgress } from 'vs/platform/progress/common/progress';
import { ExtHostChatAgentsShape, ExtHostContext, MainContext, MainThreadChatAgentsShape } from 'vs/workbench/api/common/extHost.protocol';
import { IChatAgentMetadata, IChatAgentService } from 'vs/workbench/contrib/chat/common/chatAgents';
import { IChatSlashFragment } from 'vs/workbench/contrib/chat/common/chatSlashCommands';
import { IExtHostContext, extHostNamedCustomer } from 'vs/workbench/services/extensions/common/extHostCustomers';


@extHostNamedCustomer(MainContext.MainThreadChatAgents)
export class MainThreadChatAgents implements MainThreadChatAgentsShape {

private readonly _agents = new DisposableMap<number>;
private readonly _pendingProgress = new Map<number, IProgress<IChatSlashFragment>>();
private readonly _proxy: ExtHostChatAgentsShape;

constructor(
extHostContext: IExtHostContext,
@IChatAgentService private readonly _chatAgentService: IChatAgentService
) {
this._proxy = extHostContext.getProxy(ExtHostContext.ExtHostChatAgents);
}

$unregisterAgent(handle: number): void {
this._agents.deleteAndDispose(handle);
}

dispose(): void {
this._agents.clearAndDisposeAll();
}

$registerAgent(handle: number, name: string, metadata: IChatAgentMetadata): void {
if (!this._chatAgentService.hasAgent(name)) {
// dynamic!
this._chatAgentService.registerAgentData({
id: name,
metadata: revive(metadata)
});
}

const d = this._chatAgentService.registerAgentCallback(name, async (prompt, progress, history, token) => {
const requestId = Math.random();
this._pendingProgress.set(requestId, progress);
try {
return await this._proxy.$invokeAgent(handle, requestId, prompt, { history }, token);
} finally {
this._pendingProgress.delete(requestId);
}
});
this._agents.set(handle, d);
}

async $handleProgressChunk(requestId: number, chunk: IChatSlashFragment): Promise<void> {
this._pendingProgress.get(requestId)?.report(revive(chunk));
}

$unregisterCommand(handle: number): void {
this._agents.deleteAndDispose(handle);
}
}
7 changes: 7 additions & 0 deletions src/vs/workbench/api/common/extHost.api.impl.ts
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,7 @@ import { ExtHostChatSlashCommands } from 'vs/workbench/api/common/extHostChatSla
import { ExtHostChatVariables } from 'vs/workbench/api/common/extHostChatVariables';
import { ExtHostRelatedInformation } from 'vs/workbench/api/common/extHostAiRelatedInformation';
import { ExtHostAiEmbeddingVector } from 'vs/workbench/api/common/extHostEmbeddingVector';
import { ExtHostChatAgents } from 'vs/workbench/api/common/extHostChatAgents';

export interface IExtensionRegistries {
mine: ExtensionDescriptionRegistry;
Expand Down Expand Up @@ -209,6 +210,7 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
const extHostInteractiveEditor = rpcProtocol.set(ExtHostContext.ExtHostInlineChat, new ExtHostInteractiveEditor(rpcProtocol, extHostCommands, extHostDocuments, extHostLogService));
const extHostChatProvider = rpcProtocol.set(ExtHostContext.ExtHostChatProvider, new ExtHostChatProvider(rpcProtocol, extHostLogService));
const extHostChatSlashCommands = rpcProtocol.set(ExtHostContext.ExtHostChatSlashCommands, new ExtHostChatSlashCommands(rpcProtocol, extHostChatProvider, extHostLogService));
const extHostChatAgents = rpcProtocol.set(ExtHostContext.ExtHostChatAgents, new ExtHostChatAgents(rpcProtocol, extHostChatProvider, extHostLogService));
const extHostChatVariables = rpcProtocol.set(ExtHostContext.ExtHostChatVariables, new ExtHostChatVariables(rpcProtocol));
const extHostChat = rpcProtocol.set(ExtHostContext.ExtHostChat, new ExtHostChat(rpcProtocol, extHostLogService));
const extHostAiRelatedInformation = rpcProtocol.set(ExtHostContext.ExtHostAiRelatedInformation, new ExtHostRelatedInformation(rpcProtocol));
Expand Down Expand Up @@ -1360,7 +1362,12 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
registerMappedEditsProvider(selector: vscode.DocumentSelector, provider: vscode.MappedEditsProvider) {
checkProposedApiEnabled(extension, 'mappedEditsProvider');
return extHostLanguageFeatures.registerMappedEditsProvider(extension, selector, provider);
},
registerAgent(name: string, agent: vscode.ChatAgent, metadata: vscode.ChatAgentMetadata) {
checkProposedApiEnabled(extension, 'chatAgents');
return extHostChatAgents.registerAgent(extension.identifier, name, agent, metadata);
}

};

return <typeof vscode>{
Expand Down
13 changes: 13 additions & 0 deletions src/vs/workbench/api/common/extHost.protocol.ts
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ import { IChatMessage, IChatResponseFragment, IChatResponseProviderMetadata } fr
import { IChatSlashFragment } from 'vs/workbench/contrib/chat/common/chatSlashCommands';
import { IChatRequestVariableValue, IChatVariableData } from 'vs/workbench/contrib/chat/common/chatVariables';
import { RelatedInformationResult, RelatedInformationType } from 'vs/workbench/services/aiRelatedInformation/common/aiRelatedInformation';
import { IChatAgentMetadata } from 'vs/workbench/contrib/chat/common/chatAgents';

export interface IWorkspaceData extends IStaticWorkspaceData {
folders: { uri: UriComponents; name: string; index: number }[];
Expand Down Expand Up @@ -1154,6 +1155,16 @@ export interface ExtHostChatSlashCommandsShape {
$executeCommand(handle: number, requestId: number, prompt: string, context: { history: IChatMessage[] }, token: CancellationToken): Promise<any>;
}

export interface MainThreadChatAgentsShape extends IDisposable {
$registerAgent(handle: number, name: string, metadata: IChatAgentMetadata): void;
$unregisterAgent(handle: number): void;
$handleProgressChunk(requestId: number, chunk: IChatSlashFragment): Promise<void>;
}

export interface ExtHostChatAgentsShape {
$invokeAgent(handle: number, requestId: number, prompt: string, context: { history: IChatMessage[] }, token: CancellationToken): Promise<any>;
}

export interface MainThreadChatVariablesShape extends IDisposable {
$registerVariable(handle: number, data: IChatVariableData): void;
$unregisterVariable(handle: number): void;
Expand Down Expand Up @@ -2605,6 +2616,7 @@ export const MainContext = {
MainThreadBulkEdits: createProxyIdentifier<MainThreadBulkEditsShape>('MainThreadBulkEdits'),
MainThreadChatProvider: createProxyIdentifier<MainThreadChatProviderShape>('MainThreadChatProvider'),
MainThreadChatSlashCommands: createProxyIdentifier<MainThreadChatSlashCommandsShape>('MainThreadChatSlashCommands'),
MainThreadChatAgents: createProxyIdentifier<MainThreadChatAgentsShape>('MainThreadChatAgents'),
MainThreadChatVariables: createProxyIdentifier<MainThreadChatVariablesShape>('MainThreadChatVariables'),
MainThreadClipboard: createProxyIdentifier<MainThreadClipboardShape>('MainThreadClipboard'),
MainThreadCommands: createProxyIdentifier<MainThreadCommandsShape>('MainThreadCommands'),
Expand Down Expand Up @@ -2725,6 +2737,7 @@ export const ExtHostContext = {
ExtHostInlineChat: createProxyIdentifier<ExtHostInlineChatShape>('ExtHostInlineChatShape'),
ExtHostChat: createProxyIdentifier<ExtHostChatShape>('ExtHostChat'),
ExtHostChatSlashCommands: createProxyIdentifier<ExtHostChatSlashCommandsShape>('ExtHostChatSlashCommands'),
ExtHostChatAgents: createProxyIdentifier<ExtHostChatAgentsShape>('ExtHostChatAgents'),
ExtHostChatVariables: createProxyIdentifier<ExtHostChatVariablesShape>('ExtHostChatVariables'),
ExtHostChatProvider: createProxyIdentifier<ExtHostChatProviderShape>('ExtHostChatProvider'),
ExtHostAiRelatedInformation: createProxyIdentifier<ExtHostAiRelatedInformationShape>('ExtHostAiRelatedInformation'),
Expand Down
91 changes: 91 additions & 0 deletions src/vs/workbench/api/common/extHostChatAgents.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
/*---------------------------------------------------------------------------------------------
* Copyright (c) Microsoft Corporation. All rights reserved.
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/

import { DeferredPromise, raceCancellation } from 'vs/base/common/async';
import { CancellationToken } from 'vs/base/common/cancellation';
import { IDisposable, toDisposable } from 'vs/base/common/lifecycle';
import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions';
import { ILogService } from 'vs/platform/log/common/log';
import { Progress } from 'vs/platform/progress/common/progress';
import { ExtHostChatAgentsShape, IMainContext, MainContext, MainThreadChatAgentsShape } from 'vs/workbench/api/common/extHost.protocol';
import { ExtHostChatProvider } from 'vs/workbench/api/common/extHostChatProvider';
import * as typeConvert from 'vs/workbench/api/common/extHostTypeConverters';
import { ChatMessageRole } from 'vs/workbench/api/common/extHostTypes';
import { IChatMessage } from 'vs/workbench/contrib/chat/common/chatProvider';
import type * as vscode from 'vscode';

export class ExtHostChatAgents implements ExtHostChatAgentsShape {

private static _idPool = 0;

private readonly _agents = new Map<number, { extension: ExtensionIdentifier; agent: vscode.ChatAgent }>();
private readonly _proxy: MainThreadChatAgentsShape;

constructor(
mainContext: IMainContext,
private readonly _extHostChatProvider: ExtHostChatProvider,
private readonly _logService: ILogService,
) {
this._proxy = mainContext.getProxy(MainContext.MainThreadChatAgents);
}

registerAgent(extension: ExtensionIdentifier, name: string, agent: vscode.ChatAgent, metadata: vscode.ChatAgentMetadata): IDisposable {
const handle = ExtHostChatAgents._idPool++;
this._agents.set(handle, { extension, agent });
this._proxy.$registerAgent(handle, name, metadata);

return toDisposable(() => {
this._proxy.$unregisterAgent(handle);
this._agents.delete(handle);
});
}

async $invokeAgent(handle: number, requestId: number, prompt: string, context: { history: IChatMessage[] }, token: CancellationToken): Promise<any> {
const data = this._agents.get(handle);
if (!data) {
this._logService.warn(`[CHAT](${handle}) CANNOT invoke agent because the agent is not registered`);
return;
}

let done = false;
function throwIfDone() {
if (done) {
throw new Error('Only valid while executing the command');
}
}

const commandExecution = new DeferredPromise<void>();
token.onCancellationRequested(() => commandExecution.complete());
setTimeout(() => commandExecution.complete(), 3 * 1000);
this._extHostChatProvider.allowListExtensionWhile(data.extension, commandExecution.p);

const task = data.agent(
{ role: ChatMessageRole.User, content: prompt },
{ history: context.history.map(typeConvert.ChatMessage.to) },
new Progress<vscode.ChatAgentResponse>(p => {
throwIfDone();
this._proxy.$handleProgressChunk(requestId, { content: isInteractiveProgressFileTree(p.message) ? p.message : p.message.value });
}),
token
);

try {
return await raceCancellation(Promise.resolve(task).then((v) => {
if (v && 'followUp' in v) {
const convertedFollowup = v?.followUp?.map(f => typeConvert.ChatFollowup.from(f));
return { followUp: convertedFollowup };
}
return undefined;
}), token);
} finally {
done = true;
commandExecution.complete();
}
}
}

function isInteractiveProgressFileTree(thing: unknown): thing is vscode.InteractiveProgressFileTree {
return !!thing && typeof thing === 'object' && 'treeData' in thing;
}
2 changes: 2 additions & 0 deletions src/vs/workbench/contrib/chat/browser/chat.contribution.ts
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ import { ICommandService } from 'vs/platform/commands/common/commands';
import { ChatVariablesService, IChatVariablesService } from 'vs/workbench/contrib/chat/common/chatVariables';
import { registerChatFileTreeActions } from 'vs/workbench/contrib/chat/browser/actions/chatFileTreeActions';
import { QuickChatService } from 'vs/workbench/contrib/chat/browser/chatQuick';
import { ChatAgentService, IChatAgentService } from 'vs/workbench/contrib/chat/common/chatAgents';

// Register configuration
const configurationRegistry = Registry.as<IConfigurationRegistry>(ConfigurationExtensions.Configuration);
Expand Down Expand Up @@ -253,4 +254,5 @@ registerSingleton(IChatAccessibilityService, ChatAccessibilityService, Instantia
registerSingleton(IChatWidgetHistoryService, ChatWidgetHistoryService, InstantiationType.Delayed);
registerSingleton(IChatProviderService, ChatProviderService, InstantiationType.Delayed);
registerSingleton(IChatSlashCommandService, ChatSlashCommandService, InstantiationType.Delayed);
registerSingleton(IChatAgentService, ChatAgentService, InstantiationType.Delayed);
registerSingleton(IChatVariablesService, ChatVariablesService, InstantiationType.Delayed);
Loading

0 comments on commit bc6b75e

Please sign in to comment.