From 592ba2f12903a67249a81c275331e59eb4d93e7b Mon Sep 17 00:00:00 2001 From: Alena Khineika Date: Mon, 8 Jul 2024 14:23:23 +0200 Subject: [PATCH 1/4] feat: convert copilot POC to feature utilising latest API VSCODE-550 --- package.json | 35 +- scripts/generate-constants.ts | 59 +++ scripts/generate-keyfile.ts | 28 -- src/commands/index.ts | 4 + src/editors/playgroundController.ts | 21 + src/mdbExtensionController.ts | 53 +++ src/participant/participant.ts | 386 +++++++++++++++++++ src/templates/playgroundBasicTextTemplate.ts | 16 + 8 files changed, 573 insertions(+), 29 deletions(-) create mode 100644 scripts/generate-constants.ts delete mode 100644 scripts/generate-keyfile.ts create mode 100644 src/participant/participant.ts create mode 100644 src/templates/playgroundBasicTextTemplate.ts diff --git a/package.json b/package.json index dd40df552..bc3f6ac36 100644 --- a/package.json +++ b/package.json @@ -32,6 +32,9 @@ "color": "#3D4F58", "theme": "dark" }, + "enabledApiProposals": [ + "chatVariableResolver" + ], "license": "SEE LICENSE IN LICENSE.txt", "main": "./dist/extension.js", "scripts": { @@ -40,7 +43,7 @@ "update-grammar": "ts-node ./scripts/update-grammar.ts", "precompile": "npm run clean", "compile": "npm-run-all compile:*", - "compile:keyfile": "ts-node ./scripts/generate-keyfile.ts", + "compile:keyfile": "ts-node ./scripts/generate-constants.ts", "compile:resources": "npm run update-grammar", "compile:extension": "tsc -p ./", "compile:extension-bundles": "webpack --mode development", @@ -71,11 +74,25 @@ }, "activationEvents": [ "onView:mongoDB", + "onStartupFinished", "onLanguage:json", "onLanguage:javascript", "onLanguage:plaintext" ], "contributes": { + "chatParticipants": [ + { + "id": "mongodb.participant", + "name": "MongoDB", + "description": "Ask anything about MongoDB, from writing queries to questions about your cluster.", + "commands": [ + { + "name": "query", + "description": "Ask how to write MongoDB queries or pipelines. For example, you can ask: \"Show me all the documents where the address contains the word street\"." + } + ] + } + ], "viewsContainers": { "activitybar": [ { @@ -142,6 +159,14 @@ } ], "commands": [ + { + "command": "mdb.runParticipantQuery", + "title": "Run Content Generated by the Chat Participant" + }, + { + "command": "mdb.openParticipantQueryInPlayground", + "title": "Open Generated by the Chat Participant Content In Playground" + }, { "command": "mdb.connect", "title": "MongoDB: Connect" @@ -688,6 +713,14 @@ } ], "commandPalette": [ + { + "command": "mdb.openParticipantQueryInPlayground", + "when": "false" + }, + { + "command": "mdb.runParticipantQuery", + "when": "false" + }, { "command": "mdb.disconnect", "when": "mdb.connectedToMongoDB == true" diff --git a/scripts/generate-constants.ts b/scripts/generate-constants.ts new file mode 100644 index 000000000..b48fb4175 --- /dev/null +++ b/scripts/generate-constants.ts @@ -0,0 +1,59 @@ +#! /usr/bin/env ts-node + +import ora from 'ora'; +import fs from 'fs'; +import path from 'path'; +import { resolve } from 'path'; +import { config } from 'dotenv'; +import { promisify } from 'util'; + +const writeFile = promisify(fs.writeFile); +const ROOT_DIR = path.join(__dirname, '..'); +const ui = ora('Generate constants keyfile').start(); + +config({ path: resolve(__dirname, '../.env') }); + +interface Constants { + segmentKey?: string; + useMongodbChatParticipant?: string; + chatParticipantGenericPrompt?: string; + chatParticipantQueryPrompt?: string; + chatParticipantModel?: string; +} + +const constants: Constants = {}; + +(async () => { + if (process.env.SEGMENT_KEY) { + constants.segmentKey = process.env.SEGMENT_KEY; + } + if (process.env.USE_MONGODB_CHAT_PARTICIPANT) { + constants.useMongodbChatParticipant = + process.env.USE_MONGODB_CHAT_PARTICIPANT; + } + if (process.env.CHAT_PARTICIPANT_GENERIC_PROMPT) { + constants.chatParticipantGenericPrompt = + process.env.CHAT_PARTICIPANT_GENERIC_PROMPT; + } + if (process.env.CHAT_PARTICIPANT_QUERY_PROMPT) { + constants.chatParticipantQueryPrompt = + process.env.CHAT_PARTICIPANT_QUERY_PROMPT; + } + if (process.env.CHAT_PARTICIPANT_MODEL) { + constants.chatParticipantModel = process.env.CHAT_PARTICIPANT_MODEL; + } + if (Object.keys(constants).length === 0) { + ui.warn('No constants to write'); + return; + } + + await writeFile( + `${ROOT_DIR}/constants.json`, + JSON.stringify(constants, null, 2) + ); + ui.succeed('The constants file was written'); +})().catch((error) => { + ui.fail( + `An error occurred while writing the constants file: ${error.message}` + ); +}); diff --git a/scripts/generate-keyfile.ts b/scripts/generate-keyfile.ts deleted file mode 100644 index b7a6add73..000000000 --- a/scripts/generate-keyfile.ts +++ /dev/null @@ -1,28 +0,0 @@ -#! /usr/bin/env ts-node - -import ora from 'ora'; -import fs from 'fs'; -import path from 'path'; -import { resolve } from 'path'; -import { config } from 'dotenv'; -import { promisify } from 'util'; - -const writeFile = promisify(fs.writeFile); -const ROOT_DIR = path.join(__dirname, '..'); -const ui = ora('Generate constants keyfile').start(); - -config({ path: resolve(__dirname, '../.env') }); - -(async () => { - if (process.env.SEGMENT_KEY) { - await writeFile( - `${ROOT_DIR}/constants.json`, - JSON.stringify({ segmentKey: process.env.SEGMENT_KEY }, null, 2) - ); - ui.succeed('Generated segment constants file'); - } else { - throw new Error('The Segment key is missing in environment variables'); - } -})().catch((error) => { - ui.fail(`Failed to generate segment constants file: ${error.message}`); -}); diff --git a/src/commands/index.ts b/src/commands/index.ts index ebef4dec1..61513dc2c 100644 --- a/src/commands/index.ts +++ b/src/commands/index.ts @@ -73,6 +73,10 @@ enum EXTENSION_COMMANDS { MDB_START_STREAM_PROCESSOR = 'mdb.startStreamProcessor', MDB_STOP_STREAM_PROCESSOR = 'mdb.stopStreamProcessor', MDB_DROP_STREAM_PROCESSOR = 'mdb.dropStreamProcessor', + + // Chat participant. + OPEN_PARTICIPANT_QUERY_IN_PLAYGROUND = 'mdb.openParticipantQueryInPlayground', + RUN_PARTICIPANT_QUERY = 'mdb.runParticipantQuery', } export default EXTENSION_COMMANDS; diff --git a/src/editors/playgroundController.ts b/src/editors/playgroundController.ts index 64b2594a3..79c9bb62a 100644 --- a/src/editors/playgroundController.ts +++ b/src/editors/playgroundController.ts @@ -16,6 +16,7 @@ import { DatabaseTreeItem } from '../explorer'; import type ExportToLanguageCodeLensProvider from './exportToLanguageCodeLensProvider'; import formatError from '../utils/formatError'; import type { LanguageServerController } from '../language'; +import playgroundBasicTextTemplate from '../templates/playgroundBasicTextTemplate'; import playgroundCreateIndexTemplate from '../templates/playgroundCreateIndexTemplate'; import playgroundCreateCollectionTemplate from '../templates/playgroundCreateCollectionTemplate'; import playgroundCloneDocumentTemplate from '../templates/playgroundCloneDocumentTemplate'; @@ -382,6 +383,21 @@ export default class PlaygroundController { return this._createPlaygroundFileWithContent(content); } + createPlaygroundFromParticipantQuery({ + text, + }: { + text: string; + }): Promise { + const useDefaultTemplate = !!vscode.workspace + .getConfiguration('mdb') + .get('useDefaultTemplateForPlayground'); + const content = useDefaultTemplate + ? playgroundBasicTextTemplate.replace('PLAYGROUND_CONTENT', text) + : text; + this._telemetryService.trackPlaygroundCreated('agent'); + return this._createPlaygroundFileWithContent(content); + } + createPlaygroundForCloneDocument( documentContents: string, databaseName: string, @@ -802,6 +818,11 @@ export default class PlaygroundController { return { namespace, expression }; } + async evaluateParticipantQuery({ text }: { text: string }): Promise { + this._codeToEvaluate = text; + return this._evaluatePlayground(); + } + async _transpile(): Promise { const { selectedText, importStatements, driverSyntax, builders, language } = this._exportToLanguageCodeLensProvider._exportToLanguageAddons; diff --git a/src/mdbExtensionController.ts b/src/mdbExtensionController.ts index 791d85287..bdf07b2c9 100644 --- a/src/mdbExtensionController.ts +++ b/src/mdbExtensionController.ts @@ -40,6 +40,7 @@ import WebviewController from './views/webviewController'; import { createIdFactory, generateId } from './utils/objectIdHelper'; import { ConnectionStorage } from './storage/connectionStorage'; import type StreamProcessorTreeItem from './explorer/streamProcessorTreeItem'; +import { ParticipantController } from './participant/participant'; // This class is the top-level controller for our extension. // Commands which the extensions handles are defined in the function `activate`. @@ -63,6 +64,7 @@ export default class MDBExtensionController implements vscode.Disposable { _activeConnectionCodeLensProvider: ActiveConnectionCodeLensProvider; _editDocumentCodeLensProvider: EditDocumentCodeLensProvider; _exportToLanguageCodeLensProvider: ExportToLanguageCodeLensProvider; + _participantController?: ParticipantController; constructor( context: vscode.ExtensionContext, @@ -137,6 +139,11 @@ export default class MDBExtensionController implements vscode.Disposable { telemetryService: this._telemetryService, }); this._editorsController.registerProviders(); + this._participantController = new ParticipantController({ + context, + connectionController: this._connectionController, + playgroundController: this._playgroundController, + }); } async activate(): Promise { @@ -265,6 +272,52 @@ export default class MDBExtensionController implements vscode.Disposable { this.registerEditorCommands(); this.registerTreeViewCommands(); + + // ------ CHAT PARTICIPANT ------ // + this.registerParticipantCommand( + EXTENSION_COMMANDS.OPEN_PARTICIPANT_QUERY_IN_PLAYGROUND, + () => { + if (!this._participantController) { + return Promise.resolve(false); + } + return this._playgroundController.createPlaygroundFromParticipantQuery({ + text: + this._participantController.chatResult.metadata.queryContent || '', + }); + } + ); + this.registerParticipantCommand( + EXTENSION_COMMANDS.RUN_PARTICIPANT_QUERY, + () => { + if (!this._participantController) { + return Promise.resolve(false); + } + return this._playgroundController.evaluateParticipantQuery({ + text: + this._participantController.chatResult.metadata.queryContent || '', + }); + } + ); + }; + + registerParticipantCommand = ( + command: string, + commandHandler: (...args: any[]) => Promise + ): void => { + if (!this._participantController) { + return; + } + + const commandHandlerWithTelemetry = (args: any[]): Promise => { + this._telemetryService.trackCommandRun(command); + + return commandHandler(args); + }; + + this._context.subscriptions.push( + this._participantController.participant, + vscode.commands.registerCommand(command, commandHandlerWithTelemetry) + ); }; registerCommand = ( diff --git a/src/participant/participant.ts b/src/participant/participant.ts new file mode 100644 index 000000000..4ed30988b --- /dev/null +++ b/src/participant/participant.ts @@ -0,0 +1,386 @@ +import * as vscode from 'vscode'; +import { config } from 'dotenv'; +import fs from 'fs'; +import path from 'path'; + +import { createLogger } from '../logging'; +import type { PlaygroundController } from '../editors'; +import type ConnectionController from '../connectionController'; +import EXTENSION_COMMANDS from '../commands'; + +const log = createLogger('participant'); + +interface ChatResult extends vscode.ChatResult { + metadata: { + command: string; + databaseName?: string; + collectionName?: string; + queryContent?: string; + description?: string; + }; + stream?: vscode.ChatResponseStream; +} + +interface GenAIConstants { + useMongodbChatParticipant?: string; + chatParticipantGenericPrompt?: string; + chatParticipantQueryPrompt?: string; + chatParticipantModel?: string; +} + +const PARTICIPANT_ID = 'mongodb.participant'; + +function handleEmptyQueryRequest(participantId: string): ChatResult { + log.info('Chat request participant id', participantId); + + return { + metadata: { + command: '', + }, + errorDetails: { + message: + 'Please specify a question when using this command. Usage: @MongoDB /query find documents where "name" contains "database".', + }, + }; +} + +function getRunnableContentFromString(responseContent: string) { + const matchedJSQueryContent = responseContent.match( + /```javascript((.|\n)*)```/ + ); + log.info('matchedJSQueryContent', matchedJSQueryContent); + + const queryContent = + matchedJSQueryContent && matchedJSQueryContent.length > 1 + ? matchedJSQueryContent[1] + : ''; + log.info('queryContent', queryContent); + return queryContent; +} + +export class ParticipantController { + participant: vscode.ChatParticipant; + chatResult: ChatResult; + _connectionController: ConnectionController; + _playgroundController: PlaygroundController; + + private _context: vscode.ExtensionContext; + private _useMongodbChatParticipant = false; + private _chatParticipantGenericPrompt = 'You are a MongoDB expert!'; + private _chatParticipantQueryPrompt = 'You are a MongoDB expert!'; + private _chatParticipantModel = 'gpt-3.5-turbo'; + + constructor({ + context, + connectionController, + playgroundController, + }: { + context: vscode.ExtensionContext; + connectionController: ConnectionController; + playgroundController: PlaygroundController; + }) { + this.participant = this.createParticipant(context); + this.chatResult = { metadata: { command: '' } }; + this._connectionController = connectionController; + this._playgroundController = playgroundController; + this._context = context; + + this._readConstants(); + } + + private _readConstants(): string | undefined { + config({ path: path.join(this._context.extensionPath, '.env') }); + + try { + const constantsLocation = path.join( + this._context.extensionPath, + './constants.json' + ); + // eslint-disable-next-line no-sync + const constantsFile = fs.readFileSync(constantsLocation, 'utf8'); + const constants = JSON.parse(constantsFile) as GenAIConstants; + + this._useMongodbChatParticipant = + constants.useMongodbChatParticipant === 'true'; + this._chatParticipantGenericPrompt = + constants.chatParticipantGenericPrompt || + this._chatParticipantGenericPrompt; + this._chatParticipantQueryPrompt = + constants.chatParticipantQueryPrompt || + this._chatParticipantQueryPrompt; + this._chatParticipantModel = + constants.chatParticipantModel || this._chatParticipantModel; + } catch (error) { + log.error('An error occurred while reading the constants file', error); + return; + } + } + + createParticipant(context: vscode.ExtensionContext) { + // Chat participants appear as top-level options in the chat input + // when you type `@`, and can contribute sub-commands in the chat input + // that appear when you type `/`. + const cat = vscode.chat.createChatParticipant( + PARTICIPANT_ID, + this.chatHandler.bind(this) + ); + cat.iconPath = vscode.Uri.joinPath( + context.extensionUri, + 'images', + 'mongodb.png' + ); + return cat; + } + + handleError(err: any, stream: vscode.ChatResponseStream): void { + // Making the chat request might fail because + // - model does not exist + // - user consent not given + // - quote limits exceeded + if (err instanceof vscode.LanguageModelError) { + log.error(err.message, err.code, err.cause); + if ( + err.cause instanceof Error && + err.cause.message.includes('off_topic') + ) { + stream.markdown( + vscode.l10n.t( + "I'm sorry, I can only explain computer science concepts.\n\n" + ) + ); + } + } + } + + async getChatResponseContent({ + messages, + stream, + token, + }: { + messages: vscode.LanguageModelChatMessage[]; + stream: vscode.ChatResponseStream; + token: vscode.CancellationToken; + }): Promise { + let responseContent = ''; + try { + const [model] = await vscode.lm.selectChatModels({ + vendor: 'copilot', + family: this._chatParticipantModel, + }); + if (model) { + const chatResponse = await model.sendRequest(messages, {}, token); + for await (const fragment of chatResponse.text) { + responseContent += fragment; + stream.markdown(fragment); + } + stream.markdown('\n\n'); + } + } catch (err) { + this.handleError(err, stream); + } + + return responseContent; + } + + // @MongoDB what is mongodb? + async handleGenericRequest({ + request, + context, + stream, + token, + }: { + request: vscode.ChatRequest; + context: vscode.ChatContext; + stream: vscode.ChatResponseStream; + token: vscode.CancellationToken; + }) { + const messages = [ + // eslint-disable-next-line new-cap + vscode.LanguageModelChatMessage.Assistant( + this._chatParticipantGenericPrompt + ), + ]; + + context.history.map((historyItem) => { + if ( + historyItem.participant === PARTICIPANT_ID && + historyItem instanceof vscode.ChatRequestTurn + ) { + // eslint-disable-next-line new-cap + messages.push(vscode.LanguageModelChatMessage.User(historyItem.prompt)); + } + + if ( + historyItem.participant === 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)); + } + }); + + // eslint-disable-next-line new-cap + messages.push(vscode.LanguageModelChatMessage.User(request.prompt)); + + const abortController = new AbortController(); + token.onCancellationRequested(() => { + abortController.abort(); + }); + + const responseContent = await this.getChatResponseContent({ + messages, + stream, + token, + }); + const queryContent = getRunnableContentFromString(responseContent); + + if (!queryContent || queryContent.trim().length === 0) { + return { metadata: { command: '' } }; + } + + stream.button({ + command: EXTENSION_COMMANDS.RUN_PARTICIPANT_QUERY, + title: vscode.l10n.t('▶️ Run'), + }); + + stream.button({ + command: EXTENSION_COMMANDS.OPEN_PARTICIPANT_QUERY_IN_PLAYGROUND, + title: vscode.l10n.t('Open in playground'), + }); + + return { + metadata: { + command: '', + stream, + queryContent, + }, + }; + } + + // @MongoDB /query find all documents where the "address" has the word Broadway in it. + async handleQueryRequest({ + request, + stream, + token, + }: { + request: vscode.ChatRequest; + context?: vscode.ChatContext; + stream: vscode.ChatResponseStream; + token: vscode.CancellationToken; + }) { + if (!request.prompt || request.prompt.trim().length === 0) { + return handleEmptyQueryRequest(this.participant.id); + } + + let dataService = this._connectionController.getActiveDataService(); + if (!dataService) { + stream.markdown( + "Looks like you aren't currently connected, first let's get you connected to the cluster we'd like to create this query to run against.\n\n" + ); + // We add a delay so the user can read the message. + // TODO: maybe there is better way to handle this. + // stream.button() does not awaits so we can't use it here. + // Followups do not support input so we can't use that either. + await new Promise((resolve) => setTimeout(resolve, 1000)); + const successfullyConnected = + await this._connectionController.changeActiveConnection(); + dataService = this._connectionController.getActiveDataService(); + + if (!dataService || !successfullyConnected) { + stream.markdown( + 'No connection for command provided. Please use a valid connection for running commands.\n\n' + ); + return { metadata: { command: '' } }; + } + + stream.markdown( + `Connected to "${this._connectionController.getActiveConnectionName()}".\n\n` + ); + } + + const abortController = new AbortController(); + token.onCancellationRequested(() => { + abortController.abort(); + }); + + const messages = [ + // eslint-disable-next-line new-cap + vscode.LanguageModelChatMessage.Assistant( + this._chatParticipantQueryPrompt + ), + // eslint-disable-next-line new-cap + vscode.LanguageModelChatMessage.User(request.prompt), + ]; + + const responseContent = await this.getChatResponseContent({ + messages, + stream, + token, + }); + const queryContent = getRunnableContentFromString(responseContent); + + if (!queryContent || queryContent.trim().length === 0) { + return { metadata: { command: '' } }; + } + + stream.button({ + command: EXTENSION_COMMANDS.RUN_PARTICIPANT_QUERY, + title: vscode.l10n.t('▶️ Run'), + }); + stream.button({ + command: EXTENSION_COMMANDS.OPEN_PARTICIPANT_QUERY_IN_PLAYGROUND, + title: vscode.l10n.t('Open in playground'), + }); + + return { + metadata: { + command: '', + stream, + queryContent, + }, + }; + } + + async chatHandler( + request: vscode.ChatRequest, + context: vscode.ChatContext, + stream: vscode.ChatResponseStream, + token: vscode.CancellationToken + ): Promise { + if (!this._useMongodbChatParticipant) { + stream.markdown( + vscode.l10n.t( + 'Under construction. Will be available soon. Stay tuned!\n\n' + ) + ); + return { metadata: { command: '' } }; + } + + if (request.command === 'query') { + this.chatResult = await this.handleQueryRequest({ + request, + context, + stream, + token, + }); + return this.chatResult; + } else if (request.command === 'docs') { + // TODO: Implement this. + } else if (request.command === 'schema') { + // TODO: Implement this. + } else if (request.command === 'logs') { + // TODO: Implement this. + } + + return await this.handleGenericRequest({ + request, + context, + stream, + token, + }); + } +} diff --git a/src/templates/playgroundBasicTextTemplate.ts b/src/templates/playgroundBasicTextTemplate.ts new file mode 100644 index 000000000..2fe3e8186 --- /dev/null +++ b/src/templates/playgroundBasicTextTemplate.ts @@ -0,0 +1,16 @@ +const template = `/* global use, db */ +// MongoDB Playground +// To disable this template go to Settings | MongoDB | Use Default Template For Playground. +// Make sure you are connected to enable completions and to be able to run a playground. +// Use Ctrl+Space inside a snippet or a string literal to trigger completions. +// The result of the last command run in a playground is shown on the results panel. +// By default the first 20 documents will be returned with a cursor. +// Use 'console.log()' to print to the debug output. +// For more documentation on playgrounds please refer to +// https://www.mongodb.com/docs/mongodb-vscode/playgrounds/ +// Select the database to use. +use('CURRENT_DATABASE'); +PLAYGROUND_CONTENT +`; + +export default template; From afea98164b8192c0d809f6217d2f6aefdc0a5a88 Mon Sep 17 00:00:00 2001 From: Alena Khineika Date: Wed, 10 Jul 2024 12:49:48 +0200 Subject: [PATCH 2/4] refactor: use prompts in code --- package.json | 2 +- scripts/generate-constants.ts | 59 ----------------------- scripts/generate-keyfile.ts | 28 +++++++++++ src/participant/participant.ts | 85 +++++++--------------------------- 4 files changed, 47 insertions(+), 127 deletions(-) delete mode 100644 scripts/generate-constants.ts create mode 100644 scripts/generate-keyfile.ts diff --git a/package.json b/package.json index bc3f6ac36..f60d77464 100644 --- a/package.json +++ b/package.json @@ -43,7 +43,7 @@ "update-grammar": "ts-node ./scripts/update-grammar.ts", "precompile": "npm run clean", "compile": "npm-run-all compile:*", - "compile:keyfile": "ts-node ./scripts/generate-constants.ts", + "compile:keyfile": "ts-node ./scripts/generate-keyfile.ts", "compile:resources": "npm run update-grammar", "compile:extension": "tsc -p ./", "compile:extension-bundles": "webpack --mode development", diff --git a/scripts/generate-constants.ts b/scripts/generate-constants.ts deleted file mode 100644 index b48fb4175..000000000 --- a/scripts/generate-constants.ts +++ /dev/null @@ -1,59 +0,0 @@ -#! /usr/bin/env ts-node - -import ora from 'ora'; -import fs from 'fs'; -import path from 'path'; -import { resolve } from 'path'; -import { config } from 'dotenv'; -import { promisify } from 'util'; - -const writeFile = promisify(fs.writeFile); -const ROOT_DIR = path.join(__dirname, '..'); -const ui = ora('Generate constants keyfile').start(); - -config({ path: resolve(__dirname, '../.env') }); - -interface Constants { - segmentKey?: string; - useMongodbChatParticipant?: string; - chatParticipantGenericPrompt?: string; - chatParticipantQueryPrompt?: string; - chatParticipantModel?: string; -} - -const constants: Constants = {}; - -(async () => { - if (process.env.SEGMENT_KEY) { - constants.segmentKey = process.env.SEGMENT_KEY; - } - if (process.env.USE_MONGODB_CHAT_PARTICIPANT) { - constants.useMongodbChatParticipant = - process.env.USE_MONGODB_CHAT_PARTICIPANT; - } - if (process.env.CHAT_PARTICIPANT_GENERIC_PROMPT) { - constants.chatParticipantGenericPrompt = - process.env.CHAT_PARTICIPANT_GENERIC_PROMPT; - } - if (process.env.CHAT_PARTICIPANT_QUERY_PROMPT) { - constants.chatParticipantQueryPrompt = - process.env.CHAT_PARTICIPANT_QUERY_PROMPT; - } - if (process.env.CHAT_PARTICIPANT_MODEL) { - constants.chatParticipantModel = process.env.CHAT_PARTICIPANT_MODEL; - } - if (Object.keys(constants).length === 0) { - ui.warn('No constants to write'); - return; - } - - await writeFile( - `${ROOT_DIR}/constants.json`, - JSON.stringify(constants, null, 2) - ); - ui.succeed('The constants file was written'); -})().catch((error) => { - ui.fail( - `An error occurred while writing the constants file: ${error.message}` - ); -}); diff --git a/scripts/generate-keyfile.ts b/scripts/generate-keyfile.ts new file mode 100644 index 000000000..b7a6add73 --- /dev/null +++ b/scripts/generate-keyfile.ts @@ -0,0 +1,28 @@ +#! /usr/bin/env ts-node + +import ora from 'ora'; +import fs from 'fs'; +import path from 'path'; +import { resolve } from 'path'; +import { config } from 'dotenv'; +import { promisify } from 'util'; + +const writeFile = promisify(fs.writeFile); +const ROOT_DIR = path.join(__dirname, '..'); +const ui = ora('Generate constants keyfile').start(); + +config({ path: resolve(__dirname, '../.env') }); + +(async () => { + if (process.env.SEGMENT_KEY) { + await writeFile( + `${ROOT_DIR}/constants.json`, + JSON.stringify({ segmentKey: process.env.SEGMENT_KEY }, null, 2) + ); + ui.succeed('Generated segment constants file'); + } else { + throw new Error('The Segment key is missing in environment variables'); + } +})().catch((error) => { + ui.fail(`Failed to generate segment constants file: ${error.message}`); +}); diff --git a/src/participant/participant.ts b/src/participant/participant.ts index 4ed30988b..e67fae906 100644 --- a/src/participant/participant.ts +++ b/src/participant/participant.ts @@ -1,7 +1,4 @@ import * as vscode from 'vscode'; -import { config } from 'dotenv'; -import fs from 'fs'; -import path from 'path'; import { createLogger } from '../logging'; import type { PlaygroundController } from '../editors'; @@ -21,14 +18,8 @@ interface ChatResult extends vscode.ChatResult { stream?: vscode.ChatResponseStream; } -interface GenAIConstants { - useMongodbChatParticipant?: string; - chatParticipantGenericPrompt?: string; - chatParticipantQueryPrompt?: string; - chatParticipantModel?: string; -} - -const PARTICIPANT_ID = 'mongodb.participant'; +const CHAT_PARTICIPANT_ID = 'mongodb.participant'; +const CHAT_PARTICIPANT_MODEL = 'gpt-3.5-turbo'; function handleEmptyQueryRequest(participantId: string): ChatResult { log.info('Chat request participant id', participantId); @@ -64,12 +55,6 @@ export class ParticipantController { _connectionController: ConnectionController; _playgroundController: PlaygroundController; - private _context: vscode.ExtensionContext; - private _useMongodbChatParticipant = false; - private _chatParticipantGenericPrompt = 'You are a MongoDB expert!'; - private _chatParticipantQueryPrompt = 'You are a MongoDB expert!'; - private _chatParticipantModel = 'gpt-3.5-turbo'; - constructor({ context, connectionController, @@ -83,37 +68,6 @@ export class ParticipantController { this.chatResult = { metadata: { command: '' } }; this._connectionController = connectionController; this._playgroundController = playgroundController; - this._context = context; - - this._readConstants(); - } - - private _readConstants(): string | undefined { - config({ path: path.join(this._context.extensionPath, '.env') }); - - try { - const constantsLocation = path.join( - this._context.extensionPath, - './constants.json' - ); - // eslint-disable-next-line no-sync - const constantsFile = fs.readFileSync(constantsLocation, 'utf8'); - const constants = JSON.parse(constantsFile) as GenAIConstants; - - this._useMongodbChatParticipant = - constants.useMongodbChatParticipant === 'true'; - this._chatParticipantGenericPrompt = - constants.chatParticipantGenericPrompt || - this._chatParticipantGenericPrompt; - this._chatParticipantQueryPrompt = - constants.chatParticipantQueryPrompt || - this._chatParticipantQueryPrompt; - this._chatParticipantModel = - constants.chatParticipantModel || this._chatParticipantModel; - } catch (error) { - log.error('An error occurred while reading the constants file', error); - return; - } } createParticipant(context: vscode.ExtensionContext) { @@ -121,7 +75,7 @@ export class ParticipantController { // when you type `@`, and can contribute sub-commands in the chat input // that appear when you type `/`. const cat = vscode.chat.createChatParticipant( - PARTICIPANT_ID, + CHAT_PARTICIPANT_ID, this.chatHandler.bind(this) ); cat.iconPath = vscode.Uri.joinPath( @@ -165,7 +119,7 @@ export class ParticipantController { try { const [model] = await vscode.lm.selectChatModels({ vendor: 'copilot', - family: this._chatParticipantModel, + family: CHAT_PARTICIPANT_MODEL, }); if (model) { const chatResponse = await model.sendRequest(messages, {}, token); @@ -196,14 +150,17 @@ export class ParticipantController { }) { const messages = [ // eslint-disable-next-line new-cap - vscode.LanguageModelChatMessage.Assistant( - this._chatParticipantGenericPrompt - ), + 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 === PARTICIPANT_ID && + historyItem.participant === CHAT_PARTICIPANT_ID && historyItem instanceof vscode.ChatRequestTurn ) { // eslint-disable-next-line new-cap @@ -211,7 +168,7 @@ export class ParticipantController { } if ( - historyItem.participant === PARTICIPANT_ID && + historyItem.participant === CHAT_PARTICIPANT_ID && historyItem instanceof vscode.ChatResponseTurn ) { let res = ''; @@ -309,9 +266,12 @@ export class ParticipantController { const messages = [ // eslint-disable-next-line new-cap - vscode.LanguageModelChatMessage.Assistant( - this._chatParticipantQueryPrompt - ), + 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), ]; @@ -351,15 +311,6 @@ export class ParticipantController { stream: vscode.ChatResponseStream, token: vscode.CancellationToken ): Promise { - if (!this._useMongodbChatParticipant) { - stream.markdown( - vscode.l10n.t( - 'Under construction. Will be available soon. Stay tuned!\n\n' - ) - ); - return { metadata: { command: '' } }; - } - if (request.command === 'query') { this.chatResult = await this.handleQueryRequest({ request, From b714d2205c2abcf9c5b12ac06447b8217bbc7500 Mon Sep 17 00:00:00 2001 From: Alena Khineika Date: Thu, 11 Jul 2024 11:27:53 +0200 Subject: [PATCH 3/4] fix: use uri path --- src/participant/participant.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/participant/participant.ts b/src/participant/participant.ts index e67fae906..88836bc40 100644 --- a/src/participant/participant.ts +++ b/src/participant/participant.ts @@ -79,7 +79,7 @@ export class ParticipantController { this.chatHandler.bind(this) ); cat.iconPath = vscode.Uri.joinPath( - context.extensionUri, + vscode.Uri.parse(context.extensionPath), 'images', 'mongodb.png' ); From 365024e389275a9001a620cc9ced3d395ece4ce5 Mon Sep 17 00:00:00 2001 From: Alena Khineika Date: Thu, 18 Jul 2024 12:58:16 +0200 Subject: [PATCH 4/4] refactor: pass participant to playground controller --- src/editors/playgroundController.ts | 5 +++ src/mdbExtensionController.ts | 27 ++++--------- src/participant/participant.ts | 39 ++++++++----------- .../editors/playgroundController.test.ts | 9 ++++- ...aygroundSelectedCodeActionProvider.test.ts | 27 +++++++++++++ .../language/languageServerController.test.ts | 6 +++ 6 files changed, 71 insertions(+), 42 deletions(-) diff --git a/src/editors/playgroundController.ts b/src/editors/playgroundController.ts index 79c9bb62a..187f320fa 100644 --- a/src/editors/playgroundController.ts +++ b/src/editors/playgroundController.ts @@ -45,6 +45,7 @@ import { isPlayground, getPlaygroundExtensionForTelemetry, } from '../utils/playground'; +import type { ParticipantController } from '../participant/participant'; const log = createLogger('playground controller'); @@ -133,6 +134,7 @@ export default class PlaygroundController { private _playgroundResultTextDocument?: vscode.TextDocument; private _statusView: StatusView; private _playgroundResultViewProvider: PlaygroundResultProvider; + private _participantController: ParticipantController; private _codeToEvaluate = ''; @@ -145,6 +147,7 @@ export default class PlaygroundController { activeConnectionCodeLensProvider, exportToLanguageCodeLensProvider, playgroundSelectedCodeActionProvider, + participantController, }: { connectionController: ConnectionController; languageServerController: LanguageServerController; @@ -154,6 +157,7 @@ export default class PlaygroundController { activeConnectionCodeLensProvider: ActiveConnectionCodeLensProvider; exportToLanguageCodeLensProvider: ExportToLanguageCodeLensProvider; playgroundSelectedCodeActionProvider: PlaygroundSelectedCodeActionProvider; + participantController: ParticipantController; }) { this._connectionController = connectionController; this._activeTextEditor = vscode.window.activeTextEditor; @@ -165,6 +169,7 @@ export default class PlaygroundController { this._exportToLanguageCodeLensProvider = exportToLanguageCodeLensProvider; this._playgroundSelectedCodeActionProvider = playgroundSelectedCodeActionProvider; + this._participantController = participantController; this._connectionController.addEventListener( DataServiceEventTypes.ACTIVE_CONNECTION_CHANGED, diff --git a/src/mdbExtensionController.ts b/src/mdbExtensionController.ts index bdf07b2c9..579757dac 100644 --- a/src/mdbExtensionController.ts +++ b/src/mdbExtensionController.ts @@ -64,7 +64,7 @@ export default class MDBExtensionController implements vscode.Disposable { _activeConnectionCodeLensProvider: ActiveConnectionCodeLensProvider; _editDocumentCodeLensProvider: EditDocumentCodeLensProvider; _exportToLanguageCodeLensProvider: ExportToLanguageCodeLensProvider; - _participantController?: ParticipantController; + _participantController: ParticipantController; constructor( context: vscode.ExtensionContext, @@ -107,6 +107,9 @@ export default class MDBExtensionController implements vscode.Disposable { new PlaygroundSelectedCodeActionProvider(); this._playgroundDiagnosticsCodeActionProvider = new PlaygroundDiagnosticsCodeActionProvider(); + this._participantController = new ParticipantController({ + connectionController: this._connectionController, + }); this._playgroundController = new PlaygroundController({ connectionController: this._connectionController, languageServerController: this._languageServerController, @@ -117,6 +120,7 @@ export default class MDBExtensionController implements vscode.Disposable { exportToLanguageCodeLensProvider: this._exportToLanguageCodeLensProvider, playgroundSelectedCodeActionProvider: this._playgroundSelectedCodeActionProvider, + participantController: this._participantController, }); this._editorsController = new EditorsController({ context, @@ -139,11 +143,6 @@ export default class MDBExtensionController implements vscode.Disposable { telemetryService: this._telemetryService, }); this._editorsController.registerProviders(); - this._participantController = new ParticipantController({ - context, - connectionController: this._connectionController, - playgroundController: this._playgroundController, - }); } async activate(): Promise { @@ -277,24 +276,18 @@ export default class MDBExtensionController implements vscode.Disposable { this.registerParticipantCommand( EXTENSION_COMMANDS.OPEN_PARTICIPANT_QUERY_IN_PLAYGROUND, () => { - if (!this._participantController) { - return Promise.resolve(false); - } return this._playgroundController.createPlaygroundFromParticipantQuery({ text: - this._participantController.chatResult.metadata.queryContent || '', + this._participantController._chatResult.metadata.queryContent || '', }); } ); this.registerParticipantCommand( EXTENSION_COMMANDS.RUN_PARTICIPANT_QUERY, () => { - if (!this._participantController) { - return Promise.resolve(false); - } return this._playgroundController.evaluateParticipantQuery({ text: - this._participantController.chatResult.metadata.queryContent || '', + this._participantController._chatResult.metadata.queryContent || '', }); } ); @@ -304,10 +297,6 @@ export default class MDBExtensionController implements vscode.Disposable { command: string, commandHandler: (...args: any[]) => Promise ): void => { - if (!this._participantController) { - return; - } - const commandHandlerWithTelemetry = (args: any[]): Promise => { this._telemetryService.trackCommandRun(command); @@ -315,7 +304,7 @@ export default class MDBExtensionController implements vscode.Disposable { }; this._context.subscriptions.push( - this._participantController.participant, + this._participantController.getParticipant(this._context), vscode.commands.registerCommand(command, commandHandlerWithTelemetry) ); }; diff --git a/src/participant/participant.ts b/src/participant/participant.ts index 88836bc40..d49e58cf0 100644 --- a/src/participant/participant.ts +++ b/src/participant/participant.ts @@ -1,7 +1,6 @@ import * as vscode from 'vscode'; import { createLogger } from '../logging'; -import type { PlaygroundController } from '../editors'; import type ConnectionController from '../connectionController'; import EXTENSION_COMMANDS from '../commands'; @@ -18,10 +17,10 @@ interface ChatResult extends vscode.ChatResult { stream?: vscode.ChatResponseStream; } -const CHAT_PARTICIPANT_ID = 'mongodb.participant'; -const CHAT_PARTICIPANT_MODEL = 'gpt-3.5-turbo'; +export const CHAT_PARTICIPANT_ID = 'mongodb.participant'; +export const CHAT_PARTICIPANT_MODEL = 'gpt-4'; -function handleEmptyQueryRequest(participantId: string): ChatResult { +function handleEmptyQueryRequest(participantId?: string): ChatResult { log.info('Chat request participant id', participantId); return { @@ -35,7 +34,7 @@ function handleEmptyQueryRequest(participantId: string): ChatResult { }; } -function getRunnableContentFromString(responseContent: string) { +export function getRunnableContentFromString(responseContent: string) { const matchedJSQueryContent = responseContent.match( /```javascript((.|\n)*)```/ ); @@ -50,40 +49,37 @@ function getRunnableContentFromString(responseContent: string) { } export class ParticipantController { - participant: vscode.ChatParticipant; - chatResult: ChatResult; + _participant?: vscode.ChatParticipant; + _chatResult: ChatResult; _connectionController: ConnectionController; - _playgroundController: PlaygroundController; constructor({ - context, connectionController, - playgroundController, }: { - context: vscode.ExtensionContext; connectionController: ConnectionController; - playgroundController: PlaygroundController; }) { - this.participant = this.createParticipant(context); - this.chatResult = { metadata: { command: '' } }; + this._chatResult = { metadata: { command: '' } }; this._connectionController = connectionController; - this._playgroundController = playgroundController; } createParticipant(context: vscode.ExtensionContext) { // Chat participants appear as top-level options in the chat input // when you type `@`, and can contribute sub-commands in the chat input // that appear when you type `/`. - const cat = vscode.chat.createChatParticipant( + this._participant = vscode.chat.createChatParticipant( CHAT_PARTICIPANT_ID, this.chatHandler.bind(this) ); - cat.iconPath = vscode.Uri.joinPath( + this._participant.iconPath = vscode.Uri.joinPath( vscode.Uri.parse(context.extensionPath), 'images', 'mongodb.png' ); - return cat; + return this._participant; + } + + getParticipant(context: vscode.ExtensionContext) { + return this._participant || this.createParticipant(context); } handleError(err: any, stream: vscode.ChatResponseStream): void { @@ -230,7 +226,7 @@ export class ParticipantController { token: vscode.CancellationToken; }) { if (!request.prompt || request.prompt.trim().length === 0) { - return handleEmptyQueryRequest(this.participant.id); + return handleEmptyQueryRequest(this._participant?.id); } let dataService = this._connectionController.getActiveDataService(); @@ -275,7 +271,6 @@ export class ParticipantController { // eslint-disable-next-line new-cap vscode.LanguageModelChatMessage.User(request.prompt), ]; - const responseContent = await this.getChatResponseContent({ messages, stream, @@ -312,13 +307,13 @@ export class ParticipantController { token: vscode.CancellationToken ): Promise { if (request.command === 'query') { - this.chatResult = await this.handleQueryRequest({ + this._chatResult = await this.handleQueryRequest({ request, context, stream, token, }); - return this.chatResult; + return this._chatResult; } else if (request.command === 'docs') { // TODO: Implement this. } else if (request.command === 'schema') { diff --git a/src/test/suite/editors/playgroundController.test.ts b/src/test/suite/editors/playgroundController.test.ts index 14df23617..66af1b06b 100644 --- a/src/test/suite/editors/playgroundController.test.ts +++ b/src/test/suite/editors/playgroundController.test.ts @@ -22,6 +22,7 @@ import { StorageController } from '../../../storage'; import TelemetryService from '../../../telemetry/telemetryService'; import { TEST_DATABASE_URI } from '../dbTestHelper'; import { ExtensionContextStub, LanguageServerControllerStub } from '../stubs'; +import { ParticipantController } from '../../../participant/participant'; const expect = chai.expect; @@ -57,6 +58,7 @@ suite('Playground Controller Test Suite', function () { let testPlaygroundController: PlaygroundController; let showErrorMessageStub: SinonStub; let sandbox: sinon.SinonSandbox; + let testParticipantController: ParticipantController; beforeEach(() => { sandbox = sinon.createSandbox(); @@ -84,11 +86,13 @@ suite('Playground Controller Test Suite', function () { testExportToLanguageCodeLensProvider = new ExportToLanguageCodeLensProvider(); testCodeActionProvider = new PlaygroundSelectedCodeActionProvider(); - languageServerControllerStub = new LanguageServerControllerStub( extensionContextStub, testStorageController ); + testParticipantController = new ParticipantController({ + connectionController: testConnectionController, + }); testPlaygroundController = new PlaygroundController({ connectionController: testConnectionController, languageServerController: languageServerControllerStub, @@ -98,6 +102,7 @@ suite('Playground Controller Test Suite', function () { activeConnectionCodeLensProvider: testActiveDBCodeLensProvider, exportToLanguageCodeLensProvider: testExportToLanguageCodeLensProvider, playgroundSelectedCodeActionProvider: testCodeActionProvider, + participantController: testParticipantController, }); showErrorMessageStub = sandbox.stub(vscode.window, 'showErrorMessage'); sandbox.stub(testTelemetryService, 'trackNewConnection'); @@ -354,6 +359,7 @@ suite('Playground Controller Test Suite', function () { exportToLanguageCodeLensProvider: testExportToLanguageCodeLensProvider, playgroundSelectedCodeActionProvider: testCodeActionProvider, + participantController: testParticipantController, }); expect(playgroundController._activeTextEditor).to.deep.equal( @@ -372,6 +378,7 @@ suite('Playground Controller Test Suite', function () { exportToLanguageCodeLensProvider: testExportToLanguageCodeLensProvider, playgroundSelectedCodeActionProvider: testCodeActionProvider, + participantController: testParticipantController, }); const textFromEditor = 'var x = { name: qwerty }'; const selection = { diff --git a/src/test/suite/editors/playgroundSelectedCodeActionProvider.test.ts b/src/test/suite/editors/playgroundSelectedCodeActionProvider.test.ts index 1dea18fbd..74f143249 100644 --- a/src/test/suite/editors/playgroundSelectedCodeActionProvider.test.ts +++ b/src/test/suite/editors/playgroundSelectedCodeActionProvider.test.ts @@ -13,6 +13,11 @@ import type { PlaygroundResult } from '../../../types/playgroundType'; import { ExportToLanguageMode } from '../../../types/playgroundType'; import { TEST_DATABASE_URI } from '../dbTestHelper'; import { ExtensionContextStub } from '../stubs'; +import { ParticipantController } from '../../../participant/participant'; +import ConnectionController from '../../../connectionController'; +import StatusView from '../../../views/statusView'; +import StorageController from '../../../storage/storageController'; +import TelemetryService from '../../../telemetry/telemetryService'; // eslint-disable-next-line @typescript-eslint/no-var-requires const { version } = require('../../../../package.json'); @@ -33,8 +38,29 @@ suite('Playground Selected CodeAction Provider Test Suite', function () { suite('the MongoDB playground in JS', () => { const testCodeActionProvider = new PlaygroundSelectedCodeActionProvider(); const sandbox = sinon.createSandbox(); + let testStorageController: StorageController; + let testTelemetryService: TelemetryService; + let testStatusView: StatusView; + let testConnectionController: ConnectionController; + let testParticipantController: ParticipantController; beforeEach(async () => { + testStorageController = new StorageController(extensionContextStub); + testTelemetryService = new TelemetryService( + testStorageController, + extensionContextStub + ); + testStatusView = new StatusView(extensionContextStub); + testConnectionController = new ConnectionController({ + statusView: testStatusView, + storageController: testStorageController, + telemetryService: testTelemetryService, + }); + + testParticipantController = new ParticipantController({ + connectionController: testConnectionController, + }); + sandbox.replace( mdbTestExtension.testExtensionController, '_languageServerController', @@ -73,6 +99,7 @@ suite('Playground Selected CodeAction Provider Test Suite', function () { exportToLanguageCodeLensProvider: testExportToLanguageCodeLensProvider, playgroundSelectedCodeActionProvider: testCodeActionProvider, + participantController: testParticipantController, }); const fakeOpenPlaygroundResult = sandbox.fake(); diff --git a/src/test/suite/language/languageServerController.test.ts b/src/test/suite/language/languageServerController.test.ts index b34944301..52a660a37 100644 --- a/src/test/suite/language/languageServerController.test.ts +++ b/src/test/suite/language/languageServerController.test.ts @@ -22,6 +22,7 @@ import { StorageController } from '../../../storage'; import { TEST_DATABASE_URI } from '../dbTestHelper'; import TelemetryService from '../../../telemetry/telemetryService'; import { ExtensionContextStub } from '../stubs'; +import { ParticipantController } from '../../../participant/participant'; const expect = chai.expect; @@ -60,6 +61,7 @@ suite('Language Server Controller Test Suite', () => { let languageServerControllerStub: LanguageServerController; let testPlaygroundController: PlaygroundController; + let testParticipantController: ParticipantController; const sandbox = sinon.createSandbox(); @@ -67,6 +69,9 @@ suite('Language Server Controller Test Suite', () => { languageServerControllerStub = new LanguageServerController( extensionContextStub ); + testParticipantController = new ParticipantController({ + connectionController: testConnectionController, + }); testPlaygroundController = new PlaygroundController({ connectionController: testConnectionController, languageServerController: languageServerControllerStub, @@ -76,6 +81,7 @@ suite('Language Server Controller Test Suite', () => { activeConnectionCodeLensProvider: testActiveDBCodeLensProvider, exportToLanguageCodeLensProvider: testExportToLanguageCodeLensProvider, playgroundSelectedCodeActionProvider: testCodeActionProvider, + participantController: testParticipantController, }); await languageServerControllerStub.startLanguageServer(); await testPlaygroundController._activeConnectionChanged();