310 lines
9.8 KiB
TypeScript
310 lines
9.8 KiB
TypeScript
import * as plugins from '../plugins.js';
|
|
import type { ITrimConfig, ContextMode } from './types.js';
|
|
|
|
/**
|
|
* Class responsible for trimming file contents to reduce token usage
|
|
* while preserving important information for context
|
|
*/
|
|
export class ContextTrimmer {
|
|
private config: ITrimConfig;
|
|
|
|
/**
|
|
* Create a new ContextTrimmer with the given configuration
|
|
* @param config The trimming configuration
|
|
*/
|
|
constructor(config?: ITrimConfig) {
|
|
this.config = {
|
|
removeImplementations: true,
|
|
preserveInterfaces: true,
|
|
preserveTypeDefs: true,
|
|
preserveJSDoc: true,
|
|
maxFunctionLines: 5,
|
|
removeComments: true,
|
|
removeBlankLines: true,
|
|
...config
|
|
};
|
|
}
|
|
|
|
/**
|
|
* Trim a file's contents based on the configuration
|
|
* @param filePath The path to the file
|
|
* @param content The file's contents
|
|
* @param mode The context mode to use
|
|
* @returns The trimmed file contents
|
|
*/
|
|
public trimFile(filePath: string, content: string, mode: ContextMode = 'trimmed'): string {
|
|
// If mode is 'full', return the original content
|
|
if (mode === 'full') {
|
|
return content;
|
|
}
|
|
|
|
// Process based on file type
|
|
if (filePath.endsWith('.ts') || filePath.endsWith('.tsx')) {
|
|
return this.trimTypeScriptFile(content);
|
|
} else if (filePath.endsWith('.md')) {
|
|
return this.trimMarkdownFile(content);
|
|
} else if (filePath.endsWith('.json')) {
|
|
return this.trimJsonFile(content);
|
|
}
|
|
|
|
// Default to returning the original content for unknown file types
|
|
return content;
|
|
}
|
|
|
|
/**
|
|
* Trim a TypeScript file to reduce token usage
|
|
* @param content The TypeScript file contents
|
|
* @returns The trimmed file contents
|
|
*/
|
|
private trimTypeScriptFile(content: string): string {
|
|
let result = content;
|
|
|
|
// Step 1: Preserve JSDoc comments if configured
|
|
const jsDocComments: string[] = [];
|
|
if (this.config.preserveJSDoc) {
|
|
const jsDocRegex = /\/\*\*[\s\S]*?\*\//g;
|
|
const matches = result.match(jsDocRegex) || [];
|
|
jsDocComments.push(...matches);
|
|
}
|
|
|
|
// Step 2: Remove comments if configured
|
|
if (this.config.removeComments) {
|
|
// Remove single-line comments
|
|
result = result.replace(/\/\/.*$/gm, '');
|
|
// Remove multi-line comments (except JSDoc if preserveJSDoc is true)
|
|
if (!this.config.preserveJSDoc) {
|
|
result = result.replace(/\/\*[\s\S]*?\*\//g, '');
|
|
} else {
|
|
// Only remove non-JSDoc comments
|
|
result = result.replace(/\/\*(?!\*)[\s\S]*?\*\//g, '');
|
|
}
|
|
}
|
|
|
|
// Step 3: Remove function implementations if configured
|
|
if (this.config.removeImplementations) {
|
|
// Match function and method bodies
|
|
result = result.replace(
|
|
/(\b(function|constructor|async function)\s+[\w$]*\s*\([^)]*\)\s*{)([\s\S]*?)(})/g,
|
|
(match, start, funcType, body, end) => {
|
|
// Keep function signature and opening brace, replace body with comment
|
|
return `${start} /* implementation removed */ ${end}`;
|
|
}
|
|
);
|
|
|
|
// Match arrow function bodies
|
|
result = result.replace(
|
|
/(\([^)]*\)\s*=>\s*{)([\s\S]*?)(})/g,
|
|
(match, start, body, end) => {
|
|
return `${start} /* implementation removed */ ${end}`;
|
|
}
|
|
);
|
|
|
|
// Match method declarations
|
|
result = result.replace(
|
|
/(^\s*[\w$]*\s*\([^)]*\)\s*{)([\s\S]*?)(})/gm,
|
|
(match, start, body, end) => {
|
|
return `${start} /* implementation removed */ ${end}`;
|
|
}
|
|
);
|
|
|
|
// Match class methods
|
|
result = result.replace(
|
|
/(\b(public|private|protected|static|async)?\s+[\w$]+\s*\([^)]*\)\s*{)([\s\S]*?)(})/g,
|
|
(match, start, modifier, body, end) => {
|
|
return `${start} /* implementation removed */ ${end}`;
|
|
}
|
|
);
|
|
} else if (this.config.maxFunctionLines && this.config.maxFunctionLines > 0) {
|
|
// If not removing implementations completely, limit the number of lines
|
|
// Match function and method bodies
|
|
result = result.replace(
|
|
/(\b(function|constructor|async function)\s+[\w$]*\s*\([^)]*\)\s*{)([\s\S]*?)(})/g,
|
|
(match, start, funcType, body, end) => {
|
|
return this.limitFunctionBody(start, body, end);
|
|
}
|
|
);
|
|
|
|
// Match arrow function bodies
|
|
result = result.replace(
|
|
/(\([^)]*\)\s*=>\s*{)([\s\S]*?)(})/g,
|
|
(match, start, body, end) => {
|
|
return this.limitFunctionBody(start, body, end);
|
|
}
|
|
);
|
|
|
|
// Match method declarations
|
|
result = result.replace(
|
|
/(^\s*[\w$]*\s*\([^)]*\)\s*{)([\s\S]*?)(})/gm,
|
|
(match, start, body, end) => {
|
|
return this.limitFunctionBody(start, body, end);
|
|
}
|
|
);
|
|
|
|
// Match class methods
|
|
result = result.replace(
|
|
/(\b(public|private|protected|static|async)?\s+[\w$]+\s*\([^)]*\)\s*{)([\s\S]*?)(})/g,
|
|
(match, start, modifier, body, end) => {
|
|
return this.limitFunctionBody(start, body, end);
|
|
}
|
|
);
|
|
}
|
|
|
|
// Step 4: Remove blank lines if configured
|
|
if (this.config.removeBlankLines) {
|
|
result = result.replace(/^\s*[\r\n]/gm, '');
|
|
}
|
|
|
|
// Step 5: Restore preserved JSDoc comments
|
|
if (this.config.preserveJSDoc && jsDocComments.length > 0) {
|
|
// This is a placeholder; we already preserved JSDoc comments in the regex steps
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
/**
|
|
* Limit a function body to a maximum number of lines
|
|
* @param start The function signature and opening brace
|
|
* @param body The function body
|
|
* @param end The closing brace
|
|
* @returns The limited function body
|
|
*/
|
|
private limitFunctionBody(start: string, body: string, end: string): string {
|
|
const lines = body.split('\n');
|
|
if (lines.length > this.config.maxFunctionLines!) {
|
|
const limitedBody = lines.slice(0, this.config.maxFunctionLines!).join('\n');
|
|
return `${start}${limitedBody}\n // ... (${lines.length - this.config.maxFunctionLines!} lines trimmed)\n${end}`;
|
|
}
|
|
return `${start}${body}${end}`;
|
|
}
|
|
|
|
/**
|
|
* Trim a Markdown file to reduce token usage
|
|
* @param content The Markdown file contents
|
|
* @returns The trimmed file contents
|
|
*/
|
|
private trimMarkdownFile(content: string): string {
|
|
// For markdown files, we generally want to keep most content
|
|
// but we can remove lengthy code blocks if needed
|
|
return content;
|
|
}
|
|
|
|
/**
|
|
* Trim a JSON file to reduce token usage
|
|
* @param content The JSON file contents
|
|
* @returns The trimmed file contents
|
|
*/
|
|
private trimJsonFile(content: string): string {
|
|
try {
|
|
// Parse the JSON
|
|
const json = JSON.parse(content);
|
|
|
|
// For package.json, keep only essential information
|
|
if ('name' in json && 'version' in json && 'dependencies' in json) {
|
|
const essentialKeys = [
|
|
'name', 'version', 'description', 'author', 'license',
|
|
'main', 'types', 'exports', 'type'
|
|
];
|
|
|
|
const trimmedJson: any = {};
|
|
essentialKeys.forEach(key => {
|
|
if (key in json) {
|
|
trimmedJson[key] = json[key];
|
|
}
|
|
});
|
|
|
|
// Add dependency information without versions
|
|
if ('dependencies' in json) {
|
|
trimmedJson.dependencies = Object.keys(json.dependencies).reduce((acc, dep) => {
|
|
acc[dep] = '*'; // Replace version with wildcard
|
|
return acc;
|
|
}, {} as Record<string, string>);
|
|
}
|
|
|
|
// Return the trimmed JSON
|
|
return JSON.stringify(trimmedJson, null, 2);
|
|
}
|
|
|
|
// For other JSON files, leave as is
|
|
return content;
|
|
} catch (error) {
|
|
// If there's an error parsing the JSON, return the original content
|
|
return content;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Update the trimmer configuration
|
|
* @param config The new configuration to apply
|
|
*/
|
|
public updateConfig(config: ITrimConfig): void {
|
|
this.config = {
|
|
...this.config,
|
|
...config
|
|
};
|
|
}
|
|
|
|
/**
|
|
* Trim a file based on its importance tier
|
|
* @param filePath The path to the file
|
|
* @param content The file's contents
|
|
* @param level The trimming level to apply ('none', 'light', 'aggressive')
|
|
* @returns The trimmed file contents
|
|
*/
|
|
public trimFileWithLevel(
|
|
filePath: string,
|
|
content: string,
|
|
level: 'none' | 'light' | 'aggressive'
|
|
): string {
|
|
// No trimming for essential files
|
|
if (level === 'none') {
|
|
return content;
|
|
}
|
|
|
|
// Create a temporary config based on level
|
|
const originalConfig = { ...this.config };
|
|
|
|
try {
|
|
if (level === 'light') {
|
|
// Light trimming: preserve signatures, remove only complex implementations
|
|
this.config = {
|
|
...this.config,
|
|
removeImplementations: false,
|
|
preserveInterfaces: true,
|
|
preserveTypeDefs: true,
|
|
preserveJSDoc: true,
|
|
maxFunctionLines: 10,
|
|
removeComments: false,
|
|
removeBlankLines: true
|
|
};
|
|
} else if (level === 'aggressive') {
|
|
// Aggressive trimming: remove all implementations, keep only signatures
|
|
this.config = {
|
|
...this.config,
|
|
removeImplementations: true,
|
|
preserveInterfaces: true,
|
|
preserveTypeDefs: true,
|
|
preserveJSDoc: true,
|
|
maxFunctionLines: 3,
|
|
removeComments: true,
|
|
removeBlankLines: true
|
|
};
|
|
}
|
|
|
|
// Process based on file type
|
|
let result = content;
|
|
if (filePath.endsWith('.ts') || filePath.endsWith('.tsx')) {
|
|
result = this.trimTypeScriptFile(content);
|
|
} else if (filePath.endsWith('.md')) {
|
|
result = this.trimMarkdownFile(content);
|
|
} else if (filePath.endsWith('.json')) {
|
|
result = this.trimJsonFile(content);
|
|
}
|
|
|
|
return result;
|
|
} finally {
|
|
// Restore original config
|
|
this.config = originalConfig;
|
|
}
|
|
}
|
|
} |