Skip to content

Commit

Permalink
Add tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@gagik
gagik committed Nov 5, 2024
1 parent 548f247 commit f468adf
Show file tree
Hide file tree
Showing 2 changed files with 234 additions and 3 deletions.
43 changes: 40 additions & 3 deletions src/participant/prompts/promptBase.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@ export interface PromptArgsBase {
};
context?: vscode.ChatContext;
connectionNames?: string[];
databaseName?: string;
collectionName?: string;
}

export interface UserPromptResponse {
Expand Down Expand Up @@ -163,16 +165,24 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
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[] = [];

if (!context) {
return [];
}

let previousItem:
| vscode.ChatRequestTurn
| vscode.ChatResponseTurn
| undefined = undefined;
for (const historyItem of context.history) {
if (historyItem instanceof vscode.ChatRequestTurn) {
if (
Expand All @@ -181,9 +191,25 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
) {
// 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' &&
databaseName !== undefined &&
collectionName !== undefined
) {
previousItem = historyItem;
continue;
}
}

// eslint-disable-next-line new-cap
messages.push(vscode.LanguageModelChatMessage.User(historyItem.prompt));
}
Expand All @@ -206,11 +232,21 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
'emptyRequest',
'askToConnect',
];

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 (
responseTypesToSkip.indexOf(
(historyItem.result as ChatResult)?.metadata?.intent
) > -1
responseType === 'askForNamespace' &&
databaseName !== undefined &&
collectionName !== undefined
) {
previousItem = historyItem;
continue;
}

Expand All @@ -232,6 +268,7 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
// eslint-disable-next-line new-cap
messages.push(vscode.LanguageModelChatMessage.Assistant(message));
}
previousItem = historyItem;
}

return messages;
Expand Down
194 changes: 194 additions & 0 deletions src/test/suite/participant/participant.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2418,6 +2418,200 @@ Schema:
);
});

test('removes askForNameSpace messages from history if the metadata exists', async function () {
sinon.replace(
testParticipantController._chatMetadataStore,
'getChatMetadata',
() => ({
databaseName: 'dbOne',
collectionName: 'collOne',
})
);
// The user is responding to an `askToConnect` message, so the prompt is just the
// name of the connection
const chatRequestMock = {
prompt: 'localhost',
command: 'query',
};

const userMessages = [
'find all docs by a name example',
'what other queries can be used as an example',
];

chatContextStub = {
history: [
Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
prompt: userMessages[0],
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 query within this database?',
} as vscode.MarkdownString,
},
],
command: 'query',
result: {
metadata: {
intent: 'askForNamespace',
},
},
}),
Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
prompt: 'dbOne',
command: 'query',
references: [],
participant: CHAT_PARTICIPANT_ID,
}),
Object.assign(Object.create(vscode.ChatResponseTurn.prototype), {
participant: CHAT_PARTICIPANT_ID,
response: [
{
value: {
value:
'Which collection would you like to query within dbOne?',
} as vscode.MarkdownString,
},
],
command: 'query',
result: {
metadata: {
intent: 'askForNamespace',
databaseName: 'dbOne',
collectionName: undefined,
chatId: testChatId,
},
},
}),
Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
prompt: 'collectionOne',
command: 'query',
references: [],
participant: CHAT_PARTICIPANT_ID,
}),
Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
prompt: userMessages[1],
command: 'query',
references: [],
participant: CHAT_PARTICIPANT_ID,
}),
],
};

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('does not remove askForNameSpace messages if there is no metadata', async function () {
// The user is responding to an `askToConnect` message, so the prompt is just the
// name of the connection
const chatRequestMock = {
prompt: 'localhost',
command: 'query',
};

const userMessages = [
'find all docs by a name example',
'what other queries can be used as an example',
];

chatContextStub = {
history: [
Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
prompt: userMessages[0],
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 query within this database?',
} as vscode.MarkdownString,
},
],
command: 'query',
result: {
metadata: {
intent: 'askForNamespace',
},
},
}),
Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
prompt: 'dbOne',
command: 'query',
references: [],
participant: CHAT_PARTICIPANT_ID,
}),
Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
prompt: userMessages[1],
command: 'query',
references: [],
participant: CHAT_PARTICIPANT_ID,
}),
],
};

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(6);
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
Expand Down

0 comments on commit f468adf

Please sign in to comment.