Integrate with the magic-wand extension

Author: jtpio

package.json

@@ -66,7 +66,8 @@
         "@langchain/mistralai": "^0.1.1",
         "@lumino/coreutils": "^2.1.2",
         "@lumino/polling": "^2.1.2",
-        "@lumino/signaling": "^2.1.2"
+        "@lumino/signaling": "^2.1.2",
+        "jupyterlab_magic_wand": "*"
     },
     "devDependencies": {
         "@jupyterlab/builder": "^4.0.0",
@@ -104,7 +105,16 @@
     "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
+            }
+        }
     },
     "eslintIgnore": [
         "node_modules",
@@ -196,5 +206,8 @@
             "selector-no-vendor-prefix": null,
             "value-no-vendor-prefix": null
         }
+    },
+    "resolutions": {
+        "jupyterlab_magic_wand": "file:/Users/jtp/git/jtpio/jupyterlab-magic-wand"
     }
 }

src/index.ts

@@ -19,6 +19,10 @@ import { getSettings } from './llm-models';
 import { AIProvider } from './provider';
 import { IAIProvider } from './token';
 
+import { IMagicProvider } from 'jupyterlab_magic_wand';
+import { PartialJSONValue } from '@lumino/coreutils';
+import { HumanMessage, SystemMessage } from '@langchain/core/messages';
+
 const chatPlugin: JupyterFrontEndPlugin<void> = {
   id: '@jupyterlite/ai:chat',
   description: 'LLM chat extension',
@@ -151,4 +155,69 @@ const aiProviderPlugin: JupyterFrontEndPlugin<IAIProvider> = {
   }
 };
 
-export default [chatPlugin, 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, magicProviderPlugin];

tsconfig.json

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