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

87 lines
2.7 KiB
TypeScript

import type TreeSitterParser from "web-tree-sitter";
import { getParser, languagesConfigs } from "../syntax/parser";
import { typeList } from "../syntax/typeList";
import { CompletionContext } from "../Agent";
import { PostprocessFilter, logger } from "./base";
export const supportedLanguages = Object.keys(languagesConfigs);
function findLineBegin(text: string, position: number): number {
let lastNonBlankCharPos = position - 1;
while (lastNonBlankCharPos >= 0 && text[lastNonBlankCharPos].match(/\s/)) {
lastNonBlankCharPos--;
}
if (lastNonBlankCharPos < 0) {
return 0;
}
const lineBegin = text.lastIndexOf("\n", lastNonBlankCharPos);
if (lineBegin < 0) {
return 0;
}
const line = text.slice(lineBegin + 1, position);
const indentation = line.search(/\S/);
return lineBegin + 1 + indentation;
}
function findLineEnd(text: string, position: number): number {
let firstNonBlankCharPos = position;
while (firstNonBlankCharPos < text.length && text[firstNonBlankCharPos].match(/\s/)) {
firstNonBlankCharPos++;
}
if (firstNonBlankCharPos >= text.length) {
return text.length;
}
const lineEnd = text.indexOf("\n", firstNonBlankCharPos);
if (lineEnd < 0) {
return text.length;
}
return lineEnd;
}
function findScope(node: TreeSitterParser.SyntaxNode, typeList: string[][]): TreeSitterParser.SyntaxNode {
for (const types of typeList) {
let scope = node;
while (scope) {
if (types.includes(scope.type)) {
return scope;
}
scope = scope.parent;
}
}
return node;
}
export function limitScopeBySyntax(context: CompletionContext): PostprocessFilter {
return async (input) => {
const { position, text, language, prefix, suffix } = context;
if (!supportedLanguages.includes(language)) {
return input;
}
const languageConfig = languagesConfigs[language];
const parser = await getParser(languageConfig);
const updatedText = prefix + input + suffix;
const updatedTree = parser.parse(updatedText);
const lineBegin = findLineBegin(updatedText, position);
const lineEnd = findLineEnd(updatedText, position);
const scope = findScope(updatedTree.rootNode.namedDescendantForIndex(lineBegin, lineEnd), typeList[languageConfig]);
if (scope.endIndex < position + input.length) {
logger.debug(
{
languageConfig,
text,
updatedText,
position,
lineBegin,
lineEnd,
scope: { type: scope.type, start: scope.startIndex, end: scope.endIndex },
},
"Remove content out of syntax scope",
);
return input.slice(0, scope.endIndex - position);
}
return input;
};
}