Context Management Guide

import { encode } from 'tiktoken';

class AdvancedContextManager {
  constructor(options = {}) {
    this.maxTokens = options.maxTokens || 4000;
    this.maxOutputTokens = options.maxOutputTokens || 1000;
    this.compressionThreshold = options.compressionThreshold || 0.7;
    this.messages = [];
    this.summaries = [];
  }
  
  // Accurately count tokens
  countTokens(messages) {
    const encoding = encode(JSON.stringify(messages));
    return encoding.length;
  }
  
  // Add message and automatically manage context
  async addMessage(message) {
    this.messages.push(message);
    
    const currentTokens = this.countTokens(this.messages);
    const availableTokens = this.maxTokens - this.maxOutputTokens;
    
    if (currentTokens > availableTokens * this.compressionThreshold) {
      await this.compressContext();
    }
  }
  
  // Compress context
  async compressContext() {
    // Separate system messages and conversation
    const systemMessages = this.messages.filter(m => m.role === 'system');
    const conversation = this.messages.filter(m => m.role !== 'system');
    
    // Choose compression strategy
    if (conversation.length > 10) {
      // Strategy 1: Generate summary
      const oldMessages = conversation.slice(0, -6);
      const summary = await this.generateSummary(oldMessages);
      
      this.summaries.push({
        timestamp: Date.now(),
        summary: summary,
        messageCount: oldMessages.length
      });
      
      // Reconstruct message list
      this.messages = [
        ...systemMessages,
        { 
          role: 'system', 
          content: `Previous conversation summary: ${summary}` 
        },
        ...conversation.slice(-6)
      ];
    } else {
      // Strategy 2: Simple truncation
      this.messages = [
        ...systemMessages,
        ...conversation.slice(-6)
      ];
    }
  }
  
  // Generate summary
  async generateSummary(messages) {
    const prompt = `Summarize the following conversation in 100 words:
    ${messages.map(m => `${m.role}: ${m.content}`).join('\n')}`;
    
    // Call AI to generate summary
    const response = await callAPI({
      messages: [{ role: 'user', content: prompt }],
      max_tokens: 200
    });
    
    return response.content;
  }
  
  // Get optimized context
  getOptimizedContext() {
    const tokens = this.countTokens(this.messages);
    const budget = this.maxTokens - this.maxOutputTokens;
    
    if (tokens <= budget) {
      return this.messages;
    }
    
    // Need further trimming
    return this.trimToFit(this.messages, budget);
  }
  
  // Trim to specified size
  trimToFit(messages, targetTokens) {
    let trimmed = [...messages];
    
    // Priority: system > recent > old
    while (this.countTokens(trimmed) > targetTokens && trimmed.length > 2) {
      // Remove oldest non-system message
      const nonSystemIndex = trimmed.findIndex(m => m.role !== 'system');
      if (nonSystemIndex > -1 && nonSystemIndex < trimmed.length - 2) {
        trimmed.splice(nonSystemIndex, 1);
      } else {
        break;
      }
    }
    
    return trimmed;
  }
}

// usingExample
const contextManager = new AdvancedContextManager({
  maxTokens: 4000,
  maxOutputTokens: 1000,
  compressionThreshold: 0.7
});

// Handle conversation
async function handleConversation(userInput) {
  await contextManager.addMessage({ role: 'user', content: userInput });
  
  const optimizedContext = contextManager.getOptimizedContext();
  
  const response = await callAPI({
    messages: optimizedContext,
    max_tokens: contextManager.maxOutputTokens
  });
  
  await contextManager.addMessage(response.message);
  
  return response.content;
}