feat: coding assistant UI updates for beta (#169)

feat: Add Ask Sourcery code lens
feat: Add initial chat message
refactor: drive showing messages in the UI from the binary

Author: Hellebore

package.json

@@ -229,6 +229,11 @@
 				"title": "Clear",
 				"category": "Sourcery"
 			},
+			{
+				"command": "sourcery.chat.ask",
+				"title": "Ask Sourcery",
+				"category": "Sourcery"
+			},
 			{
 				"command": "sourcery.scan.selectLanguage",
 				"title": "Select Language",
@@ -302,6 +307,11 @@
 					"command": "sourcery.scan.open",
 					"title": "Create Sourcery rule with this pattern",
 					"when": "editorLangId==python && editorHasSelection"
+				},
+				{
+					"command": "sourcery.chat.ask",
+					"title": "Ask Sourcery",
+					"when": "editorHasSelection && sourcery.features.coding_assistant"
 				}
 			],
 			"sourcery.scans": [
@@ -345,6 +355,10 @@
 					"command": "sourcery.chat.clearChat",
 					"when": "false"
 				},
+				{
+					"command": "sourcery.chat.ask",
+					"when": "sourcery.features.coding_assistant"
+				},
 				{
 					"command": "sourcery.scan.selectLanguage",
 					"when": "false"

src/ask-sourcery.ts

@@ -0,0 +1,60 @@
+import * as vscode from "vscode";
+import { Recipe, RecipeProvider } from "./recipes";
+import { ChatRequest } from "./chat";
+
+export function askSourceryCommand(recipes: Recipe[], contextRange?) {
+  showAskSourceryQuickPick(recipes).then((result: any) => {
+    vscode.commands.executeCommand("sourcery.chat.focus").then(() => {
+      let request: ChatRequest;
+      if ("id" in result) {
+        request = {
+          type: "recipe_request",
+          data: {
+            kind: "recipe_request",
+            name: result.label,
+            id: result.id,
+          },
+          context_range: contextRange,
+        };
+      } else {
+        request = {
+          type: "chat_request",
+          data: { kind: "user_message", message: result.label },
+          context_range: contextRange,
+        };
+      }
+
+      vscode.commands.executeCommand("sourcery.chat_request", request);
+    });
+  });
+}
+
+export function showAskSourceryQuickPick(recipes: Recipe[]) {
+  return new Promise((resolve) => {
+    const recipeNames = recipes.map((item) => item.name);
+    const recipeItems = recipes.map((item) => ({
+      label: item.name,
+      id: item.id,
+    }));
+
+    const quickPick = vscode.window.createQuickPick();
+    quickPick.placeholder = "Ask any question or choose one of these recipes";
+    quickPick.items = recipeItems;
+
+    quickPick.onDidAccept(() => {
+      const selection = quickPick.activeItems[0];
+      resolve(selection);
+      quickPick.hide();
+    });
+
+    quickPick.onDidChangeValue(() => {
+      // add what the user has typed to the pick list as the first item
+      if (!recipeNames.includes(quickPick.value)) {
+        const newItems = [{ label: quickPick.value }, ...recipeItems];
+        quickPick.items = newItems;
+      }
+    });
+    quickPick.onDidHide(() => quickPick.dispose());
+    quickPick.show();
+  });
+}

src/chat.ts

@@ -19,16 +19,35 @@ marked.use(
 enum ChatResultOutcome {
   Success = "success",
   Error = "error",
+  Finished = "finished",
+}
+
+enum ChatResultRole {
+  Assistant = "assistant",
+  User = "user",
 }
 
 type ChatResult = {
   outcome: ChatResultOutcome;
   textContent: string;
+  role: ChatResultRole;
+};
+
+export type ChatRequestData = {
+  kind: "user_message";
+  message: string;
+};
+
+export type RecipeRequestData = {
+  kind: "recipe_request";
+  name: string;
+  id: string;
 };
 
 export type ChatRequest = {
-  type: string;
-  data: string;
+  type: "recipe_request" | "chat_request";
+  data: ChatRequestData | RecipeRequestData;
+  context_range?: any;
 };
 
 export class ChatProvider implements vscode.WebviewViewProvider {
@@ -38,7 +57,9 @@ export class ChatProvider implements vscode.WebviewViewProvider {
 
   private _extensionUri: vscode.Uri;
 
-  private currentAssistantMessage: string = "";
+  private _currentAssistantMessage: string = "";
+
+  private _unhandledMessages: ChatResult[] = [];
 
   constructor(private _context: vscode.ExtensionContext) {
     this._extensionUri = _context.extensionUri;
@@ -71,50 +92,89 @@ export class ChatProvider implements vscode.WebviewViewProvider {
     webviewView.webview.onDidReceiveMessage(async (data: ChatRequest) => {
       switch (data.type) {
         case "chat_request": {
-          this.currentAssistantMessage = "";
           vscode.commands.executeCommand("sourcery.chat_request", data);
           break;
         }
       }
     });
+    if (this._unhandledMessages.length === 0) {
+      vscode.commands.executeCommand("sourcery.initialise_chat");
+    } else {
+      while (this._unhandledMessages.length > 0) {
+        const message = this._unhandledMessages.shift();
+        this.addResult(message);
+      }
+    }
   }
 
   public addResult(result: ChatResult) {
-    // Send the whole message we've been streamed so far to the webview,
-    // after converting from markdown to html
-
-    if (result.outcome === ChatResultOutcome.Success) {
-      this.currentAssistantMessage += result.textContent;
+    if (this._view) {
+      if (result.role === ChatResultRole.User) {
+        this.addUserResult(result);
+      } else {
+        this.addAssistantResult(result);
+      }
     } else {
-      this.currentAssistantMessage = result.textContent;
+      this._unhandledMessages.push(result);
     }
+  }
 
-    const rendered = marked(this.currentAssistantMessage, {
-      gfm: true,
-      breaks: true,
+  private addUserResult(result: ChatResult) {
+    this._view.webview.postMessage({
+      command: "add_result",
+      result: {
+        role: result.role,
+        outcome: result.outcome,
+        textContent: result.textContent,
+      },
     });
+  }
 
-    const sanitized = sanitizeHtml(rendered, {
-      allowedClasses: { span: false, code: false },
-    });
+  private addAssistantResult(result: ChatResult) {
+    if (result.outcome === ChatResultOutcome.Finished) {
+      this._currentAssistantMessage = "";
+      this._view.webview.postMessage({ command: "assistant_finished" });
+      return;
+    }
+
+    if (result.outcome === ChatResultOutcome.Success) {
+      this._currentAssistantMessage += result.textContent;
+    } else {
+      this._currentAssistantMessage = result.textContent;
+    }
+
+    let sanitized = this.renderAssistantMessage(this._currentAssistantMessage);
 
     this._view.webview.postMessage({
       command: "add_result",
-      result: { outcome: result.outcome, textContent: sanitized },
+      result: {
+        role: result.role,
+        outcome: result.outcome,
+        textContent: sanitized,
+      },
     });
   }
 
-  public executeRecipeRequest(message: string) {
-    this.currentAssistantMessage = "";
-    this._view.webview.postMessage({
-      command: "recipe_request",
-      result: message,
+  private renderAssistantMessage(message: string) {
+    // Send the whole message we've been streamed so far to the webview,
+    // after converting from markdown to html
+
+    const rendered = marked(message, {
+      gfm: true,
+      breaks: true,
+      mangle: false,
+      headerIds: false,
+    });
+
+    // Allow any classes on span and code blocks or highlightjs classes get removed
+    return sanitizeHtml(rendered, {
+      allowedClasses: { span: false, code: false },
     });
   }
 
   public clearChat() {
     this._view.webview.postMessage({ command: "clear_chat" });
-    this.currentAssistantMessage = "";
+    this._currentAssistantMessage = "";
   }
 
   private async _getHtmlForWebview(webview: vscode.Webview) {

src/extension.ts

@@ -27,12 +27,14 @@ import {
   LanguageClient,
   LanguageClientOptions,
   ServerOptions,
+  URI,
 } from "vscode-languageclient/node";
 import { getHubSrc } from "./hub";
 import { RuleInputProvider } from "./rule-search";
 import { ScanResultProvider } from "./rule-search-results";
 import { ChatProvider, ChatRequest } from "./chat";
-import { RecipeProvider } from "./recipes";
+import { Recipe, RecipeProvider } from "./recipes";
+import { askSourceryCommand } from "./ask-sourcery";
 
 function createLangServer(): LanguageClient {
   const token = workspace.getConfiguration("sourcery").get<string>("token");
@@ -172,7 +174,8 @@ function registerCommands(
   tree: ScanResultProvider,
   treeView: TreeView<TreeItem>,
   hubWebviewPanel: WebviewPanel,
-  chatProvider: ChatProvider
+  chatProvider: ChatProvider,
+  recipeProvider: RecipeProvider
 ) {
   context.subscriptions.push(
     vscode.window.registerWebviewViewProvider(
@@ -214,6 +217,13 @@ function registerCommands(
     })
   );
 
+  context.subscriptions.push(
+    commands.registerCommand("sourcery.chat.ask", (arg?) => {
+      let contextRange = arg && "start" in arg ? arg : null;
+      askSourceryCommand(recipeProvider.recipes, contextRange);
+    })
+  );
+
   context.subscriptions.push(
     commands.registerCommand("sourcery.scan.selectLanguage", () => {
       const items = ["python", "javascript"];
@@ -382,24 +392,30 @@ function registerCommands(
     )
   );
 
+  context.subscriptions.push(
+    commands.registerCommand("sourcery.initialise_chat", () => {
+      let request: ExecuteCommandParams = {
+        command: "sourcery/chat/initialise",
+        arguments: [],
+      };
+      languageClient.sendRequest(ExecuteCommandRequest.type, request);
+    })
+  );
+
   context.subscriptions.push(
     commands.registerCommand(
       "sourcery.chat_request",
       (message: ChatRequest) => {
-        const selectionLocation = getSelectionLocation();
-        const activeEditor = window.activeTextEditor;
-        let activeFile = undefined;
-        if (activeEditor) {
-          activeFile = activeEditor.document.uri;
-        }
-        const allFiles = [];
-        for (const tabGroup of vscode.window.tabGroups.all) {
-          for (const tab of tabGroup.tabs) {
-            if (tab.input instanceof vscode.TabInputText) {
-              allFiles.push(tab.input.uri);
-            }
-          }
+        // Use the editor selection unless a range was passed through in
+        // the message
+        let selectionLocation = getSelectionLocation();
+        if (message.context_range != null) {
+          selectionLocation = {
+            uri: selectionLocation.uri,
+            range: message.context_range,
+          };
         }
+        let { activeFile, allFiles } = activeFiles();
 
         let request: ExecuteCommandParams = {
           command: "sourcery/chat/request",
@@ -417,15 +433,6 @@ function registerCommands(
     )
   );
 
-  context.subscriptions.push(
-    commands.registerCommand(
-      "sourcery.recipe_request",
-      (message: ChatRequest) => {
-        chatProvider.executeRecipeRequest(message.data);
-      }
-    )
-  );
-
   context.subscriptions.push(
     commands.registerCommand(
       "sourcery.scan.rule",
@@ -561,7 +568,8 @@ export function activate(context: ExtensionContext) {
     tree,
     treeView,
     hubWebviewPanel,
-    chatProvider
+    chatProvider,
+    recipeProvider
   );
 
   showSourceryStatusBarItem(context);
@@ -588,3 +596,20 @@ function openDocument(document_path: string) {
     window.showTextDocument(doc);
   });
 }
+
+function activeFiles(): { activeFile: URI | null; allFiles: URI[] } {
+  const activeEditor = window.activeTextEditor;
+  let activeFile = null;
+  if (activeEditor) {
+    activeFile = activeEditor.document.uri;
+  }
+  const allFiles = [];
+  for (const tabGroup of vscode.window.tabGroups.all) {
+    for (const tab of tabGroup.tabs) {
+      if (tab.input instanceof vscode.TabInputText) {
+        allFiles.push(tab.input.uri);
+      }
+    }
+  }
+  return { activeFile, allFiles };
+}