Skip to content

Commit 860e08c

Browse files
committed
chore: better system prompt
1 parent f84748d commit 860e08c

File tree

1 file changed

+23
-2
lines changed

1 file changed

+23
-2
lines changed

snippets/ai/src/providers/ai-provider.ts

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -93,8 +93,24 @@ export class AiProvider {
9393
if (/[\n\r]/.test(trimmedText)) {
9494
// If the text includes a newline or carriage return, we should enter editor mode first
9595
process.stdin.unshift('.editor\n', 'utf-8');
96+
// Add backspace characters to remove indentation carried over from previous lines
97+
const lines = trimmedText.split(/\r?\n/);
98+
let result = lines[0]; // First line doesn't need any backspaces
99+
100+
for (let i = 1; i < lines.length; i++) {
101+
// Count leading whitespace from the previous line
102+
const previousLineLeadingWhitespace =
103+
lines[i - 1].match(/^\s*/)?.[0].length ?? 0;
104+
result += '\n';
105+
// Add backspace characters to remove carried-over indentation
106+
result += '\x08'.repeat(previousLineLeadingWhitespace);
107+
result += lines[i];
108+
}
109+
110+
process.stdin.unshift(result, 'utf-8');
111+
} else {
112+
process.stdin.unshift(trimmedText, 'utf-8');
96113
}
97-
process.stdin.unshift(trimmedText, 'utf-8');
98114
}
99115
/** @internal */
100116
respond(text: string) {
@@ -177,8 +193,10 @@ export class AiProvider {
177193
systemPrompt: string,
178194
{
179195
includeSampleDocs = false,
196+
expectedOutput,
180197
}: {
181198
includeSampleDocs?: boolean;
199+
expectedOutput?: GetResponseOptions['expectedOutput'];
182200
},
183201
): Promise<string> {
184202
return (
@@ -193,6 +211,9 @@ export class AiProvider {
193211
? `Sample documents from ${this.cliContext.db._name}.${this.session.collection}: ${JSON.stringify(
194212
await this.getSampleDocuments(this.session.collection),
195213
)}. Skip the use command to switch to the database, it is already set.`
214+
: '') +
215+
(expectedOutput === 'command'
216+
? `Expected output is meant to be just a mongo shell, no explanation or formatting.`
196217
: '')
197218
);
198219
}
@@ -208,7 +229,7 @@ export class AiProvider {
208229

209230
await this.processResponse(prompt, {
210231
systemPrompt: await this.getSystemPrompt(
211-
"You generate the exact mongosh aggregate command that matches the user's request",
232+
"You generate the exact mongosh aggregate command that matches the user's request. The message specifies the requirements to find and/or aggregate the documents by.",
212233
{ includeSampleDocs: true },
213234
),
214235
signal,

0 commit comments

Comments
 (0)