tsdoc/ts/context/enhanced-context.ts

import * as plugins from '../plugins.js';
import type { ContextMode, IContextResult, IFileInfo, TaskType } from './types.js';
import { ContextTrimmer } from './context-trimmer.js';
import { ConfigManager } from './config-manager.js';

/**
 * Enhanced ProjectContext that supports context optimization strategies
 */
export class EnhancedContext {
  private projectDir: string;
  private trimmer: ContextTrimmer;
  private configManager: ConfigManager;
  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());
  }
  
  /**
   * 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());
  }
  
  /**
   * 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;
  }
  
  /**
   * Gather files from the project
   * @param includePaths Optional paths to include
   * @param excludePaths Optional paths to exclude
   */
  public async gatherFiles(includePaths?: string[], excludePaths?: string[]): Promise<Record<string, plugins.smartfile.SmartFile | plugins.smartfile.SmartFile[]>> {
    const smartfilePackageJSON = await plugins.smartfile.SmartFile.fromFilePath(
      plugins.path.join(this.projectDir, 'package.json'),
      this.projectDir,
    );
    
    const smartfilesReadme = await plugins.smartfile.SmartFile.fromFilePath(
      plugins.path.join(this.projectDir, 'readme.md'),
      this.projectDir,
    );

    const smartfilesReadmeHints = await plugins.smartfile.SmartFile.fromFilePath(
      plugins.path.join(this.projectDir, 'readme.hints.md'),
      this.projectDir,
    );
    
    const smartfilesNpmextraJSON = await plugins.smartfile.SmartFile.fromFilePath(
      plugins.path.join(this.projectDir, 'npmextra.json'),
      this.projectDir,
    );
    
    // Use provided include paths or default to all TypeScript files
    const includeGlobs = includePaths?.map(path => `${path}/**/*.ts`) || ['ts*/**/*.ts'];
    
    // Get TypeScript files
    const smartfilesModPromises = includeGlobs.map(glob => 
      plugins.smartfile.fs.fileTreeToObject(this.projectDir, glob)
    );
    
    const smartfilesModArrays = await Promise.all(smartfilesModPromises);
    
    // Flatten the arrays
    const smartfilesMod: plugins.smartfile.SmartFile[] = [];
    smartfilesModArrays.forEach(array => {
      smartfilesMod.push(...array);
    });
    
    // Get test files if not excluded
    let smartfilesTest: plugins.smartfile.SmartFile[] = [];
    if (!excludePaths?.includes('test/')) {
      smartfilesTest = await plugins.smartfile.fs.fileTreeToObject(
        this.projectDir,
        'test/**/*.ts',
      );
    }
    
    return {
      smartfilePackageJSON,
      smartfilesReadme,
      smartfilesReadmeHints,
      smartfilesNpmextraJSON,
      smartfilesMod,
      smartfilesTest,
    };
  }
  
  /**
   * Convert files to context string
   * @param files The files to convert
   * @param mode The context mode to use
   */
  public async convertFilesToContext(
    files: plugins.smartfile.SmartFile[],
    mode: ContextMode = this.contextMode
  ): Promise<string> {
    // Reset context result
    this.contextResult = {
      context: '',
      tokenCount: 0,
      includedFiles: [],
      trimmedFiles: [],
      excludedFiles: [],
      tokenSavings: 0
    };
    
    let totalTokenCount = 0;
    let totalOriginalTokens = 0;
    
    // Sort files by importance (for now just a simple alphabetical sort)
    // Later this could be enhanced with more sophisticated prioritization
    const sortedFiles = [...files].sort((a, b) => a.relative.localeCompare(b.relative));
    
    const processedFiles: string[] = [];
    
    for (const smartfile of sortedFiles) {
      // Calculate original token count
      const originalContent = smartfile.contents.toString();
      const originalTokenCount = this.countTokens(originalContent);
      totalOriginalTokens += originalTokenCount;
      
      // Apply trimming based on mode
      let processedContent = originalContent;
      
      if (mode !== 'full') {
        processedContent = this.trimmer.trimFile(
          smartfile.relative,
          originalContent,
          mode
        );
      }
      
      // Calculate new token count
      const processedTokenCount = this.countTokens(processedContent);
      
      // Check if we have budget for this file
      if (totalTokenCount + processedTokenCount > this.tokenBudget) {
        // We don't have budget for this file
        this.contextResult.excludedFiles.push({
          path: smartfile.path,
          contents: originalContent,
          relativePath: smartfile.relative,
          tokenCount: originalTokenCount
        });
        continue;
      }
      
      // Format the file for context
      const formattedContent = `
====== START OF FILE ${smartfile.relative} ======
  
${processedContent}
  
====== END OF FILE ${smartfile.relative} ======
      `;
      
      processedFiles.push(formattedContent);
      totalTokenCount += processedTokenCount;
      
      // Track file in appropriate list
      const fileInfo: IFileInfo = {
        path: smartfile.path,
        contents: processedContent,
        relativePath: smartfile.relative,
        tokenCount: processedTokenCount
      };
      
      if (mode === 'full' || processedContent === originalContent) {
        this.contextResult.includedFiles.push(fileInfo);
      } else {
        this.contextResult.trimmedFiles.push(fileInfo);
        this.contextResult.tokenSavings += (originalTokenCount - processedTokenCount);
      }
    }
    
    // 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
   * @param taskType Optional task type for task-specific context
   */
  public async buildContext(taskType?: TaskType): Promise<IContextResult> {
    // Initialize if needed
    if (this.tokenBudget === 0) {
      await this.initialize();
    }
    
    // Get task-specific configuration if a task type is provided
    if (taskType) {
      const taskConfig = this.configManager.getTaskConfig(taskType);
      if (taskConfig.mode) {
        this.setContextMode(taskConfig.mode);
      }
    }
    
    // Gather files
    const taskConfig = taskType ? this.configManager.getTaskConfig(taskType) : undefined;
    const files = await this.gatherFiles(
      taskConfig?.includePaths,
      taskConfig?.excludePaths
    );
    
    // Convert files to context
    // Create an array of all files to process
    const allFiles: plugins.smartfile.SmartFile[] = [];
    
    // Add individual files
    if (files.smartfilePackageJSON) allFiles.push(files.smartfilePackageJSON as plugins.smartfile.SmartFile);
    if (files.smartfilesReadme) allFiles.push(files.smartfilesReadme as plugins.smartfile.SmartFile);
    if (files.smartfilesReadmeHints) allFiles.push(files.smartfilesReadmeHints as plugins.smartfile.SmartFile);
    if (files.smartfilesNpmextraJSON) allFiles.push(files.smartfilesNpmextraJSON as plugins.smartfile.SmartFile);
    
    // Add arrays of files
    if (files.smartfilesMod) {
      if (Array.isArray(files.smartfilesMod)) {
        allFiles.push(...files.smartfilesMod);
      } else {
        allFiles.push(files.smartfilesMod);
      }
    }
    
    if (files.smartfilesTest) {
      if (Array.isArray(files.smartfilesTest)) {
        allFiles.push(...files.smartfilesTest);
      } else {
        allFiles.push(files.smartfilesTest);
      }
    }
    
    const context = await this.convertFilesToContext(allFiles);
    
    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
    };
  }
}