From 61fed40ec142b548712790808513747577ac03f7 Mon Sep 17 00:00:00 2001 From: Rhys Howell Date: Mon, 23 Sep 2024 12:56:53 -0400 Subject: [PATCH] chore(participant): add progress indicators for asynchronous requests in chat VSCODE-613 --- src/participant/participant.ts | 173 ++++++++++++++++++--------------- 1 file changed, 96 insertions(+), 77 deletions(-) diff --git a/src/participant/participant.ts b/src/participant/participant.ts index 4098014e9..ae6806d49 100644 --- a/src/participant/participant.ts +++ b/src/participant/participant.ts @@ -449,24 +449,30 @@ export default class ParticipantController { ) as Promise; } - async getDatabasesTree({ + async renderDatabasesTree({ command, context, + stream, }: { command: ParticipantCommand; context: vscode.ChatContext; - }): Promise { + stream: vscode.ChatResponseStream; + }): Promise { const dataService = this._connectionController.getActiveDataService(); if (!dataService) { - return []; + return; } + stream.push( + new vscode.ChatResponseProgressPart('Fetching database names...') + ); + try { const databases = await dataService.listDatabases({ nameOnly: true, }); - return [ - ...databases.slice(0, MAX_MARKDOWN_LIST_LENGTH).map((db) => + databases.slice(0, MAX_MARKDOWN_LIST_LENGTH).forEach((db) => + stream.markdown( createMarkdownLink({ commandId: EXTENSION_COMMANDS.SELECT_DATABASE_WITH_PARTICIPANT, data: { @@ -478,46 +484,54 @@ export default class ParticipantController { }, name: db.name, }) - ), - ...(databases.length > MAX_MARKDOWN_LIST_LENGTH - ? [ - createMarkdownLink({ - data: { - command, - chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId( - context.history - ), - }, - commandId: EXTENSION_COMMANDS.SELECT_DATABASE_WITH_PARTICIPANT, - name: 'Show more', - }), - ] - : []), - ]; + ) + ); + if (databases.length > MAX_MARKDOWN_LIST_LENGTH) { + stream.markdown( + createMarkdownLink({ + data: { + command, + chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId( + context.history + ), + }, + commandId: EXTENSION_COMMANDS.SELECT_DATABASE_WITH_PARTICIPANT, + name: 'Show more', + }) + ); + } } catch (error) { + log.error('Unable to fetch databases:', error); + // Users can always do this manually when asked to provide a database name. - return []; + return; } } - async getCollectionTree({ + async renderCollectionsTree({ command, context, databaseName, + stream, }: { command: ParticipantCommand; databaseName: string; context: vscode.ChatContext; - }): Promise { + stream: vscode.ChatResponseStream; + }): Promise { const dataService = this._connectionController.getActiveDataService(); if (!dataService) { - return []; + return; } + stream.push( + new vscode.ChatResponseProgressPart('Fetching collection names...') + ); + try { const collections = await dataService.listCollections(databaseName); - return [ - ...collections.slice(0, MAX_MARKDOWN_LIST_LENGTH).map((coll) => + collections.slice(0, MAX_MARKDOWN_LIST_LENGTH).forEach((coll) => + stream.markdown( createMarkdownLink({ commandId: EXTENSION_COMMANDS.SELECT_COLLECTION_WITH_PARTICIPANT, data: { @@ -530,27 +544,28 @@ export default class ParticipantController { }, name: coll.name, }) - ), - ...(collections.length > MAX_MARKDOWN_LIST_LENGTH - ? [ - createMarkdownLink({ - commandId: - EXTENSION_COMMANDS.SELECT_COLLECTION_WITH_PARTICIPANT, - data: { - command, - chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId( - context.history - ), - databaseName, - }, - name: 'Show more', - }), - ] - : []), - ]; + ) + ); + if (collections.length > MAX_MARKDOWN_LIST_LENGTH) { + stream.markdown( + createMarkdownLink({ + commandId: EXTENSION_COMMANDS.SELECT_COLLECTION_WITH_PARTICIPANT, + data: { + command, + chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId( + context.history + ), + databaseName, + }, + name: 'Show more', + }) + ); + } } catch (error) { + log.error('Unable to fetch collections:', error); + // Users can always do this manually when asked to provide a collection name. - return []; + return; } } @@ -614,30 +629,26 @@ export default class ParticipantController { // we retrieve the available namespaces from the current connection. // Users can then select a value by clicking on an item in the list. if (!databaseName) { - const tree = await this.getDatabasesTree({ - command, - context, - }); stream.markdown( `What is the name of the database you would like${ command === '/query' ? ' this query' : '' } to run against?\n\n` ); - for (const item of tree) { - stream.markdown(item); - } - } else if (!collectionName) { - const tree = await this.getCollectionTree({ + await this.renderDatabasesTree({ command, - databaseName, context, + stream, }); + } else if (!collectionName) { stream.markdown( `Which collection would you like to use within ${databaseName}?\n\n` ); - for (const item of tree) { - stream.markdown(item); - } + await this.renderCollectionsTree({ + command, + databaseName, + context, + stream, + }); } return namespaceRequestChatResult({ @@ -675,12 +686,14 @@ export default class ParticipantController { collectionName, amountOfDocumentsToSample = NUM_DOCUMENTS_TO_SAMPLE, schemaFormat = 'simplified', + stream, }: { abortSignal; databaseName: string; collectionName: string; amountOfDocumentsToSample?: number; schemaFormat?: 'simplified' | 'full'; + stream: vscode.ChatResponseStream; }): Promise<{ schema?: string; sampleDocuments?: Document[]; @@ -693,6 +706,12 @@ export default class ParticipantController { }; } + stream.push( + new vscode.ChatResponseProgressPart( + 'Fetching documents and analyzing schema...' + ) + ); + try { const sampleDocuments = await dataService.sample( `${databaseName}.${collectionName}`, @@ -735,7 +754,7 @@ export default class ParticipantController { amountOfDocumentsSampled: sampleDocuments.length, }; } catch (err: any) { - log.error('Unable to fetch schema and sample documents', err); + log.error('Unable to fetch schema and sample documents:', err); throw err; } } @@ -759,7 +778,6 @@ export default class ParticipantController { // When the last message was asking for a database or collection name, // we re-ask the question. - let tree: vscode.MarkdownString[]; const databaseName = lastMessage.metadata.databaseName; if (databaseName) { stream.markdown( @@ -767,10 +785,11 @@ export default class ParticipantController { 'Please select a collection by either clicking on an item in the list or typing the name manually in the chat.' ) ); - tree = await this.getCollectionTree({ + await this.renderCollectionsTree({ command, databaseName, context, + stream, }); } else { stream.markdown( @@ -778,16 +797,13 @@ export default class ParticipantController { 'Please select a database by either clicking on an item in the list or typing the name manually in the chat.' ) ); - tree = await this.getDatabasesTree({ + await this.renderDatabasesTree({ command, context, + stream, }); } - for (const item of tree) { - stream.markdown(item); - } - return namespaceRequestChatResult({ databaseName, collectionName: undefined, @@ -841,12 +857,6 @@ export default class ParticipantController { abortController.abort(); }); - stream.push( - new vscode.ChatResponseProgressPart( - 'Fetching documents and analyzing schema...' - ) - ); - let sampleDocuments: Document[] | undefined; let amountOfDocumentsSampled: number; let schema: string | undefined; @@ -861,6 +871,7 @@ export default class ParticipantController { schemaFormat: 'full', collectionName, amountOfDocumentsToSample: DOCUMENTS_TO_SAMPLE_FOR_SCHEMA_PROMPT, + stream, })); if (!schema || amountOfDocumentsSampled === 0) { @@ -971,6 +982,7 @@ export default class ParticipantController { abortSignal: abortController.signal, databaseName, collectionName, + stream, })); } catch (e) { // When an error fetching the collection schema or sample docs occurs, @@ -978,7 +990,7 @@ export default class ParticipantController { // we do want to notify the user. stream.markdown( vscode.l10n.t( - 'An error occurred while fetching the collection schema and sample documents.\nThe generated query will not be able to reference your data.' + 'An error occurred while fetching the collection schema and sample documents.\nThe generated query will not be able to reference the shape of your data.' ) ); } @@ -1012,14 +1024,20 @@ export default class ParticipantController { async _handleDocsRequestWithChatbot({ prompt, chatId, + stream, }: { prompt: string; chatId: string; + stream: vscode.ChatResponseStream; }): Promise<{ responseContent: string; responseReferences?: Reference[]; docsChatbotMessageId: string; }> { + stream.push( + new vscode.ChatResponseProgressPart('Consulting MongoDB documentation...') + ); + let { docsChatbotConversationId } = this._chatMetadataStore.getChatMetadata(chatId) ?? {}; if (!docsChatbotConversationId) { @@ -1115,6 +1133,7 @@ export default class ParticipantController { docsResult = await this._handleDocsRequestWithChatbot({ prompt: request.prompt, chatId, + stream, }); } catch (error) { // If the docs chatbot API is not available, fall back to Copilot’s LLM and include @@ -1163,9 +1182,9 @@ export default class ParticipantController { if (!hasBeenShownWelcomeMessageAlready) { stream.markdown( vscode.l10n.t(` - Welcome to MongoDB Participant!\n\n - Interact with your MongoDB clusters and generate MongoDB-related code more efficiently with intelligent AI-powered feature, available today in the MongoDB extension.\n\n - Please see our [FAQ](https://www.mongodb.com/docs/generative-ai-faq/) for more information.\n\n`) +Welcome to MongoDB Participant!\n\n +Interact with your MongoDB clusters and generate MongoDB-related code more efficiently with intelligent AI-powered feature, available today in the MongoDB extension.\n\n +Please see our [FAQ](https://www.mongodb.com/docs/generative-ai-faq/) for more information.\n\n`) ); this._telemetryService.track(