feat(chat-ui): add smart apply to chat panel

clients/tabby-agent/package.json

@@ -42,6 +42,7 @@
     "@types/fast-levenshtein": "^0.0.4",
     "@types/fs-extra": "^11.0.1",
     "@types/glob": "^7.2.0",
+    "@types/js-levenshtein": "^1.1.3",
     "@types/mocha": "^10.0.1",
     "@types/node": "18.x",
     "@types/object-hash": "^3.0.0",
@@ -68,6 +69,7 @@
     "file-stream-rotator": "^1.0.0",
     "fs-extra": "^11.1.1",
     "glob": "^7.2.0",
+    "js-levenshtein": "^1.1.6",
     "jwt-decode": "^3.1.2",
     "lru-cache": "^9.1.1",
     "mac-ca": "^2.0.3",

clients/tabby-agent/src/chat/SmartApply.ts

@@ -0,0 +1,335 @@
+import {
+  CancellationToken,
+  Connection,
+  Location,
+  Range,
+  ShowDocumentParams,
+  TextDocuments,
+} from "vscode-languageserver";
+import type { Feature } from "../feature";
+import {
+  ChatEditDocumentTooLongError,
+  ChatEditMutexError,
+  ChatFeatureNotAvailableError,
+  ServerCapabilities,
+  SmartApplyCodeRequest,
+  SmartApplyCodeParams,
+} from "../protocol";
+import { Configurations } from "../config";
+import { TabbyApiClient } from "../http/tabbyApiClient";
+import cryptoRandomString from "crypto-random-string";
+import { getLogger } from "../logger";
+import { readResponseStream, showDocument } from "./utils";
+import { TextDocument } from "vscode-languageserver-textdocument";
+import { getSmartApplyRange } from "./SmartRange";
+import { Edit } from "./inlineEdit";
+const logger = getLogger("ChatEditProvider");
+
+export class SmartApplyFeature implements Feature {
+  private lspConnection: Connection | undefined = undefined;
+  private mutexAbortController: AbortController | undefined = undefined;
+  constructor(
+    private readonly configurations: Configurations,
+    private readonly tabbyApiClient: TabbyApiClient,
+    private readonly documents: TextDocuments<TextDocument>,
+  ) {}
+
+  initialize(connection: Connection): ServerCapabilities | Promise<ServerCapabilities> {
+    this.lspConnection = connection;
+    connection.onRequest(SmartApplyCodeRequest.type, async (params, token) => {
+      return this.provideSmartApplyEdit(params, token);
+    });
+
+    return {};
+  }
+  initialized?(): void | Promise<void> {
+    //nothing
+  }
+  shutdown?(): void | Promise<void> {
+    //nothing
+  }
+
+  async provideSmartApplyEdit(params: SmartApplyCodeParams, token: CancellationToken): Promise<boolean> {
+    logger.info("Getting document");
+    const document = this.documents.get(params.location.uri);
+    if (!document) {
+      logger.info("Document not found, returning false");
+      return false;
+    }
+    if (!this.lspConnection) {
+      logger.info("LSP connection lost.");
+      return false;
+    }
+
+    if (this.mutexAbortController && !this.mutexAbortController.signal.aborted) {
+      logger.warn("Another smart edit is already in progress");
+      throw {
+        name: "ChatEditMutexError",
+        message: "Another smart edit is already in progress",
+      } as ChatEditMutexError;
+    }
+    this.mutexAbortController = new AbortController();
+    logger.info("mutex abort status: " + (this.mutexAbortController === undefined));
+    token.onCancellationRequested(() => this.mutexAbortController?.abort());
+
+    let applyRange = getSmartApplyRange(document, params.applyCode);
+    //if cannot find range, lets use backend LLMs
+    if (!applyRange) {
+      if (!this.tabbyApiClient.isChatApiAvailable) {
+        return false;
+      }
+      applyRange = await provideSmartApplyLineRange(
+        document,
+        params.applyCode,
+        this.tabbyApiClient,
+        this.configurations,
+      );
+    }
+
+    if (!applyRange) {
+      return false;
+    }
+
+    try {
+      //reveal editor range
+      const revealEditorRangeParams: ShowDocumentParams = {
+        uri: params.location.uri,
+        selection: {
+          start: applyRange.range.start,
+          end: applyRange.range.end,
+        },
+        takeFocus: true,
+      };
+      await showDocument(revealEditorRangeParams, this.lspConnection);
+    } catch (error) {
+      logger.warn("cline not support reveal range");
+    }
+
+    try {
+      await provideSmartApplyEditLLM(
+        {
+          uri: params.location.uri,
+          range: {
+            start: applyRange.range.start,
+            end: { line: applyRange.range.end.line + 1, character: 0 },
+          },
+        },
+        params.applyCode,
+        applyRange.action === "insert" ? true : false,
+        document,
+        this.lspConnection,
+        this.tabbyApiClient,
+        this.configurations,
+        this.mutexAbortController,
+        () => {
+          this.mutexAbortController = undefined;
+        },
+      );
+      return true;
+    } catch (error) {
+      logger.error("Error applying smart edit:", error);
+      return false;
+    } finally {
+      logger.info("Resetting mutex abort controller");
+      this.mutexAbortController = undefined;
+    }
+  }
+}
+
+export async function provideSmartApplyLineRange(
+  document: TextDocument,
+  applyCodeBlock: string,
+  tabbyApiClient: TabbyApiClient,
+  configurations: Configurations,
+): Promise<{ range: Range; action: "insert" | "replace" } | undefined> {
+  if (!document) {
+    return undefined;
+  }
+  if (!tabbyApiClient.isChatApiAvailable()) {
+    throw {
+      name: "ChatFeatureNotAvailableError",
+      message: "Chat feature not available",
+    } as ChatFeatureNotAvailableError;
+  }
+
+  const documentText = document
+    .getText()
+    .split("\n")
+    .map((line, idx) => `${idx + 1} | ${line}`)
+    .join("\n");
+
+  const config = configurations.getMergedConfig();
+  const promptTemplate = config.chat.provideSmartApplyLineRange.promptTemplate;
+
+  const messages: { role: "user"; content: string }[] = [
+    {
+      role: "user",
+      content: promptTemplate.replace(/{{document}}|{{applyCode}}/g, (pattern: string) => {
+        switch (pattern) {
+          case "{{document}}":
+            return documentText;
+          case "{{applyCode}}":
+            return applyCodeBlock;
+          default:
+            return "";
+        }
+      }),
+    },
+  ];
+
+  try {
+    const readableStream = await tabbyApiClient.fetchChatStream({
+      messages,
+      model: "",
+      stream: true,
+    });
+
+    if (!readableStream) {
+      return undefined;
+    }
+
+    let response = "";
+    for await (const chunk of readableStream) {
+      response += chunk;
+    }
+
+    const regex = /<GENERATEDCODE>(.*?)<\/GENERATEDCODE>/s;
+    const match = response.match(regex);
+    if (match && match[1]) {
+      response = match[1].trim();
+    }
+
+    const range = response.split("-");
+    if (range.length !== 2) {
+      return undefined;
+    }
+
+    const startLine = parseInt(range[0] ?? "0", 10) - 1;
+    const endLine = parseInt(range[1] ?? "0", 10) - 1;
+
+    return {
+      range: {
+        start: { line: startLine < 0 ? 0 : startLine, character: 0 },
+        end: { line: endLine < 0 ? 0 : endLine, character: Number.MAX_SAFE_INTEGER },
+      },
+      action: startLine == endLine ? "insert" : "replace",
+    };
+  } catch (error) {
+    return undefined;
+  }
+}
+
+export async function provideSmartApplyEditLLM(
+  location: Location,
+  applyCode: string,
+  insertMode: boolean,
+  document: TextDocument,
+  lspConnection: Connection,
+  tabbyApiClient: TabbyApiClient,
+  configurations: Configurations,
+  mutexAbortController: AbortController,
+  onResetMutex: () => void,
+): Promise<boolean> {
+  if (!document) {
+    logger.warn("Document not found");
+    return false;
+  }
+  if (!lspConnection) {
+    logger.warn("LSP connection failed");
+    return false;
+  }
+
+  if (!tabbyApiClient.isChatApiAvailable()) {
+    throw {
+      name: "ChatFeatureNotAvailableError",
+      message: "Chat feature not available",
+    } as ChatFeatureNotAvailableError;
+  }
+
+  const config = configurations.getMergedConfig();
+  const documentText = document.getText();
+  const selection = {
+    start: document.offsetAt(location.range.start),
+    end: document.offsetAt(location.range.end),
+  };
+  const selectedDocumentText = documentText.substring(selection.start, selection.end);
+
+  logger.info("current selectedDoc: " + selectedDocumentText);
+
+  if (selection.end - selection.start > config.chat.edit.documentMaxChars) {
+    throw { name: "ChatEditDocumentTooLongError", message: "Document too long" } as ChatEditDocumentTooLongError;
+  }
+
+  const promptTemplate = config.chat.provideSmartApply.promptTemplate;
+
+  // Extract the selected text and the surrounding context
+  let documentPrefix = documentText.substring(0, selection.start);
+  let documentSuffix = documentText.substring(selection.end);
+  if (documentText.length > config.chat.edit.documentMaxChars) {
+    const charsRemain = config.chat.edit.documentMaxChars - selectedDocumentText.length;
+    if (documentPrefix.length < charsRemain / 2) {
+      documentSuffix = documentSuffix.substring(0, charsRemain - documentPrefix.length);
+    } else if (documentSuffix.length < charsRemain / 2) {
+      documentPrefix = documentPrefix.substring(documentPrefix.length - charsRemain + documentSuffix.length);
+    } else {
+      documentPrefix = documentPrefix.substring(documentPrefix.length - charsRemain / 2);
+      documentSuffix = documentSuffix.substring(0, charsRemain / 2);
+    }
+  }
+
+  const messages: { role: "user"; content: string }[] = [
+    {
+      role: "user",
+      content: promptTemplate.replace(/{{document}}|{{code}}/g, (pattern: string) => {
+        switch (pattern) {
+          case "{{document}}":
+            return selectedDocumentText;
+          case "{{code}}":
+            return applyCode || "";
+          default:
+            return "";
+        }
+      }),
+    },
+  ];
+
+  try {
+    const readableStream = await tabbyApiClient.fetchChatStream({
+      messages,
+      model: "",
+      stream: true,
+    });
+
+    if (!readableStream) {
+      return false;
+    }
+    const editId = "tabby-" + cryptoRandomString({ length: 6, type: "alphanumeric" });
+    const currentEdit: Edit = {
+      id: editId,
+      location: location,
+      languageId: document.languageId,
+      originalText: selectedDocumentText,
+      editedRange: insertMode
+        ? { start: location.range.start, end: location.range.end }
+        : { start: location.range.start, end: location.range.end },
+      editedText: "",
+      comments: "",
+      buffer: "",
+      state: "editing",
+    };
+
+    await readResponseStream(
+      readableStream,
+      lspConnection,
+      currentEdit,
+      mutexAbortController,
+      onResetMutex,
+      config.chat.edit.responseDocumentTag,
+      config.chat.edit.responseCommentTag,
+    );
+
+    return true;
+  } catch (error) {
+    return false;
+  }
+}

clients/tabby-agent/src/chat/SmartRange.ts

@@ -0,0 +1,45 @@
+import levenshtein from "js-levenshtein";
+import { Position, Range } from "vscode-languageserver-protocol";
+import { TextDocument } from "vscode-languageserver-textdocument";
+
+//return [start, end] close interval 0-based range
+export function getSmartApplyRange(
+  document: TextDocument,
+  snippet: string,
+): { range: Range; action: "insert" | "replace" } | undefined {
+  const applyRange = fuzzyApplyRange(document, snippet);
+  if (!applyRange) {
+    return undefined;
+  }
+  //insert mode
+  if (applyRange.range.start.line === applyRange.range.end.line || document.getText().trim() === "") {
+    return { range: applyRange.range, action: "insert" };
+  }
+  return { range: applyRange.range, action: "replace" };
+}
+
+export function fuzzyApplyRange(document: TextDocument, snippet: string): { range: Range; score: number } | null {
+  const lines = document.getText().split("\n");
+  const snippetLines = snippet.split("\n");
+
+  let [minDistance, index] = [Number.MAX_SAFE_INTEGER, 0];
+  for (let i = 0; i <= lines.length - snippetLines.length; i++) {
+    const window = lines.slice(i, i + snippetLines.length).join("\n");
+    const distance = levenshtein(window, snippet);
+    if (minDistance >= distance) {
+      minDistance = distance;
+      index = i;
+    }
+  }
+
+  if (minDistance === Number.MAX_SAFE_INTEGER && index === 0) {
+    return null;
+  }
+
+  const startLine = index;
+  const endLine = index + snippetLines.length - 1;
+  const start: Position = { line: startLine, character: 0 };
+  const end: Position = { line: endLine, character: lines[endLine]?.length || 0 };
+
+  return { range: { start, end }, score: minDistance };
+}