Add support for lightspeed multi-task pattern match and model Id upda…

…tes (#979) * Add support for lightspeed multi-task pattern match and model Id updates * Trigger inline suggestion if the multi-task comment pattern is matched for a logged user that has lightspeed seat assignment. * Rename modelId to modelIdOverride and update the description * update logic to trigger multi-task suggestions
ansible · Sep 25, 2023 · 354cf6f · 354cf6f
1 parent 7b5fbbe
commit 354cf6f
Show file tree

Hide file tree

Showing 10 changed files with 360 additions and 84 deletions.
diff --git a/README.md b/README.md
@@ -225,8 +225,9 @@ any level (User, Remote, Workspace and/or Folder).
 - `ansible.lightspeed.URL`: URL for Ansible Lightspeed.
 - `ansible.lightspeed.suggestions.enabled`: Enable Ansible Lightspeed with
   watsonx Code Assistant inline suggestions.
-- `ansible.lightspeed.modelId`: The model to be used for inline suggestions.
-  This setting applies only to subscribed users.
+- `ansible.lightspeed.modelIdOverride`: Model ID to override your organization's
+  default model. This setting is only applicable to commercial users with an
+  Ansible Lightspeed seat assignment.
 
 ## Data and Telemetry
 

diff --git a/package.json b/package.json
@@ -452,11 +452,10 @@
             "markdownDescription": "Enable Ansible Lightspeed with watsonx Code Assistant inline suggestions.",
             "order": 2
           },
-          "ansible.lightspeed.modelId": {
+          "ansible.lightspeed.modelIdOverride": {
             "scope": "resource",
             "type": "string",
-            "markdownDescription": "The model to be used for inline suggestions. This setting applies only to subscribed users.",
-            "hidden": true,
+            "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": 3
           }
         }

diff --git a/src/definitions/lightspeed.ts b/src/definitions/lightspeed.ts
@@ -10,7 +10,7 @@ export interface CompletionRequestParams {
   prompt: string;
   suggestionId?: string;
   metadata?: MetadataParams;
-  modelId?: string;
+  modelIdOverride?: string;
 }
 
 export enum UserAction {
@@ -57,3 +57,21 @@ export const LIGHTSPEED_STATUS_BAR_CLICK_HANDLER =
 
 export const LIGHTSPEED_CLIENT_ID = "Vu2gClkeR5qUJTUGHoFAePmBznd6RZjDdy5FW2wy";
 export const LIGHTSPEED_SERVICE_LOGIN_TIMEOUT = 120000;
+
+export type LIGHTSPEED_SUGGESTION_TYPE = "SINGLE-TASK" | "MULTI-TASK";
+
+export const tasksInPlaybookKeywords = [
+  /(?<!\S)tasks\s*:(?!\S)\s*$/,
+  /(?<!\S)block\s*:(?!\S)\s*$/,
+  /(?<!\S)rescue\s*:(?!\S)\s*$/,
+  /(?<!\S)always\s*:(?!\S)\s*$/,
+  /(?<!\S)pre_tasks\s*:(?!\S)\s*$/,
+  /(?<!\S)post_tasks\s*:(?!\S)\s*$/,
+  /(?<!\S)handlers\s*:(?!\S)\s*$/,
+];
+
+export const tasksFileKeywords = [
+  /(?<!\S)block\s*:(?!\S)\s*$/,
+  /(?<!\S)rescue\s*:(?!\S)\s*$/,
+  /(?<!\S)always\s*:(?!\S)\s*$/,
+];
diff --git a/src/extension.ts b/src/extension.ts
@@ -232,7 +232,7 @@ export async function activate(context: ExtensionContext): Promise<void> {
         pythonInterpreterManager
       );
       if (editor) {
-        lightSpeedManager.ansibleContentFeedback(
+        await lightSpeedManager.ansibleContentFeedback(
           editor.document,
           AnsibleContentUploadTrigger.TAB_CHANGE
         );
@@ -250,8 +250,8 @@ export async function activate(context: ExtensionContext): Promise<void> {
       AnsibleContentUploadTrigger.FILE_OPEN
     );
   });
-  workspace.onDidCloseTextDocument((document: vscode.TextDocument) => {
-    lightSpeedManager.ansibleContentFeedback(
+  workspace.onDidCloseTextDocument(async (document: vscode.TextDocument) => {
+    await lightSpeedManager.ansibleContentFeedback(
       document,
       AnsibleContentUploadTrigger.FILE_CLOSE
     );

diff --git a/src/features/lightspeed/base.ts b/src/features/lightspeed/base.ts
@@ -118,10 +118,10 @@ export class LightSpeedManager {
       watchRolesDirectory(this, rolePath);
     }
   }
-  public ansibleContentFeedback(
+  public async ansibleContentFeedback(
     document: vscode.TextDocument,
     trigger: AnsibleContentUploadTrigger
-  ): void {
+  ): Promise<void> {
     if (
       document.languageId !== "ansible" ||
       !this.settingsManager.settings.lightSpeedService.enabled ||
@@ -130,6 +130,10 @@ export class LightSpeedManager {
       return;
     }
 
+    if (await this.lightSpeedAuthenticationProvider.rhUserHasSeat()) {
+      return;
+    }
+
     const currentFileContent = document.getText();
     const documentUri = document.uri.toString();
     let activityId: string;

diff --git a/src/features/lightspeed/inlineSuggestions.ts b/src/features/lightspeed/inlineSuggestions.ts
@@ -16,13 +16,17 @@ import {
   IRoleContext,
   IStandaloneTaskContext,
 } from "../../interfaces/lightspeed";
-import { UserAction } from "../../definitions/lightspeed";
+import {
+  LIGHTSPEED_SUGGESTION_TYPE,
+  UserAction,
+} from "../../definitions/lightspeed";
 import { LightSpeedCommands } from "../../definitions/lightspeed";
 import {
   getIncludeVarsContext,
   getRelativePath,
   getRolePathFromPathWithinRole,
   shouldRequestInlineSuggestions,
+  shouldTriggerMultiTaskSuggestion,
 } from "./utils/data";
 import { getVarsFilesContext } from "./utils/data";
 import {
@@ -33,9 +37,9 @@ import {
 import { getAnsibleFileType, getCustomRolePaths } from "../utils/ansible";
 import { watchRolesDirectory } from "./utils/watchers";
 
-const TASK_REGEX_EP =
+const SINGLE_TASK_REGEX_EP =
   /^(?<![\s-])(?<blank>\s*)(?<list>- \s*name\s*:\s*)(?<description>\S.*)(?<end>$)/;
-
+const MULTI_TASK_REGEX_EP = /^\s*#\s*\S+.*$/;
 let suggestionId = "";
 let currentSuggestion = "";
 let inlineSuggestionData: InlineSuggestionEvent = {};
@@ -134,8 +138,12 @@ export async function getInlineSuggestionItems(
   let result: CompletionResponseParams = {
     predictions: [],
   };
-  inlineSuggestionData = {};
-  suggestionId = "";
+  const range = new vscode.Range(new vscode.Position(0, 0), currentPosition);
+  const documentContent = range.isEmpty
+    ? ""
+    : document.getText(range).trimEnd();
+
+  let suggestionMatchType: LIGHTSPEED_SUGGESTION_TYPE | undefined = undefined;
 
   let rhUserHasSeat =
     await lightSpeedManager.lightSpeedAuthenticationProvider.rhUserHasSeat();
@@ -144,41 +152,116 @@ export async function getInlineSuggestionItems(
   }
 
   const lineToExtractPrompt = document.lineAt(currentPosition.line - 1);
-  const taskMatchedPattern = lineToExtractPrompt.text.match(TASK_REGEX_EP);
+  const taskMatchedPattern =
+    lineToExtractPrompt.text.match(SINGLE_TASK_REGEX_EP);
   const currentLineText = document.lineAt(currentPosition);
-  const spacesBeforeTaskNameStart =
-    lineToExtractPrompt?.text.match(/^ +/)?.[0].length || 0;
   const spacesBeforeCursor =
     currentLineText?.text.slice(0, currentPosition.character).match(/^ +/)?.[0]
       .length || 0;
 
+  if (taskMatchedPattern) {
+    suggestionMatchType = "SINGLE-TASK";
+  } else {
+    const commentMatchedPattern =
+      lineToExtractPrompt.text.match(MULTI_TASK_REGEX_EP);
+    if (commentMatchedPattern) {
+      suggestionMatchType = "MULTI-TASK";
+    }
+  }
+
+  const spacesBeforePromptStart =
+    lineToExtractPrompt?.text.match(/^ +/)?.[0].length || 0;
+
   if (
-    !taskMatchedPattern ||
+    !suggestionMatchType ||
     !currentLineText.isEmptyOrWhitespace ||
-    spacesBeforeTaskNameStart !== spacesBeforeCursor
+    spacesBeforePromptStart !== spacesBeforeCursor
   ) {
     resetInlineSuggestionDisplayed();
     // If the user has triggered the inline suggestion by pressing the configured keys,
     // we will show an information message to the user to help them understand the
     // correct cursor position to trigger the inline suggestion.
     if (context.triggerKind === vscode.InlineCompletionTriggerKind.Invoke) {
-      if (!taskMatchedPattern || !currentLineText.isEmptyOrWhitespace) {
-        vscode.window.showInformationMessage(
-          "Cursor should be positioned on the line after the task name with the same indent as that of the task name line to trigger an inline suggestion."
-        );
-      } else if (
-        taskMatchedPattern &&
-        currentLineText.isEmptyOrWhitespace &&
-        spacesBeforeTaskNameStart !== spacesBeforeCursor
+      if (rhUserHasSeat) {
+        if (!suggestionMatchType || !currentLineText.isEmptyOrWhitespace) {
+          vscode.window.showInformationMessage(
+            "Cursor should be positioned on the line after the task name or a comment line within task context to trigger an inline suggestion."
+          );
+        } else if (
+          suggestionMatchType &&
+          currentLineText.isEmptyOrWhitespace &&
+          spacesBeforePromptStart !== spacesBeforeCursor
+        ) {
+          vscode.window.showInformationMessage(
+            `Cursor must be in column ${spacesBeforePromptStart} to trigger an inline suggestion.`
+          );
+        }
+      } else {
+        if (!taskMatchedPattern || !currentLineText.isEmptyOrWhitespace) {
+          vscode.window.showInformationMessage(
+            "Cursor should be positioned on the line after the task name with the same indent as that of the task name line to trigger an inline suggestion."
+          );
+        } else if (
+          taskMatchedPattern &&
+          currentLineText.isEmptyOrWhitespace &&
+          spacesBeforePromptStart !== spacesBeforeCursor
+        ) {
+          vscode.window.showInformationMessage(
+            `Cursor must be in column ${spacesBeforePromptStart} to trigger an inline suggestion.`
+          );
+        }
+      }
+    }
+    return [];
+  }
+
+  let parsedAnsibleDocument = undefined;
+  const documentUri = document.uri.toString();
+  const documentDirPath = pathUri.dirname(URI.parse(documentUri).path);
+  const documentFilePath = URI.parse(documentUri).path;
+  const ansibleFileType: IAnsibleFileType = getAnsibleFileType(
+    documentFilePath,
+    parsedAnsibleDocument
+  );
+
+  if (suggestionMatchType === "MULTI-TASK") {
+    if (!rhUserHasSeat) {
+      console.debug(
+        "[inline-suggestions] Multitask suggestions not supported for a non seat user."
+      );
+      return [];
+    } else {
+      if (
+        !shouldTriggerMultiTaskSuggestion(
+          documentContent,
+          spacesBeforePromptStart,
+          ansibleFileType
+        )
       ) {
-        vscode.window.showInformationMessage(
-          `Cursor must be in column ${spacesBeforeTaskNameStart} to trigger an inline suggestion.`
-        );
+        return [];
       }
     }
+  }
+
+  try {
+    parsedAnsibleDocument = yaml.parse(documentContent, {
+      keepSourceTokens: true,
+    });
+  } catch (err) {
+    vscode.window.showErrorMessage(
+      `Ansible Lightspeed expects valid YAML syntax to provide inline suggestions. Error: ${err}`
+    );
+    return [];
+  }
+  if (
+    suggestionMatchType === "SINGLE-TASK" &&
+    !shouldRequestInlineSuggestions(parsedAnsibleDocument)
+  ) {
     return [];
   }
+
   inlineSuggestionData = {};
+  suggestionId = "";
   inlineSuggestionDisplayTime = getCurrentUTCDateTime();
   const requestTime = getCurrentUTCDateTime();
 
@@ -187,7 +270,6 @@ export async function getInlineSuggestionItems(
   );
   try {
     suggestionId = uuidv4();
-    const documentUri = document.uri.toString();
     let activityId: string | undefined = undefined;
     inlineSuggestionData["suggestionId"] = suggestionId;
     inlineSuggestionData["documentUri"] = documentUri;
@@ -203,48 +285,18 @@ export async function getInlineSuggestionItems(
         lightSpeedManager.lightSpeedActivityTracker[documentUri].activityId;
     }
     inlineSuggestionData["activityId"] = activityId;
-    const range = new vscode.Range(new vscode.Position(0, 0), currentPosition);
-
-    const documentContent = range.isEmpty
-      ? ""
-      : document.getText(range).trimEnd();
-
-    let parsedAnsibleDocument = undefined;
-    try {
-      parsedAnsibleDocument = yaml.parse(documentContent, {
-        keepSourceTokens: true,
-      });
-      if (!parsedAnsibleDocument) {
-        return [];
-      }
-      // check if YAML is a list, if not it is not a valid Ansible document
-      if (
-        typeof parsedAnsibleDocument === "object" &&
-        !Array.isArray(parsedAnsibleDocument)
-      ) {
-        vscode.window.showErrorMessage(
-          "Ansible Lightspeed expects valid Ansible syntax. For playbook files it should be a list of plays and for tasks files it should be list of tasks."
-        );
-        return [];
-      }
-    } catch (err) {
-      vscode.window.showErrorMessage(
-        `Ansible Lightspeed expects valid YAML syntax to provide inline suggestions. Error: ${err}`
-      );
-      return [];
-    }
 
-    if (!shouldRequestInlineSuggestions(parsedAnsibleDocument)) {
-      return [];
-    }
     lightSpeedManager.statusBarProvider.statusBar.text =
       "$(loading~spin) Lightspeed";
     result = await requestInlineSuggest(
       documentContent,
       parsedAnsibleDocument,
       documentUri,
       activityId,
-      rhUserHasSeat
+      rhUserHasSeat,
+      documentDirPath,
+      documentFilePath,
+      ansibleFileType
     );
     lightSpeedManager.statusBarProvider.statusBar.text = "Lightspeed";
   } catch (error) {
@@ -304,15 +356,11 @@ async function requestInlineSuggest(
   parsedAnsibleDocument: any,
   documentUri: string,
   activityId: string,
-  rhUserHasSeat: boolean
+  rhUserHasSeat: boolean,
+  documentDirPath: string,
+  documentFilePath: string,
+  ansibleFileType: IAnsibleFileType
 ): Promise<CompletionResponseParams> {
-  const documentDirPath = pathUri.dirname(URI.parse(documentUri).path);
-  const documentFilePath = URI.parse(documentUri).path;
-  const ansibleFileType: IAnsibleFileType = getAnsibleFileType(
-    documentFilePath,
-    parsedAnsibleDocument
-  );
-
   const hash = crypto.createHash("sha256").update(documentUri).digest("hex");
   const completionData: CompletionRequestParams = {
     prompt: content,
@@ -323,13 +371,15 @@ async function requestInlineSuggest(
       activityId: activityId,
     },
   };
-  if (rhUserHasSeat) {
-    const modelId =
-      lightSpeedManager.settingsManager.settings.lightSpeedService.modelId;
-    if (modelId && modelId !== "") {
-      completionData.modelId = modelId;
-    }
 
+  const modelIdOverride =
+    lightSpeedManager.settingsManager.settings.lightSpeedService
+      .modelIdOverride;
+  if (modelIdOverride && modelIdOverride !== "") {
+    completionData.model_name = modelIdOverride;
+  }
+
+  if (rhUserHasSeat) {
     const additionalContext = getAdditionalContext(
       parsedAnsibleDocument,
       documentDirPath,