feat(agent): add experimental option: scope of indentation filter.

TabbyML · Oct 27, 2023 · 39ddc92 · 39ddc92
1 parent d6c1324
commit 39ddc92
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 7 deletions.
diff --git a/clients/tabby-agent/src/AgentConfig.ts b/clients/tabby-agent/src/AgentConfig.ts
@@ -20,6 +20,14 @@ export type AgentConfig = {
       manually: number;
     };
   };
+  postprocess: {
+    limitScopeByIndentation: {
+      // When completion is continuing the current line, limit the scope to:
+      // false(default): the line scope, meaning use the next indent level as the limit.
+      // true: the block scope, meaning use the current indent level as the limit.
+      keepBlockScopeWhenCompletingLine: boolean;
+    };
+  };
   logs: {
     level: "debug" | "error" | "silent";
   };
@@ -59,6 +67,11 @@ export const defaultAgentConfig: AgentConfig = {
       manually: 4000, // 4s
     },
   },
+  postprocess: {
+    limitScopeByIndentation: {
+      keepBlockScopeWhenCompletingLine: false,
+    },
+  },
   logs: {
     level: "silent",
   },

diff --git a/clients/tabby-agent/src/TabbyAgent.ts b/clients/tabby-agent/src/TabbyAgent.ts
@@ -541,7 +541,7 @@ export class TabbyAgent extends EventEmitter implements Agent {
             throw error;
           }
           // Postprocess (pre-cache)
-          completionResponse = await preCacheProcess(context, completionResponse);
+          completionResponse = await preCacheProcess(context, this.config.postprocess, completionResponse);
           if (options?.signal?.aborted) {
             throw options.signal.reason;
           }
@@ -550,7 +550,7 @@ export class TabbyAgent extends EventEmitter implements Agent {
         }
       }
       // Postprocess (post-cache)
-      completionResponse = await postCacheProcess(context, completionResponse);
+      completionResponse = await postCacheProcess(context, this.config.postprocess, completionResponse);
       if (options?.signal?.aborted) {
         throw options.signal.reason;
       }

diff --git a/clients/tabby-agent/src/postprocess/index.ts b/clients/tabby-agent/src/postprocess/index.ts
@@ -1,4 +1,5 @@
 import { CompletionContext, CompletionResponse } from "../Agent";
+import { AgentConfig } from "../AgentConfig";
 import { applyFilter } from "./base";
 import { removeRepetitiveBlocks } from "./removeRepetitiveBlocks";
 import { removeRepetitiveLines } from "./removeRepetitiveLines";
@@ -10,6 +11,7 @@ import { dropBlank } from "./dropBlank";
 
 export async function preCacheProcess(
   context: CompletionContext,
+  config: AgentConfig["postprocess"],
   response: CompletionResponse,
 ): Promise<CompletionResponse> {
   return Promise.resolve(response)
@@ -21,12 +23,13 @@ export async function preCacheProcess(
 
 export async function postCacheProcess(
   context: CompletionContext,
+  config: AgentConfig["postprocess"],
   response: CompletionResponse,
 ): Promise<CompletionResponse> {
   return Promise.resolve(response)
     .then(applyFilter(removeRepetitiveBlocks(context), context))
     .then(applyFilter(removeRepetitiveLines(context), context))
-    .then(applyFilter(limitScopeByIndentation(context), context))
+    .then(applyFilter(limitScopeByIndentation(context, config["limitScopeByIndentation"]), context))
     .then(applyFilter(dropDuplicated(context), context))
     .then(applyFilter(trimSpace(context), context))
     .then(applyFilter(dropBlank(), context));

diff --git a/clients/tabby-agent/src/postprocess/limitScopeByIndentation.ts b/clients/tabby-agent/src/postprocess/limitScopeByIndentation.ts
@@ -1,4 +1,5 @@
 import { CompletionContext } from "../Agent";
+import { AgentConfig } from "../AgentConfig";
 import { PostprocessFilter, logger } from "./base";
 import { isBlank, splitLines } from "../utils";
 
@@ -17,6 +18,7 @@ function processContext(
   lines: string[],
   prefixLines: string[],
   suffixLines: string[],
+  config: AgentConfig["postprocess"]["limitScopeByIndentation"],
 ): { indentLevelLimit: number; allowClosingLine: (closingLine: string) => boolean } {
   let allowClosingLine = false;
   let result = { indentLevelLimit: 0, allowClosingLine: (closingLine: string) => allowClosingLine };
@@ -57,7 +59,11 @@ function processContext(
   if (!isCurrentLineInCompletionBlank && !isCurrentLineInPrefixBlank) {
     // if two reference lines are contacted at current line, it is continuing uncompleted sentence
 
-    result.indentLevelLimit = referenceLineInPrefixIndent + 1; // + 1 for comparison, no matter how many spaces indent
+    if (config.keepBlockScopeWhenCompletingLine) {
+      result.indentLevelLimit = referenceLineInPrefixIndent;
+    } else {
+      result.indentLevelLimit = referenceLineInPrefixIndent + 1; // + 1 for comparison, no matter how many spaces indent
+    }
     // allow closing line if first line is opening a new indent block
     allowClosingLine = !!lines[1] && calcIndentLevel(lines[1]) > referenceLineInPrefixIndent;
   } else if (referenceLineInCompletionIndent > referenceLineInPrefixIndent) {
@@ -95,7 +101,10 @@ function processContext(
   return result;
 }
 
-export const limitScopeByIndentation: (context: CompletionContext) => PostprocessFilter = (context) => {
+export function limitScopeByIndentation(
+  context: CompletionContext,
+  config: AgentConfig["postprocess"]["limitScopeByIndentation"],
+): PostprocessFilter {
   return (input) => {
     const { prefix, suffix, prefixLines, suffixLines } = context;
     const inputLines = splitLines(input);
@@ -105,7 +114,7 @@ export const limitScopeByIndentation: (context: CompletionContext) => Postproces
         return null;
       }
     }
-    const indentContext = processContext(inputLines, prefixLines, suffixLines);
+    const indentContext = processContext(inputLines, prefixLines, suffixLines, config);
     let index;
     for (index = 1; index < inputLines.length; index++) {
       if (isBlank(inputLines[index])) {
@@ -135,4 +144,4 @@ export const limitScopeByIndentation: (context: CompletionContext) => Postproces
     }
     return input;
   };
-};
+}