feat(agent): postprocess add calculate replace range by syntax. (#765)

clients/tabby-agent/src/AgentConfig.ts

@@ -31,6 +31,10 @@ export type AgentConfig = {
         experimentalKeepBlockScopeWhenCompletingLine: boolean;
       };
     };
+    calculateReplaceRange: {
+      // Prefer to use syntax parser than bracket stack
+      experimentalSyntax: boolean;
+    };
   };
   logs: {
     level: "debug" | "error" | "silent";
@@ -76,6 +80,9 @@ export const defaultAgentConfig: AgentConfig = {
         experimentalKeepBlockScopeWhenCompletingLine: false,
       },
     },
+    calculateReplaceRange: {
+      experimentalSyntax: false,
+    },
   },
   logs: {
     level: "silent",

clients/tabby-agent/src/TabbyAgent.ts

@@ -538,7 +538,7 @@ export class TabbyAgent extends EventEmitter implements Agent {
         throw options.signal.reason;
       }
       // Calculate replace range
-      completionResponse = await calculateReplaceRange(completionResponse, context);
+      completionResponse = await calculateReplaceRange(context, this.config.postprocess, completionResponse);
       if (options?.signal?.aborted) {
         throw options.signal.reason;
       }

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

@@ -170,44 +170,46 @@ describe("postprocess", () => {
   });
 
   describe("calculateReplaceRangeByBracketStack: bad cases", () => {
-    const context = {
+    it("cannot handle the case of completion bracket stack is same with suffix but should not be replaced", () => {
-      ...documentContext`
+      const context = {
-      function clamp(n: number, max: number, min: number): number {
+        ...documentContext`
-        return Math.max(Math.min(║);
+        function clamp(n: number, max: number, min: number): number {
-      }
+          return Math.max(Math.min(║);
-      `,
+        }
-      language: "typescript",
+        `,
-    };
+        language: "typescript",
-    const response = {
+      };
-      id: "",
+      const response = {
-      choices: [
+        id: "",
-        {
+        choices: [
-          index: 0,
+          {
-          text: inline`
+            index: 0,
-                                 ├n, max), min┤
+            text: inline`
-          `,
+                                   ├n, max), min┤
-          replaceRange: {
+            `,
-            start: context.position,
+            replaceRange: {
-            end: context.position,
+              start: context.position,
+              end: context.position,
+            },
           },
-        },
+        ],
-      ],
+      };
-    };
+      const expected = {
-    const expected = {
+        id: "",
-      id: "",
+        choices: [
-      choices: [
+          {
-        {
+            index: 0,
-          index: 0,
+            text: inline`
-          text: inline`
+                                   ├n, max), min┤
-                                 ├n, max), min┤
+            `,
-          `,
+            replaceRange: {
-          replaceRange: {
+              start: context.position,
-            start: context.position,
+              end: context.position,
-            end: context.position,
+            },
           },
-        },
+        ],
-      ],
+      };
-    };
+      expect(calculateReplaceRangeByBracketStack(response, context)).not.to.deep.equal(expected);
-    expect(calculateReplaceRangeByBracketStack(response, context)).not.to.deep.equal(expected);
+    });
   });
 });

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

@@ -19,10 +19,16 @@ export function calculateReplaceRangeByBracketStack(
     }
     if (suffixText.startsWith(unpaired)) {
       choice.replaceRange.end = context.position + unpaired.length;
-      logger.trace({ context, completion: choice.text, range: choice.replaceRange, unpaired }, "Adjust replace range");
+      logger.trace(
+        { context, completion: choice.text, range: choice.replaceRange, unpaired },
+        "Adjust replace range by bracket stack",
+      );
     } else if (unpaired.startsWith(suffixText)) {
       choice.replaceRange.end = context.position + suffixText.length;
-      logger.trace({ context, completion: choice.text, range: choice.replaceRange, unpaired }, "Adjust replace range");
+      logger.trace(
+        { context, completion: choice.text, range: choice.replaceRange, unpaired },
+        "Adjust replace range by bracket stack",
+      );
     }
   }
   return response;

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

@@ -0,0 +1,213 @@
+import { expect } from "chai";
+import { documentContext, inline } from "./testUtils";
+import { calculateReplaceRangeBySyntax } from "./calculateReplaceRangeBySyntax";
+
+describe("postprocess", () => {
+  describe("calculateReplaceRangeBySyntax", () => {
+    it("should handle auto closing quotes", async () => {
+      const context = {
+        ...documentContext`
+        const hello = "║"
+        `,
+        language: "typescript",
+      };
+      const response = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                       ├hello";┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position,
+            },
+          },
+        ],
+      };
+      const expected = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                       ├hello";┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position + 1,
+            },
+          },
+        ],
+      };
+      expect(await calculateReplaceRangeBySyntax(response, context)).to.deep.equal(expected);
+    });
+
+    it("should handle auto closing quotes", async () => {
+      const context = {
+        ...documentContext`
+        let htmlMarkup = \`║\`
+        `,
+        language: "typescript",
+      };
+      const response = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                           ├<h1>\${message}</h1>\`;┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position,
+            },
+          },
+        ],
+      };
+      const expected = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                           ├<h1>\${message}</h1>\`;┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position + 1,
+            },
+          },
+        ],
+      };
+      expect(await calculateReplaceRangeBySyntax(response, context)).to.deep.equal(expected);
+    });
+
+    it("should handle multiple auto closing brackets", async () => {
+      const context = {
+        ...documentContext`
+        process.on('data', (data) => {║})
+        `,
+        language: "typescript",
+      };
+      const response = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                                      ├
+              console.log(data);
+            });┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position,
+            },
+          },
+        ],
+      };
+      const expected = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                                      ├
+              console.log(data);
+            });┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position + 2,
+            },
+          },
+        ],
+      };
+      expect(await calculateReplaceRangeBySyntax(response, context)).to.deep.equal(expected);
+    });
+
+    it("should handle multiple auto closing brackets", async () => {
+      const context = {
+        ...documentContext`
+        let mat: number[][][] = [[[║]]]
+        `,
+        language: "typescript",
+      };
+      const response = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                                   ├1, 2], [3, 4]], [[5, 6], [7, 8]]];┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position,
+            },
+          },
+        ],
+      };
+      const expected = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                                   ├1, 2], [3, 4]], [[5, 6], [7, 8]]];┤
+            `,
+            replaceRange: {
+              start: context.position,
+              end: context.position + 3,
+            },
+          },
+        ],
+      };
+      expect(await calculateReplaceRangeBySyntax(response, context)).to.deep.equal(expected);
+    });
+
+    it("should handle the bad case of calculateReplaceRangeByBracketStack", async () => {
+      const context = {
+        ...documentContext`
+      function clamp(n: number, max: number, min: number): number {
+        return Math.max(Math.min(║);
+      }
+      `,
+        language: "typescript",
+      };
+      const response = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                                 ├n, max), min┤
+          `,
+            replaceRange: {
+              start: context.position,
+              end: context.position,
+            },
+          },
+        ],
+      };
+      const expected = {
+        id: "",
+        choices: [
+          {
+            index: 0,
+            text: inline`
+                                 ├n, max), min┤
+          `,
+            replaceRange: {
+              start: context.position,
+              end: context.position,
+            },
+          },
+        ],
+      };
+      expect(await calculateReplaceRangeBySyntax(response, context)).to.deep.equal(expected);
+    });
+  });
+});

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

@@ -0,0 +1,54 @@
+import type TreeSitterParser from "web-tree-sitter";
+import { getParser, languagesConfigs } from "../syntax/parser";
+import { CompletionContext, CompletionResponse } from "../Agent";
+import { isBlank, splitLines } from "../utils";
+import { logger } from "./base";
+
+export const supportedLanguages = Object.keys(languagesConfigs);
+
+export async function calculateReplaceRangeBySyntax(
+  response: CompletionResponse,
+  context: CompletionContext,
+): Promise<CompletionResponse> {
+  const { position, prefix, suffix, prefixLines, suffixLines, language } = context;
+  if (supportedLanguages.indexOf(language) < 0) {
+    return response;
+  }
+  const languageConfig = languagesConfigs[language];
+  const parser = await getParser(languageConfig);
+  const prefixText = prefixLines[prefixLines.length - 1];
+  const suffixText = suffixLines[0]?.trimEnd() || "";
+  if (isBlank(suffixText)) {
+    return response;
+  }
+  for (const choice of response.choices) {
+    const completionText = choice.text.slice(position - choice.replaceRange.start);
+    const completionLines = splitLines(completionText);
+    let replaceLength = 0;
+    let tree = parser.parse(prefix + completionText + suffix);
+    let node = tree.rootNode.namedDescendantForIndex(prefix.length + completionText.length);
+    while (node.hasError() && replaceLength < suffixText.length) {
+      replaceLength++;
+      const row = prefixLines.length - 1 + completionLines.length - 1;
+      let column = completionLines[completionLines.length - 1].length;
+      if (completionLines.length == 1) {
+        column += prefixLines[prefixLines.length - 1].length;
+      }
+      tree.edit({
+        startIndex: prefix.length + completionText.length,
+        oldEndIndex: prefix.length + completionText.length + 1,
+        newEndIndex: prefix.length + completionText.length,
+        startPosition: { row, column },
+        oldEndPosition: { row, column: column + 1 },
+        newEndPosition: { row, column },
+      });
+      tree = parser.parse(prefix + completionText + suffix.slice(replaceLength), tree);
+      node = tree.rootNode.namedDescendantForIndex(prefix.length + completionText.length);
+    }
+    if (!node.hasError()) {
+      choice.replaceRange.end = position + replaceLength;
+      logger.trace({ context, completion: choice.text, range: choice.replaceRange }, "Adjust replace range by syntax");
+    }
+  }
+  return response;
+}

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

@@ -1,5 +1,6 @@
 import { CompletionContext, CompletionResponse } from "../Agent";
 import { AgentConfig } from "../AgentConfig";
+import { isBrowser } from "../env";
 import { applyFilter } from "./base";
 import { removeRepetitiveBlocks } from "./removeRepetitiveBlocks";
 import { removeRepetitiveLines } from "./removeRepetitiveLines";
@@ -9,6 +10,7 @@ import { trimSpace } from "./trimSpace";
 import { dropDuplicated } from "./dropDuplicated";
 import { dropBlank } from "./dropBlank";
 import { calculateReplaceRangeByBracketStack } from "./calculateReplaceRangeByBracketStack";
+import { calculateReplaceRangeBySyntax, supportedLanguages } from "./calculateReplaceRangeBySyntax";
 
 export async function preCacheProcess(
   context: CompletionContext,
@@ -37,8 +39,13 @@ export async function postCacheProcess(
 }
 
 export async function calculateReplaceRange(
-  response: CompletionResponse,
   context: CompletionContext,
+  config: AgentConfig["postprocess"],
+  response: CompletionResponse,
 ): Promise<CompletionResponse> {
-  return calculateReplaceRangeByBracketStack(response, context);
+  return isBrowser || // syntax parser is not supported in browser yet
+    !config["calculateReplaceRange"].experimentalSyntax ||
+    !supportedLanguages.includes(context.language)
+    ? calculateReplaceRangeByBracketStack(response, context)
+    : calculateReplaceRangeBySyntax(response, context);
 }

clients/tabby-agent/tests/golden.test.ts

@@ -85,6 +85,7 @@ describe("agent golden test", () => {
         experimentalSyntax: false,
         indentation: { experimentalKeepBlockScopeWhenCompletingLine: false },
       },
+      calculateReplaceRange: { experimentalSyntax: false },
     },
     logs: { level: "debug" },
     anonymousUsageTracking: { disable: true },
@@ -147,20 +148,37 @@ describe("agent golden test", () => {
   });
 
   it("updateConfig experimental", async () => {
-    requestId++;
-    const updateConfigRequest = [
-      requestId,
-      {
-        func: "updateConfig",
-        args: ["postprocess.limitScope.experimentalSyntax", true],
-      },
-    ];
-    agent.stdin.write(JSON.stringify(updateConfigRequest) + "\n");
-    await waitForResponse(requestId);
     const expectedConfig = { ...config };
-    expectedConfig.postprocess.limitScope.experimentalSyntax = true;
+    {
-    expect(output.shift()).to.deep.equal([0, { event: "configUpdated", config: expectedConfig }]);
+      requestId++;
-    expect(output.shift()).to.deep.equal([requestId, true]);
+      const updateConfigRequest = [
+        requestId,
+        {
+          func: "updateConfig",
+          args: ["postprocess.limitScope.experimentalSyntax", true],
+        },
+      ];
+      agent.stdin.write(JSON.stringify(updateConfigRequest) + "\n");
+      await waitForResponse(requestId);
+      expectedConfig.postprocess.limitScope.experimentalSyntax = true;
+      expect(output.shift()).to.deep.equal([0, { event: "configUpdated", config: expectedConfig }]);
+      expect(output.shift()).to.deep.equal([requestId, true]);
+    }
+    {
+      requestId++;
+      const updateConfigRequest = [
+        requestId,
+        {
+          func: "updateConfig",
+          args: ["postprocess.calculateReplaceRange.experimentalSyntax", true],
+        },
+      ];
+      agent.stdin.write(JSON.stringify(updateConfigRequest) + "\n");
+      await waitForResponse(requestId);
+      expectedConfig.postprocess.calculateReplaceRange.experimentalSyntax = true;
+      expect(output.shift()).to.deep.equal([0, { event: "configUpdated", config: expectedConfig }]);
+      expect(output.shift()).to.deep.equal([requestId, true]);
+    }
   });
   badCasesFiles.forEach((goldenFile) => {
     it("experimental: " + goldenFile.path, async () => {