feat(agent): format indentation if not match with editor config. (#911)

clients/tabby-agent/src/CompletionContext.ts

@@ -6,6 +6,7 @@ export type CompletionRequest = {
   language: string;
   text: string;
   position: number;
+  indentation?: string;
   clipboard?: string;
   manually?: boolean;
 };
@@ -36,6 +37,7 @@ function isAtLineEndExcludingAutoClosedChar(suffix: string) {
 export class CompletionContext {
   filepath: string;
   language: string;
+  indentation?: string;
   text: string;
   position: number;
 
@@ -57,6 +59,7 @@ export class CompletionContext {
     this.language = request.language;
     this.text = request.text;
     this.position = request.position;
+    this.indentation = request.indentation;
 
     this.prefix = request.text.slice(0, request.position);
     this.suffix = request.text.slice(request.position);

clients/tabby-agent/src/postprocess/formatIndentation.test.ts

@@ -0,0 +1,166 @@
+import { expect } from "chai";
+import { documentContext, inline } from "./testUtils";
+import { formatIndentation } from "./formatIndentation";
+
+describe("postprocess", () => {
+  describe("formatIndentation", () => {
+    it("should format indentation if first line of completion is over indented.", () => {
+      const context = {
+        ...documentContext`
+        function clamp(n: number, max: number, min: number): number {
+          ║
+        }
+        `,
+        indentation: "  ",
+        language: "typescript",
+      };
+      const completion = inline`
+          ├  return Math.max(Math.min(n, max), min);┤
+      `;
+      const expected = inline`
+          ├return Math.max(Math.min(n, max), min);┤
+      `;
+      expect(formatIndentation(context)(completion)).to.eq(expected);
+    });
+
+    it("should format indentation if first line of completion is wrongly indented.", () => {
+      const context = {
+        ...documentContext`
+        function clamp(n: number, max: number, min: number): number {
+        ║
+        }
+        `,
+        indentation: "    ",
+        language: "typescript",
+      };
+      const completion = inline`
+        ├  return Math.max(Math.min(n, max), min);┤
+      `;
+      const expected = inline`
+        ├    return Math.max(Math.min(n, max), min);┤
+      `;
+      expect(formatIndentation(context)(completion)).to.eq(expected);
+    });
+
+    it("should format indentation if completion lines is over indented.", () => {
+      const context = {
+        ...documentContext`
+        def findMax(arr):║
+        `,
+        indentation: "  ",
+        language: "python",
+      };
+      const completion = inline`
+                         ├
+            max = arr[0]
+            for i in range(1, len(arr)):
+                if arr[i] > max:
+                    max = arr[i]
+            return max
+        }┤
+      `;
+      const expected = inline`
+                         ├
+          max = arr[0]
+          for i in range(1, len(arr)):
+            if arr[i] > max:
+              max = arr[i]
+          return max
+        }┤
+      `;
+      expect(formatIndentation(context)(completion)).to.eq(expected);
+    });
+
+    it("should format indentation if completion lines is wrongly indented.", () => {
+      const context = {
+        ...documentContext`
+        def findMax(arr):║
+        `,
+        indentation: "    ",
+        language: "python",
+      };
+      const completion = inline`
+                         ├
+          max = arr[0]
+          for i in range(1, len(arr)):
+            if arr[i] > max:
+              max = arr[i]
+          return max
+        }┤
+      `;
+      const expected = inline`
+                         ├
+            max = arr[0]
+            for i in range(1, len(arr)):
+                if arr[i] > max:
+                    max = arr[i]
+            return max
+        }┤
+      `;
+      expect(formatIndentation(context)(completion)).to.eq(expected);
+    });
+
+    it("should keep it unchanged if it no indentation specified.", () => {
+      const context = {
+        ...documentContext`
+        def findMax(arr):║
+        `,
+        indentation: undefined,
+        language: "python",
+      };
+      const completion = inline`
+                          ├
+            max = arr[0]
+            for i in range(1, len(arr)):
+                if arr[i] > max:
+                    max = arr[i]
+            return max
+        }┤
+      `;
+      expect(formatIndentation(context)(completion)).to.eq(completion);
+    });
+
+    it("should keep it unchanged if there is indentation in the context.", () => {
+      const context = {
+        ...documentContext`
+        def hello():
+            return "world"
+
+        def findMax(arr):║
+        `,
+        indentation: "\t",
+        language: "python",
+      };
+      const completion = inline`
+                          ├
+            max = arr[0]
+            for i in range(1, len(arr)):
+                if arr[i] > max:
+                    max = arr[i]
+            return max
+        }┤
+      `;
+      expect(formatIndentation(context)(completion)).to.eq(completion);
+    });
+
+    it("should keep it unchanged if it is well indented.", () => {
+      const context = {
+        ...documentContext`
+        def findMax(arr):║
+        `,
+        indentation: "    ",
+        language: "python",
+      };
+      const completion = inline`
+                          ├
+            max = arr[0]
+            for i in range(1, len(arr)):
+                if arr[i] > max:
+                    max = arr[i]
+            return max
+        }┤
+      `;
+      expect(formatIndentation(context)(completion)).to.eq(completion);
+    });
+  });
+});

clients/tabby-agent/src/postprocess/formatIndentation.ts

@@ -0,0 +1,100 @@
+import { CompletionContext } from "../Agent";
+import { PostprocessFilter, logger } from "./base";
+import { isBlank, splitLines } from "../utils";
+
+function detectIndentation(lines: string[]): string | null {
+  const matches = {
+    "\t": 0,
+    "  ": 0,
+    "    ": 0,
+  };
+  for (const line of lines) {
+    if (line.match(/^\t/)) {
+      matches["\t"]++;
+    } else {
+      const spaces = line.match(/^ */)[0].length;
+      if (spaces > 0) {
+        if (spaces % 4 === 0) {
+          matches["    "]++;
+        }
+        if (spaces % 2 === 0) {
+          matches["  "]++;
+        }
+      }
+    }
+  }
+  if (matches["\t"] > 0) {
+    return "\t";
+  }
+  if (matches["  "] > matches["    "]) {
+    return "  ";
+  }
+  if (matches["    "] > 0) {
+    return "    ";
+  }
+  return null;
+}
+
+function getIndentLevel(line: string, indentation: string): number {
+  if (indentation === "\t") {
+    return line.match(/^\t*/g)[0].length;
+  } else {
+    const spaces = line.match(/^ */)[0].length;
+    return spaces / indentation.length;
+  }
+}
+
+export function formatIndentation(context: CompletionContext): PostprocessFilter {
+  return (input) => {
+    const { prefixLines, suffixLines, indentation } = context;
+    const inputLines = splitLines(input);
+
+    // if no indentation is specified
+    if (!indentation) {
+      return input;
+    }
+
+    // if there is any indentation in context, the server output should have learned from it
+    const prefixLinesForDetection = isBlank(prefixLines[prefixLines.length - 1])
+      ? prefixLines.slice(0, prefixLines.length - 1)
+      : prefixLines;
+    if (prefixLines.length > 1 && detectIndentation(prefixLinesForDetection) !== null) {
+      return input;
+    }
+    const suffixLinesForDetection = suffixLines.slice(1);
+    if (suffixLines.length > 1 && detectIndentation(suffixLinesForDetection) !== null) {
+      return input;
+    }
+
+    // if the input is well indented with specific indentation
+    const inputLinesForDetection = inputLines.map((line, index) => {
+      return index === 0 ? prefixLines[prefixLines.length - 1] + line : line;
+    });
+    const inputIndentation = detectIndentation(inputLinesForDetection);
+    if (inputIndentation === null || inputIndentation === indentation) {
+      return input;
+    }
+
+    // otherwise, do formatting
+    const formatted = inputLinesForDetection.map((line, index) => {
+      const level = getIndentLevel(inputLinesForDetection[index], inputIndentation);
+      if (level === 0) {
+        return inputLines[index];
+      }
+      const rest = line.slice(inputIndentation.length * level);
+      if (index === 0) {
+        // for first line
+        if (!isBlank(prefixLines[prefixLines.length - 1])) {
+          return inputLines[0];
+        } else {
+          return indentation.repeat(level).slice(prefixLines[prefixLines.length - 1].length) + rest;
+        }
+      } else {
+        // for next lines
+        return indentation.repeat(level) + rest;
+      }
+    });
+    logger.debug({ prefixLines, suffixLines, inputLines, formatted }, "Format indentation.");
+    return formatted.join("");
+  };
+}

clients/tabby-agent/src/postprocess/index.ts

@@ -6,6 +6,7 @@ import { removeRepetitiveBlocks } from "./removeRepetitiveBlocks";
 import { removeRepetitiveLines } from "./removeRepetitiveLines";
 import { removeLineEndsWithRepetition } from "./removeLineEndsWithRepetition";
 import { limitScope } from "./limitScope";
+import { formatIndentation } from "./formatIndentation";
 import { trimSpace } from "./trimSpace";
 import { dropDuplicated } from "./dropDuplicated";
 import { dropBlank } from "./dropBlank";
@@ -33,6 +34,7 @@ export async function postCacheProcess(
     .then(applyFilter(removeRepetitiveBlocks(context), context))
     .then(applyFilter(removeRepetitiveLines(context), context))
     .then(applyFilter(limitScope(context, config["limitScope"]), context))
+    .then(applyFilter(formatIndentation(context), context))
     .then(applyFilter(dropDuplicated(context), context))
     .then(applyFilter(trimSpace(context), context))
     .then(applyFilter(dropBlank(), context));

clients/vscode/src/TabbyCompletionProvider.ts

@@ -72,6 +72,7 @@ export class TabbyCompletionProvider extends EventEmitter implements InlineCompl
       language: document.languageId, // https://code.visualstudio.com/docs/languages/identifiers
       text: additionalContext.prefix + document.getText() + additionalContext.suffix,
       position: additionalContext.prefix.length + document.offsetAt(position),
+      indentation: this.getEditorIndentation(),
       clipboard: await env.clipboard.readText(),
       manually: context.triggerKind === InlineCompletionTriggerKind.Invoke,
     };
@@ -166,6 +167,21 @@ export class TabbyCompletionProvider extends EventEmitter implements InlineCompl
     }
   }
 
+  private getEditorIndentation(): string | undefined {
+    const editor = window.activeTextEditor;
+    if (!editor) {
+      return undefined;
+    }
+
+    const { insertSpaces, tabSize } = editor.options;
+    if (insertSpaces && typeof tabSize === "number" && tabSize > 0) {
+      return " ".repeat(tabSize);
+    } else if (!insertSpaces) {
+      return "\t";
+    }
+    return undefined;
+  }
+
   private updateConfiguration() {
     if (!workspace.getConfiguration("editor").get("inlineSuggest.enabled", true)) {
       this.triggerMode = "disabled";