Implement toggling the AI completer via statusbar item

also adds the icon for provider re-using jupyternaut icon

packages/jupyter-ai/src/components/statusbar-item.tsx

@@ -0,0 +1,92 @@
+import { Popup, showPopup } from '@jupyterlab/statusbar';
+import React from 'react';
+import { VDomModel, VDomRenderer } from '@jupyterlab/ui-components';
+import { CommandRegistry } from '@lumino/commands';
+import { MenuSvg, RankedMenu, IRankedMenu } from '@jupyterlab/ui-components';
+import { Jupyternaut } from '../icons';
+import { IJupyternautStatus } from '../tokens';
+
+/**
+ * StatusBar item to display menu for toggling the completion.
+ */
+export class JupyternautStatus
+  extends VDomRenderer<VDomModel>
+  implements IJupyternautStatus
+{
+  constructor(options: JupyternautStatus.IOptions) {
+    super(new VDomModel());
+    this._commandRegistry = options.commandRegistry;
+    this._items = [];
+
+    this.addClass('jp-mod-highlighted');
+    this.title.caption = 'Open Jupyternaut status menu';
+    this.node.addEventListener('click', this._handleClick);
+  }
+
+  addItem(item: IRankedMenu.IItemOptions): void {
+    this._items.push(item);
+  }
+
+  hasItems(): boolean {
+    return this._items.length !== 0;
+  }
+
+  /**
+   * Render the status item.
+   */
+  render(): JSX.Element | null {
+    if (!this.model) {
+      return null;
+    }
+    return <Jupyternaut top={'2px'} width={'16px'} stylesheet={'statusBar'} />;
+  }
+
+  dispose(): void {
+    this.node.removeEventListener('click', this._handleClick);
+    super.dispose();
+  }
+
+  /**
+   * Create a menu for viewing status and changing options.
+   */
+  private _handleClick = () => {
+    if (this._popup) {
+      this._popup.dispose();
+    }
+    if (this._menu) {
+      this._menu.dispose();
+    }
+    this._menu = new RankedMenu({
+      commands: this._commandRegistry,
+      renderer: MenuSvg.defaultRenderer
+    });
+    for (const item of this._items) {
+      this._menu.addItem(item);
+    }
+    this._popup = showPopup({
+      body: this._menu,
+      anchor: this,
+      align: 'left'
+    });
+  };
+
+  private _items: IRankedMenu.IItemOptions[];
+  private _commandRegistry: CommandRegistry;
+  private _menu: RankedMenu | null = null;
+  private _popup: Popup | null = null;
+}
+
+/**
+ * A namespace for JupyternautStatus statics.
+ */
+export namespace JupyternautStatus {
+  /**
+   * Options for the JupyternautStatus item.
+   */
+  export interface IOptions {
+    /**
+     * The application command registry.
+     */
+    commandRegistry: CommandRegistry;
+  }
+}

packages/jupyter-ai/src/index.ts

@@ -13,6 +13,7 @@ import { SelectionWatcher } from './selection-watcher';
 import { ChatHandler } from './chat_handler';
 import { buildErrorWidget } from './widgets/chat-error';
 import { inlineCompletionProvider } from './inline-completions';
+import { jupyternautStatus } from './status';
 
 export type DocumentTracker = IWidgetTracker<IDocumentWidget>;
 
@@ -61,4 +62,4 @@ const plugin: JupyterFrontEndPlugin<void> = {
   }
 };
 
-export default [plugin, inlineCompletionProvider];
+export default [plugin, jupyternautStatus, inlineCompletionProvider];

packages/jupyter-ai/src/inline-completions.ts

@@ -8,7 +8,7 @@ import {
   IInlineCompletionContext,
   CompletionHandler
 } from '@jupyterlab/completer';
-import type { ISettingRegistry } from '@jupyterlab/settingregistry';
+import { ISettingRegistry } from '@jupyterlab/settingregistry';
 import { Notification, showErrorMessage } from '@jupyterlab/apputils';
 import { IDisposable } from '@lumino/disposable';
 import { JSONValue, PromiseDelegate } from '@lumino/coreutils';
@@ -22,6 +22,9 @@ import { NotebookPanel } from '@jupyterlab/notebook';
 import { AiService } from './handler';
 import { DocumentWidget } from '@jupyterlab/docregistry';
 import { Signal, ISignal } from '@lumino/signaling';
+import { jupyternautIcon } from './icons';
+import { getEditor } from './selection-watcher';
+import { IJupyternautStatus } from './tokens';
 
 const SERVICE_URL = 'api/ai/completion/inline';
 
@@ -161,8 +164,22 @@ export class CompletionWebsocketHandler implements IDisposable {
   private _initialized: PromiseDelegate<void> = new PromiseDelegate<void>();
 }
 
+/**
+ * Format the language name nicely.
+ */
+function displayName(language: IEditorLanguage): string {
+  if (language.name === 'ipythongfm') {
+    return 'Markdown (IPython)';
+  }
+  if (language.name === 'ipython') {
+    return 'IPython';
+  }
+  return language.displayName ?? language.name;
+}
+
 class JupyterAIInlineProvider implements IInlineCompletionProvider {
   readonly identifier = 'jupyter-ai';
+  readonly icon = jupyternautIcon.bindprops({ width: 16, top: 1 });
 
   constructor(protected options: JupyterAIInlineProvider.IOptions) {
     options.completionHandler.modelChanged.connect(
@@ -186,8 +203,15 @@ class JupyterAIInlineProvider implements IInlineCompletionProvider {
     context: IInlineCompletionContext
   ) {
     const mime = request.mimeType ?? 'text/plain';
-    if (mime === 'text/x-ipythongfm') {
-      // Do not offer suggestions in markdown cells.
+    const language = this.options.languageRegistry.findByMIME(mime);
+    if (!language) {
+      console.warn(
+        `Could not recognise language for ${mime} - cannot complete`
+      );
+      return { items: [] };
+    }
+    if (!this.isLanguageEnabled(language?.name)) {
+      // Do not offer suggestions if disabled.
       return { items: [] };
     }
     let cellId = undefined;
@@ -202,7 +226,6 @@ class JupyterAIInlineProvider implements IInlineCompletionProvider {
       path = context.widget.context.path;
     }
     const number = ++this._counter;
-    const language = this.options.languageRegistry.findByMIME(mime);
     const result = await this.options.completionHandler.sendMessage({
       path: context.session?.path,
       mime,
@@ -236,6 +259,7 @@ class JupyterAIInlineProvider implements IInlineCompletionProvider {
   }
 
   get schema(): ISettingRegistry.IProperty {
+    const knownLanguages = this.options.languageRegistry.getLanguages();
     return {
       properties: {
         maxPrefix: {
@@ -251,6 +275,18 @@ class JupyterAIInlineProvider implements IInlineCompletionProvider {
           type: 'number',
           description:
             'At most how many suffix characters should be provided to the model.'
+        },
+        disabledLanguages: {
+          title: 'Disabled languages',
+          type: 'array',
+          items: {
+            type: 'string',
+            oneOf: knownLanguages.map(language => {
+              return { const: language.name, title: displayName(language) };
+            })
+          },
+          description:
+            'Languages for which the completions should not be shown.'
         }
       },
       default: JupyterAIInlineProvider.DEFAULT_SETTINGS as any
@@ -261,6 +297,14 @@ class JupyterAIInlineProvider implements IInlineCompletionProvider {
     this._settings = settings as unknown as JupyterAIInlineProvider.ISettings;
   }
 
+  isEnabled(): boolean {
+    return this._settings.enabled;
+  }
+
+  isLanguageEnabled(language: string) {
+    return !this._settings.disabledLanguages.includes(language);
+  }
+
   /**
    * Extract prefix from request, accounting for context window limit.
    */
@@ -310,25 +354,46 @@ namespace JupyterAIInlineProvider {
     maxPrefix: number;
     maxSuffix: number;
     debouncerDelay: number;
+    enabled: boolean;
+    disabledLanguages: string[];
   }
   export const DEFAULT_SETTINGS: ISettings = {
     maxPrefix: 10000,
     maxSuffix: 10000,
     // The debouncer delay handling is implemented upstream in JupyterLab;
     // here we just increase the default from 0, as compared to kernel history
     // the external AI models may have a token cost associated.
-    debouncerDelay: 250
+    debouncerDelay: 250,
+    enabled: true,
+    // ipythongfm means "IPython GitHub Flavoured Markdown"
+    disabledLanguages: ['ipythongfm']
   };
 }
 
+export namespace CommandIDs {
+  export const toggleCompletions = 'jupyter-ai:toggle-completions';
+  export const toggleLanguageCompletions =
+    'jupyter-ai:toggle-language-completions';
+}
+
+const INLINE_COMPLETER_PLUGIN =
+  '@jupyterlab/completer-extension:inline-completer';
+
 export const inlineCompletionProvider: JupyterFrontEndPlugin<void> = {
   id: 'jupyter_ai:inline-completions',
   autoStart: true,
-  requires: [ICompletionProviderManager, IEditorLanguageRegistry],
+  requires: [
+    ICompletionProviderManager,
+    IEditorLanguageRegistry,
+    ISettingRegistry
+  ],
+  optional: [IJupyternautStatus],
   activate: async (
     app: JupyterFrontEnd,
     manager: ICompletionProviderManager,
-    languageRegistry: IEditorLanguageRegistry
+    languageRegistry: IEditorLanguageRegistry,
+    settingRegistry: ISettingRegistry,
+    statusMenu: IJupyternautStatus | null
   ): Promise<void> => {
     if (typeof manager.registerInlineProvider === 'undefined') {
       // Gracefully short-circuit on JupyterLab 4.0 and Notebook 7.0
@@ -344,5 +409,92 @@ export const inlineCompletionProvider: JupyterFrontEndPlugin<void> = {
     });
     await completionHandler.initialize();
     manager.registerInlineProvider(provider);
+
+    const findCurrentLanguage = (): IEditorLanguage | null => {
+      const widget = app.shell.currentWidget;
+      const editor = getEditor(widget);
+      if (!editor) {
+        return null;
+      }
+      return languageRegistry.findByMIME(editor.model.mimeType);
+    };
+
+    let settings: ISettingRegistry.ISettings | null = null;
+
+    settingRegistry.pluginChanged.connect(async (_emitter, plugin) => {
+      if (plugin === INLINE_COMPLETER_PLUGIN) {
+        // Only load the settings once the plugin settings were transformed
+        settings = await settingRegistry.load(INLINE_COMPLETER_PLUGIN);
+      }
+    });
+
+    app.commands.addCommand(CommandIDs.toggleCompletions, {
+      execute: () => {
+        if (!settings) {
+          return;
+        }
+        const providers = Object.assign({}, settings.user.providers) as any;
+        const ourSettings = {
+          ...JupyterAIInlineProvider.DEFAULT_SETTINGS,
+          ...providers[provider.identifier]
+        };
+        const wasEnabled = ourSettings['enabled'];
+        providers[provider.identifier]['enabled'] = !wasEnabled;
+        settings.set('providers', providers);
+      },
+      label: 'Enable Jupyternaut Completions',
+      isToggled: () => {
+        return provider.isEnabled();
+      }
+    });
+
+    app.commands.addCommand(CommandIDs.toggleLanguageCompletions, {
+      execute: () => {
+        const language = findCurrentLanguage();
+        if (!settings || !language) {
+          return;
+        }
+        const providers = Object.assign({}, settings.user.providers) as any;
+        const ourSettings = {
+          ...JupyterAIInlineProvider.DEFAULT_SETTINGS,
+          ...providers[provider.identifier]
+        };
+        const wasDisabled = ourSettings['disabledLanguages'].includes(
+          language.name
+        );
+        const disabledList: string[] =
+          providers[provider.identifier]['disabledLanguages'];
+        if (wasDisabled) {
+          disabledList.filter(name => name !== language.name);
+        } else {
+          disabledList.push(language.name);
+        }
+        settings.set('providers', providers);
+      },
+      label: () => {
+        const language = findCurrentLanguage();
+        return language
+          ? `Enable Completions in ${displayName(language)}`
+          : 'Enable Completions for Language of Current Editor';
+      },
+      isToggled: () => {
+        const language = findCurrentLanguage();
+        return !!language && provider.isLanguageEnabled(language.name);
+      },
+      isEnabled: () => {
+        return !!findCurrentLanguage() && provider.isEnabled();
+      }
+    });
+
+    if (statusMenu) {
+      statusMenu.addItem({
+        command: CommandIDs.toggleCompletions,
+        rank: 1
+      });
+      statusMenu.addItem({
+        command: CommandIDs.toggleLanguageCompletions,
+        rank: 2
+      });
+    }
   }
 };

packages/jupyter-ai/src/selection-watcher.ts

@@ -14,7 +14,9 @@ import { getCellIndex } from './utils';
 /**
  * Gets the editor instance used by a document widget. Returns `null` if unable.
  */
-function getEditor(widget: Widget | null) {
+export function getEditor(
+  widget: Widget | null
+): CodeMirrorEditor | null | undefined {
   if (!(widget instanceof DocumentWidget)) {
     return null;
   }