tsdoc/ts/context/context-trimmer.ts

246 lines
7.9 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
};
}
}