AAP-30427: VS Code Extension settings page to take a new custom_promp…

…t field (ansible#1545)
manstis · Oct 14, 2024 · 16ff1a6 · 16ff1a6
1 parent 5c40c83
commit 16ff1a6
Show file tree

Hide file tree

Showing 7 changed files with 371 additions and 1 deletion.
diff --git a/package.json b/package.json
@@ -633,6 +633,18 @@
             "type": "string",
             "markdownDescription": "Model ID to override your organization's default model. This setting is only applicable to commercial users with an Ansible Lightspeed seat assignment.",
             "order": 4
+          },
+          "ansible.lightspeed.playbookGenerationCustomPrompt": {
+            "scope": "resource",
+            "type": "string",
+            "markdownDescription": "Custom Prompt for Playbook generation. This setting is only applicable to commercial users with an Ansible Lightspeed seat assignment.",
+            "order": 5
+          },
+          "ansible.lightspeed.playbookExplanationCustomPrompt": {
+            "scope": "resource",
+            "type": "string",
+            "markdownDescription": "Custom Prompt for Playbook explanation. This setting is only applicable to commercial users with an Ansible Lightspeed seat assignment.",
+            "order": 6
           }
         }
       },

diff --git a/src/features/lightspeed/api.ts b/src/features/lightspeed/api.ts
@@ -297,6 +297,12 @@ export class LightSpeedAPI {
       return {} as ExplanationResponseParams;
     }
     try {
+      const customPrompt =
+        lightSpeedManager.settingsManager.settings.lightSpeedService
+          .playbookExplanationCustomPrompt;
+      if (customPrompt && customPrompt.length > 0) {
+        inputData.customPrompt = customPrompt;
+      }
       const requestData = {
         ...inputData,
         metadata: { ansibleExtensionVersion: this._extensionVersion },
@@ -336,6 +342,13 @@ export class LightSpeedAPI {
       return {} as GenerationResponseParams;
     }
     try {
+      const customPrompt =
+        lightSpeedManager.settingsManager.settings.lightSpeedService
+          .playbookGenerationCustomPrompt;
+      if (customPrompt && customPrompt.length > 0) {
+        inputData.customPrompt = customPrompt;
+      }
+
       const requestData = {
         ...inputData,
         metadata: { ansibleExtensionVersion: this._extensionVersion },

diff --git a/src/features/lightspeed/playbookExplanation.ts b/src/features/lightspeed/playbookExplanation.ts
@@ -111,6 +111,13 @@ export const playbookExplanation = async (extensionUri: vscode.Uri) => {
         markdown = response.content;
         if (markdown.length === 0) {
           markdown = "### No explanation provided.";
+          const customPrompt =
+            lightSpeedManager.settingsManager.settings.lightSpeedService
+              .playbookExplanationCustomPrompt ?? "";
+          if (customPrompt.length > 0) {
+            markdown +=
+              "\n\nYou may want to consider amending your custom prompt.";
+          }
         }
         const html_snippet = marked.parse(markdown) as string;
         currentPanel.setContent(html_snippet, true);

diff --git a/src/interfaces/lightspeed.ts b/src/interfaces/lightspeed.ts
@@ -142,6 +142,7 @@ export interface ISuggestionDetails {
 
 export interface GenerationRequestParams {
   text: string;
+  customPrompt?: string;
   outline?: string;
   generationId: string;
   createOutline: boolean;
@@ -156,6 +157,7 @@ export interface GenerationResponseParams {
 
 export interface ExplanationRequestParams {
   content: string;
+  customPrompt?: string;
   explanationId: string;
 }
 

diff --git a/test/mockLightspeedServer/explanations.ts b/test/mockLightspeedServer/explanations.ts
@@ -9,6 +9,7 @@ export function explanations(
   res: any,
 ) {
   const playbook = req.body.content;
+  const customPrompt = req.body.customPrompt;
   const explanationId = req.body.explanationId
     ? req.body.explanationId
     : uuidv4();
@@ -38,8 +39,18 @@ export function explanations(
     });
   }
 
+  // Special case to replicate a broken custom prompt
+  if (customPrompt && customPrompt === "custom prompt broken") {
+    logger.info("Returning 400. Custom prompt is invalid");
+    return res.status(400).send({
+      code: "validation",
+      message: "invalid",
+      detail: { customPrompt: "custom prompt is invalid" },
+    });
+  }
+
   // cSpell: disable
-  const content = `
+  let content = `
 ## Playbook Overview and Structure
 
 This playbook creates an Azure Virtual Network (VNET) with the name "VNET_1"
@@ -83,6 +94,10 @@ the following parameters:
 `;
   // cSpell: enable
 
+  if (customPrompt) {
+    content += "\nCustom prompt explanation.";
+  }
+
   return res.send({
     content,
     format,

diff --git a/test/mockLightspeedServer/generations.ts b/test/mockLightspeedServer/generations.ts
@@ -8,6 +8,7 @@ export function generations(
   res: any,
 ) {
   const text = req.body.text;
+  const customPrompt = req.body.customPrompt;
   const createOutline = req.body.createOutline;
   const generationId = req.body.generationId ? req.body.generationId : uuidv4();
   const wizardId = req.body.wizardId;
@@ -46,6 +47,16 @@ export function generations(
     });
   }
 
+  // Special case to replicate a broken custom prompt
+  if (customPrompt && customPrompt === "custom prompt broken") {
+    logger.info("Returning 400. Custom prompt is invalid");
+    return res.status(400).send({
+      code: "validation",
+      message: "invalid",
+      detail: { customPrompt: "custom prompt is invalid" },
+    });
+  }
+
   // cSpell: disable
   let outline: string | undefined = `1. Create VNET named VNET_1
 2. Create VNET named VNET_2
@@ -98,6 +109,10 @@ export function generations(
     outline += "\n4. Some extra step.";
   }
 
+  if (customPrompt) {
+    outline += "\n5. Custom prompt step.";
+  }
+
   return res.send({
     playbook,
     outline,