First draft for an API to add models from other VSCode extensions.

Exports functions "addExtensionModel" and
"removeExtensionModel" through the VSCode extension API.

`addExtensionModel` accepts a `CustomLLM` object, as well as
callbacks to the third-party extension to be notified
when the model is added and removed.

Added field `extensionModels` to `ContinueConfig`, which holds the `CustomLLM` model specifications for models added
by extensions. This field is not serialized to `config.json`. Config
loading routines are adadpted to construct `CustomLLMClass`
instances from all models declared in `extensionModels` and
add them to `models`.

Also reworked `vscode/src/loadConfig.ts:ConfigHandler` to
keep track of extension models on the VSCode extension side
and improve configuration update event handling generally.

core/config/load.ts

@@ -115,4 +115,18 @@ async function intermediateToFinalConfig(
   };
 }
 
-export { intermediateToFinalConfig, serializedToIntermediateConfig };
+function injectExtensionModelsToFinalConfig(config: ContinueConfig, extensionModels: readonly CustomLLM[]) {
+  var configWithExtensionModels = {...config};
+  configWithExtensionModels.extensionModels = [...extensionModels];
+
+  var models = [...config.models];
+  for (var modelDescription of extensionModels) {
+    const model = new CustomLLMClass(modelDescription);
+    models.push(model);
+  }
+  configWithExtensionModels.models = models;
+
+  return configWithExtensionModels;
+}
+
+export { injectExtensionModelsToFinalConfig, intermediateToFinalConfig, serializedToIntermediateConfig };

core/index.d.ts

@@ -1,3 +1,7 @@
+/**
+ * 2024-02 Modified by Lukas Prediger, Copyright (c) 2023 CSC - IT Center for Science Ltd.
+ */
+
 export interface ChunkWithoutID {
   content: string;
   startLine: number;
@@ -586,6 +590,7 @@ export interface Config {
 export interface ContinueConfig {
   allowAnonymousTelemetry?: boolean;
   models: ILLM[];
+  extensionModels?: CustomLLM[];
   systemMessage?: string;
   completionOptions?: BaseCompletionOptions;
   slashCommands?: SlashCommand[];

extensions/vscode/src/activation/activate.ts

@@ -1,3 +1,7 @@
+/**
+ * 2024-02 Modified by Lukas Prediger, Copyright (c) 2023 CSC - IT Center for Science Ltd.
+ */
+
 import { getTsConfigPath } from "core/util/paths";
 import * as fs from "fs";
 import path from "path";
@@ -86,7 +90,8 @@ export async function activateExtension(context: vscode.ExtensionContext) {
   setupInlineTips(context);
   showRefactorMigrationMessage();
 
-  ideProtocolClient = new IdeProtocolClient(context);
+  ideProtocolClient = new IdeProtocolClient();
+  context.subscriptions.push(ideProtocolClient);
 
   // Register Continue GUI as sidebar webview, and beginning a new session
   const provider = new ContinueGUIWebviewViewProvider();
@@ -157,4 +162,10 @@ export async function activateExtension(context: vscode.ExtensionContext) {
   } catch (e) {
     console.log("Error adding .continueignore file icon: ", e);
   }
+
+  const extensionApi = {
+    "addExtensionModel": ideProtocolClient.addExtensionModel,
+    "removeExtensionModel": ideProtocolClient.removeExtensionModel,
+  };
+  return extensionApi;
 }

extensions/vscode/src/continueIdeClient.ts

@@ -1,4 +1,8 @@
-import { FileEdit, RangeInFile } from "core";
+/**
+ * 2024-02 Modified by Lukas Prediger, Copyright (c) 2023 CSC - IT Center for Science Ltd.
+ */
+
+import { CustomLLM, FileEdit, RangeInFile } from "core";
 import { getConfigJsonPath, getDevDataFilePath } from "core/util/paths";
 import { readFileSync, writeFileSync } from "fs";
 import * as path from "path";
@@ -20,22 +24,26 @@ import {
   openEditorAndRevealRange,
   uriFromFilePath,
 } from "./util/vscode";
+import { VsCodeIde } from "./ideProtocol";
 const util = require("util");
 const exec = util.promisify(require("child_process").exec);
 
 const continueVirtualDocumentScheme = "continue";
 
-class IdeProtocolClient {
+class IdeProtocolClient implements vscode.Disposable {
   private static PREVIOUS_BRANCH_FOR_WORKSPACE_DIR: { [dir: string]: string } =
     {};
 
-  private readonly context: vscode.ExtensionContext;
+  private disposable: vscode.Disposable;
 
-  constructor(context: vscode.ExtensionContext) {
-    this.context = context;
+  constructor() {
+    const extensionModelsChangeSubscription = configHandler.onConfigChanged(
+      e => { this.configUpdate(e.config); },
+      this
+    );
 
     // Listen for file saving
-    vscode.workspace.onDidSaveTextDocument((event) => {
+    const configFileChangeSubscription = vscode.workspace.onDidSaveTextDocument((event) => {
       const filepath = event.uri.fsPath;
 
       if (
@@ -45,10 +53,7 @@ class IdeProtocolClient {
         filepath.endsWith(".continue\\config.ts") ||
         filepath.endsWith(".continuerc.json")
       ) {
-        const config = readFileSync(getConfigJsonPath(), "utf8");
-        const configJson = JSON.parse(config);
-        this.configUpdate(configJson);
-        configHandler.reloadConfig();
+        configHandler.reloadConfig(new VsCodeIde());
       } else if (
         filepath.endsWith(".continueignore") ||
         filepath.endsWith(".gitignore")
@@ -96,14 +101,21 @@ class IdeProtocolClient {
         return uri.query;
       }
     })();
-    context.subscriptions.push(
+
+    this.disposable = vscode.Disposable.from(
+      extensionModelsChangeSubscription,
+      configFileChangeSubscription,
       vscode.workspace.registerTextDocumentContentProvider(
         continueVirtualDocumentScheme,
         documentContentProvider
       )
     );
   }
 
+  dispose() {
+    this.disposable.dispose();
+  }
+
   visibleMessages: Set<string> = new Set();
 
   configUpdate(config: any) {
@@ -113,6 +125,26 @@ class IdeProtocolClient {
     });
   }
 
+  addExtensionModel(customLLM: CustomLLM, modelAddedCallback?: () => void, modelRemovedCallback?: () => void): vscode.Disposable {
+    var eventSubscription = configHandler.onExtensionModelsChange(e => {
+      if (modelAddedCallback && e.added?.find(addedModelTitle => addedModelTitle === customLLM.options?.title)) {
+        modelAddedCallback();
+      }
+
+      if (modelRemovedCallback && e.removed?.find(removedModelTitle => removedModelTitle === customLLM.options?.title)) {
+        modelRemovedCallback();
+      }
+    });
+
+    configHandler.addExtensionModel(customLLM);
+
+    return eventSubscription;
+  }
+
+  removeExtensionModel(customLLMTitle: string) {
+    configHandler.removeExtensionModel(customLLMTitle);
+  }
+
   async gotoDefinition(
     filepath: string,
     position: vscode.Position

extensions/vscode/src/debugPanel.ts

@@ -1,3 +1,7 @@
+/**
+ * 2024-02 Modified by Lukas Prediger, Copyright (c) 2023 CSC - IT Center for Science Ltd.
+ */
+
 import { ContextItemId, DiffLine, FileEdit, ModelDescription } from "core";
 import { indexDocs } from "core/indexing/docs";
 import TransformersJsEmbeddingsProvider from "core/indexing/embeddings/TransformersJsEmbeddingsProvider";
@@ -448,6 +452,10 @@ export function getSidebarContent(
           ideProtocolClient.logDevData(data.tableName, data.data);
           break;
         }
+        case "onLoad":
+          const config = await configHandler.loadConfig(ide);
+          ideProtocolClient.configUpdate(config);
+          break;
         case "addModel": {
           const model = data.model;
           const config = readFileSync(getConfigJsonPath(), "utf8");
@@ -461,7 +469,8 @@ export function getSidebarContent(
             2
           );
           writeFileSync(getConfigJsonPath(), newConfigString);
-          ideProtocolClient.configUpdate(configJson);
+          // ideProtocolClient.configUpdate(configJson);
+          await configHandler.reloadConfig(ide);
 
           ideProtocolClient.openFile(getConfigJsonPath());
 
@@ -502,13 +511,17 @@ export function getSidebarContent(
           break;
         }
         case "deleteModel": {
-          const configJson = editConfigJson((config) => {
-            config.models = config.models.filter(
-              (m: any) => m.title !== data.title
-            );
-            return config;
-          });
-          ideProtocolClient.configUpdate(configJson);
+          // if the model is an extension model, we let configHandler handle the removal
+          if (!configHandler.removeExtensionModel(data.title)) {
+            // otherwise, we need to remove it from the config JSON and reload the config manually
+            const configJson = editConfigJson((config) => {
+              config.models = config.models.filter(
+                (m: any) => m.title !== data.title
+              );
+              return config;
+            });
+            await configHandler.reloadConfig(ide);
+          }
           break;
         }
         case "addOpenAIKey": {
@@ -522,7 +535,7 @@ export function getSidebarContent(
             });
             return config;
           });
-          ideProtocolClient.configUpdate(configJson);
+          await configHandler.reloadConfig(ide);
           break;
         }
         case "llmStreamComplete": {

extensions/vscode/src/extension.ts

@@ -1,5 +1,7 @@
 /**
  * This is the entry point for the extension.
+ * 
+ * 2024-02 Modified by Lukas Prediger, Copyright (c) 2023 CSC - IT Center for Science Ltd.
  */
 
 import * as vscode from "vscode";
@@ -31,7 +33,8 @@ async function dynamicImportAndActivate(context: vscode.ExtensionContext) {
 
   const { activateExtension } = await import("./activation/activate");
   try {
-    await activateExtension(context);
+    const ownApi = await activateExtension(context);
+    return ownApi;
   } catch (e) {
     console.log("Error activating extension: ", e);
     vscode.window
@@ -52,7 +55,7 @@ async function dynamicImportAndActivate(context: vscode.ExtensionContext) {
 }
 
 export function activate(context: vscode.ExtensionContext) {
-  dynamicImportAndActivate(context);
+  return dynamicImportAndActivate(context);
 }
 
 export function deactivate() {