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/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 { EportToPlaygroundPrompt } from './prompts/exportToPlayground';
 import type { ChatResult } from './constants';
 import {
   askToConnectChatResult,
@@ -1201,6 +1202,66 @@ export default class ParticipantController {
     });
   }
 
+  async exportCodeToPlayground(): Promise<boolean> {
+    const activeTextEditor = vscode.window.activeTextEditor;
+    if (!activeTextEditor) {
+      void vscode.window.showErrorMessage('Active editor not found.');
+      return Promise.resolve(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) => {
+          const messages = EportToPlaygroundPrompt.buildMessages(code);
+          return await this.getChatResponseContent({
+            messages,
+            token,
+          });
+        }
+      );
+
+      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 (progressResult) {
+          await vscode.commands.executeCommand(
+            EXTENSION_COMMANDS.OPEN_PARTICIPANT_QUERY_IN_PLAYGROUND,
+            {
+              runnableContent,
+            }
+          );
+        }
+      }
+
+      return Promise.resolve(true);
+    } catch (error) {
+      log.error(
+        'Exporting code to a playground with cancel modal failed',
+        error
+      );
+      return Promise.resolve(false);
+    }
+  }
+
   async chatHandler(
     ...args: [
       vscode.ChatRequest,

src/participant/prompts/exportToPlayground.ts

@@ -0,0 +1,32 @@
+import * as vscode from 'vscode';
+
+export class EportToPlaygroundPrompt {
+  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 achieve this by converting user's code written in any proggramming language to the MongoDB Shell syntax.
+Take a user promt 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 \`\`\`.
+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(prompt: string): vscode.LanguageModelChatMessage {
+    // eslint-disable-next-line new-cap
+    return vscode.LanguageModelChatMessage.User(prompt);
+  }
+
+  static buildMessages(prompt: string): vscode.LanguageModelChatMessage[] {
+    const messages = [
+      EportToPlaygroundPrompt.getAssistantPrompt(),
+      EportToPlaygroundPrompt.getUserPrompt(prompt),
+    ];
+
+    return messages;
+  }
+}

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

@@ -26,6 +26,8 @@ import {
 } from '../../../storage/storageController';
 import type { LoadedConnection } from '../../../storage/connectionStorage';
 import { ChatMetadataStore } from '../../../participant/chatMetadata';
+import { getFullRange } from '../suggestTestHelpers';
+import { isPlayground } from '../../../utils/playground';
 
 // The Copilot's model in not available in tests,
 // therefore we need to mock its methods and returning values.
@@ -1279,6 +1281,105 @@ Schema:
           expect(properties.error_name).to.equal('Docs Chatbot API Issue');
         });
       });
+
+      suite('export to playground', function () {
+        beforeEach(async function () {
+          await vscode.commands.executeCommand(
+            'workbench.action.files.newUntitledFile'
+          );
+        });
+
+        afterEach(async function () {
+          await vscode.commands.executeCommand(
+            'workbench.action.closeActiveEditor'
+          );
+        });
+
+        test('exports all code to a playground', async function () {
+          this.timeout(20000);
+          const editor = vscode.window.activeTextEditor;
+          if (!editor) {
+            throw new Error('Window active text editor is undefined');
+          }
+
+          const testDocumentUri = editor.document.uri;
+          const edit = new vscode.WorkspaceEdit();
+          const code = `
+  InsertOneResult result = collection.insertOne(new Document()
+    .append("_id", new ObjectId())
+    .append("title", "Ski Bloopers")
+    .append("genres", Arrays.asList("Documentary", "Comedy")));
+  System.out.println("Success! Inserted document id: " + result.getInsertedId());
+`;
+          edit.replace(testDocumentUri, getFullRange(editor.document), code);
+          await vscode.workspace.applyEdit(edit);
+          sendRequestStub.onCall(0).resolves({
+            text: [
+              '```javascript\n' +
+                'db.collection.insertOne({\n' +
+                '_id: new ObjectId(),\n' +
+                'title: "Ski Bloopers",' +
+                'genres: ["Documentary", "Comedy"]' +
+                '});\n' +
+                'print("Success! Inserted document id: " + result.insertedId);' +
+                '```',
+            ],
+          });
+          await testParticipantController.exportCodeToPlayground();
+          const messages = sendRequestStub.firstCall.args[0];
+          expect(messages[1].content).to.include('System.out.println');
+          expect(
+            isPlayground(vscode.window.activeTextEditor?.document.uri)
+          ).to.be.eql(true);
+          expect(vscode.window.activeTextEditor?.document.getText()).to.include(
+            'Inserted document id'
+          );
+        });
+
+        test('exports selected lines of code to a playground', async function () {
+          const editor = vscode.window.activeTextEditor;
+          if (!editor) {
+            throw new Error('Window active text editor is undefined');
+          }
+
+          const testDocumentUri = editor.document.uri;
+          const edit = new vscode.WorkspaceEdit();
+          const code = `
+  InsertOneResult result = collection.insertOne(new Document()
+    .append("_id", new ObjectId())
+    .append("title", "Ski Bloopers")
+    .append("genres", Arrays.asList("Documentary", "Comedy")));
+  System.out.println("Success! Inserted document id: " + result.getInsertedId());
+`;
+          edit.replace(testDocumentUri, getFullRange(editor.document), code);
+          await vscode.workspace.applyEdit(edit);
+          const position = editor.selection.active;
+          const startPosition = position.with(0, 0);
+          const endPosition = position.with(3, 63);
+          const newSelection = new vscode.Selection(startPosition, endPosition);
+          editor.selection = newSelection;
+          sendRequestStub.onCall(0).resolves({
+            text: [
+              '```javascript\n' +
+                'db.collection.insertOne({\n' +
+                '_id: new ObjectId(),\n' +
+                'title: "Ski Bloopers",' +
+                'genres: ["Documentary", "Comedy"]' +
+                '});\n' +
+                '```',
+            ],
+          });
+          await testParticipantController.exportCodeToPlayground();
+          const messages = sendRequestStub.firstCall.args[0];
+          expect(messages[1].content).to.not.include('System.out.println');
+          expect(
+            isPlayground(vscode.window.activeTextEditor?.document.uri)
+          ).to.be.eql(true);
+          expect(
+            vscode.window.activeTextEditor?.document.getText()
+          ).to.not.include('Inserted document id');
+        });
+      });
     });
   });