Agentic-Compaction-and-Pipleline-by-GLM-5/agent-system/core/context-manager.ts

/**
 * Context Compaction Module
 * 
 * Manages conversation context by implementing multiple compaction strategies:
 * - Sliding window (keep recent messages)
 * - Summarization (compress older messages)
 * - Priority-based retention (keep important messages)
 * - Semantic clustering (group related messages)
 */

import { TokenCounter, TokenBudget } from './token-counter';
import { ConversationSummarizer, SummaryResult, ConversationTurn } from './summarizer';

export interface CompactionResult {
  messages: ConversationTurn[];
  originalTokenCount: number;
  newTokenCount: number;
  tokensSaved: number;
  compressionRatio: number;
  strategy: CompactionStrategy;
  summaryAdded: boolean;
  removedCount: number;
}

export type CompactionStrategy = 
  | 'sliding-window'
  | 'summarize-old'
  | 'priority-retention'
  | 'hybrid';

export interface CompactionConfig {
  maxTokens: number;
  targetTokens: number;
  strategy: CompactionStrategy;
  preserveRecentCount: number;
  preserveSystemMessage: boolean;
  priorityKeywords: string[];
  summaryMaxTokens: number;
  triggerThreshold: number; // Percentage (0-100) of maxTokens to trigger compaction
}

export interface MessagePriority {
  message: ConversationTurn;
  index: number;
  priority: number;
  tokens: number;
  reasons: string[];
}

const DEFAULT_CONFIG: CompactionConfig = {
  maxTokens: 120000,
  targetTokens: 80000,
  strategy: 'hybrid',
  preserveRecentCount: 6,
  preserveSystemMessage: true,
  priorityKeywords: ['important', 'critical', 'decision', 'todo', 'remember'],
  summaryMaxTokens: 2000,
  triggerThreshold: 80
};

/**
 * ContextCompactor - Manages conversation context compaction
 */
export class ContextCompactor {
  private tokenCounter: TokenCounter;
  private summarizer: ConversationSummarizer;
  private config: CompactionConfig;
  private lastCompaction: Date | null = null;
  private compactionHistory: CompactionResult[] = [];

  constructor(config: Partial<CompactionConfig> = {}) {
    this.config = { ...DEFAULT_CONFIG, ...config };
    this.tokenCounter = new TokenCounter(this.config.maxTokens);
    this.summarizer = new ConversationSummarizer();
  }

  /**
   * Check if compaction is needed
   */
  needsCompaction(messages: ConversationTurn[]): boolean {
    const tokenCount = this.tokenCounter.countConversation(messages).total;
    const threshold = this.config.maxTokens * (this.config.triggerThreshold / 100);
    return tokenCount >= threshold;
  }

  /**
   * Get current token budget status
   */
  getBudget(messages: ConversationTurn[]): TokenBudget {
    const tokenCount = this.tokenCounter.countConversation(messages).total;
    return this.tokenCounter.getBudget(tokenCount);
  }

  /**
   * Compact the conversation using the configured strategy
   */
  async compact(messages: ConversationTurn[]): Promise<CompactionResult> {
    const originalTokenCount = this.tokenCounter.countConversation(messages).total;

    // Check if compaction is needed
    if (originalTokenCount < this.config.targetTokens) {
      return {
        messages,
        originalTokenCount,
        newTokenCount: originalTokenCount,
        tokensSaved: 0,
        compressionRatio: 1,
        strategy: this.config.strategy,
        summaryAdded: false,
        removedCount: 0
      };
    }

    let result: CompactionResult;

    switch (this.config.strategy) {
      case 'sliding-window':
        result = this.slidingWindowCompaction(messages, originalTokenCount);
        break;
      case 'summarize-old':
        result = await this.summarizeOldCompaction(messages, originalTokenCount);
        break;
      case 'priority-retention':
        result = this.priorityRetentionCompaction(messages, originalTokenCount);
        break;
      case 'hybrid':
      default:
        result = await this.hybridCompaction(messages, originalTokenCount);
        break;
    }

    // Record compaction
    this.lastCompaction = new Date();
    this.compactionHistory.push(result);

    return result;
  }

  /**
   * Sliding window compaction - keep only recent messages
   */
  private slidingWindowCompaction(
    messages: ConversationTurn[],
    originalTokenCount: number
  ): CompactionResult {
    const result: ConversationTurn[] = [];
    
    // Preserve system message if configured
    if (this.config.preserveSystemMessage && messages[0]?.role === 'system') {
      result.push(messages[0]);
    }

    // Add recent messages
    const startIndex = Math.max(
      this.config.preserveSystemMessage && messages[0]?.role === 'system' ? 1 : 0,
      messages.length - this.config.preserveRecentCount
    );

    for (let i = startIndex; i < messages.length; i++) {
      result.push(messages[i]);
    }

    const newTokenCount = this.tokenCounter.countConversation(result).total;

    return {
      messages: result,
      originalTokenCount,
      newTokenCount,
      tokensSaved: originalTokenCount - newTokenCount,
      compressionRatio: newTokenCount / originalTokenCount,
      strategy: 'sliding-window',
      summaryAdded: false,
      removedCount: messages.length - result.length
    };
  }

  /**
   * Summarize old messages compaction
   */
  private async summarizeOldCompaction(
    messages: ConversationTurn[],
    originalTokenCount: number
  ): Promise<CompactionResult> {
    const result: ConversationTurn[] = [];
    
    // Preserve system message
    if (this.config.preserveSystemMessage && messages[0]?.role === 'system') {
      result.push(messages[0]);
    }

    // Find cutoff point
    const cutoffIndex = messages.length - this.config.preserveRecentCount;
    
    if (cutoffIndex > 1) {
      // Get messages to summarize
      const toSummarize = messages.slice(
        this.config.preserveSystemMessage && messages[0]?.role === 'system' ? 1 : 0,
        cutoffIndex
      );

      // Create summary
      const summaryResult = await this.summarizer.summarize(toSummarize, {
        maxSummaryTokens: this.config.summaryMaxTokens
      });

      // Add summary as a system message
      result.push({
        role: 'system',
        content: `[Previous Conversation Summary]\n${summaryResult.summary}\n\nKey Points:\n${summaryResult.keyPoints.map(p => `- ${p}`).join('\n')}`,
        metadata: { 
          type: 'compaction-summary',
          originalMessageCount: toSummarize.length,
          createdAt: new Date().toISOString()
        }
      });
    }

    // Add recent messages
    for (let i = Math.max(cutoffIndex, 0); i < messages.length; i++) {
      result.push(messages[i]);
    }

    const newTokenCount = this.tokenCounter.countConversation(result).total;

    return {
      messages: result,
      originalTokenCount,
      newTokenCount,
      tokensSaved: originalTokenCount - newTokenCount,
      compressionRatio: newTokenCount / originalTokenCount,
      strategy: 'summarize-old',
      summaryAdded: cutoffIndex > 1,
      removedCount: messages.length - result.length + (cutoffIndex > 1 ? 1 : 0)
    };
  }

  /**
   * Priority-based retention compaction
   */
  private priorityRetentionCompaction(
    messages: ConversationTurn[],
    originalTokenCount: number
  ): CompactionResult {
    // Calculate priorities for all messages
    const priorities = this.calculateMessagePriorities(messages);

    // Sort by priority (descending)
    priorities.sort((a, b) => b.priority - a.priority);

    // Select messages until we hit target tokens
    const selected: ConversationTurn[] = [];
    let currentTokens = 0;
    const selectedIndices = new Set<number>();

    // Always include system message if configured
    if (this.config.preserveSystemMessage && messages[0]?.role === 'system') {
      selected.push(messages[0]);
      selectedIndices.add(0);
      currentTokens += this.tokenCounter.countMessage(messages[0]);
    }

    // Always include recent messages (high priority)
    const recentStart = Math.max(
      this.config.preserveSystemMessage && messages[0]?.role === 'system' ? 1 : 0,
      messages.length - this.config.preserveRecentCount
    );

    for (let i = recentStart; i < messages.length; i++) {
      if (!selectedIndices.has(i)) {
        selected.push(messages[i]);
        selectedIndices.add(i);
        currentTokens += this.tokenCounter.countMessage(messages[i]);
      }
    }

    // Add high-priority messages until target is reached
    for (const mp of priorities) {
      if (selectedIndices.has(mp.index)) continue;
      if (currentTokens + mp.tokens > this.config.targetTokens) break;

      selected.push(mp.message);
      selectedIndices.add(mp.index);
      currentTokens += mp.tokens;
    }

    // Sort selected messages by original order
    selected.sort((a, b) => {
      const aIdx = messages.indexOf(a);
      const bIdx = messages.indexOf(b);
      return aIdx - bIdx;
    });

    const newTokenCount = this.tokenCounter.countConversation(selected).tokens;

    return {
      messages: selected,
      originalTokenCount,
      newTokenCount,
      tokensSaved: originalTokenCount - newTokenCount,
      compressionRatio: newTokenCount / originalTokenCount,
      strategy: 'priority-retention',
      summaryAdded: false,
      removedCount: messages.length - selected.length
    };
  }

  /**
   * Hybrid compaction - combines multiple strategies
   */
  private async hybridCompaction(
    messages: ConversationTurn[],
    originalTokenCount: number
  ): Promise<CompactionResult> {
    const result: ConversationTurn[] = [];
    
    // Preserve system message
    if (this.config.preserveSystemMessage && messages[0]?.role === 'system') {
      result.push(messages[0]);
    }

    const priorities = this.calculateMessagePriorities(messages);
    
    // Identify important messages to keep
    const importantIndices = new Set<number>();
    for (const mp of priorities) {
      if (mp.priority >= 7) { // High priority threshold
        importantIndices.add(mp.index);
      }
    }

    // Find cutoff for summarization
    const cutoffIndex = messages.length - this.config.preserveRecentCount;
    
    // Summarize middle section if needed
    const middleStart = this.config.preserveSystemMessage && messages[0]?.role === 'system' ? 1 : 0;
    const middleEnd = cutoffIndex;
    
    const middleMessages = messages.slice(middleStart, middleEnd)
      .filter((_, idx) => !importantIndices.has(middleStart + idx));
    
    if (middleMessages.length > 3) {
      const summaryResult = await this.summarizer.summarize(middleMessages, {
        maxSummaryTokens: this.config.summaryMaxTokens
      });

      result.push({
        role: 'system',
        content: `[Context Summary]\n${summaryResult.summary}`,
        metadata: {
          type: 'compaction-summary',
          originalMessageCount: middleMessages.length
        }
      });
    }

    // Add important messages from the middle section
    for (let i = middleStart; i < middleEnd; i++) {
      if (importantIndices.has(i)) {
        result.push(messages[i]);
      }
    }

    // Add recent messages
    for (let i = cutoffIndex; i < messages.length; i++) {
      result.push(messages[i]);
    }

    // Sort by original order
    result.sort((a, b) => messages.indexOf(a) - messages.indexOf(b));

    const newTokenCount = this.tokenCounter.countConversation(result).tokens;

    return {
      messages: result,
      originalTokenCount,
      newTokenCount,
      tokensSaved: originalTokenCount - newTokenCount,
      compressionRatio: newTokenCount / originalTokenCount,
      strategy: 'hybrid',
      summaryAdded: middleMessages.length > 3,
      removedCount: messages.length - result.length + (middleMessages.length > 3 ? 1 : 0)
    };
  }

  /**
   * Calculate priority scores for messages
   */
  private calculateMessagePriorities(messages: ConversationTurn[]): MessagePriority[] {
    return messages.map((msg, index) => {
      let priority = 5; // Base priority
      const reasons: string[] = [];

      // System messages are high priority
      if (msg.role === 'system') {
        priority += 3;
        reasons.push('System message');
      }

      // Recent messages are higher priority
      const recency = index / messages.length;
      priority += recency * 2;
      if (recency > 0.7) reasons.push('Recent message');

      // Check for priority keywords
      const content = msg.content.toLowerCase();
      for (const keyword of this.config.priorityKeywords) {
        if (content.includes(keyword.toLowerCase())) {
          priority += 1;
          reasons.push(`Contains "${keyword}"`);
        }
      }

      // User questions might be important
      if (msg.role === 'user' && content.includes('?')) {
        priority += 0.5;
        reasons.push('User question');
      }

      // Code blocks might be important
      if (content.includes('```')) {
        priority += 1;
        reasons.push('Contains code');
      }

      // Decisions or confirmations
      if (content.match(/(yes|no|agree|decided|confirmed|done)/i)) {
        priority += 0.5;
        reasons.push('Potential decision');
      }

      return {
        message: msg,
        index,
        priority: Math.min(10, Math.max(1, priority)),
        tokens: this.tokenCounter.countMessage(msg),
        reasons
      };
    });
  }

  /**
   * Get compaction history
   */
  getHistory(): CompactionResult[] {
    return [...this.compactionHistory];
  }

  /**
   * Get statistics about compactions
   */
  getStats(): {
    totalCompactions: number;
    totalTokensSaved: number;
    averageCompressionRatio: number;
    lastCompaction: Date | null;
  } {
    if (this.compactionHistory.length === 0) {
      return {
        totalCompactions: 0,
        totalTokensSaved: 0,
        averageCompressionRatio: 0,
        lastCompaction: null
      };
    }

    const totalTokensSaved = this.compactionHistory.reduce(
      (sum, c) => sum + c.tokensSaved, 0
    );
    
    const avgRatio = this.compactionHistory.reduce(
      (sum, c) => sum + c.compressionRatio, 0
    ) / this.compactionHistory.length;

    return {
      totalCompactions: this.compactionHistory.length,
      totalTokensSaved,
      averageCompressionRatio: avgRatio,
      lastCompaction: this.lastCompaction
    };
  }
}

/**
 * ConversationContextManager - High-level context management
 */
export class ConversationContextManager {
  private compactor: ContextCompactor;
  private messages: ConversationTurn[] = [];
  private summary: string | null = null;

  constructor(config: Partial<CompactionConfig> = {}) {
    this.compactor = new ContextCompactor(config);
  }

  /**
   * Add a message to the context
   */
  addMessage(message: ConversationTurn): void {
    this.messages.push({
      ...message,
      timestamp: message.timestamp || new Date()
    });
  }

  /**
   * Get all messages, with optional compaction
   */
  async getMessages(): Promise<ConversationTurn[]> {
    if (this.compactor.needsCompaction(this.messages)) {
      const result = await this.compactor.compact(this.messages);
      this.messages = result.messages;
      return this.messages;
    }
    return this.messages;
  }

  /**
   * Force compaction
   */
  async forceCompact(): Promise<CompactionResult> {
    const result = await this.compactor.compact(this.messages);
    this.messages = result.messages;
    return result;
  }

  /**
   * Get current token count
   */
  getTokenCount(): number {
    return this.compactor['tokenCounter'].countConversation(this.messages).total;
  }

  /**
   * Clear the context
   */
  clear(): void {
    this.messages = [];
    this.summary = null;
  }

  /**
   * Get context stats
   */
  getStats() {
    return {
      messageCount: this.messages.length,
      tokenCount: this.getTokenCount(),
      budget: this.compactor.getBudget(this.messages),
      compactionStats: this.compactor.getStats()
    };
  }
}

// Default instance
export const defaultCompactor = new ContextCompactor();
export const defaultContextManager = new ConversationContextManager();