feat(participant): export to a playground VSCODE-574

package.json

@@ -282,6 +282,10 @@
           "dark": "images/dark/play.svg"
         }
       },
+      {
+        "command": "mdb.exportCodeToPlayground",
+        "title": "Export Code to Playground"
+      },
       {
         "command": "mdb.exportToPython",
         "title": "MongoDB: Export To Python 3"
@@ -747,6 +751,17 @@
           "when": "mdb.isPlayground == true"
         }
       ],
+      "mdb.copilot": [
+        {
+          "command": "mdb.exportCodeToPlayground"
+        }
+      ],
+      "editor/context": [
+        {
+          "submenu": "mdb.copilot",
+          "group": "1_main@2"
+        }
+      ],
       "commandPalette": [
         {
           "command": "mdb.selectDatabaseWithParticipant",
@@ -948,6 +963,10 @@
           "command": "mdb.runPlayground",
           "when": "false"
         },
+        {
+          "command": "mdb.exportCodeToPlayground",
+          "when": "false"
+        },
         {
           "command": "mdb.createIndexFromTreeView",
           "when": "false"
@@ -994,6 +1013,12 @@
         }
       ]
     },
+    "submenus": [
+      {
+        "id": "mdb.copilot",
+        "label": "MongoDB Copilot"
+      }
+    ],
     "keybindings": [
       {
         "command": "mdb.runSelectedPlaygroundBlocks",

src/commands/index.ts

@@ -15,6 +15,7 @@ enum EXTENSION_COMMANDS {
   MDB_RUN_SELECTED_PLAYGROUND_BLOCKS = 'mdb.runSelectedPlaygroundBlocks',
   MDB_RUN_ALL_PLAYGROUND_BLOCKS = 'mdb.runAllPlaygroundBlocks',
   MDB_RUN_ALL_OR_SELECTED_PLAYGROUND_BLOCKS = 'mdb.runPlayground',
+  MDB_EXPORT_CODE_TO_PLAYGROUND = 'mdb.exportCodeToPlayground',
 
   MDB_FIX_THIS_INVALID_INTERACTIVE_SYNTAX = 'mdb.fixThisInvalidInteractiveSyntax',
   MDB_FIX_ALL_INVALID_INTERACTIVE_SYNTAX = 'mdb.fixAllInvalidInteractiveSyntax',

src/editors/playgroundController.ts

@@ -492,33 +492,31 @@ export default class PlaygroundController {
     }
 
     try {
-      const progressResult = await vscode.window.withProgress(
+      return await vscode.window.withProgress(
         {
           location: ProgressLocation.Notification,
           title: 'Running MongoDB playground...',
           cancellable: true,
         },
-        async (progress, token) => {
+        async (progress, token): Promise<ShellEvaluateResult> => {
           token.onCancellationRequested(() => {
-            // If a user clicked the cancel button terminate all playground scripts.
+            // If the user clicks the cancel button,
+            // terminate all processes running on the language server.
             this._languageServerController.cancelAll();
-
-            return { result: undefined };
           });
 
-          // Run all playground scripts.
-          const result: ShellEvaluateResult = await this._evaluate(
-            codeToEvaluate
-          );
-
-          return result;
+          return Promise.race([
+            this._evaluate(codeToEvaluate),
+            new Promise<{ result: undefined }>((resolve) =>
+              token.onCancellationRequested(() => {
+                resolve({ result: undefined });
+              })
+            ),
+          ]);
         }
       );
-
-      return progressResult;
     } catch (error) {
       log.error('Evaluating playground with cancel modal failed', error);
-
       return { result: undefined };
     }
   }

src/mdbExtensionController.ts

@@ -233,6 +233,9 @@ export default class MDBExtensionController implements vscode.Disposable {
       EXTENSION_COMMANDS.MDB_RUN_ALL_OR_SELECTED_PLAYGROUND_BLOCKS,
       () => this._playgroundController.runAllOrSelectedPlaygroundBlocks()
     );
+    this.registerCommand(EXTENSION_COMMANDS.MDB_EXPORT_CODE_TO_PLAYGROUND, () =>
+      this._participantController.exportCodeToPlayground()
+    );
 
     this.registerCommand(
       EXTENSION_COMMANDS.MDB_FIX_THIS_INVALID_INTERACTIVE_SYNTAX,

src/participant/participant.ts

@@ -10,6 +10,7 @@ import EXTENSION_COMMANDS from '../commands';
 import type { StorageController } from '../storage';
 import { StorageVariables } from '../storage';
 import { GenericPrompt, isPromptEmpty } from './prompts/generic';
+import { ExportToPlaygroundPrompt } from './prompts/exportToPlayground';
 import type { ChatResult } from './constants';
 import {
   askToConnectChatResult,
@@ -42,6 +43,7 @@ import {
 } from '../telemetry/telemetryService';
 import { DocsChatbotAIService } from './docsChatbotAIService';
 import type TelemetryService from '../telemetry/telemetryService';
+import formatError from '../utils/formatError';
 
 const log = createLogger('participant');
 
@@ -132,7 +134,7 @@ export default class ParticipantController {
     return this._participant;
   }
 
-  handleError(err: any, command: string): never {
+  handleError(err: any, command: string): void {
     let errorCode: string | undefined;
     let errorName: ParticipantErrorTypes;
     // Making the chat request might fail because
@@ -168,9 +170,6 @@ export default class ParticipantController {
         error_name: errorName,
       }
     );
-
-    // Re-throw other errors so they show up in the UI.
-    throw err;
   }
 
   /**
@@ -1201,6 +1200,79 @@ export default class ParticipantController {
     });
   }
 
+  async exportCodeToPlayground(): Promise<boolean> {
+    const activeTextEditor = vscode.window.activeTextEditor;
+    if (!activeTextEditor) {
+      await vscode.window.showErrorMessage('Active editor not found.');
+      return false;
+    }
+
+    const sortedSelections = Array.from(activeTextEditor.selections).sort(
+      (a, b) => a.start.compareTo(b.start)
+    );
+    const selectedText = sortedSelections
+      .map((selection) => activeTextEditor.document.getText(selection))
+      .join('\n');
+    const code =
+      selectedText || activeTextEditor.document.getText().trim() || '';
+    try {
+      const progressResult = await vscode.window.withProgress(
+        {
+          location: vscode.ProgressLocation.Notification,
+          title: 'Exporting code to a playground...',
+          cancellable: true,
+        },
+        async (progress, token): Promise<string | undefined> => {
+          token.onCancellationRequested(async () => {
+            await vscode.window.showInformationMessage(
+              'The running export to a playground operation was canceled.'
+            );
+          });
+
+          const messages = ExportToPlaygroundPrompt.buildMessages(code);
+          return Promise.race([
+            this.getChatResponseContent({
+              messages,
+              token,
+            }),
+            new Promise<undefined>((resolve) =>
+              token.onCancellationRequested(() => {
+                resolve(undefined);
+              })
+            ),
+          ]);
+        }
+      );
+
+      if (progressResult?.includes("Sorry, I can't assist with that.")) {
+        void vscode.window.showErrorMessage("Sorry, I can't assist with that.");
+        return Promise.resolve(false);
+      }
+
+      if (progressResult) {
+        const runnableContent = getRunnableContentFromString(progressResult);
+        if (runnableContent) {
+          await vscode.commands.executeCommand(
+            EXTENSION_COMMANDS.OPEN_PARTICIPANT_QUERY_IN_PLAYGROUND,
+            {
+              runnableContent,
+            }
+          );
+        }
+      }
+
+      return Promise.resolve(true);
+    } catch (error) {
+      this.handleError(error, 'exportToPlayground');
+      await vscode.window.showErrorMessage(
+        `An error occurred exporting to a playground: ${
+          formatError(error).message
+        }`
+      );
+      return Promise.resolve(false);
+    }
+  }
+
   async chatHandler(
     ...args: [
       vscode.ChatRequest,
@@ -1249,6 +1321,8 @@ Please see our [FAQ](https://www.mongodb.com/docs/generative-ai-faq/) for more i
       }
     } catch (e) {
       this.handleError(e, request.command || 'generic');
+      // Re-throw other errors so they show up in the UI.
+      throw e;
     }
   }
 

src/participant/prompts/exportToPlayground.ts

@@ -0,0 +1,31 @@
+import * as vscode from 'vscode';
+
+export class ExportToPlaygroundPrompt {
+  static getAssistantPrompt(): vscode.LanguageModelChatMessage {
+    const prompt = `You are a MongoDB expert.
+Your task is to help the user build MongoDB queries and aggregation pipelines that perform their task.
+You convert user's code written in any programming language to the MongoDB Shell syntax.
+Take a user prompt as an input string and translate it to the MongoDB Shell language.
+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 \`\`\`.
+Respond in MongoDB shell syntax using the \`\`\`javascript code block syntax.`;
+
+    // eslint-disable-next-line new-cap
+    return vscode.LanguageModelChatMessage.Assistant(prompt);
+  }
+
+  static getUserPrompt(prompt: string): vscode.LanguageModelChatMessage {
+    // eslint-disable-next-line new-cap
+    return vscode.LanguageModelChatMessage.User(prompt);
+  }
+
+  static buildMessages(prompt: string): vscode.LanguageModelChatMessage[] {
+    const messages = [
+      ExportToPlaygroundPrompt.getAssistantPrompt(),
+      ExportToPlaygroundPrompt.getUserPrompt(prompt),
+    ];
+
+    return messages;
+  }
+}