From 1a46740e60b693e22b5704b6a1615636a643d225 Mon Sep 17 00:00:00 2001 From: Rhys Howell Date: Mon, 23 Sep 2024 16:54:11 -0400 Subject: [PATCH 1/2] chore(participant): add handlers for request failed, remove extra abort controller definitions --- src/participant/constants.ts | 7 ++ src/participant/docsChatbotAIService.ts | 15 +++- src/participant/participant.ts | 83 ++++++++++++------- .../participant/docsChatbotAIService.test.ts | 30 ++++++- 4 files changed, 101 insertions(+), 34 deletions(-) diff --git a/src/participant/constants.ts b/src/participant/constants.ts index b03e66274..49e283d05 100644 --- a/src/participant/constants.ts +++ b/src/participant/constants.ts @@ -10,6 +10,7 @@ export type ParticipantResponseType = | 'docs' | 'generic' | 'emptyRequest' + | 'cancelledRequest' | 'askToConnect' | 'askForNamespace'; @@ -49,6 +50,12 @@ export function namespaceRequestChatResult({ }; } +export function createCancelledRequestChatResult( + history: ReadonlyArray +): ChatResult { + return createChatResult('cancelledRequest', history); +} + function createChatResult( intent: ParticipantResponseType, history: ReadonlyArray diff --git a/src/participant/docsChatbotAIService.ts b/src/participant/docsChatbotAIService.ts index 4b0c8b1b8..86a5d2dfd 100644 --- a/src/participant/docsChatbotAIService.ts +++ b/src/participant/docsChatbotAIService.ts @@ -58,29 +58,37 @@ export class DocsChatbotAIService { uri, method, body, + signal, headers, }: { uri: string; method: string; + signal?: AbortSignal; body?: string; headers?: { [key: string]: string }; }): Promise { - return fetch(uri, { + return await fetch(uri, { headers: { origin: this._serverBaseUri, 'User-Agent': `mongodb-vscode/${version}`, ...headers, }, method, + signal, ...(body && { body }), }); } - async createConversation(): Promise { + async createConversation({ + signal, + }: { + signal: AbortSignal; + }): Promise { const uri = this.getUri('/conversations'); const res = await this._fetch({ uri, method: 'POST', + signal, }); let data; @@ -113,9 +121,11 @@ export class DocsChatbotAIService { async addMessage({ conversationId, message, + signal, }: { conversationId: string; message: string; + signal: AbortSignal; }): Promise { const uri = this.getUri(`/conversations/${conversationId}/messages`); const res = await this._fetch({ @@ -123,6 +133,7 @@ export class DocsChatbotAIService { method: 'POST', body: JSON.stringify({ message }), headers: { 'Content-Type': 'application/json' }, + signal, }); let data; diff --git a/src/participant/participant.ts b/src/participant/participant.ts index 09cc3707c..1a4dd80c5 100644 --- a/src/participant/participant.ts +++ b/src/participant/participant.ts @@ -20,6 +20,7 @@ import { queryRequestChatResult, docsRequestChatResult, schemaRequestChatResult, + createCancelledRequestChatResult, } from './constants'; import { QueryPrompt } from './prompts/query'; import { COL_NAME_ID, DB_NAME_ID, NamespacePrompt } from './prompts/namespace'; @@ -241,10 +242,6 @@ export default class ParticipantController { context, }); - const abortController = new AbortController(); - token.onCancellationRequested(() => { - abortController.abort(); - }); const responseContent = await this.getChatResponseContent({ messages, token, @@ -667,20 +664,32 @@ export default class ParticipantController { return askToConnectChatResult(context.history); } + _handleCancelledRequest({ + context, + stream, + }: { + context: vscode.ChatContext; + stream: vscode.ChatResponseStream; + }): ChatResult { + stream.markdown('\nRequest cancelled.'); + + return createCancelledRequestChatResult(context.history); + } + // The sample documents returned from this are simplified (strings and arrays shortened). // The sample documents are only returned when a user has the setting enabled. async _fetchCollectionSchemaAndSampleDocuments({ - abortSignal, databaseName, collectionName, amountOfDocumentsToSample = NUM_DOCUMENTS_TO_SAMPLE, schemaFormat = 'simplified', + token, }: { - abortSignal; databaseName: string; collectionName: string; amountOfDocumentsToSample?: number; schemaFormat?: 'simplified' | 'full'; + token: vscode.CancellationToken; }): Promise<{ schema?: string; sampleDocuments?: Document[]; @@ -693,6 +702,11 @@ export default class ParticipantController { }; } + const abortController = new AbortController(); + token.onCancellationRequested(() => { + abortController.abort(); + }); + try { const sampleDocuments = await dataService.sample( `${databaseName}.${collectionName}`, @@ -702,7 +716,7 @@ export default class ParticipantController { }, { promoteValues: false }, { - abortSignal, + abortSignal: abortController.signal, } ); @@ -836,10 +850,12 @@ export default class ParticipantController { }); } - const abortController = new AbortController(); - token.onCancellationRequested(() => { - abortController.abort(); - }); + if (token.isCancellationRequested) { + return this._handleCancelledRequest({ + context, + stream, + }); + } stream.push( new vscode.ChatResponseProgressPart( @@ -856,11 +872,11 @@ export default class ParticipantController { amountOfDocumentsSampled, // There can be fewer than the amount we attempt to sample. schema, } = await this._fetchCollectionSchemaAndSampleDocuments({ - abortSignal: abortController.signal, databaseName, schemaFormat: 'full', collectionName, amountOfDocumentsToSample: DOCUMENTS_TO_SAMPLE_FOR_SCHEMA_PROMPT, + token, })); if (!schema || amountOfDocumentsSampled === 0) { @@ -958,19 +974,21 @@ export default class ParticipantController { }); } - const abortController = new AbortController(); - token.onCancellationRequested(() => { - abortController.abort(); - }); + if (token.isCancellationRequested) { + return this._handleCancelledRequest({ + context, + stream, + }); + } let schema: string | undefined; let sampleDocuments: Document[] | undefined; try { ({ schema, sampleDocuments } = await this._fetchCollectionSchemaAndSampleDocuments({ - abortSignal: abortController.signal, databaseName, collectionName, + token, })); } catch (e) { // When an error fetching the collection schema or sample docs occurs, @@ -1012,9 +1030,11 @@ export default class ParticipantController { async _handleDocsRequestWithChatbot({ prompt, chatId, + token, }: { prompt: string; chatId: string; + token: vscode.CancellationToken; }): Promise<{ responseContent: string; responseReferences?: Reference[]; @@ -1022,9 +1042,14 @@ export default class ParticipantController { }> { let { docsChatbotConversationId } = this._chatMetadataStore.getChatMetadata(chatId) ?? {}; + const abortController = new AbortController(); + token.onCancellationRequested(() => { + abortController.abort(); + }); if (!docsChatbotConversationId) { - const conversation = - await this._docsChatbotAIService.createConversation(); + const conversation = await this._docsChatbotAIService.createConversation({ + signal: abortController.signal, + }); docsChatbotConversationId = conversation._id; this._chatMetadataStore.setChatMetadata(chatId, { docsChatbotConversationId, @@ -1035,6 +1060,7 @@ export default class ParticipantController { const response = await this._docsChatbotAIService.addMessage({ message: prompt, conversationId: docsChatbotConversationId, + signal: abortController.signal, }); log.info('Docs chatbot message sent', { @@ -1067,10 +1093,6 @@ export default class ParticipantController { context, }); - const abortController = new AbortController(); - token.onCancellationRequested(() => { - abortController.abort(); - }); const responseContent = await this.getChatResponseContent({ messages, token, @@ -1096,11 +1118,7 @@ export default class ParticipantController { vscode.CancellationToken ] ): Promise { - const [request, context, stream, token] = args; - const abortController = new AbortController(); - token.onCancellationRequested(() => { - abortController.abort(); - }); + const [request, context, stream] = args; const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId( context.history @@ -1115,11 +1133,20 @@ export default class ParticipantController { docsResult = await this._handleDocsRequestWithChatbot({ prompt: request.prompt, chatId, + token: args[3], }); } catch (error) { // If the docs chatbot API is not available, fall back to Copilot’s LLM and include // the MongoDB documentation link for users to go to our documentation site directly. log.error(error); + + if (args[3].isCancellationRequested) { + return this._handleCancelledRequest({ + context, + stream: args[2], + }); + } + this._telemetryService.track( TelemetryEventTypes.PARTICIPANT_RESPONSE_FAILED, { diff --git a/src/test/suite/participant/docsChatbotAIService.test.ts b/src/test/suite/participant/docsChatbotAIService.test.ts index 827cf6841..e0f97e7ff 100644 --- a/src/test/suite/participant/docsChatbotAIService.test.ts +++ b/src/test/suite/participant/docsChatbotAIService.test.ts @@ -28,7 +28,9 @@ suite('DocsChatbotAIService Test Suite', function () { }), }); global.fetch = fetchStub; - const conversation = await docsChatbotAIService.createConversation(); + const conversation = await docsChatbotAIService.createConversation({ + signal: new AbortController().signal, + }); expect(conversation._id).to.be.eql('650b4b260f975ef031016c8a'); }); @@ -42,13 +44,28 @@ suite('DocsChatbotAIService Test Suite', function () { global.fetch = fetchStub; try { - await docsChatbotAIService.createConversation(); + await docsChatbotAIService.createConversation({ + signal: new AbortController().signal, + }); expect.fail('It must fail with the server error'); } catch (error) { expect((error as Error).message).to.include('Internal server error'); } }); + test('throws when aborted', async () => { + try { + const abortController = new AbortController(); + abortController.abort(); + await docsChatbotAIService.createConversation({ + signal: abortController.signal, + }); + expect.fail('It must fail with the server error'); + } catch (error) { + expect((error as Error).message).to.include('This operation was aborted'); + } + }); + test('throws on bad requests', async () => { const fetchStub = sinon.stub().resolves({ status: 400, @@ -59,7 +76,9 @@ suite('DocsChatbotAIService Test Suite', function () { global.fetch = fetchStub; try { - await docsChatbotAIService.createConversation(); + await docsChatbotAIService.createConversation({ + signal: new AbortController().signal, + }); expect.fail('It must fail with the bad request error'); } catch (error) { expect((error as Error).message).to.include('Bad request'); @@ -76,7 +95,9 @@ suite('DocsChatbotAIService Test Suite', function () { global.fetch = fetchStub; try { - await docsChatbotAIService.createConversation(); + await docsChatbotAIService.createConversation({ + signal: new AbortController().signal, + }); expect.fail('It must fail with the rate limited error'); } catch (error) { expect((error as Error).message).to.include('Rate limited'); @@ -95,6 +116,7 @@ suite('DocsChatbotAIService Test Suite', function () { await docsChatbotAIService.addMessage({ conversationId: '650b4b260f975ef031016c8a', message: 'what is mongosh?', + signal: new AbortController().signal, }); expect.fail('It must fail with the timeout error'); } catch (error) { From 4e17634b87ec3b99b6f71ab6d73fd80717034410 Mon Sep 17 00:00:00 2001 From: Rhys Howell Date: Tue, 24 Sep 2024 08:24:26 -0400 Subject: [PATCH 2/2] fixup: better variable decompose, remove extra await/async --- src/participant/docsChatbotAIService.ts | 4 ++-- src/participant/participant.ts | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/participant/docsChatbotAIService.ts b/src/participant/docsChatbotAIService.ts index 86a5d2dfd..47858280f 100644 --- a/src/participant/docsChatbotAIService.ts +++ b/src/participant/docsChatbotAIService.ts @@ -54,7 +54,7 @@ export class DocsChatbotAIService { return `${this._serverBaseUri}api/${MONGODB_DOCS_CHATBOT_API_VERSION}${path}`; } - async _fetch({ + _fetch({ uri, method, body, @@ -67,7 +67,7 @@ export class DocsChatbotAIService { body?: string; headers?: { [key: string]: string }; }): Promise { - return await fetch(uri, { + return fetch(uri, { headers: { origin: this._serverBaseUri, 'User-Agent': `mongodb-vscode/${version}`, diff --git a/src/participant/participant.ts b/src/participant/participant.ts index 1a4dd80c5..cb2b44f25 100644 --- a/src/participant/participant.ts +++ b/src/participant/participant.ts @@ -1118,7 +1118,7 @@ export default class ParticipantController { vscode.CancellationToken ] ): Promise { - const [request, context, stream] = args; + const [request, context, stream, token] = args; const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId( context.history @@ -1133,17 +1133,17 @@ export default class ParticipantController { docsResult = await this._handleDocsRequestWithChatbot({ prompt: request.prompt, chatId, - token: args[3], + token, }); } catch (error) { // If the docs chatbot API is not available, fall back to Copilot’s LLM and include // the MongoDB documentation link for users to go to our documentation site directly. log.error(error); - if (args[3].isCancellationRequested) { + if (token.isCancellationRequested) { return this._handleCancelledRequest({ context, - stream: args[2], + stream, }); }