feat(agent): added adaptive completion debouncing. (#389)

* feat(agent): added adaptive completion debouncing.

* chore: bump tabby-agent version to v0.1.0-dev.

clients/tabby-agent/package.json

@@ -1,6 +1,6 @@
 {
   "name": "tabby-agent",
-  "version": "0.0.1",
+  "version": "0.1.0-dev",
   "description": "Generic client agent for Tabby AI coding assistant IDE extensions.",
   "repository": "https://github.com/TabbyML/tabby",
   "main": "./dist/index.js",

clients/tabby-agent/src/Agent.ts

@@ -20,8 +20,6 @@ export type CompletionRequest = {
   text: string;
   position: number;
   manually?: boolean;
-  maxPrefixLines?: number;
-  maxSuffixLines?: number;
 };
 
 export type CompletionResponse = ApiCompletionResponse;
@@ -116,11 +114,14 @@ export interface AgentFunction {
   waitForAuthToken(code: string): CancelablePromise<any>;
 
   /**
+   * Provide completions for the given request. This method is debounced, calling it before the previous
+   * call is resolved will cancel the previous call. The debouncing interval is automatically calculated
+   * or can be set in the config.
    * @param request
    * @returns
    * @throws Error if agent is not initialized
    */
-  getCompletions(request: CompletionRequest): CancelablePromise<CompletionResponse>;
+  provideCompletions(request: CompletionRequest): CancelablePromise<CompletionResponse>;
 
   /**
    * @param event

clients/tabby-agent/src/AgentConfig.ts

@@ -7,12 +7,18 @@ export type AgentConfig = {
     requestTimeout: number;
   };
   completion: {
+    prompt: {
+      maxPrefixLines: number;
+      maxSuffixLines: number;
+    };
+    debounce: {
+      mode: "adaptive" | "fixed";
+      interval: number;
+    };
     timeout: {
       auto: number;
       manually: number;
     };
-    maxPrefixLines: number;
-    maxSuffixLines: number;
   };
   logs: {
     level: "debug" | "error" | "silent";
@@ -39,12 +45,18 @@ export const defaultAgentConfig: AgentConfig = {
     requestTimeout: 30000, // 30s
   },
   completion: {
+    prompt: {
+      maxPrefixLines: 20,
+      maxSuffixLines: 20,
+    },
+    debounce: {
+      mode: "adaptive",
+      interval: 250, // ms
+    },
     timeout: {
       auto: 5000, // 5s
       manually: 30000, // 30s
     },
-    maxPrefixLines: 20,
-    maxSuffixLines: 20,
   },
   logs: {
     level: "silent",

clients/tabby-agent/src/CompletionDebounce.ts

@@ -0,0 +1,114 @@
+import { CancelablePromise } from "./generated";
+import { CompletionRequest } from "./Agent";
+import { AgentConfig } from "./AgentConfig";
+import { rootLogger } from "./logger";
+import { splitLines } from "./utils";
+
+function clamp(min: number, max: number, value: number): number {
+  return Math.max(min, Math.min(max, value));
+}
+
+export class CompletionDebounce {
+  private readonly logger = rootLogger.child({ component: "CompletionDebounce" });
+  private ongoing: CancelablePromise<any> | null = null;
+  private lastCalledTimeStamp = 0;
+
+  private baseInterval = 200; // ms
+  private calledIntervalHistory: number[] = [];
+
+  private options = {
+    baseIntervalSlideWindowAvg: {
+      minSize: 20,
+      maxSize: 100,
+      min: 100,
+      max: 400,
+    },
+    adaptiveRate: {
+      min: 1.5,
+      max: 3.0,
+    },
+    contextScoreWeights: {
+      triggerCharacter: 0.5,
+      noSuffixInCurrentLine: 0.4,
+      noSuffix: 0.1,
+    },
+    requestDelay: {
+      min: 100, // ms
+      max: 1000,
+    },
+  };
+
+  debounce(
+    request: CompletionRequest,
+    config: AgentConfig["completion"]["debounce"],
+    responseTime: number,
+  ): CancelablePromise<any> {
+    if (request.manually) {
+      return this.renewPromise(0);
+    }
+    if (config.mode === "fixed") {
+      return this.renewPromise(config.interval);
+    }
+    const now = Date.now();
+    this.updateBaseInterval(now - this.lastCalledTimeStamp);
+    this.lastCalledTimeStamp = now;
+    const contextScore = this.calcContextScore(request);
+    const adaptiveRate =
+      this.options.adaptiveRate.max - (this.options.adaptiveRate.max - this.options.adaptiveRate.min) * contextScore;
+    const expectedLatency = adaptiveRate * this.baseInterval;
+    const delay = clamp(this.options.requestDelay.min, this.options.requestDelay.max, expectedLatency - responseTime);
+    return this.renewPromise(delay);
+  }
+
+  private renewPromise(delay: number): CancelablePromise<any> {
+    if (this.ongoing) {
+      this.ongoing.cancel();
+    }
+    this.ongoing = new CancelablePromise<any>((resolve, reject, onCancel) => {
+      const timer = setTimeout(
+        () => {
+          resolve(true);
+        },
+        Math.min(delay, 0x7fffffff),
+      );
+      onCancel(() => {
+        clearTimeout(timer);
+      });
+    });
+    return this.ongoing;
+  }
+
+  private updateBaseInterval(interval: number) {
+    if (interval > this.options.baseIntervalSlideWindowAvg.max) {
+      return;
+    }
+    this.calledIntervalHistory.push(interval);
+    if (this.calledIntervalHistory.length > this.options.baseIntervalSlideWindowAvg.maxSize) {
+      this.calledIntervalHistory.shift();
+    }
+    if (this.calledIntervalHistory.length > this.options.baseIntervalSlideWindowAvg.minSize) {
+      const avg = this.calledIntervalHistory.reduce((a, b) => a + b, 0) / this.calledIntervalHistory.length;
+      this.baseInterval = clamp(
+        this.options.baseIntervalSlideWindowAvg.min,
+        this.options.baseIntervalSlideWindowAvg.max,
+        avg,
+      );
+    }
+  }
+
+  // return score in [0, 1], 1 means the context has a high chance to accept the completion
+  private calcContextScore(request: CompletionRequest): number {
+    let score = 0;
+    const weights = this.options.contextScoreWeights;
+    const triggerCharacter = request.text[request.position - 1] ?? "";
+    score += triggerCharacter.match(/^\W*$/) ? weights.triggerCharacter : 0;
+
+    const suffix = request.text.slice(request.position) ?? "";
+    const currentLineInSuffix = splitLines(suffix)[0] ?? "";
+    score += currentLineInSuffix.match(/^\W*$/) ? weights.noSuffixInCurrentLine : 0;
+    score += suffix.match(/^\W*$/) ? weights.noSuffix : 0;
+
+    score = clamp(0, 1, score);
+    return score;
+  }
+}

clients/tabby-agent/src/TabbyAgent.ts

@@ -19,6 +19,7 @@ import {
 import { Auth } from "./Auth";
 import { AgentConfig, PartialAgentConfig, defaultAgentConfig, userAgentConfig } from "./AgentConfig";
 import { CompletionCache } from "./CompletionCache";
+import { CompletionDebounce } from "./CompletionDebounce";
 import { DataStore } from "./dataStore";
 import { postprocess, preCacheProcess } from "./postprocess";
 import { rootLogger, allLoggers } from "./logger";
@@ -46,6 +47,7 @@ export class TabbyAgent extends EventEmitter implements Agent {
   private auth: Auth;
   private dataStore: DataStore | null = null;
   private completionCache: CompletionCache = new CompletionCache();
+  private CompletionDebounce: CompletionDebounce = new CompletionDebounce();
   static readonly tryConnectInterval = 1000 * 30; // 30s
   private tryingConnectTimer: ReturnType<typeof setInterval> | null = null;
   private completionResponseStats: ResponseStats = new ResponseStats(completionResponseTimeStatsStrategy);
@@ -261,8 +263,8 @@ export class TabbyAgent extends EventEmitter implements Agent {
 
   private createSegments(request: CompletionRequest): { prefix: string; suffix: string } {
     // max lines in prefix and suffix configurable
-    const maxPrefixLines = request.maxPrefixLines ?? this.config.completion.maxPrefixLines;
-    const maxSuffixLines = request.maxSuffixLines ?? this.config.completion.maxSuffixLines;
+    const maxPrefixLines = this.config.completion.prompt.maxPrefixLines;
+    const maxSuffixLines = this.config.completion.prompt.maxSuffixLines;
     const prefix = request.text.slice(0, request.position);
     const prefixLines = splitLines(prefix);
     const suffix = request.text.slice(request.position);
@@ -379,7 +381,7 @@ export class TabbyAgent extends EventEmitter implements Agent {
     );
   }
 
-  public getCompletions(request: CompletionRequest): CancelablePromise<CompletionResponse> {
+  public provideCompletions(request: CompletionRequest): CancelablePromise<CompletionResponse> {
     if (this.status === "notInitialized") {
       return cancelable(Promise.reject("Agent is not initialized"), () => {});
     }
@@ -387,15 +389,19 @@ export class TabbyAgent extends EventEmitter implements Agent {
     return cancelable(
       Promise.resolve(null)
         // From cache
-        .then((response: CompletionResponse | null) => {
+        .then(async (response: CompletionResponse | null) => {
           if (response) return response;
           if (this.completionCache.has(request)) {
             this.logger.debug({ request }, "Completion cache hit");
+            const debounce = this.CompletionDebounce.debounce(request, this.config.completion.debounce, 0);
+            cancelableList.push(debounce);
+            await debounce;
             return this.completionCache.get(request);
           }
+          return null;
         })
         // From api
-        .then((response: CompletionResponse | null) => {
+        .then(async (response: CompletionResponse | null) => {
           if (response) return response;
           const segments = this.createSegments(request);
           if (isBlank(segments.prefix)) {
@@ -405,6 +411,13 @@ export class TabbyAgent extends EventEmitter implements Agent {
               choices: [],
             };
           }
+          const debounce = this.CompletionDebounce.debounce(
+            request,
+            this.config.completion.debounce,
+            this.completionResponseStats.stats()["averageResponseTime"],
+          );
+          cancelableList.push(debounce);
+          await debounce;
           const apiRequest = this.callApi(
             this.api.v1.completion,
             {
@@ -417,17 +430,13 @@ export class TabbyAgent extends EventEmitter implements Agent {
             },
           );
           cancelableList.push(apiRequest);
-          return apiRequest
-            .then((response) => {
-              return preCacheProcess(request, response);
-            })
-            .then((response) => {
-              this.completionCache.set(request, response);
-              return response;
-            });
+          let res = await apiRequest;
+          res = await preCacheProcess(request, res);
+          this.completionCache.set(request, res);
+          return res;
         })
         // Postprocess
-        .then((response: CompletionResponse | null) => {
+        .then(async (response: CompletionResponse | null) => {
           return postprocess(request, response);
         }),
       () => {

clients/tabby-agent/src/index.ts

@@ -6,6 +6,11 @@ export {
   AgentEvent,
   StatusChangedEvent,
   ConfigUpdatedEvent,
+  AuthRequiredEvent,
+  NewIssueEvent,
+  AgentIssue,
+  SlowCompletionResponseTimeIssue,
+  HighCompletionTimeoutRateIssue,
   CompletionRequest,
   CompletionResponse,
   LogEventRequest,

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

@@ -10,8 +10,6 @@ export function documentContext(strings): PostprocessContext {
     language: null,
     text: doc.replace(/║/, ""),
     position: doc.indexOf("║"),
-    maxPrefixLines: 20,
-    maxSuffixLines: 20,
   });
 }
 

clients/vscode/src/TabbyCompletionProvider.ts

@@ -14,21 +14,13 @@ import {
 import { CompletionResponse, CancelablePromise } from "tabby-agent";
 import { agent } from "./agent";
 import { notifications } from "./notifications";
-import { sleep } from "./utils";
 
 export class TabbyCompletionProvider implements InlineCompletionItemProvider {
-  private uuid = Date.now();
-  private latestTimestamp: number = 0;
   private pendingCompletion: CancelablePromise<CompletionResponse> | null = null;
 
   // User Settings
   private enabled: boolean = true;
 
-  // These settings will be move to tabby-agent
-  private suggestionDelay: number = 150;
-  private maxPrefixLines: number = 20;
-  private maxSuffixLines: number = 20;
-
   constructor() {
     this.updateConfiguration();
     workspace.onDidChangeConfiguration((event) => {
@@ -47,14 +39,6 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider {
       return emptyResponse;
     }
 
-    const currentTimestamp = Date.now();
-    this.latestTimestamp = currentTimestamp;
-
-    await sleep(this.suggestionDelay);
-    if (currentTimestamp < this.latestTimestamp) {
-      return emptyResponse;
-    }
-
     const replaceRange = this.calculateReplaceRange(document, position);
 
     if (this.pendingCompletion) {
@@ -67,12 +51,10 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider {
       text: document.getText(),
       position: document.offsetAt(position),
       manually: context.triggerKind === InlineCompletionTriggerKind.Invoke,
-      maxPrefixLines: this.maxPrefixLines,
-      maxSuffixLines: this.maxSuffixLines,
     };
-    this.pendingCompletion = agent().getCompletions(request);
+    this.pendingCompletion = agent().provideCompletions(request);
 
-    const completion = await this.pendingCompletion.catch((_: Error) => {
+    const completion = await this.pendingCompletion.catch((e: Error) => {
       return null;
     });
     this.pendingCompletion = null;

clients/vscode/src/utils.ts

@@ -1,3 +0,0 @@
-export function sleep(milliseconds: number) {
-  return new Promise((r) => setTimeout(r, milliseconds));
-}