[WIP] Elixir code completion

Dockerfile

@@ -6,7 +6,7 @@ FROM ${BASE_IMAGE} AS build
 
 RUN apt-get update && apt-get upgrade -y && \
     apt-get install --no-install-recommends -y \
-        build-essential git && \
+        wget build-essential git && \
     apt-get purge -y --auto-remove -o APT::AutoRemove::RecommendsImportant=false && \
     apt-get clean -y && \
     rm -rf /var/lib/apt/lists/*
@@ -28,6 +28,8 @@ RUN mix local.hex --force && \
 # Build for production
 ENV MIX_ENV=prod
 
+ENV XLA_TARGET=cuda120
+
 # Install mix dependencies
 COPY mix.exs mix.lock ./
 COPY config config

assets/js/hooks/cell_editor.js

@@ -28,6 +28,7 @@ const CellEditor = {
           revision,
           this.props.language,
           this.props.intellisense,
+          this.props.copilot,
           this.props.readOnly,
           code_markers,
           doctest_reports
@@ -80,6 +81,11 @@ const CellEditor = {
         "data-intellisense",
         parseBoolean
       ),
+      copilot: getAttributeOrThrow(
+        this.el,
+        "data-copilot",
+        parseBoolean
+      ),
       readOnly: getAttributeOrThrow(this.el, "data-read-only", parseBoolean),
     };
   },

assets/js/hooks/cell_editor/live_editor.js

@@ -11,6 +11,9 @@ import Doctest from "./live_editor/doctest";
 import { initVimMode } from "monaco-vim";
 import { EmacsExtension, unregisterKey } from "monaco-emacs";
 
+import debounce from 'lodash/debounce';
+
+
 /**
  * Mounts cell source editor with real-time collaboration mechanism.
  */
@@ -24,6 +27,7 @@ class LiveEditor {
     revision,
     language,
     intellisense,
+    copilot,
     readOnly,
     codeMarkers,
     doctestReports
@@ -34,6 +38,7 @@ class LiveEditor {
     this.source = source;
     this.language = language;
     this.intellisense = intellisense;
+    this.copilot = copilot;
     this.readOnly = readOnly;
     this._onMount = [];
     this._onChange = [];
@@ -80,6 +85,10 @@ class LiveEditor {
       this._setupIntellisense();
     }
 
+    if (this.copilot) {
+      this._setupInlineCopilot();
+    }
+
     this.editorClient.setEditorAdapter(new MonacoEditorAdapter(this.editor));
 
     this.editor.onDidFocusEditorWidget(() => {
@@ -638,6 +647,156 @@ class LiveEditor {
     });
   }
 
+  _setupInlineCopilot() {
+
+    const settings = settingsStore.get();
+
+    this.copilotHandlerByRef = {};
+
+    // There's an issue in monaco that means the commands will only ever be executed
+    // on the most recently added editor. To avoid this, we can use addAction
+    // as described here: https://github.com/microsoft/monaco-editor/issues/2947
+    this.editor.addAction({
+      id: "copilot_completion",
+      label: "Copilot completion",
+      keybindings: [monaco.KeyMod.WinCtrl | monaco.KeyCode.Space],
+      run: () => {
+        console.log("Keyboard trigger!")
+        this.editor.trigger("copilot", "editor.action.inlineSuggest.trigger");
+      }
+    });
+
+
+    // TODO insert client-side completion cache to avoid needlessly hitting server
+    this.editor.getModel().__getCopilotCompletionItems__ = (model, position, context, token) => {
+
+      console.log("Copilot completion items provider", context, token)
+      const contextBeforeCursor = model.getValueInRange(new monaco.Range(1, 1, position.lineNumber, position.column));
+      const contextAfterCursor = model.getValueInRange(new monaco.Range(position.lineNumber, position.column, model.getLineCount(), model.getLineMaxColumn(model.getLineCount())));
+
+      // When triggered automatically, we don't show suggestions when the cursor is in the middle of a 
+      // word or at the beginning of the editor
+      if (context.triggerKind == monaco.languages.InlineCompletionTriggerKind.Automatic) {
+        console.log("Triggered automatically")
+        if (!/\S/.test(contextBeforeCursor)) {
+          console.log("Context before cursor consists entirely of whitespace")
+          return null
+        }
+        if (/^\S/.test(contextAfterCursor)) {
+          console.log("Bailing out because cursor is in the middle of some text")
+          return null;
+        }
+      } else {
+        console.log("Triggered manually")
+      }
+
+      // not sure if we need this, yet
+      if (context.selectedSuggestionInfo !== undefined) {
+        console.debug("Autocomplete widget is visible");
+      }
+
+      return this._asyncCopilotRequest("completion", {
+        context_before_cursor: contextBeforeCursor,
+        context_after_cursor: contextAfterCursor
+      }).then((response) => {
+        if (response.error) {
+          console.error("Copilot completion failed", response.error)
+          return
+        }
+
+        if (!response.completions) {
+          console.warn("No copilot completion items received")
+          return
+        }
+        const items = response.completions.map((completion) => {
+          return {
+            insertText: completion,
+            range: new monaco.Range(
+              position.lineNumber,
+              position.column,
+              position.lineNumber,
+              position.column
+            )
+          }
+        });
+        console.log("items", items)
+        return {
+          items
+        };
+      });
+    }
+
+    // Not sure if there is a way to get monaco to debounce inline completions.
+    // We only want to show code completion items after certain delay (unless triggered via keyboard)
+    // Adapted from here: https://blog.smithers.dev/posts/debounce-with-promise/
+    // TODO is here the right place for the debouncing? this code also gets called when manually triggering the completion, right?
+    function debouncePromise(debounceDelay, func) {
+      let promiseResolverRef = {
+        current: function () {}
+      };
+
+      // TODO don't debounce when triggered by keyboard shortcut
+      var debouncedFunc = debounce(function (...args) {
+        let promiseResolverSnapshot = promiseResolverRef.current;
+        let returnedPromise = func(...args)
+        if (returnedPromise == null) {
+          return;
+        }
+        returnedPromise.then(function (...args) {
+          console.log("then called")
+          if (promiseResolverSnapshot === promiseResolverRef.current) {
+            promiseResolverRef.current(...args);
+          }
+        });
+      }, debounceDelay);
+
+      return function (...args) {
+        return new Promise(function (resolve) {
+          promiseResolverRef.current({
+            items: []
+          });
+          promiseResolverRef.current = resolve;
+          debouncedFunc(...args);
+        });
+      };
+    }
+    this.editor.getModel().__getCopilotCompletionItemsDebounced__ = debouncePromise(300, this.editor.getModel().__getCopilotCompletionItems__)
+
+
+    this.hook.handleEvent("copilot_response", ({
+      ref,
+      response
+    }) => {
+      const handler = this.copilotHandlerByRef[ref];
+      if (handler) {
+        handler(response);
+        delete this.copilotHandlerByRef[ref];
+      }
+    });
+  }
+
+  _asyncCopilotRequest(type, props) {
+    return new Promise((resolve, reject) => {
+      this.hook.pushEvent(
+        "copilot_request",
+        { cell_id: this.cellId, type, ...props },
+        ({ ref }) => {
+          if (ref) {
+            this.copilotHandlerByRef[ref] = (response) => {
+              if (response) {
+                resolve(response);
+              } else {
+                reject(null);
+              }
+            };
+          } else {
+            reject(null);
+          }
+        }
+      );
+    });
+  }
+
   /**
    * Sets Monaco editor mode via monaco-emacs or monaco-vim packages.
    */

assets/js/hooks/cell_editor/live_editor/monaco.js

@@ -227,6 +227,29 @@ settingsStore.getAndSubscribe((settings) => {
       },
     }
   );
+
+  window.monaco = monaco // so i can access it in the dev tools
+  monaco.languages.registerInlineCompletionsProvider(
+    "elixir",
+    {
+      freeInlineCompletions: (completions) => {
+        // TODO?
+        return;
+      },
+      provideInlineCompletions: (model, position, context, token) => {
+        let getItems = model.__getCopilotCompletionItems__;
+        if (context.triggerKind == monaco.languages.InlineCompletionTriggerKind.Automatic) {
+          getItems = model.__getCopilotCompletionItemsDebounced__;
+        }
+
+        if (getItems) {
+          return getItems(model, position, context, token);
+        } else {
+          return null;
+        }
+      }
+    }
+  );
 });
 
 monaco.languages.registerHoverProvider("elixir", {

config/config.exs

@@ -37,6 +37,40 @@ config :livebook,
   within_iframe: false,
   allowed_uri_schemes: []
 
+# TODO keen for some advice how best to structure this configuration
+# feels a bit jank
+config :livebook, Livebook.Copilot,
+  enabled: true,
+  # backend: Livebook.Copilot.HuggingfaceBackend,
+  # backend_config: %{
+  #   model: "deepseek-coder-1.3b"
+  # }
+
+  backend: Livebook.Copilot.DummyBackend,
+  backend_config: %{
+    model: "echo"
+  }
+
+# backend: Livebook.Copilot.BumblebeeBackend,
+# backend_config: %{
+#   model: "gpt2"
+# }
+
+# backend: Livebook.Copilot.LlamaCppHttpBackend,
+# backend_config: %{
+#   model: "codellama-7b"
+# }
+
+# backend: Livebook.Copilot.Openai,
+# backend_config: %{
+#   api_key: System.get_env("OPENAI_API_KEY"),
+#   model: "gpt-4-1106-preview"
+# }
+
+config :nx,
+  default_backend: EXLA.Backend,
+  client: :host
+
 # Import environment specific config. This must remain at the bottom
 # of this file so it overrides the configuration defined above.
 import_config "#{Mix.env()}.exs"

config/prod.exs

@@ -15,6 +15,24 @@ config :livebook, :iframe_port, 8081
 # Set log level to warning by default to reduce output
 config :logger, level: :warning
 
+config :livebook, Livebook.Copilot,
+  enabled: true,
+  backend: Livebook.Copilot.BumblebeeBackend,
+  backend_config: %{
+    model: "deepseek-coder-1.3b",
+    client: :cuda
+  }
+
+# backend_config: %{
+#   model: "gpt2",
+#   client: :host
+# }
+
+config :nx,
+  default_backend: EXLA.Backend,
+  device: :cuda,
+  client: :cuda
+
 # ## SSL Support
 #
 # To get SSL working, you will need to add the `https` key

lib/livebook/application.ex

@@ -41,7 +41,13 @@ defmodule Livebook.Application do
         # Start the registry for managing unique connections
         {Registry, keys: :unique, name: Livebook.HubsRegistry},
         # Start the supervisor dynamically managing connections
-        {DynamicSupervisor, name: Livebook.HubsSupervisor, strategy: :one_for_one}
+        {DynamicSupervisor, name: Livebook.HubsSupervisor, strategy: :one_for_one},
+
+        # TODO no idea if this is the best way to do this ... please help
+        # My thinking is that this way we can load the model in copilot.ex for the first time
+        # without blocking anything
+        {DynamicSupervisor,
+         name: Livebook.Copilot.BumblebeeServingSupervisor, strategy: :one_for_one}
       ] ++
         if serverless?() do
           []