Skip to content

Commit f468adf

Browse files
gagikgagik
committed
Add tests
1 parent 548f247 commit f468adf

File tree

2 files changed

+234
-3
lines changed

2 files changed

+234
-3
lines changed

‎src/participant/prompts/promptBase.ts‎

Lines changed: 40 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@ export interface PromptArgsBase {
1313
};
1414
context?: vscode.ChatContext;
1515
connectionNames?: string[];
16+
databaseName?: string;
17+
collectionName?: string;
1618
}
1719

1820
export interface UserPromptResponse {
@@ -163,16 +165,24 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
163165
protected getHistoryMessages({
164166
connectionNames,
165167
context,
168+
databaseName,
169+
collectionName,
166170
}: {
167171
connectionNames?: string[]; // Used to scrape the connecting messages from the history.
168172
context?: vscode.ChatContext;
173+
databaseName?: string;
174+
collectionName?: string;
169175
}): vscode.LanguageModelChatMessage[] {
170176
const messages: vscode.LanguageModelChatMessage[] = [];
171177

172178
if (!context) {
173179
return [];
174180
}
175181

182+
let previousItem:
183+
| vscode.ChatRequestTurn
184+
| vscode.ChatResponseTurn
185+
| undefined = undefined;
176186
for (const historyItem of context.history) {
177187
if (historyItem instanceof vscode.ChatRequestTurn) {
178188
if (
@@ -181,9 +191,25 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
181191
) {
182192
// When the message is empty or a connection name then we skip it.
183193
// It's probably going to be the response to the connect step.
194+
previousItem = historyItem;
184195
continue;
185196
}
186197

198+
if (previousItem instanceof vscode.ChatResponseTurn) {
199+
const responseIntent = (previousItem.result as ChatResult).metadata
200+
.intent;
201+
202+
// If the namespace is already known, skip responses to prompts asking for it.
203+
if (
204+
responseIntent === 'askForNamespace' &&
205+
databaseName !== undefined &&
206+
collectionName !== undefined
207+
) {
208+
previousItem = historyItem;
209+
continue;
210+
}
211+
}
212+
187213
// eslint-disable-next-line new-cap
188214
messages.push(vscode.LanguageModelChatMessage.User(historyItem.prompt));
189215
}
@@ -206,11 +232,21 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
206232
'emptyRequest',
207233
'askToConnect',
208234
];
235+
236+
const responseType = (historyItem.result as ChatResult)?.metadata
237+
?.intent;
238+
if (responseTypesToSkip.includes(responseType)) {
239+
previousItem = historyItem;
240+
continue;
241+
}
242+
243+
// If the namespace is already known, skip including prompts asking for it.
209244
if (
210-
responseTypesToSkip.indexOf(
211-
(historyItem.result as ChatResult)?.metadata?.intent
212-
) > -1
245+
responseType === 'askForNamespace' &&
246+
databaseName !== undefined &&
247+
collectionName !== undefined
213248
) {
249+
previousItem = historyItem;
214250
continue;
215251
}
216252

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

237274
return messages;

‎src/test/suite/participant/participant.test.ts‎

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2418,6 +2418,200 @@ Schema:
24182418
);
24192419
});
24202420

2421+
test('removes askForNameSpace messages from history if the metadata exists', async function () {
2422+
sinon.replace(
2423+
testParticipantController._chatMetadataStore,
2424+
'getChatMetadata',
2425+
() => ({
2426+
databaseName: 'dbOne',
2427+
collectionName: 'collOne',
2428+
})
2429+
);
2430+
// The user is responding to an `askToConnect` message, so the prompt is just the
2431+
// name of the connection
2432+
const chatRequestMock = {
2433+
prompt: 'localhost',
2434+
command: 'query',
2435+
};
2436+
2437+
const userMessages = [
2438+
'find all docs by a name example',
2439+
'what other queries can be used as an example',
2440+
];
2441+
2442+
chatContextStub = {
2443+
history: [
2444+
Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
2445+
prompt: userMessages[0],
2446+
command: 'query',
2447+
references: [],
2448+
participant: CHAT_PARTICIPANT_ID,
2449+
}),
2450+
Object.assign(Object.create(vscode.ChatResponseTurn.prototype), {
2451+
participant: CHAT_PARTICIPANT_ID,
2452+
response: [
2453+
{
2454+
value: {
2455+
value:
2456+
'Which database would you like to query within this database?',
2457+
} as vscode.MarkdownString,
2458+
},
2459+
],
2460+
command: 'query',
2461+
result: {
2462+
metadata: {
2463+
intent: 'askForNamespace',
2464+
},
2465+
},
2466+
}),
2467+
Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
2468+
prompt: 'dbOne',
2469+
command: 'query',
2470+
references: [],
2471+
participant: CHAT_PARTICIPANT_ID,
2472+
}),
2473+
Object.assign(Object.create(vscode.ChatResponseTurn.prototype), {
2474+
participant: CHAT_PARTICIPANT_ID,
2475+
response: [
2476+
{
2477+
value: {
2478+
value:
2479+
'Which collection would you like to query within dbOne?',
2480+
} as vscode.MarkdownString,
2481+
},
2482+
],
2483+
command: 'query',
2484+
result: {
2485+
metadata: {
2486+
intent: 'askForNamespace',
2487+
databaseName: 'dbOne',
2488+
collectionName: undefined,
2489+
chatId: testChatId,
2490+
},
2491+
},
2492+
}),
2493+
Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
2494+
prompt: 'collectionOne',
2495+
command: 'query',
2496+
references: [],
2497+
participant: CHAT_PARTICIPANT_ID,
2498+
}),
2499+
Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
2500+
prompt: userMessages[1],
2501+
command: 'query',
2502+
references: [],
2503+
participant: CHAT_PARTICIPANT_ID,
2504+
}),
2505+
],
2506+
};
2507+
2508+
const { messages, stats } = await Prompts.query.buildMessages({
2509+
context: chatContextStub,
2510+
request: chatRequestMock,
2511+
collectionName: 'people',
2512+
connectionNames: ['localhost', 'atlas'],
2513+
databaseName: 'prod',
2514+
sampleDocuments: [],
2515+
});
2516+
2517+
expect(messages.length).to.equal(4);
2518+
expect(messages[0].role).to.equal(
2519+
vscode.LanguageModelChatMessageRole.Assistant
2520+
);
2521+
2522+
// We don't expect history because we're removing the askForConnect message as well
2523+
// as the user response to it. Therefore the actual user prompt should be the first
2524+
// message that we supplied in the history.
2525+
expect(messages[1].role).to.equal(
2526+
vscode.LanguageModelChatMessageRole.User
2527+
);
2528+
2529+
expect(
2530+
messages.slice(1, 3).map((message) => getMessageContent(message))
2531+
).to.deep.equal(userMessages);
2532+
2533+
expect(stats.command).to.equal('query');
2534+
});
2535+
2536+
test('does not remove askForNameSpace messages if there is no metadata', async function () {
2537+
// The user is responding to an `askToConnect` message, so the prompt is just the
2538+
// name of the connection
2539+
const chatRequestMock = {
2540+
prompt: 'localhost',
2541+
command: 'query',
2542+
};
2543+
2544+
const userMessages = [
2545+
'find all docs by a name example',
2546+
'what other queries can be used as an example',
2547+
];
2548+
2549+
chatContextStub = {
2550+
history: [
2551+
Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
2552+
prompt: userMessages[0],
2553+
command: 'query',
2554+
references: [],
2555+
participant: CHAT_PARTICIPANT_ID,
2556+
}),
2557+
Object.assign(Object.create(vscode.ChatResponseTurn.prototype), {
2558+
participant: CHAT_PARTICIPANT_ID,
2559+
response: [
2560+
{
2561+
value: {
2562+
value:
2563+
'Which database would you like to query within this database?',
2564+
} as vscode.MarkdownString,
2565+
},
2566+
],
2567+
command: 'query',
2568+
result: {
2569+
metadata: {
2570+
intent: 'askForNamespace',
2571+
},
2572+
},
2573+
}),
2574+
Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
2575+
prompt: 'dbOne',
2576+
command: 'query',
2577+
references: [],
2578+
participant: CHAT_PARTICIPANT_ID,
2579+
}),
2580+
Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
2581+
prompt: userMessages[1],
2582+
command: 'query',
2583+
references: [],
2584+
participant: CHAT_PARTICIPANT_ID,
2585+
}),
2586+
],
2587+
};
2588+
2589+
const { messages, stats } = await Prompts.query.buildMessages({
2590+
context: chatContextStub,
2591+
request: chatRequestMock,
2592+
connectionNames: ['localhost', 'atlas'],
2593+
sampleDocuments: [],
2594+
// @ts-expect-error Forcing undefined for the purpose of test
2595+
databaseName: undefined,
2596+
// @ts-expect-error Forcing undefined for the purpose of test
2597+
collectionName: undefined,
2598+
});
2599+
2600+
expect(messages.length).to.equal(6);
2601+
expect(messages[0].role).to.equal(
2602+
vscode.LanguageModelChatMessageRole.Assistant
2603+
);
2604+
2605+
// We don't expect history because we're removing the askForConnect message as well
2606+
// as the user response to it. Therefore the actual user prompt should be the first
2607+
// message that we supplied in the history.
2608+
expect(messages[1].role).to.equal(
2609+
vscode.LanguageModelChatMessageRole.User
2610+
);
2611+
2612+
expect(stats.command).to.equal('query');
2613+
});
2614+
24212615
test('removes askForConnect messages from history', async function () {
24222616
// The user is responding to an `askToConnect` message, so the prompt is just the
24232617
// name of the connection

0 commit comments

Comments
 (0)