Magic wand support

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,7 +11,7 @@ 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';
@@ -23,6 +23,9 @@ 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 +192,80 @@ 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 (magicContext: IMagicProvider.IMagicContext) => {
+        const { codeInput, cellId, content } = magicContext;
+        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. Do not format the response using backticks or code block delimiters, just give the code that will be inserted into the cell directly.';
+
+        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
+                }
+              },
+              {
+                name: 'show-diff',
+                args: {
+                  cell_id: cellId,
+                  original_source: codeInput
+                  // TODO
+                  // diff: {}
+                }
+              }
+            ]
+          }
+        });
+        // TODO:
+        console.log('MAGIC PROVIDER');
+        console.log(cellId, codeInput, content);
+      }
+    };
+  }
+};
+
+export default [
+  chatPlugin,
+  aiProviderPlugin,
+  autocompletionRegistryPlugin,
+  magicProviderPlugin
+];