update cloud support

package.json

@@ -148,6 +148,11 @@
           "description": "The directory to use for storing user data associated with this extension. Set to an empty string, `os.homedir()` will be used."
         },
         "firecoder.cloud.endpoint": {
+          "type": "string",
+          "default": "https://llm-api.firecoder.cc/v1",
+          "description": ""
+        },
+        "firecoder.cloud.apiToken": {
           "type": "string",
           "default": "",
           "description": ""
@@ -191,6 +196,8 @@
   },
   "dependencies": {
     "@grafana/faro-core": "^1.3.5",
-    "@grafana/faro-web-sdk": "^1.3.5"
+    "@grafana/faro-web-sdk": "^1.3.5",
+    "@langchain/community": "^0.0.27",
+    "langchain": "^0.1.17"
   }
-}
+}

src/common/chat/cloudChat.ts

@@ -0,0 +1,52 @@
+import { ChatOpenAI } from "@langchain/openai";
+import { StringOutputParser } from "langchain/schema/output_parser";
+import {
+  AIMessage,
+  HumanMessage,
+  SystemMessage,
+} from "@langchain/core/messages";
+import { HistoryMessage } from "../prompt";
+import { configuration } from "../utils/configuration";
+
+type Parameters = {
+  temperature: number;
+  n_predict: number;
+};
+
+export const sendChatRequestCloud = async (
+  history: HistoryMessage[],
+  parameters: Parameters
+) => {
+  const apiKey = configuration.get("cloud.apiToken");
+
+  const model = new ChatOpenAI({
+    maxRetries: 0,
+    openAIApiKey: apiKey,
+    configuration: {
+      baseURL: configuration.get("cloud.endpoint"),
+    },
+    temperature: parameters.temperature,
+    maxTokens: parameters.n_predict,
+  });
+
+  const parser = new StringOutputParser();
+
+  const messages = history.map((message) => {
+    switch (message.role) {
+      case "ai":
+        return new AIMessage(message.content);
+      case "user":
+        return new HumanMessage(message.content);
+      case "system":
+        return new SystemMessage(message.content);
+      default:
+        return new HumanMessage(message.content);
+    }
+  });
+
+  const stream = await model.pipe(parser).stream(messages, {
+    maxConcurrency: 1,
+  });
+
+  return stream;
+};

src/common/chat/index.ts

@@ -1,11 +1,16 @@
 import * as vscode from "vscode";
 import { randomUUID } from "crypto";
-import { HistoryMessage, getPromptChat } from "../prompt/promptChat";
+import { ChatMessage, getPromptChat } from "../prompt/promptChat";
 import Logger from "../logger";
-import { sendChatRequest } from "./localChat";
-import { servers } from "../server";
+import { sendChatRequestLocal as sendChatRequestLocal } from "./localChat";
 import { configuration } from "../utils/configuration";
 import statusBar from "../statusBar";
+import { sendChatRequestCloud } from "./cloudChat";
+import {
+  getHighlightedTextDetails,
+  humanMessageWithCodePrompt,
+  systemTemplate,
+} from "./utils/useHighlightedTextAsContext";
 
 const logCompletion = () => {
   const uuid = randomUUID();
@@ -17,37 +22,55 @@ const logCompletion = () => {
   };
 };
 
-export async function* chat(history: HistoryMessage[]) {
+export async function* chat(
+  history: ChatMessage[],
+  config?: {
+    provideHighlightedText?: boolean;
+  }
+) {
   const loggerCompletion = logCompletion();
 
   loggerCompletion.info("Chat: started");
 
-  const prompt = await getPromptChat(history);
-
   const parameters = {
     n_predict: 4096,
     stop: [],
     temperature: 0.7,
   };
 
-  const serverUrl = configuration.get("cloud.use")
-    ? configuration.get("cloud.endpoint")
-    : servers["chat-medium"].serverUrl;
-
   const { stopTask } = statusBar.startTask();
-
   try {
     Logger.info(`Start request;`, {
       component: "chat",
       sendTelemetry: true,
     });
 
-    yield* sendChatRequest(
-      prompt,
-      parameters,
-      loggerCompletion.uuid(),
-      serverUrl
-    );
+    if (config?.provideHighlightedText) {
+      const highlighted = getHighlightedTextDetails();
+
+      if (highlighted !== null && history.length === 1) {
+        const firstMessage = history.shift();
+        history.unshift({
+          role: "user",
+          content: await humanMessageWithCodePrompt.format({
+            highlightedCode: highlighted.content,
+            question: firstMessage?.content,
+          }),
+        });
+      }
+      history.unshift({
+        role: "system",
+        content: systemTemplate,
+      });
+    }
+
+    if (configuration.get("cloud.use")) {
+      yield* await sendChatRequestCloud(history, parameters);
+    } else {
+      const prompt = getPromptChat(history);
+
+      yield* sendChatRequestLocal(prompt, parameters, loggerCompletion.uuid());
+    }
 
     loggerCompletion.info("Request: finished");
   } catch (error) {

src/common/chat/localChat.ts

@@ -1,6 +1,7 @@
 import * as vscode from "vscode";
 import { randomUUID } from "crypto";
 import Logger from "../logger";
+import { servers } from "../server";
 
 const logCompletion = (uuid = randomUUID() as string) => {
   return {
@@ -33,13 +34,13 @@ const defualtParameters = {
   slot_id: -1,
 };
 
-export async function* sendChatRequest(
+export async function* sendChatRequestLocal(
   prompt: string,
   parameters: Record<string, any>,
-  uuid: string,
-  url: string
+  uuid: string
 ) {
   const loggerCompletion = logCompletion(uuid);
+  const url = servers["chat-medium"].serverUrl;
 
   const parametersForCompletion = {
     ...defualtParameters,
@@ -51,14 +52,12 @@ export async function* sendChatRequest(
 
     const startTime = performance.now();
 
-    let content = "";
     let timings;
     for await (const chunk of llama(prompt, parametersForCompletion, { url })) {
       // @ts-ignore
       if (chunk.data) {
         // @ts-ignore
-        content += chunk.data.content;
-        yield content;
+        yield chunk.data.content;
       }
 
       // @ts-ignore

src/common/chat/utils/useHighlightedTextAsContext.ts

@@ -0,0 +1,74 @@
+import { PromptTemplate } from "langchain/prompts";
+import * as vscode from "vscode";
+
+export const getHighlightedTextDetails = (): {
+  content: string;
+  startLine: number;
+  endLine: number;
+  filePath: string;
+} | null => {
+  // Get the active vscode text editor instance.
+  let editor = vscode.window.activeTextEditor;
+
+  if (!editor) {
+    return null; // If no editor is open, we'll return null.
+  }
+
+  const selection = editor?.selection;
+
+  if (selection && !selection.isEmpty) {
+    // Create a new Range object with the start and end of selected text in vscode document.
+    let rangeSelectionText: vscode.Range = new vscode.Range(
+      selection.start.line,
+      selection.start.character,
+      selection.end.line,
+      selection.end.character
+    );
+
+    // Get the text from selected lines in document using Range object we created above.
+    let content = editor?.document.getText(rangeSelectionText) || "";
+
+    return {
+      filePath: (editor && editor.document.uri.fsPath) || "",
+      startLine: selection.start.line, // line number where the selected text starts from
+      endLine: selection.end.line, // line number where the selected text ends to
+      content: content, // The actual highlighted code/text in string format
+    };
+  } else {
+    return null;
+  }
+};
+
+export const systemTemplate = `
+You are an AI Assistant who\`s a coding expert. Answer on below question.
+
+Input format:
+- The provided code is the code that the user selected for this question, usually, it is relative to the question.
+- User question.
+
+Instructions:
+- Answer in short form.
+- Don't answer, if you don't know the answer.
+- Don't explain why you did this change.
+- Only answer what the user asked.
+- If the code is not relative to the question, don't use it for the answer.
+- If the answer expects any code examples, please provide examples. 
+
+Rules for code in response:
+- Suggest only changes. 
+- Provide only the necessary code.
+- Write as less as possible comments.
+`;
+
+export const humanMessageWithCodePrompt = new PromptTemplate({
+  inputVariables: ["highlightedCode", "question"],
+  template: `Please provide an answer on this question using a provided code.
+
+Provided Code: 
+\`\`\`
+{highlightedCode}
+\`\`\`
+
+Question: {question}
+`,
+});

src/common/completion/index.ts

@@ -1,7 +1,8 @@
 import * as vscode from "vscode";
 import { randomUUID } from "crypto";
 import { getPromptCompletion } from "../prompt";
-import { abortInterval, delay } from "../utils/intervals";
+import { abortInterval, delay } from "./utils/intervals";
+import { getAdditionalDocuments } from "./utils/getAdditionalDocuments";
 import Logger from "../logger";
 import { sendCompletionRequest } from "./localCompletion";
 import { servers } from "../server";
@@ -23,6 +24,7 @@ export const getInlineCompletionProvider = (
 ) => {
   let maxToken = 50;
   let expectedTime = 1000;
+
   const provider: vscode.InlineCompletionItemProvider = {
     provideInlineCompletionItems: async (
       document,
@@ -32,15 +34,19 @@ export const getInlineCompletionProvider = (
     ) => {
       const triggerAuto =
         context.triggerKind === vscode.InlineCompletionTriggerKind.Automatic;
+
       const currentInlineSuggestModeAuto = state.workspace.get(
         "inlineSuggestModeAuto"
       );
+
       if (currentInlineSuggestModeAuto !== true && triggerAuto === true) {
         return [];
       }
+
       if (context.selectedCompletionInfo) {
         return [];
       }
+
       const loggerCompletion = logCompletion();
 
       loggerCompletion.info("Completion: started");
@@ -50,57 +56,39 @@ export const getInlineCompletionProvider = (
         loggerCompletion.info("Completion: canceled");
         return [];
       }
+
       const { abortController, requestFinish } = abortInterval(token);
 
       const modelType = triggerAuto
         ? configuration.get("completion.autoMode")
         : configuration.get("completion.manuallyMode");
 
-      const serverUrl = configuration.get("cloud.use")
-        ? configuration.get("cloud.endpoint")
-        : servers[modelType].serverUrl;
-
-      let additionalDocuments: vscode.TextDocument[] = [];
-
-      if (configuration.get("experimental.useopentabs")) {
-        const additionalDocumentsUri = vscode.window.tabGroups.all
-          .map((group) => group.tabs)
-          .flat()
-          .filter((tab) => !(tab.group.isActive && tab.isActive))
-          .filter(
-            (tab) =>
-              tab.input &&
-              "uri" in (tab.input as any) &&
-              ((tab.input as any).uri as vscode.Uri).scheme === "file"
-          )
-          .map((tab) => (tab.input as any).uri as vscode.Uri);
-        additionalDocuments = await Promise.all(
-          additionalDocumentsUri.map((uri) =>
-            vscode.workspace.openTextDocument(uri)
-          )
-        );
-      }
+      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 && !configuration.get("cloud.use") ? maxToken : 2000,
+        maxTokenExpect: triggerAuto ? maxToken : 2000,
         url: serverUrl,
       });
 
-      const parameters =
-        triggerAuto && !configuration.get("cloud.use")
-          ? {
-              n_predict: 128,
-              stop: ["\n"],
-            }
-          : {
-              n_predict: 512,
-              stop: [],
-              temperature: 0.5,
-            };
+      const parameters = triggerAuto
+        ? {
+            n_predict: 128,
+            stop: ["\n"],
+          }
+        : {
+            n_predict: 512,
+            stop: [],
+            temperature: 0.5,
+          };
 
       try {
         Logger.info(