From 6ffb553ff022fad9894d13813f4e054be7c9dd2e Mon Sep 17 00:00:00 2001 From: krassowski <5832902+krassowski@users.noreply.github.com> Date: Sun, 10 Dec 2023 13:05:55 +0000 Subject: [PATCH] Move frontend completions code to `/completions` --- .../jupyter-ai/src/completions/handler.ts | 161 +++++++++++++++++ packages/jupyter-ai/src/completions/index.ts | 1 + packages/jupyter-ai/src/completions/plugin.ts | 151 ++++++++++++++++ .../provider.ts} | 171 +----------------- packages/jupyter-ai/src/completions/types.ts | 67 +++++++ packages/jupyter-ai/src/handler.ts | 52 ------ packages/jupyter-ai/src/index.ts | 2 +- 7 files changed, 389 insertions(+), 216 deletions(-) create mode 100644 packages/jupyter-ai/src/completions/handler.ts create mode 100644 packages/jupyter-ai/src/completions/index.ts create mode 100644 packages/jupyter-ai/src/completions/plugin.ts rename packages/jupyter-ai/src/{inline-completions.ts => completions/provider.ts} (71%) create mode 100644 packages/jupyter-ai/src/completions/types.ts diff --git a/packages/jupyter-ai/src/completions/handler.ts b/packages/jupyter-ai/src/completions/handler.ts new file mode 100644 index 000000000..9517a0d29 --- /dev/null +++ b/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 { + 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 { + 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 { + return this._modelChanged; + } + + /** + * Signal emitted when a new chunk of completion is streamed. + */ + get streamed(): ISignal { + 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 { + if (this.isDisposed) { + return; + } + const promise = new PromiseDelegate(); + 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(this); + private _streamed = new Signal(this); + private _initialized: PromiseDelegate = new PromiseDelegate(); +} diff --git a/packages/jupyter-ai/src/completions/index.ts b/packages/jupyter-ai/src/completions/index.ts new file mode 100644 index 000000000..598272900 --- /dev/null +++ b/packages/jupyter-ai/src/completions/index.ts @@ -0,0 +1 @@ +export { inlineCompletionProvider } from './plugin'; diff --git a/packages/jupyter-ai/src/completions/plugin.ts b/packages/jupyter-ai/src/completions/plugin.ts new file mode 100644 index 000000000..b7286a785 --- /dev/null +++ b/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 = { + 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 => { + 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 + }); + } + } +}; diff --git a/packages/jupyter-ai/src/inline-completions.ts b/packages/jupyter-ai/src/completions/provider.ts similarity index 71% rename from packages/jupyter-ai/src/inline-completions.ts rename to packages/jupyter-ai/src/completions/provider.ts index caa87a3a6..4deb714c7 100644 --- a/packages/jupyter-ai/src/inline-completions.ts +++ b/packages/jupyter-ai/src/completions/provider.ts @@ -11,181 +11,26 @@ import { } from '@jupyterlab/completer'; import { ISettingRegistry } from '@jupyterlab/settingregistry'; import { Notification, showErrorMessage } from '@jupyterlab/apputils'; -import { IDisposable } from '@lumino/disposable'; import { JSONValue, PromiseDelegate } from '@lumino/coreutils'; -import { ServerConnection } from '@jupyterlab/services'; -import { URLExt } from '@jupyterlab/coreutils'; import { IEditorLanguageRegistry, IEditorLanguage } from '@jupyterlab/codemirror'; import { NotebookPanel } from '@jupyterlab/notebook'; -import { AiService } from './handler'; +import { AiCompleterService as AiService } from './types'; import { DocumentWidget } from '@jupyterlab/docregistry'; -import { Signal, ISignal } from '@lumino/signaling'; -import { jupyternautIcon } from './icons'; -import { getEditor } from './selection-watcher'; -import { IJupyternautStatus } from './tokens'; +import { jupyternautIcon } from '../icons'; +import { getEditor } from '../selection-watcher'; +import { IJupyternautStatus } from '../tokens'; +import { CompletionWebsocketHandler } from './handler'; -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 { - 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 { - 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 { - return this._modelChanged; - } - - /** - * Signal emitted when a new chunk of completion is streamed. - */ - get streamed(): ISignal { - 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 { - if (this.isDisposed) { - return; - } - const promise = new PromiseDelegate(); - 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(this); - private _streamed = new Signal(this); - private _initialized: PromiseDelegate = new PromiseDelegate(); -} - /** * Format the language name nicely. */ -function displayName(language: IEditorLanguage): string { +export function displayName(language: IEditorLanguage): string { if (language.name === 'ipythongfm') { return 'Markdown (IPython)'; } @@ -195,7 +40,7 @@ function displayName(language: IEditorLanguage): string { return language.displayName ?? language.name; } -class JupyterAIInlineProvider implements IInlineCompletionProvider { +export class JupyterAIInlineProvider implements IInlineCompletionProvider { readonly identifier = 'jupyter-ai'; readonly icon = jupyternautIcon.bindprops({ width: 16, top: 1 }); @@ -429,7 +274,7 @@ class JupyterAIInlineProvider implements IInlineCompletionProvider { private _counter = 0; } -namespace JupyterAIInlineProvider { +export namespace JupyterAIInlineProvider { export interface IOptions { completionHandler: CompletionWebsocketHandler; languageRegistry: IEditorLanguageRegistry; diff --git a/packages/jupyter-ai/src/completions/types.ts b/packages/jupyter-ai/src/completions/types.ts new file mode 100644 index 000000000..477c6d74b --- /dev/null +++ b/packages/jupyter-ai/src/completions/types.ts @@ -0,0 +1,67 @@ +import type { + IInlineCompletionList, + IInlineCompletionItem +} from '@jupyterlab/completer'; + +import { ServerConnection } from '@jupyterlab/services'; + +export namespace AiCompleterService { + /** + * The instantiation options for a data registry handler. + */ + export interface IOptions { + serverSettings?: ServerConnection.ISettings; + } + + export type ConnectionMessage = { + type: 'connection'; + client_id: string; + }; + + export type InlineCompletionRequest = { + number: number; + path?: string; + /* The model has to complete given prefix */ + prefix: string; + /* The model may consider the following suffix */ + suffix: string; + mime: string; + /* Whether to stream the response (if streaming is supported by the model) */ + stream: boolean; + language?: string; + cell_id?: string; + }; + + export type CompletionError = { + type: string; + traceback: string; + }; + + export type InlineCompletionReply = { + /** + * Type for this message can be skipped (`inline_completion` is presumed default). + **/ + type?: 'inline_completion'; + list: IInlineCompletionList; + reply_to: number; + error?: CompletionError; + }; + + export type InlineCompletionStreamChunk = { + type: 'stream'; + response: IInlineCompletionItem; + reply_to: number; + done: boolean; + }; + + export type InlineCompletionModelChanged = { + type: 'model_changed'; + model: string; + }; + + export type CompleterMessage = + | InlineCompletionReply + | ConnectionMessage + | InlineCompletionStreamChunk + | InlineCompletionModelChanged; +} diff --git a/packages/jupyter-ai/src/handler.ts b/packages/jupyter-ai/src/handler.ts index 5758c7275..ec91c1f3f 100644 --- a/packages/jupyter-ai/src/handler.ts +++ b/packages/jupyter-ai/src/handler.ts @@ -2,11 +2,6 @@ import { URLExt } from '@jupyterlab/coreutils'; import { ServerConnection } from '@jupyterlab/services'; -import type { - IInlineCompletionList, - IInlineCompletionItem -} from '@jupyterlab/completer'; - const API_NAMESPACE = 'api/ai'; /** @@ -60,47 +55,6 @@ export namespace AiService { prompt: string; }; - export type InlineCompletionRequest = { - number: number; - path?: string; - /* The model has to complete given prefix */ - prefix: string; - /* The model may consider the following suffix */ - suffix: string; - mime: string; - /* Whether to stream the response (if streaming is supported by the model) */ - stream: boolean; - language?: string; - cell_id?: string; - }; - - export type CompletionError = { - type: string; - traceback: string; - }; - - export type InlineCompletionReply = { - /** - * Type for this message can be skipped (`inline_completion` is presumed default). - **/ - type?: 'inline_completion'; - list: IInlineCompletionList; - reply_to: number; - error?: CompletionError; - }; - - export type InlineCompletionStreamChunk = { - type: 'stream'; - response: IInlineCompletionItem; - reply_to: number; - done: boolean; - }; - - export type InlineCompletionModelChanged = { - type: 'model_changed'; - model: string; - }; - export type Collaborator = { username: string; initials: string; @@ -146,12 +100,6 @@ export namespace AiService { | ConnectionMessage | ClearMessage; - export type CompleterMessage = - | InlineCompletionReply - | ConnectionMessage - | InlineCompletionStreamChunk - | InlineCompletionModelChanged; - export type ChatHistory = { messages: ChatMessage[]; }; diff --git a/packages/jupyter-ai/src/index.ts b/packages/jupyter-ai/src/index.ts index f7528f35f..03f8e2eb9 100644 --- a/packages/jupyter-ai/src/index.ts +++ b/packages/jupyter-ai/src/index.ts @@ -12,7 +12,7 @@ import { buildChatSidebar } from './widgets/chat-sidebar'; import { SelectionWatcher } from './selection-watcher'; import { ChatHandler } from './chat_handler'; import { buildErrorWidget } from './widgets/chat-error'; -import { inlineCompletionProvider } from './inline-completions'; +import { inlineCompletionProvider } from './completions'; import { jupyternautStatus } from './status'; export type DocumentTracker = IWidgetTracker;