From 63c7da4f969b9ee603a7e491b7f0528ec01c205a Mon Sep 17 00:00:00 2001 From: Zhiming Ma Date: Mon, 13 Nov 2023 12:53:20 +0800 Subject: [PATCH] feat(agent): postprocess add calculate replace range by syntax. (#765) --- clients/tabby-agent/src/AgentConfig.ts | 7 + clients/tabby-agent/src/TabbyAgent.ts | 2 +- ...alculateReplaceRangeByBracketStack.test.ts | 76 ++++--- .../calculateReplaceRangeByBracketStack.ts | 10 +- .../calculateReplaceRangeBySyntax.test.ts | 213 ++++++++++++++++++ .../calculateReplaceRangeBySyntax.ts | 54 +++++ clients/tabby-agent/src/postprocess/index.ts | 11 +- clients/tabby-agent/tests/golden.test.ts | 44 ++-- 8 files changed, 362 insertions(+), 55 deletions(-) create mode 100644 clients/tabby-agent/src/postprocess/calculateReplaceRangeBySyntax.test.ts create mode 100644 clients/tabby-agent/src/postprocess/calculateReplaceRangeBySyntax.ts diff --git a/clients/tabby-agent/src/AgentConfig.ts b/clients/tabby-agent/src/AgentConfig.ts index ecc2796..629db80 100644 --- a/clients/tabby-agent/src/AgentConfig.ts +++ b/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", diff --git a/clients/tabby-agent/src/TabbyAgent.ts b/clients/tabby-agent/src/TabbyAgent.ts index ef0e935..98fe2f1 100644 --- a/clients/tabby-agent/src/TabbyAgent.ts +++ b/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; } diff --git a/clients/tabby-agent/src/postprocess/calculateReplaceRangeByBracketStack.test.ts b/clients/tabby-agent/src/postprocess/calculateReplaceRangeByBracketStack.test.ts index f73b8a7..93008d7 100644 --- a/clients/tabby-agent/src/postprocess/calculateReplaceRangeByBracketStack.test.ts +++ b/clients/tabby-agent/src/postprocess/calculateReplaceRangeByBracketStack.test.ts @@ -170,44 +170,46 @@ describe("postprocess", () => { }); 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, + it("cannot handle the case of completion bracket stack is same with suffix but should not be replaced", () => { + 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, + ], + }; + 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); + ], + }; + 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 index f136992..49598ae 100644 --- a/clients/tabby-agent/src/postprocess/calculateReplaceRangeByBracketStack.ts +++ b/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; diff --git a/clients/tabby-agent/src/postprocess/calculateReplaceRangeBySyntax.test.ts b/clients/tabby-agent/src/postprocess/calculateReplaceRangeBySyntax.test.ts new file mode 100644 index 0000000..ef43ecd --- /dev/null +++ b/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` + ├

\${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(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); + }); + }); +}); diff --git a/clients/tabby-agent/src/postprocess/calculateReplaceRangeBySyntax.ts b/clients/tabby-agent/src/postprocess/calculateReplaceRangeBySyntax.ts new file mode 100644 index 0000000..1dffb6b --- /dev/null +++ b/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 { + 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; +} diff --git a/clients/tabby-agent/src/postprocess/index.ts b/clients/tabby-agent/src/postprocess/index.ts index e558909..0814c81 100644 --- a/clients/tabby-agent/src/postprocess/index.ts +++ b/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 { - 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); } diff --git a/clients/tabby-agent/tests/golden.test.ts b/clients/tabby-agent/tests/golden.test.ts index 75526c7..9b7d394 100644 --- a/clients/tabby-agent/tests/golden.test.ts +++ b/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 }]); - expect(output.shift()).to.deep.equal([requestId, true]); + { + requestId++; + 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 () => {