don't collapse constructor function

This commit is contained in:
bobo.yang 2024-04-08 22:56:48 +08:00
parent 0a3b15552a
commit 37d256a0a1
3 changed files with 38 additions and 19 deletions

View File

@ -37,6 +37,9 @@ export async function collapseFile(
if (funcBody.start === funcBody.end) {
continue;
}
if (func.name === "__init__" || func.name === "constructor") {
continue;
}
let bodyStartLine = funcBody.start.row;
let bodyEndLine = funcBody.end.row;

View File

@ -20,7 +20,8 @@ export interface FunctionRange {
body: {
start: { row: number, column: number },
end: { row: number, column: number }
}
},
name: string
}
export async function findFunctionRanges(filepath: string, node: Parser.SyntaxNode): Promise<FunctionRange[]> {
@ -41,7 +42,7 @@ export async function findFunctionRanges(filepath: string, node: Parser.SyntaxNo
// find functionNode through tag name
const functionNode = match.captures.find((capture) => capture.name === "function")?.node;
const bodyNode = match.captures.find((capture) => capture.name === "function.body")?.node;
// const nameNode = match.captures.find((capture) => capture.name === "function.name")?.node;
const nameNode = match.captures.find((capture) => capture.name === "function.name")?.node;
if (!functionNode ||!bodyNode) {
return [];
}
@ -66,7 +67,8 @@ export async function findFunctionRanges(filepath: string, node: Parser.SyntaxNo
row: bodyNode.endPosition.row,
column: bodyNode.endPosition.column,
},
}
},
name: nameNode?.text?? "",
};
return results;
}) ?? []

View File

@ -99,6 +99,9 @@ export async function collapseCodeBlock(functions: FunctionRange[], filepath: st
if (funcBody.start === funcBody.end) {
continue;
}
if (func.name === "__init__" || func.name === "constructor") {
continue;
}
let bodyStartLine = funcBody.start.row;
let bodyEndLine = funcBody.end.row;
@ -269,9 +272,13 @@ export async function symbolDefinesContext(filePath: string, fileContent: string
const functions = await findFunctionRanges(filePath, ast.rootNode);
// remove function that contain line, column
const filteredFunctions = functions.filter(f => {
let filteredFunctions = functions.filter(f => {
return!(f.define.start.row <= line && f.define.end.row >= line);
});
// remove function with name __init__ and constructor
filteredFunctions = filteredFunctions.filter(f => {
return f.name!== '__init__' && f.name!== 'constructor';
});
// collect matched ast nodes
const importTypeNodes: Parser.SyntaxNode[] = [];
@ -425,8 +432,6 @@ export async function createContextCallDefine( filepath: string, fileContent: st
export async function createPrompt(filePath: string, fileContent: string, line: number, column: number, posoffset: number, recentEdits: RecentEdit[]) {
const commentPrefix = await getCommentPrefix(filePath);
let recentEditContext = await createRecentEditContext(recentEdits, filePath);
const symbolDefines: { filepath: string, codeblock: string }[] = await symbolDefinesContext(filePath, fileContent, line, column);
let { prefix, suffix } = await currentFileContext(filePath, fileContent, line, column);
@ -455,23 +460,32 @@ export async function createPrompt(filePath: string, fileContent: string, line:
}
let symbolContext = "";
for (const symbolDefine of symbolDefines ) {
const countSymboleToken = countTokens(symbolDefine.codeblock);
if (tokenCount + countSymboleToken > CONTEXT_LIMITED_SIZE) {
break;
if (tokenCount < CONTEXT_LIMITED_SIZE) {
const symbolDefines: { filepath: string, codeblock: string }[] = await symbolDefinesContext(filePath, fileContent, line, column);
for (const symbolDefine of symbolDefines ) {
const countSymboleToken = countTokens(symbolDefine.codeblock);
if (tokenCount + countSymboleToken > CONTEXT_LIMITED_SIZE) {
break;
}
tokenCount += countSymboleToken;
symbolContext += `${commentPrefix}${symbolDefine.filepath}\n\n`;
symbolContext += `${symbolDefine.codeblock}\n\n\n\n`;
}
tokenCount += countSymboleToken;
symbolContext += `${commentPrefix}${symbolDefine.filepath}\n\n`;
symbolContext += `${symbolDefine.codeblock}\n\n\n\n`;
}
const countRecentToken = countTokens(recentEditContext);
if (tokenCount + countRecentToken < CONTEXT_LIMITED_SIZE) {
tokenCount += countRecentToken;
} else {
recentEditContext = "";
let recentEditContext = "";
if (tokenCount < CONTEXT_LIMITED_SIZE) {
recentEditContext = await createRecentEditContext(recentEdits, filePath);
const countRecentToken = countTokens(recentEditContext);
if (tokenCount + countRecentToken < CONTEXT_LIMITED_SIZE) {
tokenCount += countRecentToken;
} else {
recentEditContext = "";
}
}
logger.channel()?.info("Complete token:", tokenCount);
const prompt = "<fim_prefix>" + recentEditContext + symbolContext + callDefContext + `${commentPrefix}${filePath}\n\n` + prefix + "<fim_suffix>" + suffix + "<fim_middle>";