Integrate with the magic-wand extension

Author: jtpio

package.json

@@ -70,6 +70,7 @@
         "@lumino/signaling": "^2.1.2",
         "@mui/icons-material": "^5.11.0",
         "@mui/material": "^5.11.0",
+        "jupyterlab_magic_wand": "^0.4.0",
         "react": "^18.2.0",
         "react-dom": "^18.2.0"
     },
@@ -110,6 +111,15 @@
     "jupyterlab": {
         "extension": true,
         "outputDir": "jupyterlite_ai/labextension",
-        "schemaDir": "schema"
+        "schemaDir": "schema",
+        "disabledExtensions": [
+            "jupyterlab_magic_wand:magic-provider"
+        ],
+        "sharedPackages": {
+            "jupyterlab_magic_wand": {
+                "bundled": false,
+                "singleton": true
+            }
+        }
     }
 }

src/index.ts

@@ -11,18 +11,22 @@ import {
   JupyterFrontEnd,
   JupyterFrontEndPlugin
 } from '@jupyterlab/application';
-import { ReactWidget, IThemeManager } from '@jupyterlab/apputils';
+import { IThemeManager, ReactWidget } from '@jupyterlab/apputils';
 import { ICompletionProviderManager } from '@jupyterlab/completer';
 import { INotebookTracker } from '@jupyterlab/notebook';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import { ISettingRegistry } from '@jupyterlab/settingregistry';
 
+import { PartialJSONValue } from '@lumino/coreutils';
 import { ChatHandler } from './chat-handler';
 import { getSettings } from './llm-models';
 import { AIProvider } from './provider';
 import { renderSlashCommandOption } from './slash-commands';
 import { IAIProvider } from './token';
 
+import { HumanMessage, SystemMessage } from '@langchain/core/messages';
+import { IMagicProvider } from 'jupyterlab_magic_wand';
+
 const autocompletionRegistryPlugin: JupyterFrontEndPlugin<IAutocompletionRegistry> =
   {
     id: '@jupyterlite/ai:autocompletion-registry',
@@ -189,4 +193,74 @@ const aiProviderPlugin: JupyterFrontEndPlugin<IAIProvider> = {
   }
 };
 
-export default [chatPlugin, autocompletionRegistryPlugin, aiProviderPlugin];
+const magicProviderPlugin: JupyterFrontEndPlugin<IMagicProvider> = {
+  id: '@jupyterlite/ai:magic-provider',
+  autoStart: true,
+  requires: [IAIProvider],
+  provides: IMagicProvider,
+  activate: (app: JupyterFrontEnd, aiProvider: IAIProvider): IMagicProvider => {
+    console.log('@jupyterlite/ai magic provider plugin activated');
+    const events = app.serviceManager.events;
+
+    return {
+      magic: async (
+        cellId: string,
+        codeInput: string,
+        content: PartialJSONValue | undefined
+      ) => {
+        const trimmedPrompt = codeInput.trim();
+
+        // TODO: taken from jupyterlab-magic-wand
+        const PROMPT =
+          'The input below came from a code cell in Jupyter. If the input does not look like code, but instead a prompt, write code based on the prompt. Then, update the code to make it more efficient, add code comments, and respond with only the code and comments. ';
+
+        const messages = [
+          new SystemMessage(PROMPT),
+          new HumanMessage(trimmedPrompt)
+        ];
+
+        const response = await aiProvider.chatModel?.invoke(messages);
+        if (!response) {
+          return;
+        }
+
+        const source = response.content;
+
+        events.emit({
+          schema_id: 'https://events.jupyter.org/jupyter_ai/magic_button/v1',
+          version: '1',
+          data: {
+            agent: 'Magic Button Agent',
+            input: codeInput,
+            // @ts-expect-error: TODO
+            context: {
+              cell_id: cellId,
+              content
+            },
+            messages: [source],
+            commands: [
+              {
+                name: 'update-cell-source',
+                args: {
+                  cell_id: cellId,
+                  cell_type: 'code',
+                  source: source
+                }
+              }
+            ]
+          }
+        });
+        // TODO:
+        console.log('MAGIC PROVIDER');
+        console.log(cellId, codeInput, content);
+      }
+    };
+  }
+};
+
+export default [
+  chatPlugin,
+  aiProviderPlugin,
+  autocompletionRegistryPlugin,
+  magicProviderPlugin
+];

tsconfig.json

@@ -16,6 +16,7 @@
     "outDir": "lib",
     "rootDir": "src",
     "strict": true,
+    "skipLibCheck": true,
     "strictNullChecks": true,
     "target": "ES2018"
   },