Code completion for jupyter lab/notebook

[liyongsea]

Implementing #290, a second iteration from #378

In this PR, we leverage the inline completer jupyterlab/jupyterlab#15160 to provide copilot like code completion to jupyter ai.
This plugin currently support Starcoder, which is a open source code completion model finetuned on jupyter notebooks

Here is the expected behavior

stream_demo.mp4

[liyongsea]

As jupyterlab/jupyterlab#15160 is not merge nor released yet. There might still be modification on this PR. But you can take a quick look to see if the overall implementation logic make sense

[Wzixiao]

Hi everyone,

I am currently working on this PR that requires simultaneous development with JupyterLab core. I would like to ensure a smooth integration process and have a query regarding dependency management.

Would it be suitable to reference JupyterLab's internal @jupyterlab/completer package directly within my feature's package.json through a file path like so:

"@jupyterlab/completer": "file:../../../jupyterlab/packages/completer"

Or is there a better way you folks recommend?

Thanks!

[Wzixiao]

Code changes

Main details:

We complete this PR mainly in the following steps:

1. Load an extension that only focuses on code completion into jupyter-lab. (index.ts)

2. Binding keyboard events to the cell's editor has been used to send inline completion requests and invoke inline completion accept function.(inline-completion-handler.ts)

3. Use provider to accept requests from the upstream API. (bigcode-Inline-completion-provider.ts)

   • Get cell code.(utils/cell-context.ts)

   • Send request.(utils/bigcode-request.ts)

4. In addition, there is also a class to obtain user settings. This class is used with many files, so the following process does not show the usage logic of this class.(utils/code-completion-context-store.ts)

Major code changes or additions to code files:

packages/jupyter-ai/src
|-- bigcode-Inline-completion-provider.ts
|
|-- [components]
|   └─- [settings]
|       └─- bigcode-settings.ts
|
|-- [contexts]
|   └─- code-completion-context-store.tsx  # Used as a global context for code completion
|
|-- index.ts
|
|-- inline-completion-handler.ts
|
|-- [types]
|   └─- cell.ts  # This cell is used as logic for code completion
|
|-- [utils]  # Tool classes for obtaining cell methods and other tools
|   |-- bigcode-request.ts
|   |-- cell-context.ts
|   |-- instance.ts
|   └─- keyboard.ts
|
└─- [widgets]
    └─- bigcode-sidebar.tsx

Simple code structure diagram

index.ts

// Add inline code completion extension
const inlineCompletionPlugin: JupyterFrontEndPlugin<void> = {
  id: 'jupyter_ai:plugin:inline-completion',
  description: 'Adding an inline completion provider suggestion comes from jupyter-ai.',
  requires: [ICompletionProviderManager],
  optional: [ILayoutRestorer],
  autoStart: true,
  activate: (
    app: JupyterFrontEnd,
    completionManager: ICompletionProviderManager,
    restorer: ILayoutRestorer | null,
    translator: ITranslator | null
  ): void: { 
    // This extension will complete the loading of the sidebar and the user's keyboard request
    pass... 
  }
};

const plugins: JupyterFrontEndPlugin<void>[] = [..., inlineCompletionPlugin];

export default plugins;

inline-completion-handler.ts

/**
 * The main function to handle code completion on keydown events.
 * It initializes the keydown handlers after ensuring that the notebook is fully loaded.
 * @param {JupyterFrontEnd} app - The JupyterFrontEnd application instance.
 * @returns {Promise<void>}
 */
export const handleCodeCompletionKeyDown = async (
  app: JupyterFrontEnd,
  completionManager: ICompletionProviderManager,
  bigcodeInlineCompletionProvider: BigcodeInlineCompletionProvider
): Promise<void> => {
  // Wait for the notebook to finish initializing
  await app.start();
  initializeKeyDownHandlers();
};

// Generates a keydown extension for handling various keypress events.
const generateKeyDownExtension = (
  app: JupyterFrontEnd,
  completionManager: ICompletionProviderManager,
  bigcodeInlineCompletionProvider: BigcodeInlineCompletionProvider
): Extension => {
  return Prec.highest(
    keymap.of([
      {
        any: (view: EditorView, event: KeyboardEvent) => {
          // Only the main logic is written here(code changes).

          // Get the request shortcut key set by the user
          const parsedShortcut = parseKeyboardEventToShortcut(event);

          // If the keyboard combination pressed by the user is a set shortcut key, "invoke" function at inline-completion is performed.
          if (parsedShortcut === CodeCompletionContextStore.shortcutStr) {
              completionManager.inline?.invoke(app.shell.currentWidget?.id);
              return true;
          }

          // When enter is pressed, "accept" function is performed.
          if (event.code === 'Enter' && ...) {
              completionManager.inline?.accept(app.shell.currentWidget?.id);
              return true;
          }

          return false;
        }
      }
    ])
  );
};

/**
 * Initializes keydown event handlers for the JupyterFrontEnd application.
 * This function sets up listeners for changes in the current widget and mounts the editor accordingly.
 * @param {JupyterFrontEnd} app - The JupyterFrontEnd application instance.
 */
const initializeKeyDownHandlers = (
  app: JupyterFrontEnd,
  completionManager: ICompletionProviderManager,
  bigcodeInlineCompletionProvider: BigcodeInlineCompletionProvider
) => {
  const extension = generateKeyDownExtension();
  // Inject this extension into the cell's editor when switching cells
};

type/cell.ts

export type ICellType = null | 'markdown' | 'code' | 'output';

// This cell is used as logic for code completion
export interface ICell {
  content: string;
  type: ICellType;
}

utils/cell-context.ts

import { ICell, ICellType } from '../types/cell';

/**
 * Fetches the content of the notebook up to the current cursor position.
 *
 * @param widget - The notebook panel widget.
 * @returns An array of cell contents up to the cursor, or null.
 */
export const retrieveNotebookContentUntilCursor = (
  widget: NotebookPanel
): ICell[] | null;

utils/bigcode-request.ts

import { makeObservable, computed } from 'mobx';

export type BigCodeServiceStreamResponseItem = {
  token: {
    id: number;
    text: string;
    logprob: number;
    special: boolean;
  };
  generated_text: string | null;
  details: null;
};

export type BigCodeServiceNotStreamResponse = {
  generated_text: string;
}[];    
    
class Bigcode {
  // prompt for request
  private _prompt: string;

  constructor(private store: IGlobalStore) {
    makeObservable(this);
    this._prompt = '';
  }

  @computed get bigcodeUrl() {
    return this.store.bigcodeUrl;
  }

  @computed get accessToken() {
    return this.store.accessToken;
  }

  @computed get maxTokens() {
    return this.store.maxPromptToken;
  }

  get prompt() {
    return this._prompt;
  }

  // For stream requests
  async send(stream: true): Promise<ReadableStream<Uint8Array>>;
  // For not stream requests
  async send(stream: false): Promise<BigCodeServiceNotStreamResponse>;

  // Construct prompt using "ICell" in types/cells type
  constructContinuationPrompt(context: ICell[] | null): string;
}

bigcode-Inline-completion-provider.ts

import { retrieveNotebookContentUntilCursor } from './utils/cell-context';

import bigcodeRequestInstance, { BigCodeServiceStreamResponseItem } from './utils/bigcode-request';

export class BigcodeInlineCompletionProvider
  implements IInlineCompletionProvider
{

  // Used to record information from the last user request
  private _lastRequestInfo: {
    insertText: string;
    cellCode: string;
  } = {
    insertText: '',
    cellCode: ''
  };
  private _requesting = false;
  // Change it in "fetch" and then use this field in "stream" to call a different function
  private _requestMode: InlineCompletionTriggerKind = 0;
  // Whether to stop the stream immediately, When the stream is in progress, the user requests to change this field to true, and then interrupts the stream request.
  private _streamStop = false;
  // Whether to finish for request
  private _finish = false;
  // Debounce use
  private _timeoutId: number | null = null;
  // Debounce use
  private _callCounter = 0;
    
  // Construct the next request prompt
  private constructContinuationPrompt(
    context: IInlineCompletionContext
  ): string {
    if (context.widget instanceof NotebookPanel) {
      const widget = context.widget as NotebookPanel;
      // Get the code in the cell before the current mouse pointer(including the current cell)
      const notebookCellContent = retrieveNotebookContentUntilCursor(widget);

      // Construct prompt
      bigcodeRequestInstance.constructContinuationPrompt(notebookCellContent);
      return bigcodeRequestInstance.prompt;
    }

    return '';
  }

  /**
   * This function implements the fetch of IInlineCompletionProvider.
   * Its main logic is to call different handlers based on context.triggerKind and return stream=true and token for the upper API to call the completionHandler function.
   * 
   */
  async fetch(
    request: CompletionHandler.IRequest,
    context: IInlineCompletionContext
  ): Promise<IInlineCompletionList<IInlineCompletionItem>>{
    // shortCutCompletionHandler() if context.triggerKind else autoCompletionHandler()
  }

  /**
   * Handle requests when the user presses the keyboard
   */
  async shortCutCompletionHandler(
    request: CompletionHandler.IRequest,
    context: IInlineCompletionContext
  ): Promise<IInlineCompletionList<IInlineCompletionItem>>{
    const prompt = this.constructContinuationPrompt()
    ...
    return {
        items: [
            {
                token: prompt,
                isIncomplete: true, // All requests in the "provider" are completed by the streaming function
                insertText: ''
            }
        ]
    }
  }

  /**
   * Handle requests when upstream api automatic events occur,
   * This function will determine whether to make an automatic request based on debounceAutoRequest.
   * And this function will also filter the ghost text when the user presses the same start point as the ghost text when displaying the ghost text (no request will be made)
   */
  async autoCompletionHandler(
    request: CompletionHandler.IRequest,
    context: IInlineCompletionContext
  ): Promise<IInlineCompletionList<IInlineCompletionItem>> {
    // if debounceAutoRequest() === '<auto_stream>': 
    //     The function return value is also the same as shortCutCompletionHandler
    //     pass...
    // else:
    //     pass...
  }

  /**
   * Post-debounce strategy, the last request within the specified time returns <auto_stream>
   */
  debounceAutoRequest(): Promise<'<auto_stream>' | '<debounce>'>;

  /**
   * If the request ends, need to call this function to write the status for "accept" use
   */
  setRequestFinish(error: boolean): void {
    this._requesting = false;
    this._streamStop = false;
    this._finish = !error;
  }

  /**
   * When the user executes the "accept" function, this function will be called to clear the status.
   */
  clearState(): void {
    this._streamStop = true;
    this._finish = false;
    this._requesting = false;
    this._lastRequestInfo = {
      insertText: '',
      cellCode: ''
    };
  }
    
  /**
   * This function implements the stream of IInlineCompletionProvider.
   * The following different functions are called according to user configuration and automatic triggering.
   */
  async *stream(
    token: string
  ): AsyncGenerator<{ response: IInlineCompletionItem }, undefined, unknown>;

  private async *completionStream(
    token: string
  ): AsyncGenerator<{ response: IInlineCompletionItem }, undefined, unknown>;

  private async *keypressCompletionStream(
    token: string
  ): AsyncGenerator<{ response: IInlineCompletionItem }, undefined, unknown>{
    const reponseData = await bigcodeRequestInstance.fetchStream(true);
    // Parse responseData and return IInlineCompletionItem structure(stream)
  }

  private async *automaticCompletionStream(
    token: string
  ): AsyncGenerator<{ response: IInlineCompletionItem }, undefined, unknown>{
    const reponseData = await bigcodeRequestInstance.fetchStream(false);
    // Parse responseData and return IInlineCompletionItem structure(only once)
  }

  private async *mockCompletionStream(
    token: string
  ): AsyncGenerator<{ response: IInlineCompletionItem }, undefined, unknown>;

  private async *mockKeypressCompletionStream(
    token: string
  ): AsyncGenerator<{ response: IInlineCompletionItem }, undefined, unknown>;

  private async *mockAutomaticCompletionStream(
    token: string
  ): AsyncGenerator<{ response: IInlineCompletionItem }, undefined, unknown>;
}

contexts/code-completion-context-store.ts

import { makeObservable, observable, action } from 'mobx';

class CodeCompletionContextStore {
  /**
   * Whether to enable code completion function.
   */
  @observable enableCodeCompletion = false;

  /**
   * Observable huggingface token for authentication purposes.
   */
  @observable accessToken = '';

  /**
   * Observable URL for the BigCode service.
   */
  @observable bigcodeUrl = '';

  /**
   * Whether simulation testing is enabled (without using the real API)
   */
  @observable enableMockTest = false;

  /**
   * Observable string representing the shortcut key combination for triggering code completion.
   * Default is set to 'Ctrl + Space'.
   */
  @observable shortcutStr = 'Ctrl + Space';

  /**
   * Maximum prompt tokens when requested
   */
  @observable maxPromptToken = 400;

  /**
   * Maximum response tokens when requested
   */
  @observable maxResponseToken = 20;

  constructor() {
    makeObservable(this);

    const dataPersistenceStr = localStorage.getItem(
      '@jupyterlab-ai/CodeCompletionState'
    );

    if (dataPersistenceStr) {
      const dataPersistence: IGlobalStore = JSON.parse(dataPersistenceStr);
      this.enableCodeCompletion = dataPersistence.enableCodeCompletion;
      this.bigcodeUrl = dataPersistence.bigcodeUrl;
      this.shortcutStr = dataPersistence.shortcutStr;
      this.maxPromptToken = dataPersistence.maxPromptToken;
      this.maxResponseToken = dataPersistence.maxResponseToken;
      this.enableMockTest = dataPersistence.enableMockTest;
    }
  }

  // This function is called every time the settings change
  saveDataToLoaclStorage() {
    // Do not store sensitive information
    localStorage.setItem(
      '@jupyterlab-ai/CodeCompletionState',
      JSON.stringify({
        enableCodeCompletion: this.enableCodeCompletion,
        bigcodeUrl: this.bigcodeUrl,
        shortcutStr: this.shortcutStr,
        maxPromptToken: this.maxPromptToken,
        maxResponseToken: this.maxResponseToken,
        enableMockTest: this.enableMockTest
      })
    );
  }
}

[krassowski]

I have a few concerns on this PR:

• the credentials handling on the frontend (this should be saved once to the backend using existing JupyterLab settings system, and ideally configured by environment variable rather than stored and used on the frontend
• the custom keyboard handling code; this should be handled by the JupyterLab shortcuts to allow users to change shortcuts in the global settings
• the special-cased configuration which would be better implemented as JSON schema with the form handled by rjsf
• the introduction of new dependencies (mobx) and experimental decorators

In other words, while this PR does a great job in offering a PoC it does not integrate with the existing jupyterlab nor jupyter-ai systems to detriment of the user.

I will open a pull request proposing an infrastructure for supporting multiple providers/models for inline completions, which should help to inform how to move forward with this PR afterwards.

[dlqqq]

Thank you very much for starting the conversation on this feature. This has been implemented in #582, and is now available in Jupyter AI v2.10.0 with JupyterLab 4.1+. 🎉

[nick-youngblut]

Thank you very much for starting the conversation on this feature. This has been implemented in #582, and is now available in Jupyter AI v2.10.0 with JupyterLab 4.1+. 🎉

Where is this feature in the docs? I don't see it in the README or in the ReadTheDocs main page.

[JasonWeill]

@nick-youngblut The code completion feature is documented in the JupyterLab docs here: https://jupyterlab.readthedocs.io/en/latest/user/completer.html

To request improvements to Jupyter AI's README or docs, please open a new Jupyter AI issue. Thanks!

[krassowski]

And if you are looking for developer documentation for completer in jupyter-ai it is here: https://jupyter-ai.readthedocs.io/en/latest/developers/index.html#custom-completion-providers

Indeed a showcase of this feature in the docs could be useful, but in the latest version it should be more discoverable since it is mentioned and linked in the chat settings panel.

[nick-youngblut]

Thanks @JasonWeill!

I will submit a issue, given that I think users would greatly benefit from at least some mention of the completion feature in the Jupyter-AI docs, and not just in the Jupyter Lab docs.

Also, it's not clear (to a naive user) if the completions feature is compatible with a GitHub Copilot license. My institute has GitHub copilot licenses, which works great with the Jupyter Notebook extension in VS Code (very easy to setup); however, ssh-remote in VS Code kills all running kernels if the ssh connection is lost (unlike what can be done with running Jupyter directly). Due to this "killed-kernels" issue, large-scale data analysis is not practical with VS Code + Jupyter, since one must re-run their notebooks if they lose their ssh connection.

So, we would like to use Jupyter Lab directly (running on a Slurm cluster), instead of running inside VS Code, but it appears that our GitHub Copilot licenses will then be a waste of money, if we cannot use GitHub Copilot when running Jupyter outside of VS Code.

[krassowski]

It would be lovely if Copilot provided integration with JupyterLab. Make sure to upvote https://github.com/orgs/community/discussions/63345