Add the toolbar to websocket chat

jupyterlab · Jul 18, 2024 · 5bee9f2 · 5bee9f2
1 parent b1d557b
commit 5bee9f2
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 2 deletions.
diff --git a/packages/jupyterlab-ws-chat/package.json b/packages/jupyterlab-ws-chat/package.json
@@ -57,6 +57,7 @@
     "@jupyter/chat": "^0.2.0",
     "@jupyterlab/apputils": "^4.3.0",
     "@jupyterlab/coreutils": "^6.2.0",
+    "@jupyterlab/notebook": "^4.2.0",
     "@jupyterlab/rendermime": "^4.2.0",
     "@jupyterlab/services": "^7.2.0",
     "@jupyterlab/settingregistry": "^4.2.0",

diff --git a/packages/jupyterlab-ws-chat/schema/chat.json b/packages/jupyterlab-ws-chat/schema/chat.json
@@ -20,6 +20,12 @@
       "type": "boolean",
       "default": true,
       "readOnly": false
+    },
+    "enableCodeToolbar": {
+      "description": "Whether to enable or not the code toolbar.",
+      "type": "boolean",
+      "default": true,
+      "readOnly": false
     }
   },
   "additionalProperties": false

diff --git a/packages/jupyterlab-ws-chat/src/index.ts b/packages/jupyterlab-ws-chat/src/index.ts
@@ -4,6 +4,7 @@
  */
 
 import {
+  ActiveCellManager,
   AutocompletionRegistry,
   IAutocompletionRegistry,
   buildChatSidebar,
@@ -15,6 +16,7 @@ import {
   JupyterFrontEndPlugin
 } from '@jupyterlab/application';
 import { ReactWidget, IThemeManager } from '@jupyterlab/apputils';
+import { INotebookTracker } from '@jupyterlab/notebook';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import { ISettingRegistry } from '@jupyterlab/settingregistry';
 
@@ -49,6 +51,7 @@ const chat: JupyterFrontEndPlugin<void> = {
   optional: [
     IAutocompletionRegistry,
     ILayoutRestorer,
+    INotebookTracker,
     ISettingRegistry,
     IThemeManager
   ],
@@ -57,31 +60,44 @@ const chat: JupyterFrontEndPlugin<void> = {
     rmRegistry: IRenderMimeRegistry,
     autocompletionRegistry: IAutocompletionRegistry,
     restorer: ILayoutRestorer | null,
+    notebookTracker: INotebookTracker,
     settingsRegistry: ISettingRegistry | null,
     themeManager: IThemeManager | null
   ) => {
+    // Create an active cell manager for code toolbar.
+    const activeCellManager = new ActiveCellManager({
+      tracker: notebookTracker,
+      shell: app.shell
+    });
+
     /**
      * Initialize chat handler, open WS connection
      */
-    const chatHandler = new WebSocketHandler({ commands: app.commands });
+    const chatHandler = new WebSocketHandler({
+      commands: app.commands,
+      activeCellManager
+    });
 
     /**
      * Load the settings.
      */
     let sendWithShiftEnter = false;
     let stackMessages = true;
     let unreadNotifications = true;
+    let enableCodeToolbar = true;
     function loadSetting(setting: ISettingRegistry.ISettings): void {
       // Read the settings and convert to the correct type
       sendWithShiftEnter = setting.get('sendWithShiftEnter')
         .composite as boolean;
       stackMessages = setting.get('stackMessages').composite as boolean;
       unreadNotifications = setting.get('unreadNotifications')
         .composite as boolean;
+      enableCodeToolbar = setting.get('enableCodeToolbar').composite as boolean;
       chatHandler.config = {
         sendWithShiftEnter,
         stackMessages,
-        unreadNotifications
+        unreadNotifications,
+        enableCodeToolbar
       };
     }