Sample for using Jupyter API to run Python code with callbacks

jupyter-kernel-chat-execution-sample/.eslintrc.js

@@ -0,0 +1,23 @@
+/**@type {import('eslint').Linter.Config} */
+// eslint-disable-next-line no-undef
+module.exports = {
+	root: true,
+	parser: '@typescript-eslint/parser',
+	plugins: [
+		'@typescript-eslint',
+	],
+	extends: [
+		'eslint:recommended',
+		'plugin:@typescript-eslint/recommended',
+	],
+	rules: {
+		'semi': [
+			2,
+			"always"
+		],
+		'@typescript-eslint/no-unused-vars': 0,
+		'@typescript-eslint/no-explicit-any': 0,
+		'@typescript-eslint/explicit-module-boundary-types': 0,
+		'@typescript-eslint/no-non-null-assertion': 0,
+	}
+};

jupyter-kernel-chat-execution-sample/.vscode/extensions.json

@@ -0,0 +1,6 @@
+{
+  "recommendations": [
+    "dbaeumer.vscode-eslint",
+	"ms-toolsai.jupyter"
+  ]
+}

jupyter-kernel-chat-execution-sample/.vscode/launch.json

@@ -0,0 +1,22 @@
+// A launch configuration that compiles the extension and then opens it inside a new window
+// Use IntelliSense to learn about possible attributes.
+// Hover to view descriptions of existing attributes.
+// For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+{
+  "configurations": [
+    {
+      "args": [
+        "--extensionDevelopmentPath=${workspaceFolder}"
+      ],
+      "name": "Run Extension",
+      "outFiles": [
+        "${workspaceFolder}/out/**/*.js"
+      ],
+      "preLaunchTask": "npm: compile",
+      "request": "launch",
+      "trace": true,
+      "type": "extensionHost"
+    }
+  ],
+  "version": "0.2.0"
+}

jupyter-kernel-chat-execution-sample/.vscode/settings.json

@@ -0,0 +1,6 @@
+{
+  "search.exclude": {
+    "out": true
+  },
+  "typescript.tsc.autoDetect": "off"
+}

jupyter-kernel-chat-execution-sample/.vscode/tasks.json

@@ -0,0 +1,13 @@
+// See https://go.microsoft.com/fwlink/?LinkId=733558
+// for the documentation about the tasks.json format
+{
+	"version": "2.0.0",
+	"tasks": [
+		{
+			"type": "npm",
+			"script": "compile",
+			"problemMatcher": ["$tsc"],
+			"group": "build"
+		}
+	]
+}

jupyter-kernel-chat-execution-sample/.vscodeignore

@@ -0,0 +1,12 @@
+.vscode/**
+.vscode-test/**
+out/test/**
+src/**
+.gitignore
+.yarnrc
+vsc-extension-quickstart.md
+**/tsconfig.json
+**/.eslintrc.json
+**/*.map
+**/*.ts
+**/*.tsbuildinfo

jupyter-kernel-chat-execution-sample/README.md

@@ -0,0 +1,28 @@
+# Jupyter Server Provider Sample
+
+This is a very simple extension sample demonstrating the use of the Jupyter Extension API allowing other extensions to execute code against Jupyter Kernels.
+
+- The sample lists finds kernels associated with notebooks that are currently open in the workspace.
+- The sample the filters the kernels by language, focusing on Python kernels.
+- Upon selecting a Python kernel, code selected by the user is executed against the selected kernel
+- The output is displayed in an output panel.
+- The sample demonstrates the ability to retrieve outputs of various mime types, including streamed output.
+
+## Running this sample
+
+ 1. `cd jupyter-kernel-execution-sample`
+ 1. `code .`: Open the folder in VS Code
+ 1. Run `npm install` in terminal to install the dependencies
+ 1. Run the `Run Extension` target in the Debug View. This will:
+	- Start a task `npm: watch` to compile the code
+	- Run the extension in a new VS Code window
+ 1. Open a Jupyter Notebook and select a Python kernel and execute some code.
+ 1. Select the command `Jupyter Kernel API: Execute code against a Python Kernel`
+ 1. Select the a Kernel and then select the Code to execute.
+ 1. Watch the output panel for outputs returned by the kernel.
+
+### Notes:
+
+1. Make use of the `language` property of the kernel to ensure the language of the code matches the kernel.
+2. `getKernel` API can can return `undefined` if the user does not grant the extension access to the kernel.
+3. Access to kernels for each extension is managed via the command `Manage Access To Jupyter Kernels`.

jupyter-kernel-chat-execution-sample/package.json

@@ -0,0 +1,53 @@
+{
+	"name": "jupyter-kernel-execution-sample",
+	"displayName": "Jupyter Kernel Execution Sample",
+	"description": "Sample extension using Jupyter API to execute code against the Python Kernel",
+	"publisher": "vscode-samples",
+	"version": "0.0.1",
+	"engines": {
+		"vscode": "^1.87.0"
+	},
+	"categories": [
+		"Other"
+	],
+	"activationEvents": [
+		"onNotebook:jupyter-notebook",
+		"onNotebook:interactive"
+	],
+	"extensionDependencies": [
+		"ms-toolsai.jupyter"
+	],
+	"main": "./out/extension.js",
+	"contributes": {
+		"commands": [
+			{
+				"command": "jupyterKernelExecution.executePythonCode",
+				"title": "Execute code against a Python Kernel",
+				"category": "Jupyter Kernel API"
+			}
+		],
+		"menus": {
+			"commandPalette": [
+				{
+					"command": "jupyterKernelExecution.executePythonCode",
+					"title": "Execute code against a Python Kernel"
+				}
+			]
+		}
+	},
+	"scripts": {
+		"vscode:prepublish": "npm run compile",
+		"compile": "tsc -b",
+		"lint": "eslint src --ext ts",
+		"watch": "tsc -b --watch"
+	},
+	"devDependencies": {
+		"@types/node": "14.x",
+		"@types/vscode": "^1.82.0",
+		"@typescript-eslint/eslint-plugin": "^6.7.0",
+		"@typescript-eslint/parser": "^6.7.0",
+		"@vscode/jupyter-extension": "^1.1.0",
+		"eslint": "^7.27.0",
+		"typescript": "^5.2.2"
+	}
+}

jupyter-kernel-chat-execution-sample/src/extension.ts

@@ -0,0 +1,127 @@
+import {
+	CancellationTokenSource,
+	CancellationToken,
+	Disposable,
+	ExtensionContext,
+	NotebookCellOutputItem,
+	OutputChannel,
+	QuickPickItem,
+	commands,
+	extensions,
+	window,
+	workspace,
+} from 'vscode';
+import { Jupyter, Kernel } from '@vscode/jupyter-extension';
+import path = require('path');
+import { TextDecoder } from 'util';
+
+export function activate(context: ExtensionContext) {
+	const jupyterExt = extensions.getExtension<Jupyter>('ms-toolsai.jupyter');
+	if (!jupyterExt) {
+		throw new Error('Jupyter Extension not installed');
+	}
+	if (!jupyterExt.isActive) {
+		jupyterExt.activate();
+	}
+	const output = window.createOutputChannel('Jupyter Kernel Execution');
+	context.subscriptions.push(output);
+	context.subscriptions.push(
+		commands.registerCommand('jupyterKernelExecution.executePythonCode', async () => {
+			const notebook = window.activeNotebookEditor?.notebook;
+			if (!notebook) {
+				window.showErrorMessage('No active notebook');
+				return;
+			}
+			const kernel = await jupyterExt.exports.kernels.getKernel(notebook.uri);
+			if (!kernel) {
+				window.showErrorMessage('No active Kernel');
+				return;
+			}
+			if (kernel.language !== 'python') {
+				window.showErrorMessage('Please select a Python kernel');
+				return;
+			}
+			const token = new CancellationTokenSource();
+			const result = await executePythonCode(kernel, token.token).finally(() =>
+				token.dispose()
+			);
+			if (result) {
+				window.showInformationMessage(
+					`Result Status = ${result.status}, Output = ${result.output}`
+				);
+			} else {
+				window.showErrorMessage('No result');
+			}
+		})
+	);
+}
+
+const textMimes = [NotebookCellOutputItem.text(''), NotebookCellOutputItem.stdout('')];
+const errorMimes = [
+	NotebookCellOutputItem.error(new Error('')),
+	NotebookCellOutputItem.stderr(''),
+];
+
+type ProcessingResult = {
+	status: 'ok' | 'error';
+	output: string;
+};
+const textDecoder = new TextDecoder();
+async function executePythonCode(
+	kernel: Kernel,
+	token: CancellationToken
+): Promise<ProcessingResult | undefined> {
+	const code = `
+from vscode import chat
+
+def step1():
+    print("Started Step1")
+    chat.send_message("generatePandaSummary", {"name": "df", "shape": [10, 5]}, step2)
+
+def step2(data):
+    print(f"Inside Step2, got data {data['summary']}")
+    chat.send_message("generatePlot", {"name": "df", "type": "histogram", }, step3)
+
+def step3(data):
+    print(f"Inside Step3, got data {data['output']}")
+    import IPython.display
+    display({"application/vnd.custom.extension.message": {"status":"ok", "output": data["output"]}}, raw=True)
+
+step1()
+`;
+
+	const handlers: Record<string, (...data: any[]) => Promise<any>> = {
+		generatePandaSummary: async (_data: {
+			df: string;
+			shape: [rows: number, columns: number];
+		}) => {
+			await new Promise((resolve) => setTimeout(resolve, 1000));
+			return { summary: 'Panda Summary' };
+		},
+		generatePlot: async (_data: {
+			df: string;
+			type: 'histogram' | 'line' | 'scatter';
+		}) => {
+			await new Promise((resolve) => setTimeout(resolve, 1000));
+			return { output: 'df.plot(...)' };
+		},
+	};
+	for await (const output of kernel.executeChatCode(code, handlers, token)) {
+		output.items.forEach((item) => {
+			if (textMimes.some((mime) => mime.mime === item.mime)) {
+				console.log(textDecoder.decode(item.data));
+			}
+			if (errorMimes.some((mime) => mime.mime === item.mime)) {
+				const errorMessage = textDecoder.decode(item.data);
+				console.error(errorMessage);
+				throw new Error(errorMessage);
+			}
+		});
+		const result = output.items.find(
+			(item) => item.mime === 'application/vnd.custom.extension.message'
+		);
+		if (result) {
+			return JSON.parse(textDecoder.decode(result.data));
+		}
+	}
+}

jupyter-kernel-chat-execution-sample/src/types.d.ts

@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { CancellationToken } from 'vscode';
+
+declare module '@vscode/jupyter-extension' {
+    /**
+     * Represents a Jupyter Kernel.
+     */
+    export interface Kernel {
+        /**
+         * Executes code specific to chat in the Python kernel without affecting the execution count & execution history.
+         * Supports the ability to handle messages from the kernel.
+         *
+         * @param code Code to be executed.
+         * @param handlers Callbacks for handling messages from the kernel.
+         * @param token Triggers the cancellation of the execution.
+         * @returns Async iterable of outputs, that completes when the execution is complete.
+         */
+        executeChatCode(
+            code: string,
+            handlers: Record<string, (...data: any[]) => Promise<any>>,
+            token: CancellationToken
+        ): AsyncIterable<Output>;
+    }
+}

jupyter-kernel-chat-execution-sample/tsconfig.json

@@ -0,0 +1,17 @@
+{
+	"compilerOptions": {
+		"module": "commonjs",
+		"target": "es2020",
+		"lib": [
+			"es2020"
+		],
+		"outDir": "out",
+		"sourceMap": true,
+		"strict": true,
+		"rootDir": "src"
+	},
+	"exclude": [
+		"node_modules",
+		".vscode-test"
+	]
+}