Code completion for jupyter lab/notebook using StarCoder

.gitignore

@@ -132,3 +132,5 @@ dev.sh
 .jupyter_ystore.db
 
 .yarn
+
+.venv

README.md

@@ -94,6 +94,17 @@ that include variable values.
 
 ![Sample with code interpolation and markdown output](./docs/source/_static/sample-markdown.png)
 
+
+## Code completion
+
+In a cell, press {Short cut configured in the bigcode sidebar} to trigger
+
+![Sample of pressing a shortcut key in code completion](./docs/source/_static/code-completion-setting-key-down.jpg)
+
+Press "enter" to accpet the proposition.
+
+![Sample of pressing enter after a successful request in code completion](./docs/source/_static/code-completion-setting-completion.jpg)
+
 ## JupyterLab extension
 
 The Jupyter AI extension for JupyterLab offers a native UI that enables multiple users

docs/source/users/index.md

@@ -743,3 +743,43 @@ The `--region-name` parameter is set to the [AWS region code](https://docs.aws.a
 The `--request-schema` parameter is the JSON object the endpoint expects as input, with the prompt being substituted into any value that matches the string literal `"<prompt>"`. For example, the request schema `{"text_inputs":"<prompt>"}` will submit a JSON object with the prompt stored under the `text_inputs` key.
 
 The `--response-path` option is a [JSONPath](https://goessner.net/articles/JsonPath/index.html) string that retrieves the language model's output from the endpoint's JSON response. For example, if your endpoint returns an object with the schema `{"generated_texts":["<output>"]}`, its response path is `generated_texts.[0]`.
+
+
+## Code completion
+
+### Usage
+
+We first need to configure the huggingface token and short cut
+
+![](../_static/code-completion-setting.jpg)
+
+In a cell, press {The short cut configured in the image above} to trigger
+
+![](../_static/code-completion-setting-key-down.jpg)
+
+Press "enter" to accpet the proposition.
+
+![](../_static/code-completion-setting-completion.jpg)
+
+
+### Configure
+
+Bigcode service url:
+
+- The endpoint URL of the Bigcode service. This is the base URL that the code completion system will connect to when requesting code completions.
+
+Huggingface Access Token:
+
+- Your personal access token for Huggingface's service. This token authenticates and authorizes your requests to the Huggingface API.
+
+Short cut for completion:
+
+- A shortcut or key combination that, when pressed, triggers the code completion feature.
+
+Max prompt tokens:
+
+- The maximum number of tokens that can be sent in the prompt to the code completion API.
+
+Max response tokens:
+
+- Set 'Max Response Tokens' judiciously; a higher value can slow response times, and since we don't support streaming, you'll await the entire response.

examples/code-completion/sklearn_digits.ipynb

@@ -0,0 +1,83 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "dcbb2940-4d47-45ac-923a-f80eed1d26f5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# load Digits Dataset from sklearn\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f3a60c45-244b-49f2-8527-519c5a2557a6",
+   "metadata": {},
+   "source": [
+    "Print to show there are 1797 images (8 by 8 images for a dimensionality of 64)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0c159867-a499-4181-9713-d2558f1a258d",
+   "metadata": {},
+   "source": [
+    "Print to show there are 1797 labels (integers from 0–9)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e216707c-1ea7-446a-ac69-cd013e76b5f2",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fe688a3e-29cb-464d-9eef-7e17ce4b9dee",
+   "metadata": {},
+   "source": [
+    "Showing the Images and the Labels (Digits Dataset)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f2f5505f-32b1-4e22-8878-39f20e71335d",
+   "metadata": {},
+   "source": [
+    "This section is really just to show what the images and labels look like. It usually helps to visualize your data to see what you are working with."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4328fb40-b40c-4fde-8ef4-6a5b38c5dd10",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

packages/jupyter-ai/package.json

@@ -73,6 +73,8 @@
     "@jupyterlab/ui-components": "^4",
     "@mui/icons-material": "^5.11.0",
     "@mui/material": "^5.11.0",
+    "mobx": "^6.10.0",
+    "mobx-react": "^9.0.0",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
     "react-markdown": "^8.0.6",
@@ -103,7 +105,7 @@
     "stylelint-config-standard": "~24.0.0",
     "stylelint-prettier": "^2.0.0",
     "ts-jest": "^29",
-    "typescript": "~4.3.0",
+    "typescript": "5.1.6",
     "y-protocols": "^1.0.5"
   },
   "sideEffects": [

packages/jupyter-ai/src/bigcode/bigcode-code-completion.ts

@@ -0,0 +1,188 @@
+import { JupyterFrontEnd } from '@jupyterlab/application';
+import { EditorView } from '@codemirror/view';
+import { getNotebookContentUntilCursor } from '../utils/cell-context';
+import {
+  sendToBigCode,
+  processCompletionResult,
+  constructContinuationPrompt
+} from '../utils/bigcode-request';
+import {
+  insertAndHighlightCode,
+  removeTextStatus,
+  moveCursorToEnd,
+  replaceText
+} from '../utils/cell-modification';
+import {
+  addLoadingAnimation,
+  updateAnimation,
+  removeLoadingAnimation
+} from '../utils/animation';
+import { ICell } from '../types/cell';
+
+import CodeCompletionContextStore from '../contexts/code-completion-context-store';
+
+/**
+ * An object to track the state of the request and the visibility of the result.
+ * @property {boolean} loading - Indicates if a request is currently being made.
+ * @property {boolean} viewResult - Indicates if the result of a request is currently being displayed.
+ */
+const requestState = {
+  loading: false,
+  viewResult: false
+};
+
+/**
+ * Checks if the provided context is empty or not.
+ * @param {ICell[]} context - An array of cells representing the context.
+ * @returns {boolean} Returns true if the context is empty or undefined, otherwise false.
+ */
+const isContextEmpty = (context: ICell[]): boolean => {
+  if (!context || context.length === 0) {
+    return true; // If the context is empty or undefined, we assume that the user did not write code
+  }
+
+  const combinedContent = context.reduce((acc, cell) => acc + cell.content, '');
+  return combinedContent.trim() === '';
+};
+
+/**
+ * Handles the success response of the code completion request.
+ * @param {JupyterFrontEnd} app - The JupyterFrontEnd application instance.
+ * @param {EditorView} view - The editor view instance.
+ * @param {{ generated_text: string }[]} result - The result array containing generated text.
+ */
+const requestSuccess = (
+  app: JupyterFrontEnd,
+  view: EditorView,
+  result: { generated_text: string }[]
+) => {
+  const resultCode = processCompletionResult(result);
+  requestState.viewResult = true;
+
+  if (resultCode === '') {
+    requestState.viewResult = false;
+    CodeCompletionContextStore.setCodeOnRequest('');
+  } else {
+    insertAndHighlightCode(
+      app,
+      CodeCompletionContextStore.codeOnRequest,
+      resultCode
+    );
+  }
+
+  updateAnimation(view, 'success');
+  requestState.loading = false;
+};
+
+/**
+ * Handles the failure of the code completion request.
+ * @param {EditorView} view - The editor view instance.
+ */
+const requestFailed = (view: EditorView) => {
+  CodeCompletionContextStore.setCodeOnRequest('');
+
+  updateAnimation(view, 'failed');
+  requestState.viewResult = false;
+  requestState.loading = false;
+};
+
+/**
+ * Initiates code completion based on the provided context.
+ * @param {JupyterFrontEnd} app - The JupyterFrontEnd application instance.
+ * @param {EditorView} view - The editor view instance.
+ * @returns {boolean} Indicates whether the function completed successfully.
+ */
+export const codeCompletion = (
+  app: JupyterFrontEnd,
+  view: EditorView
+): boolean => {
+  const context = getNotebookContentUntilCursor(app);
+
+  if (!context || isContextEmpty(context)) {
+    updateAnimation(view, 'failed');
+    console.error('continueWriting() => context is null');
+    return false;
+  }
+
+  if (requestState.loading || requestState.viewResult) {
+    console.error('continueWriting() => request is running');
+    return true;
+  }
+
+  console.debug('continueWriting() => context: ', context);
+
+  requestState.loading = true;
+  removeLoadingAnimation(view);
+  addLoadingAnimation(view);
+  CodeCompletionContextStore.setCodeOnRequest(
+    context[context.length - 1].content
+  );
+  const prompt = constructContinuationPrompt(
+    context,
+    CodeCompletionContextStore.maxPromptToken
+  );
+
+  console.debug('codeCompletion() => prompt: ', prompt);
+
+  sendToBigCode(prompt, CodeCompletionContextStore.maxResponseToken)
+    .then(result => {
+      console.debug('continueWriting() => state is success, result: ', result);
+      requestSuccess(app, view, result);
+    })
+    .catch(async err => {
+      console.error(await err);
+      requestFailed(view);
+    });
+
+  return true;
+};
+
+/**
+ * Removes custom colors from the text in the editor.
+ * @param {EditorView} view - The editor view instance.
+ * @returns {boolean} Indicates whether the function completed successfully.
+ */
+export const removeColor = (view: EditorView): boolean => {
+  if (
+    CodeCompletionContextStore.codeOnRequest === '' &&
+    !requestState.viewResult
+  ) {
+    console.debug(
+      'removeColor() => No request is running or no continuation code is showing'
+    );
+    return false;
+  }
+
+  requestState.viewResult = false;
+  removeTextStatus(view);
+  CodeCompletionContextStore.setCodeOnRequest('');
+  moveCursorToEnd(view);
+  console.debug('removeColor() => remove code customize color');
+  return true;
+};
+
+/**
+ * Handles any keypress events when a request is loading or a result is being displayed.
+ * @param {EditorView} view - The editor view instance.
+ * @returns {boolean} Indicates whether the function completed successfully.
+ */
+export const handleAnyKeyPress = (view: EditorView): boolean => {
+  if (requestState.loading) {
+    console.debug('handleAnyKeyPress() => request is loading');
+    return true;
+  }
+
+  if (
+    CodeCompletionContextStore.codeOnRequest !== '' ||
+    requestState.viewResult
+  ) {
+    console.debug(
+      'handleAnyKeyPress() => code for the user to cancel the display'
+    );
+    removeTextStatus(view);
+    replaceText(view, CodeCompletionContextStore.codeOnRequest);
+    CodeCompletionContextStore.setCodeOnRequest('');
+    requestState.viewResult = false;
+  }
+  return false;
+};