Display advice in empty editor to use Mito AI

mito-ai/package.json

@@ -57,8 +57,15 @@
     "watch:labextension": "jupyter labextension watch ."
   },
   "dependencies": {
-    "@jupyterlab/application": "^4.2.4",
-    "@jupyterlab/apputils": "^4.3.4",
+    "@codemirror/state": "^6.2.0",
+    "@codemirror/view": "^6.9.6",
+    "@jupyterlab/application": "^4.1.0",
+    "@jupyterlab/apputils": "^4.2.0",
+    "@jupyterlab/cells": "^4.1.0",
+    "@jupyterlab/codeeditor": "^4.1.0",
+    "@jupyterlab/codemirror": "^4.1.0",
+    "@lumino/commands": "^2.0.0",
+    "@lumino/coreutils": "^2.0.0",
     "@lumino/widgets": "^2.5.0",
     "@types/react": "^18.0.26",
     "@types/react-dom": "^18.0.10",

mito-ai/src/Extensions/AiChat/AiChatPlugin.ts

@@ -1,103 +1,117 @@
 import {
   ILayoutRestorer,
   JupyterFrontEnd,
-  JupyterFrontEndPlugin,
+  JupyterFrontEndPlugin
 } from '@jupyterlab/application';
 import { ICommandPalette, WidgetTracker } from '@jupyterlab/apputils';
 import { INotebookTracker } from '@jupyterlab/notebook';
 import { buildChatWidget } from './ChatWidget';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import { IVariableManager } from '../VariableManager/VariableManagerPlugin';
 import { COMMAND_MITO_AI_OPEN_CHAT } from '../../commands';
-
+import { IChatTracker } from './token';
 
 /**
  * Initialization data for the mito-ai extension.
  */
-const AiChatPlugin: JupyterFrontEndPlugin<void> = {
+const AiChatPlugin: JupyterFrontEndPlugin<WidgetTracker> = {
   id: 'mito_ai:plugin',
   description: 'AI chat for JupyterLab',
   autoStart: true,
-  requires: [INotebookTracker, ICommandPalette, IRenderMimeRegistry, IVariableManager],
+  requires: [
+    INotebookTracker,
+    ICommandPalette,
+    IRenderMimeRegistry,
+    IVariableManager
+  ],
   optional: [ILayoutRestorer],
+  provides: IChatTracker,
   activate: (
-    app: JupyterFrontEnd, 
-    notebookTracker: INotebookTracker, 
-    palette: ICommandPalette, 
+    app: JupyterFrontEnd,
+    notebookTracker: INotebookTracker,
+    palette: ICommandPalette,
     rendermime: IRenderMimeRegistry,
     variableManager: IVariableManager,
     restorer: ILayoutRestorer | null
   ) => {
-
     // Define a widget creator function,
     // then call it to make a new widget
     const newWidget = () => {
       // Create a blank content widget inside of a MainAreaWidget
-      const chatWidget = buildChatWidget(app, notebookTracker, rendermime, variableManager)
-      return chatWidget
-    }
+      const chatWidget = buildChatWidget(
+        app,
+        notebookTracker,
+        rendermime,
+        variableManager
+      );
+      return chatWidget;
+    };
 
     let widget = newWidget();
 
     // Add an application command
     app.commands.addCommand(COMMAND_MITO_AI_OPEN_CHAT, {
       label: 'Your friendly Python Expert chat bot',
       execute: () => {
-
         // In order for the widget to be accessible, the widget must be:
         // 1. Created
         // 2. Added to the widget tracker
-        // 3. Attatched to the frontend 
+        // 3. Attatched to the frontend
 
         // Step 1: Create the widget if its not already created
         if (!widget || widget.isDisposed) {
           widget = newWidget();
         }
 
-        // Step 2: Add the widget to the widget tracker if 
+        // Step 2: Add the widget to the widget tracker if
         // its not already there
         if (!tracker.has(widget)) {
           tracker.add(widget);
         }
 
-        // Step 3: Attatch the widget to the app if its not 
+        // Step 3: Attatch the widget to the app if its not
         // already there
         if (!widget.isAttached) {
           app.shell.add(widget, 'left', { rank: 2000 });
         }
-         
-        // Now that the widget is potentially accessible, activating the 
+
+        // Now that the widget is potentially accessible, activating the
         // widget opens the taskpane
         app.shell.activateById(widget.id);
-        
+
         // Set focus on the chat input
-        const chatInput: HTMLTextAreaElement | null = widget.node.querySelector('.chat-input');
+        const chatInput: HTMLTextAreaElement | null =
+          widget.node.querySelector('.chat-input');
         chatInput?.focus();
       }
     });
 
     app.commands.addKeyBinding({
       command: COMMAND_MITO_AI_OPEN_CHAT,
       keys: ['Accel E'],
-      selector: 'body',
+      selector: 'body'
     });
 
     app.shell.add(widget, 'left', { rank: 2000 });
 
     // Add the command to the palette.
-    palette.addItem({ command: COMMAND_MITO_AI_OPEN_CHAT, category: 'AI Chat' });
+    palette.addItem({
+      command: COMMAND_MITO_AI_OPEN_CHAT,
+      category: 'AI Chat'
+    });
 
     // Track and restore the widget state
-    let tracker = new WidgetTracker({
+    const tracker = new WidgetTracker({
       namespace: widget.id
     });
 
     if (restorer) {
       restorer.add(widget, 'mito_ai');
     }
+
+    // This allows us to force plugin load order
+    return tracker;
   }
 };
 
 export default AiChatPlugin;
-
-

mito-ai/src/Extensions/AiChat/token.ts

@@ -0,0 +1,7 @@
+import type { WidgetTracker } from '@jupyterlab/apputils';
+import { Token } from '@lumino/coreutils';
+
+export const IChatTracker = new Token<WidgetTracker>(
+  'mito-ai/IChatTracker',
+  'Widget tracker for the chat sidebar.'
+);

mito-ai/src/Extensions/emptyCell/emptyCell.ts

@@ -0,0 +1,93 @@
+import { Facet, type Extension } from '@codemirror/state';
+import {
+  Decoration,
+  EditorView,
+  ViewPlugin,
+  WidgetType,
+  type DecorationSet,
+  type ViewUpdate
+} from '@codemirror/view';
+
+/**
+ * Theme for the advice widget.
+ */
+const adviceTheme = EditorView.baseTheme({
+  '& .cm-mito-advice': {
+    color: 'var(--jp-ui-font-color3)',
+    fontStyle: 'italic'
+  },
+  '& .cm-mito-advice > kbd': {
+    borderRadius: '3px',
+    borderStyle: 'solid',
+    borderWidth: '1px',
+    fontSize: 'calc(var(--jp-code-font-size) - 2px)',
+    padding: '3px 5px',
+    verticalAlign: 'middle'
+  }
+});
+
+/**
+ * A facet that stores the chat shortcut.
+ */
+const adviceText = Facet.define<string, string>({
+  combine: values => (values.length ? values[values.length - 1] : '')
+});
+
+class AdviceWidget extends WidgetType {
+  constructor(readonly advice: string) {
+    super();
+  }
+
+  eq(other: AdviceWidget) {
+    return false;
+  }
+
+  toDOM() {
+    const wrap = document.createElement('span');
+    wrap.className = 'cm-mito-advice';
+    wrap.innerHTML = this.advice;
+    return wrap;
+  }
+}
+
+function shouldDisplayAdvice(view: EditorView): DecorationSet {
+  const shortcut = view.state.facet(adviceText);
+  const widgets = [];
+  if (view.hasFocus && view.state.doc.length == 0) {
+    const deco = Decoration.widget({
+      widget: new AdviceWidget(shortcut),
+      side: 1
+    });
+    widgets.push(deco.range(0));
+  }
+  return Decoration.set(widgets);
+}
+
+export const showAdvice = ViewPlugin.fromClass(
+  class {
+    decorations: DecorationSet;
+
+    constructor(view: EditorView) {
+      this.decorations = shouldDisplayAdvice(view);
+    }
+
+    update(update: ViewUpdate) {
+      const advice = update.view.state.facet(adviceText);
+      if (
+        update.docChanged ||
+        update.focusChanged ||
+        advice !== update.startState.facet(adviceText)
+      ) {
+        this.decorations = shouldDisplayAdvice(update.view);
+      }
+    }
+  },
+  {
+    decorations: v => v.decorations,
+    provide: () => [adviceTheme]
+  }
+);
+
+export function advicePlugin(options: { advice?: string } = {}): Extension {
+  return [adviceText.of(options.advice ?? ''), showAdvice];
+}

mito-ai/src/Extensions/emptyCell/index.ts

@@ -0,0 +1,72 @@
+import type {
+  JupyterFrontEnd,
+  JupyterFrontEndPlugin
+} from '@jupyterlab/application';
+import type { WidgetTracker } from '@jupyterlab/apputils';
+import type { ICellModel } from '@jupyterlab/cells';
+import { IEditorMimeTypeService } from '@jupyterlab/codeeditor';
+import {
+  EditorExtensionRegistry,
+  IEditorExtensionRegistry
+} from '@jupyterlab/codemirror';
+import { CommandRegistry } from '@lumino/commands';
+import { COMMAND_MITO_AI_OPEN_CHAT } from '../../commands';
+import { IChatTracker } from '../AiChat/token';
+import { advicePlugin } from './emptyCell';
+
+export const localPrompt: JupyterFrontEndPlugin<void> = {
+  id: 'mito-ai:empty-editor-advice',
+  description: 'Display a default message in an empty code cell.',
+  autoStart: true,
+  requires: [IChatTracker, IEditorExtensionRegistry],
+  activate: (
+    app: JupyterFrontEnd,
+    // Trick to ensure the chat plugin is available and loaded before this one
+    // so that the keybinding can be properly resolved.
+    tracker: WidgetTracker,
+    extensions: IEditorExtensionRegistry
+  ): void => {
+    const keyBindings = app.commands.keyBindings.find(
+      b => b.command === COMMAND_MITO_AI_OPEN_CHAT
+    );
+    const pythonAdvice = `Start writing python or Press ${CommandRegistry.formatKeystroke(
+      keyBindings?.keys[0] ?? 'Accel E'
+    )
+      .split(/[\+\s]/)
+      .map(s => `<kbd>${s}</kbd>`)
+      .join(' + ')} to ask Mito AI to write code for you.`;
+    extensions.addExtension({
+      name: 'mito-ai:empty-editor-advice',
+      factory: options => {
+        let advice = ''; // Default advice
+        // Add the advice only for cells (not for file editor)
+        if (options.inline) {
+          let guessedMimeType = options.model.mimeType;
+          if (
+            guessedMimeType === IEditorMimeTypeService.defaultMimeType &&
+            (options.model as ICellModel).type === 'code'
+          ) {
+            // Assume the kernel is not yet ready and will be a Python one.
+            // FIXME it will be better to deal with model.mimeTypeChanged signal
+            // but this is gonna be hard.
+            guessedMimeType = 'text/x-ipython';
+          }
+          // Tune the advice with the mimetype
+          switch (guessedMimeType) {
+            case 'text/x-ipython': // Python code cell
+              advice = pythonAdvice;
+              break;
+            case 'text/x-ipythongfm': // Jupyter Markdown cell
+              advice = 'Start writing markdown.';
+              break;
+          }
+        }
+        return EditorExtensionRegistry.createImmutableExtension(
+          advicePlugin({
+            advice
+          })
+        );
+      }
+    });
+  }
+};

mito-ai/src/index.ts

@@ -1,14 +1,15 @@
-
 import AiChatPlugin from './Extensions/AiChat/AiChatPlugin';
 import VariableManagerPlugin from './Extensions/VariableManager/VariableManagerPlugin';
 import ErrorMimeRendererPlugin from './Extensions/ErrorMimeRenderer/ErrorMimeRendererPlugin';
 import CellToolbarButtonsPlugin from './Extensions/CellToolbarButtons/CellToolbarButtonsPlugin';
+import { localPrompt } from './Extensions/emptyCell';
 
-// This is the main entry point to the mito-ai extension. It must export all of the top level 
+// This is the main entry point to the mito-ai extension. It must export all of the top level
 // extensions that we want to load.
 export default [
-    AiChatPlugin, 
-    ErrorMimeRendererPlugin, 
-    VariableManagerPlugin,
-    CellToolbarButtonsPlugin
+  AiChatPlugin,
+  ErrorMimeRendererPlugin,
+  VariableManagerPlugin,
+  CellToolbarButtonsPlugin,
+  localPrompt
 ];