diff --git a/src/participant/prompts/promptBase.ts b/src/participant/prompts/promptBase.ts index 364284240..0a02a5bab 100644 --- a/src/participant/prompts/promptBase.ts +++ b/src/participant/prompts/promptBase.ts @@ -13,6 +13,8 @@ export interface PromptArgsBase { }; context?: vscode.ChatContext; connectionNames?: string[]; + databaseName?: string; + collectionName?: string; } export interface UserPromptResponse { @@ -163,9 +165,13 @@ export abstract class PromptBase { protected getHistoryMessages({ connectionNames, context, + databaseName, + collectionName, }: { connectionNames?: string[]; // Used to scrape the connecting messages from the history. context?: vscode.ChatContext; + databaseName?: string; + collectionName?: string; }): vscode.LanguageModelChatMessage[] { const messages: vscode.LanguageModelChatMessage[] = []; @@ -173,6 +179,13 @@ export abstract class PromptBase { return []; } + let previousItem: + | vscode.ChatRequestTurn + | vscode.ChatResponseTurn + | undefined = undefined; + + const namespaceIsKnown = + databaseName !== undefined && collectionName !== undefined; for (const historyItem of context.history) { if (historyItem instanceof vscode.ChatRequestTurn) { if ( @@ -181,9 +194,21 @@ export abstract class PromptBase { ) { // When the message is empty or a connection name then we skip it. // It's probably going to be the response to the connect step. + previousItem = historyItem; continue; } + if (previousItem instanceof vscode.ChatResponseTurn) { + const responseIntent = (previousItem.result as ChatResult).metadata + ?.intent; + + // If the namespace is already known, skip responses to prompts asking for it. + if (responseIntent === 'askForNamespace' && namespaceIsKnown) { + previousItem = historyItem; + continue; + } + } + // eslint-disable-next-line new-cap messages.push(vscode.LanguageModelChatMessage.User(historyItem.prompt)); } @@ -206,11 +231,17 @@ export abstract class PromptBase { 'emptyRequest', 'askToConnect', ]; - if ( - responseTypesToSkip.indexOf( - (historyItem.result as ChatResult)?.metadata?.intent - ) > -1 - ) { + + const responseType = (historyItem.result as ChatResult)?.metadata + ?.intent; + if (responseTypesToSkip.includes(responseType)) { + previousItem = historyItem; + continue; + } + + // If the namespace is already known, skip including prompts asking for it. + if (responseType === 'askForNamespace' && namespaceIsKnown) { + previousItem = historyItem; continue; } @@ -232,6 +263,7 @@ export abstract class PromptBase { // eslint-disable-next-line new-cap messages.push(vscode.LanguageModelChatMessage.Assistant(message)); } + previousItem = historyItem; } return messages; diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts index f5209b8a1..f6c05bbec 100644 --- a/src/test/suite/participant/participant.test.ts +++ b/src/test/suite/participant/participant.test.ts @@ -35,6 +35,10 @@ import { createMarkdownLink } from '../../../participant/markdown'; import EXTENSION_COMMANDS from '../../../commands'; import { getContentLength } from '../../../participant/prompts/promptBase'; import { ParticipantErrorTypes } from '../../../participant/participantErrorTypes'; +import { + createChatRequestTurn, + createChatResponseTurn, +} from './participantHelpers'; // The Copilot's model in not available in tests, // therefore we need to mock its methods and returning values. @@ -1073,33 +1077,26 @@ suite('Participant Controller Test Suite', function () { chatContextStub = { history: [ - { - prompt: 'find all docs by a name example', - command: 'query', - references: [], - participant: CHAT_PARTICIPANT_ID, - } as vscode.ChatRequestTurn, - Object.assign( - Object.create(vscode.ChatResponseTurn.prototype), - { - participant: CHAT_PARTICIPANT_ID, - response: [ - { - value: { - value: - 'Which database would you like this query to run against? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n', - } as vscode.MarkdownString, - }, - ], - command: 'query', - result: { - metadata: { - intent: 'askForNamespace', - chatId: firstChatId, - }, - }, - } as vscode.ChatResponseTurn + createChatRequestTurn( + '/query', + 'find all docs by a name example' ), + createChatResponseTurn('/query', { + response: [ + { + value: { + value: + 'Which database would you like this query to run against? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n', + } as vscode.MarkdownString, + }, + ], + result: { + metadata: { + intent: 'askForNamespace', + chatId: firstChatId, + }, + }, + }), ], }; @@ -1150,61 +1147,44 @@ suite('Participant Controller Test Suite', function () { }); chatContextStub = { history: [ - Object.assign(Object.create(vscode.ChatRequestTurn.prototype), { - prompt: 'find all docs by a name example', - command: 'query', - references: [], - participant: CHAT_PARTICIPANT_ID, - }), - Object.assign( - Object.create(vscode.ChatResponseTurn.prototype), - { - participant: CHAT_PARTICIPANT_ID, - response: [ - { - value: { - value: - 'Which database would you like to this query to run against? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n', - } as vscode.MarkdownString, - }, - ], - command: 'query', - result: { - metadata: { - intent: 'askForNamespace', - }, - }, - } + createChatRequestTurn( + '/query', + 'find all docs by a name example' ), - Object.assign(Object.create(vscode.ChatRequestTurn.prototype), { - prompt: 'dbOne', - command: 'query', - references: [], - participant: CHAT_PARTICIPANT_ID, + createChatResponseTurn('/query', { + response: [ + { + value: { + value: + 'Which database would you like to this query to run against? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n', + } as vscode.MarkdownString, + }, + ], + result: { + metadata: { + intent: 'askForNamespace', + }, + }, }), - Object.assign( - Object.create(vscode.ChatResponseTurn.prototype), - { - participant: CHAT_PARTICIPANT_ID, - response: [ - { - value: { - value: - 'Which collection would you like to query within dbOne? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n', - } as vscode.MarkdownString, - }, - ], - command: 'query', - result: { - metadata: { - intent: 'askForNamespace', - databaseName: 'dbOne', - collectionName: 'collOne', - chatId: firstChatId, - }, + createChatRequestTurn('/query', 'dbOne'), + createChatResponseTurn('/query', { + response: [ + { + value: { + value: + 'Which collection would you like to query within dbOne? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n', + } as vscode.MarkdownString, }, - } - ), + ], + result: { + metadata: { + intent: 'askForNamespace', + databaseName: 'dbOne', + collectionName: 'collOne', + chatId: firstChatId, + }, + }, + }), ], }; await invokeChatHandler(chatRequestMock); @@ -1240,33 +1220,26 @@ suite('Participant Controller Test Suite', function () { }; chatContextStub = { history: [ - { - prompt: 'find all docs by a name example', - command: 'query', - references: [], - participant: CHAT_PARTICIPANT_ID, - } as vscode.ChatRequestTurn, - Object.assign( - Object.create(vscode.ChatResponseTurn.prototype), - { - participant: CHAT_PARTICIPANT_ID, - response: [ - { - value: { - value: - 'Which database would you like this query to run against? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n', - } as vscode.MarkdownString, - }, - ], - command: 'query', - result: { - metadata: { - intent: 'askForNamespace', - chatId: 'pineapple', - }, - }, - } as vscode.ChatResponseTurn + createChatRequestTurn( + '/query', + 'find all docs by a name example' ), + createChatResponseTurn('/query', { + response: [ + { + value: { + value: + 'Which database would you like this query to run against? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n', + } as vscode.MarkdownString, + }, + ], + result: { + metadata: { + intent: 'askForNamespace', + chatId: 'pineapple', + }, + }, + }), ], }; const chatResult = await invokeChatHandler(chatRequestMock); @@ -1371,31 +1344,23 @@ suite('Participant Controller Test Suite', function () { }; chatContextStub = { history: [ - Object.assign(Object.create(vscode.ChatRequestTurn.prototype), { - prompt: - 'how do I make a find request vs favorite_fruits.pineapple?', - command: 'query', - references: [], - participant: CHAT_PARTICIPANT_ID, - }), - Object.assign( - Object.create(vscode.ChatResponseTurn.prototype), - { - participant: CHAT_PARTICIPANT_ID, - response: [ - { - value: 'some code', - }, - ], - command: 'query', - result: { - metadata: { - intent: 'query', - chatId: 'abc', - }, - }, - } + createChatRequestTurn( + '/query', + 'how do I make a find request vs favorite_fruits.pineapple?' ), + createChatResponseTurn('/query', { + response: [ + { + value: { value: 'some code' } as vscode.MarkdownString, + }, + ], + result: { + metadata: { + intent: 'query', + chatId: 'abc', + }, + }, + }), ], }; await invokeChatHandler(chatRequestMock); @@ -2034,12 +1999,10 @@ Schema: chatContextStub = { history: [ - Object.assign(Object.create(vscode.ChatRequestTurn.prototype), { - prompt: 'give me the count of all people in the prod database', - command: 'query', - references: [], - participant: CHAT_PARTICIPANT_ID, - }), + createChatRequestTurn( + '/query', + 'give me the count of all people in the prod database' + ), ], }; const { messages, stats } = await Prompts.query.buildMessages({ @@ -2176,6 +2139,120 @@ Schema: ); }); + suite('with askForNameSpace', function () { + const userMessages = [ + 'find all docs by a name example', + 'what other queries can be used as an example', + ]; + + const chatRequestMock = { + prompt: 'localhost', + command: 'query', + }; + + beforeEach(function () { + chatContextStub = { + history: [ + createChatRequestTurn('/query', userMessages[0]), + createChatResponseTurn('/query', { + participant: CHAT_PARTICIPANT_ID, + response: [ + { + value: { + value: + 'Which database would you like to query within this database?', + } as vscode.MarkdownString, + }, + ], + result: { + metadata: { + intent: 'askForNamespace', + }, + }, + }), + createChatRequestTurn('/query', 'dbOne'), + createChatResponseTurn('/query', { + participant: CHAT_PARTICIPANT_ID, + response: [ + { + value: { + value: + 'Which collection would you like to query within dbOne?', + } as vscode.MarkdownString, + }, + ], + result: { + metadata: { + intent: 'askForNamespace', + databaseName: 'dbOne', + collectionName: undefined, + chatId: testChatId, + }, + }, + }), + createChatRequestTurn('/query', 'collectionOne'), + createChatRequestTurn('/query', userMessages[1]), + ], + }; + }); + + test('does not include askForNameSpace messages in history if the metadata exists', async function () { + const { messages, stats } = await Prompts.query.buildMessages({ + context: chatContextStub, + request: chatRequestMock, + collectionName: 'people', + connectionNames: ['localhost', 'atlas'], + databaseName: 'prod', + sampleDocuments: [], + }); + + expect(messages.length).to.equal(4); + expect(messages[0].role).to.equal( + vscode.LanguageModelChatMessageRole.Assistant + ); + + // We don't expect history because we're removing the askForConnect message as well + // as the user response to it. Therefore the actual user prompt should be the first + // message that we supplied in the history. + expect(messages[1].role).to.equal( + vscode.LanguageModelChatMessageRole.User + ); + + expect( + messages.slice(1, 3).map((message) => getMessageContent(message)) + ).to.deep.equal(userMessages); + + expect(stats.command).to.equal('query'); + }); + + test('includes askForNameSpace messages in history if there is no metadata', async function () { + const { messages, stats } = await Prompts.query.buildMessages({ + context: chatContextStub, + request: chatRequestMock, + connectionNames: ['localhost', 'atlas'], + sampleDocuments: [], + // @ts-expect-error Forcing undefined for the purpose of test + databaseName: undefined, + // @ts-expect-error Forcing undefined for the purpose of test + collectionName: undefined, + }); + + expect(messages.length).to.equal(8); + expect(messages[0].role).to.equal( + vscode.LanguageModelChatMessageRole.Assistant + ); + + // We don't expect history because we're removing the askForConnect message as well + // as the user response to it. Therefore the actual user prompt should be the first + // message that we supplied in the history. + expect(messages[1].role).to.equal( + vscode.LanguageModelChatMessageRole.User + ); + + expect(stats.command).to.equal('query'); + }); + }); + test('removes askForConnect messages from history', async function () { // The user is responding to an `askToConnect` message, so the prompt is just the // name of the connection @@ -2190,13 +2267,8 @@ Schema: chatContextStub = { history: [ - Object.assign(Object.create(vscode.ChatRequestTurn.prototype), { - prompt: expectedPrompt, - command: 'query', - references: [], - participant: CHAT_PARTICIPANT_ID, - }), - Object.assign(Object.create(vscode.ChatResponseTurn.prototype), { + createChatRequestTurn('/query', expectedPrompt), + createChatResponseTurn('/query', { participant: CHAT_PARTICIPANT_ID, response: [ { @@ -2216,7 +2288,6 @@ Schema: } as vscode.MarkdownString, }, ], - command: 'query', result: { metadata: { intent: 'askToConnect', @@ -2270,32 +2341,20 @@ Schema: chatContextStub = { history: [ - Object.assign(Object.create(vscode.ChatRequestTurn.prototype), { - prompt: 'give me the count of all people in the prod database', - command: 'query', - references: [], - participant: CHAT_PARTICIPANT_ID, - }), - Object.assign(Object.create(vscode.ChatRequestTurn.prototype), { - prompt: 'some disallowed message', - command: 'query', - references: [], - participant: CHAT_PARTICIPANT_ID, - }), - Object.assign(Object.create(vscode.ChatResponseTurn.prototype), { + createChatRequestTurn( + '/query', + 'give me the count of all people in the prod database' + ), + createChatRequestTurn('/query', 'some disallowed message'), + createChatResponseTurn('/query', { result: { errorDetails: { message: ParticipantErrorTypes.FILTERED, }, + metadata: {}, }, - response: [], - participant: CHAT_PARTICIPANT_ID, - }), - Object.assign(Object.create(vscode.ChatRequestTurn.prototype), { - prompt: 'ok message', - references: [], - participant: CHAT_PARTICIPANT_ID, }), + createChatRequestTurn(undefined, 'ok message'), ], }; const { messages } = await Prompts.generic.buildMessages({ diff --git a/src/test/suite/participant/participantHelpers.ts b/src/test/suite/participant/participantHelpers.ts new file mode 100644 index 000000000..b66d810d4 --- /dev/null +++ b/src/test/suite/participant/participantHelpers.ts @@ -0,0 +1,50 @@ +import { CHAT_PARTICIPANT_ID } from '../../../participant/constants'; +import * as vscode from 'vscode'; +import type { ParticipantCommand } from '../../../participant/participant'; + +export function createChatRequestTurn( + command: ParticipantCommand | undefined, + prompt: vscode.ChatRequestTurn['prompt'] = 'some prompt', + options: { + participant?: vscode.ChatRequestTurn['participant']; + references?: vscode.ChatRequestTurn['references']; + } = { + participant: CHAT_PARTICIPANT_ID, + references: [], + } +): vscode.ChatRequestTurn { + const { participant = CHAT_PARTICIPANT_ID, references = [] } = options; + + return Object.assign(Object.create(vscode.ChatRequestTurn.prototype), { + prompt, + command: command?.substring(1), + references, + participant, + }); +} + +export function createChatResponseTurn( + command: ParticipantCommand, + options: { + response?: vscode.ChatResponseTurn['response']; + result?: vscode.ChatResponseTurn['result']; + participant?: string; + } = { + response: [], + result: {}, + participant: CHAT_PARTICIPANT_ID, + } +): vscode.ChatRequestTurn { + const { + response = [], + result = {}, + participant = CHAT_PARTICIPANT_ID, + } = options; + + return Object.assign(Object.create(vscode.ChatResponseTurn.prototype), { + participant, + response, + command: command.substring(1), + result, + }); +}