diff --git a/src/participant/constants.ts b/src/participant/constants.ts new file mode 100644 index 000000000..e4aac61e0 --- /dev/null +++ b/src/participant/constants.ts @@ -0,0 +1 @@ +export const CHAT_PARTICIPANT_ID = 'mongodb.participant'; diff --git a/src/participant/participant.ts b/src/participant/participant.ts index 3b4b527da..44992c24b 100644 --- a/src/participant/participant.ts +++ b/src/participant/participant.ts @@ -3,6 +3,9 @@ import * as vscode from 'vscode'; import { createLogger } from '../logging'; import type ConnectionController from '../connectionController'; import EXTENSION_COMMANDS from '../commands'; +import { GenericPrompt } from './prompts/generic'; +import { CHAT_PARTICIPANT_ID } from './constants'; +import { QueryPrompt } from './prompts/query'; const log = createLogger('participant'); @@ -17,7 +20,6 @@ interface ChatResult extends vscode.ChatResult { stream?: vscode.ChatResponseStream; } -export const CHAT_PARTICIPANT_ID = 'mongodb.participant'; export const CHAT_PARTICIPANT_MODEL = 'gpt-4o'; export function getRunnableContentFromString(responseContent: string) { @@ -148,41 +150,11 @@ export class ParticipantController { stream: vscode.ChatResponseStream; token: vscode.CancellationToken; }) { - const messages = [ - // eslint-disable-next-line new-cap - vscode.LanguageModelChatMessage.Assistant(`You are a MongoDB expert! - You create MongoDB queries and aggregation pipelines, - and you are very good at it. The user will provide the basis for the query. - Keep your response concise. Respond with markdown, code snippets are possible with '''javascript. - You can imagine the schema, collection, and database name. - Respond in MongoDB shell syntax using the '''javascript code style.`), - ]; - - context.history.map((historyItem) => { - if ( - historyItem.participant === CHAT_PARTICIPANT_ID && - historyItem instanceof vscode.ChatRequestTurn - ) { - // eslint-disable-next-line new-cap - messages.push(vscode.LanguageModelChatMessage.User(historyItem.prompt)); - } - - if ( - historyItem.participant === CHAT_PARTICIPANT_ID && - historyItem instanceof vscode.ChatResponseTurn - ) { - let res = ''; - for (const fragment of historyItem.response) { - res += fragment; - } - // eslint-disable-next-line new-cap - messages.push(vscode.LanguageModelChatMessage.Assistant(res)); - } + const messages = GenericPrompt.buildMessages({ + request, + context, }); - // eslint-disable-next-line new-cap - messages.push(vscode.LanguageModelChatMessage.User(request.prompt)); - const abortController = new AbortController(); token.onCancellationRequested(() => { abortController.abort(); @@ -222,11 +194,12 @@ export class ParticipantController { // @MongoDB /query find all documents where the "address" has the word Broadway in it. async handleQueryRequest({ request, + context, stream, token, }: { request: vscode.ChatRequest; - context?: vscode.ChatContext; + context: vscode.ChatContext; stream: vscode.ChatResponseStream; token: vscode.CancellationToken; }) { @@ -265,17 +238,11 @@ export class ParticipantController { abortController.abort(); }); - const messages = [ - // eslint-disable-next-line new-cap - vscode.LanguageModelChatMessage.Assistant(`You are a MongoDB expert! - You create MongoDB queries and aggregation pipelines, - and you are very good at it. The user will provide the basis for the query. - Keep your response concise. Respond with markdown, code snippets are possible with '''javascript. - You can imagine the schema, collection, and database name. - Respond in MongoDB shell syntax using the '''javascript code style.`), - // eslint-disable-next-line new-cap - vscode.LanguageModelChatMessage.User(request.prompt), - ]; + const messages = QueryPrompt.buildMessages({ + context, + request, + }); + const responseContent = await this.getChatResponseContent({ messages, stream, @@ -320,9 +287,9 @@ export class ParticipantController { }); return this._chatResult; } else if (request.command === 'docs') { - // TODO: Implement this. + // TODO(VSCODE-570): Implement this. } else if (request.command === 'schema') { - // TODO: Implement this. + // TODO(VSCODE-571): Implement this. } return await this.handleGenericRequest({ diff --git a/src/participant/prompts/generic.ts b/src/participant/prompts/generic.ts new file mode 100644 index 000000000..bde4604f0 --- /dev/null +++ b/src/participant/prompts/generic.ts @@ -0,0 +1,41 @@ +import * as vscode from 'vscode'; + +import { getHistoryMessages } from './history'; + +export class GenericPrompt { + static getSystemPrompt(): vscode.LanguageModelChatMessage { + const prompt = `You are a MongoDB expert. +Your task is to help the user craft MongoDB queries and aggregation pipelines that perform their task. +Keep your response concise. +You should suggest queries that are performant and correct. +Respond with markdown, suggest code in a Markdown code block that begins with \'\'\'javascript and ends with \`\`\`. +You can imagine the schema, collection, and database name. +Respond in MongoDB shell syntax using the \'\'\'javascript code block syntax.`; + + // eslint-disable-next-line new-cap + return vscode.LanguageModelChatMessage.Assistant(prompt); + } + + static getUserPrompt( + request: vscode.ChatRequest + ): vscode.LanguageModelChatMessage { + // eslint-disable-next-line new-cap + return vscode.LanguageModelChatMessage.User(request.prompt); + } + + static buildMessages({ + context, + request, + }: { + request: vscode.ChatRequest; + context: vscode.ChatContext; + }): vscode.LanguageModelChatMessage[] { + const messages = [ + GenericPrompt.getSystemPrompt(), + ...getHistoryMessages({ context }), + GenericPrompt.getUserPrompt(request), + ]; + + return messages; + } +} diff --git a/src/participant/prompts/history.ts b/src/participant/prompts/history.ts new file mode 100644 index 000000000..907838b7e --- /dev/null +++ b/src/participant/prompts/history.ts @@ -0,0 +1,35 @@ +import * as vscode from 'vscode'; + +import { CHAT_PARTICIPANT_ID } from '../constants'; + +export function getHistoryMessages({ + context, +}: { + context: vscode.ChatContext; +}): vscode.LanguageModelChatMessage[] { + const messages: vscode.LanguageModelChatMessage[] = []; + + context.history.map((historyItem) => { + if ( + historyItem.participant === CHAT_PARTICIPANT_ID && + historyItem instanceof vscode.ChatRequestTurn + ) { + // eslint-disable-next-line new-cap + messages.push(vscode.LanguageModelChatMessage.User(historyItem.prompt)); + } + + if ( + historyItem.participant === CHAT_PARTICIPANT_ID && + historyItem instanceof vscode.ChatResponseTurn + ) { + let res = ''; + for (const fragment of historyItem.response) { + res += fragment; + } + // eslint-disable-next-line new-cap + messages.push(vscode.LanguageModelChatMessage.Assistant(res)); + } + }); + + return messages; +} diff --git a/src/participant/prompts/query.ts b/src/participant/prompts/query.ts new file mode 100644 index 000000000..7b27ce2de --- /dev/null +++ b/src/participant/prompts/query.ts @@ -0,0 +1,41 @@ +import * as vscode from 'vscode'; + +import { getHistoryMessages } from './history'; + +export class QueryPrompt { + static getSystemPrompt(): vscode.LanguageModelChatMessage { + const prompt = `You are a MongoDB expert. +Your task is to help the user craft MongoDB queries and aggregation pipelines that perform their task. +Keep your response concise. +You should suggest queries that are performant and correct. +Respond with markdown, suggest code in a Markdown code block that begins with \'\'\'javascript and ends with \`\`\`. +You can imagine the schema, collection, and database name. +Respond in MongoDB shell syntax using the \'\'\'javascript code block syntax.`; + + // eslint-disable-next-line new-cap + return vscode.LanguageModelChatMessage.Assistant(prompt); + } + + static getUserPrompt( + request: vscode.ChatRequest + ): vscode.LanguageModelChatMessage { + // eslint-disable-next-line new-cap + return vscode.LanguageModelChatMessage.User(request.prompt); + } + + static buildMessages({ + context, + request, + }: { + request: vscode.ChatRequest; + context: vscode.ChatContext; + }): vscode.LanguageModelChatMessage[] { + const messages = [ + QueryPrompt.getSystemPrompt(), + ...getHistoryMessages({ context }), + QueryPrompt.getUserPrompt(request), + ]; + + return messages; + } +}