Move frontend completions code to /completions

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

@@ -0,0 +1,161 @@
+import { IDisposable } from '@lumino/disposable';
+import { PromiseDelegate } from '@lumino/coreutils';
+import { ServerConnection } from '@jupyterlab/services';
+import { URLExt } from '@jupyterlab/coreutils';
+import { AiCompleterService as AiService } from './types';
+import { Signal, ISignal } from '@lumino/signaling';
+
+const SERVICE_URL = 'api/ai/completion/inline';
+
+type StreamChunk = AiService.InlineCompletionStreamChunk;
+
+export class CompletionWebsocketHandler implements IDisposable {
+  /**
+   * The server settings used to make API requests.
+   */
+  readonly serverSettings: ServerConnection.ISettings;
+
+  /**
+   * Create a new completion handler.
+   */
+  constructor(options: AiService.IOptions = {}) {
+    this.serverSettings =
+      options.serverSettings ?? ServerConnection.makeSettings();
+  }
+
+  /**
+   * Initializes the WebSocket connection to the completion backend. Promise is
+   * resolved when server acknowledges connection and sends the client ID. This
+   * must be awaited before calling any other method.
+   */
+  public async initialize(): Promise<void> {
+    await this._initialize();
+  }
+
+  /**
+   * Sends a message across the WebSocket. Promise resolves to the message ID
+   * when the server sends the same message back, acknowledging receipt.
+   */
+  public sendMessage(
+    message: AiService.InlineCompletionRequest
+  ): Promise<AiService.InlineCompletionReply> {
+    return new Promise(resolve => {
+      this._socket?.send(JSON.stringify(message));
+      this._replyForResolver[message.number] = resolve;
+    });
+  }
+
+  /**
+   * Signal emitted when completion AI model changes.
+   */
+  get modelChanged(): ISignal<CompletionWebsocketHandler, string> {
+    return this._modelChanged;
+  }
+
+  /**
+   * Signal emitted when a new chunk of completion is streamed.
+   */
+  get streamed(): ISignal<CompletionWebsocketHandler, StreamChunk> {
+    return this._streamed;
+  }
+
+  /**
+   * Whether the completion handler is disposed.
+   */
+  get isDisposed(): boolean {
+    return this._isDisposed;
+  }
+
+  /**
+   * Dispose the completion handler.
+   */
+  dispose(): void {
+    if (this.isDisposed) {
+      return;
+    }
+    this._isDisposed = true;
+
+    // Clean up socket.
+    const socket = this._socket;
+    if (socket) {
+      this._socket = null;
+      socket.onopen = () => undefined;
+      socket.onerror = () => undefined;
+      socket.onmessage = () => undefined;
+      socket.onclose = () => undefined;
+      socket.close();
+    }
+  }
+
+  private _onMessage(message: AiService.CompleterMessage): void {
+    switch (message.type) {
+      case 'model_changed': {
+        this._modelChanged.emit(message.model);
+        break;
+      }
+      case 'connection': {
+        this._initialized.resolve();
+        break;
+      }
+      case 'stream': {
+        this._streamed.emit(message);
+        break;
+      }
+      default: {
+        if (message.reply_to in this._replyForResolver) {
+          this._replyForResolver[message.reply_to](message);
+          delete this._replyForResolver[message.reply_to];
+        } else {
+          console.warn('Unhandled message', message);
+        }
+        break;
+      }
+    }
+  }
+
+  /**
+   * Dictionary mapping message IDs to Promise resolvers.
+   */
+  private _replyForResolver: Record<
+    number,
+    (value: AiService.InlineCompletionReply) => void
+  > = {};
+
+  private _onClose(e: CloseEvent, reject: any) {
+    reject(new Error('Inline completion websocket disconnected'));
+    console.error('Inline completion websocket disconnected');
+    // only attempt re-connect if there was an abnormal closure
+    // WebSocket status codes defined in RFC 6455: https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
+    if (e.code === 1006) {
+      const delaySeconds = 1;
+      console.info(`Will try to reconnect in ${delaySeconds} s.`);
+      setTimeout(async () => await this._initialize(), delaySeconds * 1000);
+    }
+  }
+
+  private async _initialize(): Promise<void> {
+    if (this.isDisposed) {
+      return;
+    }
+    const promise = new PromiseDelegate<void>();
+    this._initialized = promise;
+    console.log(
+      'Creating a new websocket connection for inline completions...'
+    );
+    const { token, WebSocket, wsUrl } = this.serverSettings;
+    const url =
+      URLExt.join(wsUrl, SERVICE_URL) +
+      (token ? `?token=${encodeURIComponent(token)}` : '');
+
+    const socket = (this._socket = new WebSocket(url));
+    socket.onclose = e => this._onClose(e, promise.reject);
+    socket.onerror = e => promise.reject(e);
+    socket.onmessage = msg => msg.data && this._onMessage(JSON.parse(msg.data));
+  }
+
+  private _isDisposed = false;
+  private _socket: WebSocket | null = null;
+  private _modelChanged = new Signal<CompletionWebsocketHandler, string>(this);
+  private _streamed = new Signal<CompletionWebsocketHandler, StreamChunk>(this);
+  private _initialized: PromiseDelegate<void> = new PromiseDelegate<void>();
+}

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

@@ -0,0 +1 @@
+export { inlineCompletionProvider } from './plugin';

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

@@ -0,0 +1,151 @@
+import {
+  JupyterFrontEnd,
+  JupyterFrontEndPlugin
+} from '@jupyterlab/application';
+import {
+  ICompletionProviderManager,
+} from '@jupyterlab/completer';
+import { ISettingRegistry } from '@jupyterlab/settingregistry';
+import {
+  IEditorLanguageRegistry,
+  IEditorLanguage
+} from '@jupyterlab/codemirror';
+import { getEditor } from '../selection-watcher';
+import { IJupyternautStatus } from '../tokens';
+import { displayName, JupyterAIInlineProvider } from './provider';
+import { CompletionWebsocketHandler } from './handler';
+
+export namespace CommandIDs {
+  /**
+   * Command to toggle completions globally.
+   */
+  export const toggleCompletions = 'jupyter-ai:toggle-completions';
+    /**
+   * Command to toggle completions for specific language.
+   */
+  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,
+    ISettingRegistry
+  ],
+  optional: [IJupyternautStatus],
+  activate: async (
+    app: JupyterFrontEnd,
+    manager: ICompletionProviderManager,
+    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
+      console.warn(
+        'Inline completions are only supported in JupyterLab 4.1+ and Jupyter Notebook 7.1+'
+      );
+      return;
+    }
+    const completionHandler = new CompletionWebsocketHandler();
+    const provider = new JupyterAIInlineProvider({
+      completionHandler,
+      languageRegistry
+    });
+    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
+      });
+    }
+  }
+};