refactor(agent): move calcReplaceRange to a postprocess filter. (#1138)

* refactor(agent): Refactor calcReplaceRange to a postprocess filter. Fallback to non-syntax based filter when syntax error. * fix: remove console.log.
TabbyML · Dec 28, 2023 · a18ff03 · a18ff03
1 parent d6174d0
commit a18ff03
Show file tree

Hide file tree

Showing 14 changed files with 325 additions and 383 deletions.
diff --git a/clients/tabby-agent/src/TabbyAgent.ts b/clients/tabby-agent/src/TabbyAgent.ts
@@ -27,7 +27,7 @@ import { configFile } from "./configFile";
 import { CompletionCache } from "./CompletionCache";
 import { CompletionDebounce } from "./CompletionDebounce";
 import { CompletionContext } from "./CompletionContext";
-import { preCacheProcess, postCacheProcess, calculateReplaceRange } from "./postprocess";
+import { preCacheProcess, postCacheProcess } from "./postprocess";
 import { rootLogger, allLoggers } from "./logger";
 import { AnonymousUsageLogger } from "./AnonymousUsageLogger";
 import { CompletionProviderStats, CompletionProviderStatsEntry } from "./CompletionProviderStats";
@@ -556,11 +556,6 @@ export class TabbyAgent extends EventEmitter implements Agent {
       if (signal.aborted) {
         throw signal.reason;
       }
-      // Calculate replace range
-      completionResponse = await calculateReplaceRange(context, this.config.postprocess, completionResponse);
-      if (signal.aborted) {
-        throw signal.reason;
-      }
     } catch (error) {
       if (isCanceledError(error) || isTimeoutError(error)) {
         if (stats) {

diff --git a/clients/tabby-agent/src/postprocess/base.ts b/clients/tabby-agent/src/postprocess/base.ts
@@ -1,13 +1,20 @@
-import { CompletionResponse, CompletionContext } from "../CompletionContext";
+import { CompletionResponse, CompletionResponseChoice, CompletionContext } from "../CompletionContext";
 import { rootLogger } from "../logger";
 
-export type PostprocessFilter = (item: string, context: CompletionContext) => string | null | Promise<string | null>;
+type PostprocessFilterBase<T extends string | CompletionResponseChoice> = (
+  input: T,
+  context: CompletionContext,
+) => T | null | Promise<T | null>;
+
+export type PostprocessFilter = PostprocessFilterBase<string>;
+export type PostprocessChoiceFilter = PostprocessFilterBase<CompletionResponseChoice>;
 
 export const logger = rootLogger.child({ component: "Postprocess" });
 
 declare global {
   interface Array<T> {
     distinct(identity?: (x: T) => any): Array<T>;
+    mapAsync<U>(callbackfn: (value: T, index: number, array: T[]) => U | Promise<U>, thisArg?: any): Promise<U[]>;
   }
 }
 
@@ -17,22 +24,43 @@ if (!Array.prototype.distinct) {
   };
 }
 
+if (!Array.prototype.mapAsync) {
+  Array.prototype.mapAsync = async function <T, U>(
+    this: T[],
+    callbackfn: (value: T, index: number, array: T[]) => U | Promise<U>,
+    thisArg?: any,
+  ): Promise<U[]> {
+    return await Promise.all(this.map((item, index) => callbackfn.call(thisArg, item, index, this)));
+  };
+}
+
 export function applyFilter(
   filter: PostprocessFilter,
   context: CompletionContext,
+): (response: CompletionResponse) => Promise<CompletionResponse> {
+  return applyChoiceFilter(async (choice) => {
+    const replaceLength = context.position - choice.replaceRange.start;
+    const input = choice.text.slice(replaceLength);
+    const filtered = await filter(input, context);
+    choice.text = choice.text.slice(0, replaceLength) + (filtered ?? "");
+    return choice;
+  }, context);
+}
+
+export function applyChoiceFilter(
+  choiceFilter: PostprocessChoiceFilter,
+  context: CompletionContext,
 ): (response: CompletionResponse) => Promise<CompletionResponse> {
   return async (response: CompletionResponse) => {
     response.choices = (
-      await Promise.all(
-        response.choices.map(async (choice) => {
-          const replaceLength = context.position - choice.replaceRange.start;
-          const filtered = await filter(choice.text.slice(replaceLength), context);
-          choice.text = choice.text.slice(0, replaceLength) + (filtered ?? "");
-          return choice;
-        }),
-      )
+      await response.choices.mapAsync(async (choice) => {
+        return await choiceFilter(choice, context);
+      })
     )
-      .filter((choice) => !!choice.text)
+      .filter<CompletionResponseChoice>((choice): choice is NonNullable<CompletionResponseChoice> => {
+        // Filter out empty choices.
+        return !!choice && !!choice.text;
+      })
       .distinct((choice) => choice.text);
     return response;
   };

diff --git a/clients/tabby-agent/src/postprocess/calculateReplaceRange.ts b/clients/tabby-agent/src/postprocess/calculateReplaceRange.ts
@@ -0,0 +1,24 @@
+import { AgentConfig } from "../AgentConfig";
+import { isBrowser } from "../env";
+import { PostprocessChoiceFilter, logger } from "./base";
+import { calculateReplaceRangeByBracketStack } from "./calculateReplaceRangeByBracketStack";
+import { calculateReplaceRangeBySyntax } from "./calculateReplaceRangeBySyntax";
+
+export function calculateReplaceRange(
+  config: AgentConfig["postprocess"]["calculateReplaceRange"],
+): PostprocessChoiceFilter {
+  return async (choice, context) => {
+    const preferSyntaxParser =
+      !isBrowser && // syntax parser is not supported in browser yet
+      config.experimentalSyntax;
+
+    if (preferSyntaxParser) {
+      try {
+        return await calculateReplaceRangeBySyntax(choice, context);
+      } catch (error) {
+        logger.debug({ error }, "Failed to calculate replace range by syntax parser");
+      }
+    }
+    return calculateReplaceRangeByBracketStack(choice, context);
+  };
+}
diff --git a/clients/tabby-agent/src/postprocess/calculateReplaceRangeByBracketStack.test.ts b/clients/tabby-agent/src/postprocess/calculateReplaceRangeByBracketStack.test.ts
@@ -11,37 +11,27 @@ describe("postprocess", () => {
         `,
         language: "typescript",
       };
-      const response = {
-        id: "",
-        choices: [
-          {
-            index: 0,
-            text: inline`
+      const choice = {
+        index: 0,
+        text: inline`
                        ├hello";┤
-            `,
-            replaceRange: {
-              start: context.position,
-              end: context.position,
-            },
-          },
-        ],
+        `,
+        replaceRange: {
+          start: context.position,
+          end: context.position,
+        },
       };
       const expected = {
-        id: "",
-        choices: [
-          {
-            index: 0,
-            text: inline`
+        index: 0,
+        text: inline`
                        ├hello";┤
-            `,
-            replaceRange: {
-              start: context.position,
-              end: context.position + 1,
-            },
-          },
-        ],
+        `,
+        replaceRange: {
+          start: context.position,
+          end: context.position + 1,
+        },
       };
-      expect(calculateReplaceRangeByBracketStack(response, context)).to.deep.equal(expected);
+      expect(calculateReplaceRangeByBracketStack(choice, context)).to.deep.equal(expected);
     });
 
     it("should handle auto closing quotes", () => {
@@ -51,37 +41,27 @@ describe("postprocess", () => {
         `,
         language: "typescript",
       };
-      const response = {
-        id: "",
-        choices: [
-          {
-            index: 0,
-            text: inline`
+      const choice = {
+        index: 0,
+        text: inline`
                            ├<h1>\${message}</h1>\`;┤
-            `,
-            replaceRange: {
-              start: context.position,
-              end: context.position,
-            },
-          },
-        ],
+        `,
+        replaceRange: {
+          start: context.position,
+          end: context.position,
+        },
       };
       const expected = {
-        id: "",
-        choices: [
-          {
-            index: 0,
-            text: inline`
+        index: 0,
+        text: inline`
                            ├<h1>\${message}</h1>\`;┤
-            `,
-            replaceRange: {
-              start: context.position,
-              end: context.position + 1,
-            },
-          },
-        ],
+        `,
+        replaceRange: {
+          start: context.position,
+          end: context.position + 1,
+        },
       };
-      expect(calculateReplaceRangeByBracketStack(response, context)).to.deep.equal(expected);
+      expect(calculateReplaceRangeByBracketStack(choice, context)).to.deep.equal(expected);
     });
 
     it("should handle multiple auto closing brackets", () => {
@@ -91,41 +71,31 @@ describe("postprocess", () => {
         `,
         language: "typescript",
       };
-      const response = {
-        id: "",
-        choices: [
-          {
-            index: 0,
-            text: inline`
+      const choice = {
+        index: 0,
+        text: inline`
                                       ├
-              console.log(data);
-            });┤
-            `,
-            replaceRange: {
-              start: context.position,
-              end: context.position,
-            },
-          },
-        ],
+          console.log(data);
+        });┤
+        `,
+        replaceRange: {
+          start: context.position,
+          end: context.position,
+        },
       };
       const expected = {
-        id: "",
-        choices: [
-          {
-            index: 0,
-            text: inline`
+        index: 0,
+        text: inline`
                                       ├
-              console.log(data);
-            });┤
-            `,
-            replaceRange: {
-              start: context.position,
-              end: context.position + 2,
-            },
-          },
-        ],
+          console.log(data);
+        });┤
+        `,
+        replaceRange: {
+          start: context.position,
+          end: context.position + 2,
+        },
       };
-      expect(calculateReplaceRangeByBracketStack(response, context)).to.deep.equal(expected);
+      expect(calculateReplaceRangeByBracketStack(choice, context)).to.deep.equal(expected);
     });
 
     it("should handle multiple auto closing brackets", () => {
@@ -135,37 +105,27 @@ describe("postprocess", () => {
         `,
         language: "typescript",
       };
-      const response = {
-        id: "",
-        choices: [
-          {
-            index: 0,
-            text: inline`
+      const choice = {
+        index: 0,
+        text: inline`
                                    ├1, 2], [3, 4]], [[5, 6], [7, 8]]];┤
-            `,
-            replaceRange: {
-              start: context.position,
-              end: context.position,
-            },
-          },
-        ],
+        `,
+        replaceRange: {
+          start: context.position,
+          end: context.position,
+        },
       };
       const expected = {
-        id: "",
-        choices: [
-          {
-            index: 0,
-            text: inline`
+        index: 0,
+        text: inline`
                                    ├1, 2], [3, 4]], [[5, 6], [7, 8]]];┤
-            `,
-            replaceRange: {
-              start: context.position,
-              end: context.position + 3,
-            },
-          },
-        ],
+        `,
+        replaceRange: {
+          start: context.position,
+          end: context.position + 3,
+        },
       };
-      expect(calculateReplaceRangeByBracketStack(response, context)).to.deep.equal(expected);
+      expect(calculateReplaceRangeByBracketStack(choice, context)).to.deep.equal(expected);
     });
   });
 
@@ -179,37 +139,27 @@ describe("postprocess", () => {
         `,
         language: "typescript",
       };
-      const response = {
-        id: "",
-        choices: [
-          {
-            index: 0,
-            text: inline`
+      const choice = {
+        index: 0,
+        text: inline`
                                    ├n, max), min┤
-            `,
-            replaceRange: {
-              start: context.position,
-              end: context.position,
-            },
-          },
-        ],
+        `,
+        replaceRange: {
+          start: context.position,
+          end: context.position,
+        },
       };
       const expected = {
-        id: "",
-        choices: [
-          {
-            index: 0,
-            text: inline`
+        index: 0,
+        text: inline`
                                    ├n, max), min┤
-            `,
-            replaceRange: {
-              start: context.position,
-              end: context.position,
-            },
-          },
-        ],
+        `,
+        replaceRange: {
+          start: context.position,
+          end: context.position,
+        },
       };
-      expect(calculateReplaceRangeByBracketStack(response, context)).not.to.deep.equal(expected);
+      expect(calculateReplaceRangeByBracketStack(choice, context)).not.to.deep.equal(expected);
     });
   });
 });