Skip to content

Commit

Permalink
feat: support picking related files for chat editing in Add Files pic…
Browse files Browse the repository at this point in the history
…ker (#233817)

* feat: support picking related files for chat editing in Add Files picker

* fix: try prioritizing related files when generating file completions, drop timeout to not hold up completions too much
  • Loading branch information
joyceerhl authored Nov 14, 2024
1 parent 33c552e commit a067314
Show file tree
Hide file tree
Showing 7 changed files with 96 additions and 29 deletions.
5 changes: 3 additions & 2 deletions src/vs/workbench/api/browser/mainThreadChatAgents2.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ import { IChatWidgetService } from '../../contrib/chat/browser/chat.js';
import { ChatInputPart } from '../../contrib/chat/browser/chatInputPart.js';
import { AddDynamicVariableAction, IAddDynamicVariableContext } from '../../contrib/chat/browser/contrib/chatDynamicVariables.js';
import { ChatAgentLocation, IChatAgentHistoryEntry, IChatAgentImplementation, IChatAgentRequest, IChatAgentService } from '../../contrib/chat/common/chatAgents.js';
import { IChatEditingService } from '../../contrib/chat/common/chatEditingService.js';
import { IChatEditingService, IChatRelatedFileProviderMetadata } from '../../contrib/chat/common/chatEditingService.js';
import { ChatRequestAgentPart } from '../../contrib/chat/common/chatParserTypes.js';
import { ChatRequestParser } from '../../contrib/chat/common/chatRequestParser.js';
import { IChatContentInlineReference, IChatContentReference, IChatFollowup, IChatProgress, IChatService, IChatTask, IChatWarningMessage } from '../../contrib/chat/common/chatService.js';
Expand Down Expand Up @@ -349,8 +349,9 @@ export class MainThreadChatAgents2 extends Disposable implements MainThreadChatA
this._chatParticipantDetectionProviders.deleteAndDispose(handle);
}

$registerRelatedFilesProvider(handle: number): void {
$registerRelatedFilesProvider(handle: number, metadata: IChatRelatedFileProviderMetadata): void {
this._chatRelatedFilesProviders.set(handle, this._chatEditingService.registerRelatedFilesProvider(handle, {
description: metadata.description,
provideRelatedFiles: async (request, token) => {
return (await this._proxy.$provideRelatedFiles(handle, request, token))?.map((v) => ({ uri: URI.from(v.uri), description: v.description })) ?? [];
}
Expand Down
4 changes: 2 additions & 2 deletions src/vs/workbench/api/common/extHost.protocol.ts
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ import { IRevealOptions, ITreeItem, IViewBadge } from '../../common/views.js';
import { CallHierarchyItem } from '../../contrib/callHierarchy/common/callHierarchy.js';
import { ChatAgentLocation, IChatAgentMetadata, IChatAgentRequest, IChatAgentResult, IChatWelcomeMessageContent } from '../../contrib/chat/common/chatAgents.js';
import { ICodeMapperRequest, ICodeMapperResult } from '../../contrib/chat/common/chatCodeMapperService.js';
import { IChatRelatedFile, IChatRequestDraft } from '../../contrib/chat/common/chatEditingService.js';
import { IChatRelatedFile, IChatRelatedFileProviderMetadata as IChatRelatedFilesProviderMetadata, IChatRequestDraft } from '../../contrib/chat/common/chatEditingService.js';
import { IChatProgressHistoryResponseContent } from '../../contrib/chat/common/chatModel.js';
import { IChatContentInlineReference, IChatFollowup, IChatProgress, IChatResponseErrorDetails, IChatTask, IChatTaskDto, IChatUserActionEvent, IChatVoteAction } from '../../contrib/chat/common/chatService.js';
import { IChatRequestVariableValue, IChatVariableData, IChatVariableResolverProgress } from '../../contrib/chat/common/chatVariables.js';
Expand Down Expand Up @@ -1275,7 +1275,7 @@ export interface MainThreadChatAgentsShape2 extends IDisposable {
$registerAgent(handle: number, extension: ExtensionIdentifier, id: string, metadata: IExtensionChatAgentMetadata, dynamicProps: IDynamicChatAgentProps | undefined): void;
$registerChatParticipantDetectionProvider(handle: number): void;
$unregisterChatParticipantDetectionProvider(handle: number): void;
$registerRelatedFilesProvider(handle: number): void;
$registerRelatedFilesProvider(handle: number, metadata: IChatRelatedFilesProviderMetadata): void;
$unregisterRelatedFilesProvider(handle: number): void;
$registerAgentCompletionsProvider(handle: number, id: string, triggerCharacters: string[]): void;
$unregisterAgentCompletionsProvider(handle: number, id: string): void;
Expand Down
2 changes: 1 addition & 1 deletion src/vs/workbench/api/common/extHostChatAgents2.ts
Original file line number Diff line number Diff line change
Expand Up @@ -369,7 +369,7 @@ export class ExtHostChatAgents2 extends Disposable implements ExtHostChatAgentsS
registerRelatedFilesProvider(extension: IExtensionDescription, provider: vscode.ChatRelatedFilesProvider, metadata: vscode.ChatRelatedFilesProviderMetadata): vscode.Disposable {
const handle = ExtHostChatAgents2._relatedFilesProviderIdPool++;
this._relatedFilesProviders.set(handle, new ExtHostRelatedFilesProvider(extension, provider));
this._proxy.$registerRelatedFilesProvider(handle);
this._proxy.$registerRelatedFilesProvider(handle, metadata);
return toDisposable(() => {
this._relatedFilesProviders.delete(handle);
this._proxy.$unregisterRelatedFilesProvider(handle);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import { Codicon } from '../../../../../base/common/codicons.js';
import { KeyCode, KeyMod } from '../../../../../base/common/keyCodes.js';
import { Schemas } from '../../../../../base/common/network.js';
import { isElectron } from '../../../../../base/common/platform.js';
import { dirname } from '../../../../../base/common/resources.js';
import { compare } from '../../../../../base/common/strings.js';
import { ThemeIcon } from '../../../../../base/common/themables.js';
import { URI } from '../../../../../base/common/uri.js';
Expand Down Expand Up @@ -63,7 +64,7 @@ export function registerChatContextActions() {
/**
* We fill the quickpick with these types, and enable some quick access providers
*/
type IAttachmentQuickPickItem = ICommandVariableQuickPickItem | IQuickAccessQuickPickItem | IToolQuickPickItem | IImageQuickPickItem | IVariableQuickPickItem | IOpenEditorsQuickPickItem | ISearchResultsQuickPickItem | IScreenShotQuickPickItem;
type IAttachmentQuickPickItem = ICommandVariableQuickPickItem | IQuickAccessQuickPickItem | IToolQuickPickItem | IImageQuickPickItem | IVariableQuickPickItem | IOpenEditorsQuickPickItem | ISearchResultsQuickPickItem | IScreenShotQuickPickItem | IRelatedFilesQuickPickItem;

/**
* These are the types that we can get out of the quick pick
Expand Down Expand Up @@ -110,6 +111,19 @@ function isScreenshotQuickPickItem(obj: unknown): obj is IScreenShotQuickPickIte
&& (obj as IScreenShotQuickPickItem).kind === 'screenshot');
}

function isRelatedFileQuickPickItem(obj: unknown): obj is IRelatedFilesQuickPickItem {
return (
typeof obj === 'object'
&& (obj as IRelatedFilesQuickPickItem).kind === 'related-files'
);
}

interface IRelatedFilesQuickPickItem extends IQuickPickItem {
kind: 'related-files';
id: string;
label: string;
}

interface IImageQuickPickItem extends IQuickPickItem {
kind: 'image';
id: string;
Expand Down Expand Up @@ -384,7 +398,7 @@ export class AttachContextAction extends Action2 {
`:${item.range.startLineNumber}`);
}

private async _attachContext(widget: IChatWidget, commandService: ICommandService, clipboardService: IClipboardService, editorService: IEditorService, labelService: ILabelService, viewsService: IViewsService, chatEditingService: IChatEditingService | undefined, hostService: IHostService, isInBackground?: boolean, ...picks: IChatContextQuickPickItem[]) {
private async _attachContext(widget: IChatWidget, quickInputService: IQuickInputService, commandService: ICommandService, clipboardService: IClipboardService, editorService: IEditorService, labelService: ILabelService, viewsService: IViewsService, chatEditingService: IChatEditingService | undefined, hostService: IHostService, isInBackground?: boolean, ...picks: IChatContextQuickPickItem[]) {
const toAttach: IChatRequestVariableEntry[] = [];
for (const pick of picks) {
if (isISymbolQuickPickItem(pick) && pick.symbol) {
Expand Down Expand Up @@ -462,6 +476,38 @@ export class AttachContextAction extends Action2 {
});
}
}
} else if (isRelatedFileQuickPickItem(pick)) {
// Get all provider results and show them in a second tier picker
const chatSessionId = widget.viewModel?.sessionId;
if (!chatSessionId || !chatEditingService) {
continue;
}
const relatedFiles = await chatEditingService.getRelatedFiles(chatSessionId, widget.getInput(), CancellationToken.None);
if (!relatedFiles) {
continue;
}
const attachments = widget.attachmentModel.getAttachmentIDs();
const itemsPromise = chatEditingService.getRelatedFiles(chatSessionId, widget.getInput(), CancellationToken.None)
.then((files) => (files ?? []).reduce<((IQuickPickItem & { value: URI }) | IQuickPickSeparator)[]>((acc, cur) => {
acc.push({ type: 'separator', label: cur.group });
const workingSet = chatEditingService.currentEditingSessionObs.get()?.workingSet;
for (const file of cur.files) {
acc.push({
type: 'item',
label: labelService.getUriBasenameLabel(file.uri),
description: labelService.getUriLabel(dirname(file.uri), { relative: true }),
detail: file.description,
value: file.uri,
disabled: workingSet?.has(file.uri) || attachments.has(this._getFileContextId({ resource: file.uri })),
picked: true
});
}
return acc;
}, []));
const selectedFiles = await quickInputService.pick(itemsPromise, { placeHolder: localize('relatedFiles', 'Add related files to your working set based on your editing prompt'), canPickMany: true });
for (const file of selectedFiles ?? []) {
chatEditingService?.currentEditingSessionObs.get()?.addFileToWorkingSet(file.value);
}
} else if (isScreenshotQuickPickItem(pick)) {
const blob = await hostService.getScreenshot();
if (blob) {
Expand Down Expand Up @@ -654,6 +700,14 @@ export class AttachContextAction extends Action2 {
});
}
} else if (context.showFilesOnly) {
if (chatEditingService?.hasRelatedFilesProviders() && (widget.getInput() || chatEditingService.currentEditingSessionObs.get()?.workingSet.size)) {
quickPickItems.push({
kind: 'related-files',
id: 'related-files',
label: localize('chatContext.relatedFiles', 'Related Files'),
iconClass: ThemeIcon.asClassName(Codicon.sparkle),
});
}
if (editorService.editors.filter(e => e instanceof FileEditorInput || e instanceof DiffEditorInput || e instanceof UntitledTextEditorInput).length > 0) {
quickPickItems.push({
kind: 'open-editors',
Expand Down Expand Up @@ -698,7 +752,7 @@ export class AttachContextAction extends Action2 {
if (!clipboardService) {
return;
}
this._attachContext(widget, commandService, clipboardService, editorService, labelService, viewsService, chatEditingService, hostService, isBackgroundAccept, item);
this._attachContext(widget, quickInputService, commandService, clipboardService, editorService, labelService, viewsService, chatEditingService, hostService, isBackgroundAccept, item);
if (isQuickChat(widget)) {
quickChatService.open();
}
Expand Down
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 { compareBy, delta } from '../../../../../base/common/arrays.js';
import { coalesce, compareBy, delta } from '../../../../../base/common/arrays.js';
import { AsyncIterableSource } from '../../../../../base/common/async.js';
import { CancellationToken, CancellationTokenSource } from '../../../../../base/common/cancellation.js';
import { Codicon } from '../../../../../base/common/codicons.js';
Expand Down Expand Up @@ -387,14 +387,18 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
return editors;
}

hasRelatedFilesProviders(): boolean {
return this._chatRelatedFilesProviders.size > 0;
}

registerRelatedFilesProvider(handle: number, provider: IChatRelatedFilesProvider): IDisposable {
this._chatRelatedFilesProviders.set(handle, provider);
return toDisposable(() => {
this._chatRelatedFilesProviders.delete(handle);
});
}

async getRelatedFiles(chatSessionId: string, prompt: string, token: CancellationToken): Promise<readonly IChatRelatedFile[] | undefined> {
async getRelatedFiles(chatSessionId: string, prompt: string, token: CancellationToken): Promise<{ group: string; files: IChatRelatedFile[] }[] | undefined> {
const currentSession = this._currentSessionObs.get();
if (!currentSession || chatSessionId !== currentSession.chatSessionId) {
return undefined;
Expand All @@ -404,18 +408,17 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
const providers = Array.from(this._chatRelatedFilesProviders.values());
const result = await Promise.all(providers.map(async provider => {
try {
return provider.provideRelatedFiles({ prompt, files: currentWorkingSet }, token);
const relatedFiles = await provider.provideRelatedFiles({ prompt, files: currentWorkingSet }, token);
if (relatedFiles?.length) {
return { group: provider.description, files: relatedFiles };
}
return undefined;
} catch (e) {
return undefined;
}
}));

return result.reduce<IChatRelatedFile[]>((acc, cur) => {
if (cur) {
acc.push(...cur);
}
return acc;
}, []);
return coalesce(result);
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -524,6 +524,20 @@ class BuiltinDynamicCompletions extends Disposable {
const seen = new ResourceSet();
const len = result.suggestions.length;

// RELATED FILES
if (widget.location === ChatAgentLocation.EditingSession && widget.viewModel && this._chatEditingService.currentEditingSessionObs.get()?.chatSessionId === widget.viewModel?.sessionId) {
const relatedFiles = (await raceTimeout(this._chatEditingService.getRelatedFiles(widget.viewModel.sessionId, widget.getInput(), token), 1000)) ?? [];
for (const relatedFileGroup of relatedFiles) {
for (const relatedFile of relatedFileGroup.files) {
if (seen.has(relatedFile.uri)) {
continue;
}
seen.add(relatedFile.uri);
result.suggestions.push(makeFileCompletionItem(relatedFile.uri, relatedFile.description));
}
}
}

// HISTORY
// always take the last N items
for (const item of this.historyService.getHistory()) {
Expand Down Expand Up @@ -582,17 +596,6 @@ class BuiltinDynamicCompletions extends Disposable {
}
}

// RELATED FILES
if (widget.location === ChatAgentLocation.EditingSession && widget.viewModel && this._chatEditingService.currentEditingSessionObs.get()?.chatSessionId === widget.viewModel?.sessionId) {
for (const relatedFile of (await raceTimeout(this._chatEditingService.getRelatedFiles(widget.viewModel.sessionId, widget.getInput(), token), 2000)) ?? []) {
if (seen.has(relatedFile.uri)) {
continue;
}
seen.add(relatedFile.uri);
result.suggestions.push(makeFileCompletionItem(relatedFile.uri, relatedFile.description));
}
}

// mark results as incomplete because further typing might yield
// in more search results
result.incomplete = true;
Expand Down
8 changes: 7 additions & 1 deletion src/vs/workbench/contrib/chat/common/chatEditingService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -42,21 +42,27 @@ export interface IChatEditingService {
getSnapshotUri(requestId: string, uri: URI): URI | undefined;
restoreSnapshot(requestId: string | undefined): Promise<void>;

hasRelatedFilesProviders(): boolean;
registerRelatedFilesProvider(handle: number, provider: IChatRelatedFilesProvider): IDisposable;
getRelatedFiles(chatSessionId: string, prompt: string, token: CancellationToken): Promise<readonly IChatRelatedFile[] | undefined>;
getRelatedFiles(chatSessionId: string, prompt: string, token: CancellationToken): Promise<{ group: string; files: IChatRelatedFile[] }[] | undefined>;
}

export interface IChatRequestDraft {
readonly prompt: string;
readonly files: readonly URI[];
}

export interface IChatRelatedFileProviderMetadata {
readonly description: string;
}

export interface IChatRelatedFile {
readonly uri: URI;
readonly description: string;
}

export interface IChatRelatedFilesProvider {
readonly description: string;
provideRelatedFiles(chatRequest: IChatRequestDraft, token: CancellationToken): Promise<IChatRelatedFile[] | undefined>;
}

Expand Down

0 comments on commit a067314

Please sign in to comment.