332 lines
10 KiB
TypeScript
332 lines
10 KiB
TypeScript
import * as plugins from '../plugins.js';
|
|
import type { ContextMode, IContextResult, IFileInfo, TaskType, IFileMetadata } from './types.js';
|
|
import { ContextTrimmer } from './context-trimmer.js';
|
|
import { ConfigManager } from './config-manager.js';
|
|
import { LazyFileLoader } from './lazy-file-loader.js';
|
|
import { ContextCache } from './context-cache.js';
|
|
import { ContextAnalyzer } from './context-analyzer.js';
|
|
|
|
/**
|
|
* Enhanced ProjectContext that supports context optimization strategies
|
|
*/
|
|
export class EnhancedContext {
|
|
private projectDir: string;
|
|
private trimmer: ContextTrimmer;
|
|
private configManager: ConfigManager;
|
|
private lazyLoader: LazyFileLoader;
|
|
private cache: ContextCache;
|
|
private analyzer: ContextAnalyzer;
|
|
private contextMode: ContextMode = 'trimmed';
|
|
private tokenBudget: number = 190000; // Default for o4-mini
|
|
private contextResult: IContextResult = {
|
|
context: '',
|
|
tokenCount: 0,
|
|
includedFiles: [],
|
|
trimmedFiles: [],
|
|
excludedFiles: [],
|
|
tokenSavings: 0
|
|
};
|
|
|
|
/**
|
|
* Create a new EnhancedContext
|
|
* @param projectDirArg The project directory
|
|
*/
|
|
constructor(projectDirArg: string) {
|
|
this.projectDir = projectDirArg;
|
|
this.configManager = ConfigManager.getInstance();
|
|
this.trimmer = new ContextTrimmer(this.configManager.getTrimConfig());
|
|
this.lazyLoader = new LazyFileLoader(projectDirArg);
|
|
this.cache = new ContextCache(projectDirArg, this.configManager.getCacheConfig());
|
|
this.analyzer = new ContextAnalyzer(
|
|
projectDirArg,
|
|
this.configManager.getPrioritizationWeights(),
|
|
this.configManager.getTierConfig()
|
|
);
|
|
}
|
|
|
|
/**
|
|
* Initialize the context builder
|
|
*/
|
|
public async initialize(): Promise<void> {
|
|
await this.configManager.initialize(this.projectDir);
|
|
this.tokenBudget = this.configManager.getMaxTokens();
|
|
this.trimmer.updateConfig(this.configManager.getTrimConfig());
|
|
await this.cache.init();
|
|
}
|
|
|
|
/**
|
|
* Set the context mode
|
|
* @param mode The context mode to use
|
|
*/
|
|
public setContextMode(mode: ContextMode): void {
|
|
this.contextMode = mode;
|
|
}
|
|
|
|
/**
|
|
* Set the token budget
|
|
* @param maxTokens The maximum tokens to use
|
|
*/
|
|
public setTokenBudget(maxTokens: number): void {
|
|
this.tokenBudget = maxTokens;
|
|
}
|
|
|
|
/**
|
|
* Convert files to context with smart analysis and prioritization
|
|
* @param metadata - File metadata to analyze
|
|
* @param taskType - Task type for context-aware prioritization
|
|
* @param mode - Context mode to use
|
|
* @returns Context string
|
|
*/
|
|
public async convertFilesToContextWithAnalysis(
|
|
metadata: IFileMetadata[],
|
|
taskType: TaskType,
|
|
mode: ContextMode = this.contextMode
|
|
): Promise<string> {
|
|
// Reset context result
|
|
this.contextResult = {
|
|
context: '',
|
|
tokenCount: 0,
|
|
includedFiles: [],
|
|
trimmedFiles: [],
|
|
excludedFiles: [],
|
|
tokenSavings: 0
|
|
};
|
|
|
|
// Analyze files for smart prioritization
|
|
const analysis = await this.analyzer.analyze(metadata, taskType, []);
|
|
|
|
// Sort files by importance score (highest first)
|
|
const sortedAnalysis = [...analysis.files].sort(
|
|
(a, b) => b.importanceScore - a.importanceScore
|
|
);
|
|
|
|
// Filter out excluded tier
|
|
const relevantFiles = sortedAnalysis.filter(f => f.tier !== 'excluded');
|
|
|
|
let totalTokenCount = 0;
|
|
let totalOriginalTokens = 0;
|
|
const processedFiles: string[] = [];
|
|
|
|
// Load files with cache support
|
|
for (const fileAnalysis of relevantFiles) {
|
|
try {
|
|
// Check cache first
|
|
let contents: string;
|
|
let originalTokenCount: number;
|
|
|
|
const cached = await this.cache.get(fileAnalysis.path);
|
|
if (cached) {
|
|
contents = cached.contents;
|
|
originalTokenCount = cached.tokenCount;
|
|
} else {
|
|
// Load file
|
|
const fileData = await plugins.smartfile.fs.toStringSync(fileAnalysis.path);
|
|
contents = fileData;
|
|
originalTokenCount = this.countTokens(contents);
|
|
|
|
// Cache it
|
|
await this.cache.set({
|
|
path: fileAnalysis.path,
|
|
contents,
|
|
tokenCount: originalTokenCount,
|
|
mtime: Date.now(),
|
|
cachedAt: Date.now()
|
|
});
|
|
}
|
|
|
|
totalOriginalTokens += originalTokenCount;
|
|
|
|
// Apply tier-based trimming
|
|
let processedContent = contents;
|
|
let trimLevel: 'none' | 'light' | 'aggressive' = 'light';
|
|
|
|
if (fileAnalysis.tier === 'essential') {
|
|
trimLevel = 'none';
|
|
} else if (fileAnalysis.tier === 'important') {
|
|
trimLevel = 'light';
|
|
} else if (fileAnalysis.tier === 'optional') {
|
|
trimLevel = 'aggressive';
|
|
}
|
|
|
|
// Apply trimming based on mode and tier
|
|
if (mode !== 'full' && trimLevel !== 'none') {
|
|
const relativePath = plugins.path.relative(this.projectDir, fileAnalysis.path);
|
|
processedContent = this.trimmer.trimFileWithLevel(
|
|
relativePath,
|
|
contents,
|
|
trimLevel
|
|
);
|
|
}
|
|
|
|
// Calculate token count
|
|
const processedTokenCount = this.countTokens(processedContent);
|
|
|
|
// Check token budget
|
|
if (totalTokenCount + processedTokenCount > this.tokenBudget) {
|
|
// We don't have budget for this file
|
|
const relativePath = plugins.path.relative(this.projectDir, fileAnalysis.path);
|
|
this.contextResult.excludedFiles.push({
|
|
path: fileAnalysis.path,
|
|
contents,
|
|
relativePath,
|
|
tokenCount: originalTokenCount,
|
|
importanceScore: fileAnalysis.importanceScore
|
|
});
|
|
continue;
|
|
}
|
|
|
|
// Format the file for context
|
|
const relativePath = plugins.path.relative(this.projectDir, fileAnalysis.path);
|
|
const formattedContent = `
|
|
====== START OF FILE ${relativePath} ======
|
|
|
|
${processedContent}
|
|
|
|
====== END OF FILE ${relativePath} ======
|
|
`;
|
|
|
|
processedFiles.push(formattedContent);
|
|
totalTokenCount += processedTokenCount;
|
|
|
|
// Track file in appropriate list
|
|
const fileInfo: IFileInfo = {
|
|
path: fileAnalysis.path,
|
|
contents: processedContent,
|
|
relativePath,
|
|
tokenCount: processedTokenCount,
|
|
importanceScore: fileAnalysis.importanceScore
|
|
};
|
|
|
|
if (trimLevel === 'none' || processedContent === contents) {
|
|
this.contextResult.includedFiles.push(fileInfo);
|
|
} else {
|
|
this.contextResult.trimmedFiles.push(fileInfo);
|
|
this.contextResult.tokenSavings += (originalTokenCount - processedTokenCount);
|
|
}
|
|
} catch (error) {
|
|
console.warn(`Failed to process file ${fileAnalysis.path}:`, error.message);
|
|
}
|
|
}
|
|
|
|
// Join all processed files
|
|
const context = processedFiles.join('\n');
|
|
|
|
// Update context result
|
|
this.contextResult.context = context;
|
|
this.contextResult.tokenCount = totalTokenCount;
|
|
|
|
return context;
|
|
}
|
|
|
|
/**
|
|
* Build context for the project using smart analysis
|
|
* @param taskType Task type for context-aware prioritization (defaults to 'description')
|
|
*/
|
|
public async buildContext(taskType?: TaskType): Promise<IContextResult> {
|
|
// Initialize if needed
|
|
if (this.tokenBudget === 0) {
|
|
await this.initialize();
|
|
}
|
|
|
|
// Smart context building always requires a task type for optimal prioritization
|
|
// Default to 'description' if not provided
|
|
const effectiveTaskType = taskType || 'description';
|
|
|
|
// Get task-specific configuration
|
|
const taskConfig = this.configManager.getTaskConfig(effectiveTaskType);
|
|
if (taskConfig.mode) {
|
|
this.setContextMode(taskConfig.mode);
|
|
}
|
|
|
|
// Build globs for scanning
|
|
const includeGlobs = taskConfig?.includePaths?.map(p => `${p}/**/*.ts`) || [
|
|
'ts/**/*.ts',
|
|
'ts*/**/*.ts'
|
|
];
|
|
|
|
// Add config files
|
|
const configGlobs = [
|
|
'package.json',
|
|
'readme.md',
|
|
'readme.hints.md',
|
|
'npmextra.json'
|
|
];
|
|
|
|
// Scan files for metadata (fast, doesn't load contents)
|
|
const metadata = await this.lazyLoader.scanFiles([...configGlobs, ...includeGlobs]);
|
|
|
|
// Use smart analyzer to build context with intelligent prioritization
|
|
await this.convertFilesToContextWithAnalysis(metadata, effectiveTaskType, this.contextMode);
|
|
|
|
return this.contextResult;
|
|
}
|
|
|
|
/**
|
|
* Update the context with git diff information for commit tasks
|
|
* @param gitDiff The git diff to include
|
|
*/
|
|
public updateWithGitDiff(gitDiff: string): IContextResult {
|
|
// If we don't have a context yet, return empty result
|
|
if (!this.contextResult.context) {
|
|
return this.contextResult;
|
|
}
|
|
|
|
// Add git diff to context
|
|
const diffSection = `
|
|
====== GIT DIFF ======
|
|
|
|
${gitDiff}
|
|
|
|
====== END GIT DIFF ======
|
|
`;
|
|
|
|
const diffTokenCount = this.countTokens(diffSection);
|
|
|
|
// Update context and token count
|
|
this.contextResult.context += diffSection;
|
|
this.contextResult.tokenCount += diffTokenCount;
|
|
|
|
return this.contextResult;
|
|
}
|
|
|
|
/**
|
|
* Count tokens in a string
|
|
* @param text The text to count tokens for
|
|
* @param model The model to use for token counting
|
|
*/
|
|
public countTokens(text: string, model: string = 'gpt-3.5-turbo'): number {
|
|
try {
|
|
// Use the gpt-tokenizer library to count tokens
|
|
const tokens = plugins.gptTokenizer.encode(text);
|
|
return tokens.length;
|
|
} catch (error) {
|
|
console.error('Error counting tokens:', error);
|
|
// Provide a rough estimate if tokenization fails
|
|
return Math.ceil(text.length / 4);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Get the context result
|
|
*/
|
|
public getContextResult(): IContextResult {
|
|
return this.contextResult;
|
|
}
|
|
|
|
/**
|
|
* Get the token count for the current context
|
|
*/
|
|
public getTokenCount(): number {
|
|
return this.contextResult.tokenCount;
|
|
}
|
|
|
|
/**
|
|
* Get both the context string and its token count
|
|
*/
|
|
public getContextWithTokenCount(): { context: string; tokenCount: number } {
|
|
return {
|
|
context: this.contextResult.context,
|
|
tokenCount: this.contextResult.tokenCount
|
|
};
|
|
}
|
|
} |