feat(cloud): auto completion (#39)

gespispace · web-flow · commit 4e39df707c9e · 2024-04-25T06:48:06.000-04:00
diff --git a/.vscode/launch.json b/.vscode/launch.json
@@ -6,12 +6,16 @@
   "version": "0.2.0",
   "configurations": [
     {
-      "name": "Run Extension",
+      "name": "Watch & Launch Extension",
       "type": "extensionHost",
       "request": "launch",
       "args": ["--extensionDevelopmentPath=${workspaceFolder}"],
-      "outFiles": ["${workspaceFolder}/out/**/*.js"],
-      "preLaunchTask": "${defaultBuildTask}"
+      "outFiles": ["${workspaceFolder}/dist/*.js"],
+      "preLaunchTask": "npm: watch",
+      "debugWebviews": true,
+      "runtimeExecutable": "${execPath}",
+      "smartStep": true,
+      "sourceMaps": true
     },
     {
       "name": "Extension Tests",
diff --git a/.vscode/tasks.json b/.vscode/tasks.json
@@ -16,6 +16,15 @@
         "kind": "build",
         "isDefault": true
       }
+    },
+    {
+      "type": "npm",
+      "script": "watch",
+      "group": {
+        "kind": "build",
+        "isDefault": true
+      },
+      "isBackground": true
     }
   ]
 }
diff --git a/src/common/completion/cloudCompletion.ts b/src/common/completion/cloudCompletion.ts
@@ -0,0 +1,39 @@
+import { OpenAI } from "@langchain/openai";
+import { configuration } from "../utils/configuration";
+import Logger from "../logger";
+
+type Parameters = {
+  temperature: number;
+  n_predict: number;
+  controller?: AbortController;
+};
+
+export const sendCompletionsRequestCloud = async (
+  prompt: string,
+  parameters: Parameters
+) => {
+  const apiKey = configuration.get("cloud.apiToken");
+
+  const model = new OpenAI({
+    maxRetries: 0,
+    openAIApiKey: apiKey,
+    configuration: {
+      baseURL: configuration.get("cloud.endpoint"),
+    },
+    stop: ["<|file_separator|>"],
+    temperature: parameters.temperature,
+    maxTokens: parameters.n_predict,
+  });
+  try {
+    const response = await model.invoke(prompt, {
+      configurable: {
+        signal: parameters.controller,
+      },
+      maxConcurrency: 1,
+    });
+
+    return response;
+  } catch (error) {
+    Logger.error(error);
+  }
+};
diff --git a/src/common/completion/index.ts b/src/common/completion/index.ts
@@ -4,10 +4,12 @@ import { getPromptCompletion } from "../prompt";
 import { abortInterval, delay } from "./utils/intervals";
 import { getAdditionalDocuments } from "./utils/getAdditionalDocuments";
 import Logger from "../logger";
-import { sendCompletionRequest } from "./localCompletion";
+import { sendCompletionRequestLocal } from "./localCompletion";
 import { servers } from "../server";
 import { configuration } from "../utils/configuration";
 import { state } from "../utils/state";
+import { sendCompletionsRequestCloud } from "./cloudCompletion";
+import statusBar from "../statusBar";
 
 const logCompletion = () => {
   const uuid = randomUUID();
@@ -39,10 +41,12 @@ export const getInlineCompletionProvider = (
         "inlineSuggestModeAuto"
       );
 
+      // If the current mode is not auto and the trigger is automatic, don't suggest any completions.
       if (currentInlineSuggestModeAuto !== true && triggerAuto === true) {
         return [];
       }
 
+      // If there is a selected completion item, don't suggest any completions.
       if (context.selectedCompletionInfo) {
         return [];
       }
@@ -52,42 +56,14 @@ export const getInlineCompletionProvider = (
       loggerCompletion.info("Completion: started");
 
       const cancelled = await delay(250, token);
+
       if (cancelled) {
         loggerCompletion.info("Completion: canceled");
         return [];
       }
 
       const { abortController, requestFinish } = abortInterval(token);
-
-      const modelType = triggerAuto
-        ? configuration.get("completion.autoMode")
-        : configuration.get("completion.manuallyMode");
-
-      const serverUrl = servers[modelType].serverUrl;
-
-      const additionalDocuments: vscode.TextDocument[] = configuration.get(
-        "experimental.useopentabs"
-      )
-        ? await getAdditionalDocuments()
-        : [];
-
-      const prompt = await getPromptCompletion({
-        activeDocument: document,
-        additionalDocuments: additionalDocuments,
-        position: position,
-        maxTokenExpect: triggerAuto ? maxToken : 2000,
-      });
-
-      const parameters = triggerAuto
-        ? {
-            n_predict: 128,
-            stop: ["\n", "<|file_separator|>"],
-          }
-        : {
-            n_predict: 512,
-            stop: ["<|file_separator|>"],
-            temperature: 0.5,
-          };
+      const { stopTask } = statusBar.startTask();
 
       try {
         Logger.info(
@@ -97,49 +73,106 @@ export const getInlineCompletionProvider = (
             sendTelemetry: true,
           }
         );
+        if (configuration.get("cloud.use")) {
+          const prompt = await getPromptCompletion({
+            activeDocument: document,
+            additionalDocuments: await getAdditionalDocuments(),
+            position: position,
+            maxTokenExpect: 3300,
+          });
+          const completion = await sendCompletionsRequestCloud(prompt, {
+            n_predict: 512,
+            temperature: 0.5,
+            controller: abortController,
+          });
 
-        const completion = await sendCompletionRequest(
-          prompt,
-          parameters,
-          abortController,
-          loggerCompletion.uuid(),
-          serverUrl
-        );
+          if (completion === "" || completion === undefined) {
+            return [];
+          }
 
-        if (completion === null) {
-          return [];
-        }
+          if (token.isCancellationRequested) {
+            loggerCompletion.info(
+              "Request: canceled by new completion cancel token"
+            );
 
-        if (token.isCancellationRequested) {
-          loggerCompletion.info(
-            "Request: canceled by new completion cancel token"
-          );
+            return [];
+          }
 
-          return [];
-        }
-        if (triggerAuto) {
-          maxToken *= expectedTime / completion.timing;
-        }
+          return [
+            {
+              insertText: completion,
+              range: new vscode.Range(position, position),
+            },
+          ];
+        } else {
+          const additionalDocuments: vscode.TextDocument[] = configuration.get(
+            "experimental.useopentabs"
+          )
+            ? await getAdditionalDocuments()
+            : [];
+
+          const prompt = await getPromptCompletion({
+            activeDocument: document,
+            additionalDocuments: additionalDocuments,
+            position: position,
+            maxTokenExpect: triggerAuto ? maxToken : 2000,
+          });
+
+          const parameters = triggerAuto
+            ? {
+                n_predict: 128,
+                stop: ["\n", "<|file_separator|>"],
+              }
+            : {
+                n_predict: 512,
+                stop: ["<|file_separator|>"],
+                temperature: 0.5,
+              };
+          const modelType = triggerAuto
+            ? configuration.get("completion.autoMode")
+            : configuration.get("completion.manuallyMode");
+          const completion = await sendCompletionRequestLocal(
+            prompt,
+            parameters,
+            abortController,
+            loggerCompletion.uuid(),
+            servers[modelType].serverUrl
+          );
 
-        loggerCompletion.info(`maxToken: ${maxToken}`);
+          if (completion === null) {
+            return [];
+          }
 
-        loggerCompletion.info("Request: finished");
+          if (token.isCancellationRequested) {
+            loggerCompletion.info(
+              "Request: canceled by new completion cancel token"
+            );
 
-        return [
-          {
-            insertText: completion.content,
-            range: new vscode.Range(position, position),
-          },
-        ];
+            return [];
+          }
+          if (triggerAuto) {
+            maxToken *= expectedTime / completion.timing;
+          }
+          loggerCompletion.info("Request: finished");
+
+          return [
+            {
+              insertText: completion.content,
+              range: new vscode.Range(position, position),
+            },
+          ];
+        }
       } catch (error) {
         const Error = error as Error;
         Logger.error(error);
 
         const errorMessage = Error.message;
         vscode.window.showErrorMessage(errorMessage);
       } finally {
+        stopTask();
         requestFinish();
       }
+      return [];
     },
   };
 
diff --git a/src/common/completion/localCompletion.ts b/src/common/completion/localCompletion.ts
@@ -1,7 +1,6 @@
 import * as vscode from "vscode";
 import { randomUUID } from "crypto";
 import Logger from "../logger";
-import statusBar from "../statusBar";
 
 const logCompletion = (uuid = randomUUID() as string) => {
   return {
@@ -36,15 +35,13 @@ const defualtParameters = {
   slot_id: -1,
 };
 
-export const sendCompletionRequest = async (
+export const sendCompletionRequestLocal = async (
   prompt: string,
   parameters: Record<string, any>,
   abortController: AbortController,
   uuid: string,
   url: string
 ) => {
-  const { stopTask } = statusBar.startTask();
-
   const loggerCompletion = logCompletion(uuid);
 
   const parametersForCompletion = {
@@ -120,7 +117,5 @@ export const sendCompletionRequest = async (
     const errorMessage = Error.message;
     vscode.window.showErrorMessage(errorMessage);
     return null;
-  } finally {
-    stopTask();
   }
 };
diff --git a/src/common/prompt/promptCompletion.ts b/src/common/prompt/promptCompletion.ts
@@ -25,10 +25,10 @@ const tokenize = async (text: string): Promise<number> => {
 
 const getTextNormalized = (text: string) => {
   return text
-    .replace("<|fim_prefix|>", "")
-    .replace("<|fim_middle|>", "")
-    .replace("<|fim_suffix|>", "")
-    .replace("<|file_separator|>", "");
+    .replaceAll("<|fim_prefix|>", "")
+    .replaceAll("<|fim_middle|>", "")
+    .replaceAll("<|fim_suffix|>", "")
+    .replaceAll("<|file_separator|>", "");
 };
 
 const spliteDocumentByPosition = (
@@ -161,7 +161,7 @@ export const getPromptCompletion = async ({
 }) => {
   const start = performance.now();
 
-  const maxTokenHardLimit = 6000;
+  const maxTokenHardLimit = 4000;
   const maxToken =
     maxTokenExpect > maxTokenHardLimit ? maxTokenHardLimit : maxTokenExpect;
 
@@ -209,7 +209,7 @@ export const getPromptCompletion = async ({
 
   const activeDocumentFileName =
     additionalDocumentsText === ""
-      ? ""
+      ? "<|fim_prefix|>"
       : "\n" +
         "/" +
         getRelativePath(activeDocument.uri) +
diff --git a/src/extension.ts b/src/extension.ts

Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,15 @@`
`16`	`16`	`"kind": "build",`
`17`	`17`	`"isDefault": true`
`18`	`18`	`}`
	`19`	`+ },`
	`20`	`+ {`
	`21`	`+ "type": "npm",`
	`22`	`+ "script": "watch",`
	`23`	`+ "group": {`
	`24`	`+ "kind": "build",`
	`25`	`+ "isDefault": true`
	`26`	`+ },`
	`27`	`+ "isBackground": true`
`19`	`28`	`}`
`20`	`29`	`]`
`21`	`30`	`}`