From f2ea57bdd649b53c3c2e0414007114cd53593e52 Mon Sep 17 00:00:00 2001 From: Zhiming Ma Date: Sun, 12 Nov 2023 00:42:00 +0800 Subject: [PATCH] refactor(agent): extract calculateReplaceRange and add unit test. (#764) * refactor(agent): extract calculateReplaceRange. * test(agent): add unit test for calculateReplaceRangeByBracketStack. --- clients/tabby-agent/src/TabbyAgent.ts | 45 +--- ...alculateReplaceRangeByBracketStack.test.ts | 213 ++++++++++++++++++ .../calculateReplaceRangeByBracketStack.ts | 29 +++ clients/tabby-agent/src/postprocess/index.ts | 8 + .../tests/bad_cases/2-typescript.ts | 3 + 5 files changed, 259 insertions(+), 39 deletions(-) create mode 100644 clients/tabby-agent/src/postprocess/calculateReplaceRangeByBracketStack.test.ts create mode 100644 clients/tabby-agent/src/postprocess/calculateReplaceRangeByBracketStack.ts create mode 100644 clients/tabby-agent/tests/bad_cases/2-typescript.ts diff --git a/clients/tabby-agent/src/TabbyAgent.ts b/clients/tabby-agent/src/TabbyAgent.ts index 95c9a2ef467d..b9338e8ec96e 100644 --- a/clients/tabby-agent/src/TabbyAgent.ts +++ b/clients/tabby-agent/src/TabbyAgent.ts @@ -5,14 +5,7 @@ import { deepmerge } from "deepmerge-ts"; import { getProperty, setProperty, deleteProperty } from "dot-prop"; import createClient from "openapi-fetch"; import { paths as TabbyApi } from "./types/tabbyApi"; -import { - isBlank, - abortSignalFromAnyOf, - findUnpairedAutoClosingChars, - HttpError, - isTimeoutError, - isCanceledError, -} from "./utils"; +import { isBlank, abortSignalFromAnyOf, HttpError, isTimeoutError, isCanceledError } from "./utils"; import type { Agent, AgentStatus, @@ -32,7 +25,7 @@ import { CompletionCache } from "./CompletionCache"; import { CompletionDebounce } from "./CompletionDebounce"; import { CompletionContext } from "./CompletionContext"; import { DataStore } from "./dataStore"; -import { preCacheProcess, postCacheProcess } from "./postprocess"; +import { preCacheProcess, postCacheProcess, calculateReplaceRange } from "./postprocess"; import { rootLogger, allLoggers } from "./logger"; import { AnonymousUsageLogger } from "./AnonymousUsageLogger"; import { CompletionProviderStats, CompletionProviderStatsEntry } from "./CompletionProviderStats"; @@ -291,35 +284,6 @@ export class TabbyAgent extends EventEmitter implements Agent { return { prefix, suffix }; } - private calculateReplaceRange(response: CompletionResponse, context: CompletionContext): CompletionResponse { - const { suffixLines } = context; - const suffixText = suffixLines[0]?.trimEnd() || ""; - if (isBlank(suffixText)) { - return response; - } - for (const choice of response.choices) { - const completionText = choice.text.slice(context.position - choice.replaceRange.start); - const unpaired = findUnpairedAutoClosingChars(completionText); - if (isBlank(unpaired)) { - continue; - } - if (suffixText.startsWith(unpaired)) { - choice.replaceRange.end = context.position + unpaired.length; - this.logger.trace( - { context, completion: choice.text, range: choice.replaceRange, unpaired }, - "Adjust replace range", - ); - } else if (unpaired.startsWith(suffixText)) { - choice.replaceRange.end = context.position + suffixText.length; - this.logger.trace( - { context, completion: choice.text, range: choice.replaceRange, unpaired }, - "Adjust replace range", - ); - } - } - return response; - } - public async initialize(options: AgentInitOptions): Promise { if (options.clientProperties) { const { user: userProp, session: sessionProp } = options.clientProperties; @@ -574,7 +538,10 @@ export class TabbyAgent extends EventEmitter implements Agent { throw options.signal.reason; } // Calculate replace range - completionResponse = this.calculateReplaceRange(completionResponse, context); + completionResponse = await calculateReplaceRange(completionResponse, context); + if (options?.signal?.aborted) { + throw options.signal.reason; + } } catch (error) { if (isCanceledError(error) || isTimeoutError(error)) { if (stats) { diff --git a/clients/tabby-agent/src/postprocess/calculateReplaceRangeByBracketStack.test.ts b/clients/tabby-agent/src/postprocess/calculateReplaceRangeByBracketStack.test.ts new file mode 100644 index 000000000000..f73b8a739acc --- /dev/null +++ b/clients/tabby-agent/src/postprocess/calculateReplaceRangeByBracketStack.test.ts @@ -0,0 +1,213 @@ +import { expect } from "chai"; +import { documentContext, inline } from "./testUtils"; +import { calculateReplaceRangeByBracketStack } from "./calculateReplaceRangeByBracketStack"; + +describe("postprocess", () => { + describe("calculateReplaceRangeByBracketStack", () => { + it("should handle auto closing quotes", () => { + 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(calculateReplaceRangeByBracketStack(response, context)).to.deep.equal(expected); + }); + + it("should handle auto closing quotes", () => { + const context = { + ...documentContext` + let htmlMarkup = \`║\` + `, + language: "typescript", + }; + const response = { + id: "", + choices: [ + { + index: 0, + text: inline` + ├

\${message}

\`;┤ + `, + replaceRange: { + start: context.position, + end: context.position, + }, + }, + ], + }; + const expected = { + id: "", + choices: [ + { + index: 0, + text: inline` + ├

\${message}

\`;┤ + `, + replaceRange: { + start: context.position, + end: context.position + 1, + }, + }, + ], + }; + expect(calculateReplaceRangeByBracketStack(response, context)).to.deep.equal(expected); + }); + + it("should handle multiple auto closing brackets", () => { + 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(calculateReplaceRangeByBracketStack(response, context)).to.deep.equal(expected); + }); + + it("should handle multiple auto closing brackets", () => { + 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(calculateReplaceRangeByBracketStack(response, context)).to.deep.equal(expected); + }); + }); + + describe("calculateReplaceRangeByBracketStack: bad cases", () => { + 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(calculateReplaceRangeByBracketStack(response, context)).not.to.deep.equal(expected); + }); +}); diff --git a/clients/tabby-agent/src/postprocess/calculateReplaceRangeByBracketStack.ts b/clients/tabby-agent/src/postprocess/calculateReplaceRangeByBracketStack.ts new file mode 100644 index 000000000000..f136992dda30 --- /dev/null +++ b/clients/tabby-agent/src/postprocess/calculateReplaceRangeByBracketStack.ts @@ -0,0 +1,29 @@ +import { CompletionContext, CompletionResponse } from "../Agent"; +import { isBlank, findUnpairedAutoClosingChars } from "../utils"; +import { logger } from "./base"; + +export function calculateReplaceRangeByBracketStack( + response: CompletionResponse, + context: CompletionContext, +): CompletionResponse { + const { suffixLines } = context; + const suffixText = suffixLines[0]?.trimEnd() || ""; + if (isBlank(suffixText)) { + return response; + } + for (const choice of response.choices) { + const completionText = choice.text.slice(context.position - choice.replaceRange.start); + const unpaired = findUnpairedAutoClosingChars(completionText); + if (isBlank(unpaired)) { + continue; + } + if (suffixText.startsWith(unpaired)) { + choice.replaceRange.end = context.position + unpaired.length; + logger.trace({ context, completion: choice.text, range: choice.replaceRange, unpaired }, "Adjust replace range"); + } 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"); + } + } + return response; +} diff --git a/clients/tabby-agent/src/postprocess/index.ts b/clients/tabby-agent/src/postprocess/index.ts index 14e8c47f8890..ba0ba3b509e5 100644 --- a/clients/tabby-agent/src/postprocess/index.ts +++ b/clients/tabby-agent/src/postprocess/index.ts @@ -8,6 +8,7 @@ import { limitScopeByIndentation } from "./limitScopeByIndentation"; import { trimSpace } from "./trimSpace"; import { dropDuplicated } from "./dropDuplicated"; import { dropBlank } from "./dropBlank"; +import { calculateReplaceRangeByBracketStack } from "./calculateReplaceRangeByBracketStack"; export async function preCacheProcess( context: CompletionContext, @@ -34,3 +35,10 @@ export async function postCacheProcess( .then(applyFilter(trimSpace(context), context)) .then(applyFilter(dropBlank(), context)); } + +export async function calculateReplaceRange( + response: CompletionResponse, + context: CompletionContext, +): Promise { + return calculateReplaceRangeByBracketStack(response, context); +} diff --git a/clients/tabby-agent/tests/bad_cases/2-typescript.ts b/clients/tabby-agent/tests/bad_cases/2-typescript.ts new file mode 100644 index 000000000000..b6765df53dbf --- /dev/null +++ b/clients/tabby-agent/tests/bad_cases/2-typescript.ts @@ -0,0 +1,3 @@ +function clamp(n: number, max: number, min: number): number { + return Math.max(Math.min(⏩⏭n, max), min⏮⏪); +} \ No newline at end of file