LLM memory management

lementation provides a production-grade MemoryManager class. It handles token counting, sliding window management, vector retrieval integration, and context assembly with priority injection.

import { createHash } from 'crypto';

// Interfaces for type safety and extensibility
interface MemoryEntry {
  id: string;
  content: string;
  timestamp: number;
  priority: 'critical' | 'high' | 'normal' | 'low';
  tokens?: number;
}

interface MemoryConfig {
  maxTokens: number;
  windowSize: number; // Number of recent turns to keep
  embeddingModel: string;
  vectorStore: VectorStoreInterface;
  summarizer: SummarizerInterface;
}

interface VectorStoreInterface {
  query(query: string, topK: number): Promise<MemoryEntry[]>;
  insert(entry: MemoryEntry): Promise<void>;
}

interface SummarizerInterface {
  summarize(entries: MemoryEntry[]): Promise<string>;
}

// Token counter mock (use tiktoken in production)
const countTokens = (text: string): number => Math.ceil(text.length / 4);

export class LLMMemoryManager {
  private config: MemoryConfig;
  private workingMemory: MemoryEntry[] = [];
  private cache: Map<string, string> = new Map();

  constructor(config: MemoryConfig) {
    this.config = config;
  }

  /**
   * Adds new interaction to memory.
   * Triggers compression if window exceeds limits.
   */
  async add(role: 'user' | 'assistant' | 'system', content: string): Promise<void> {
    const entry: MemoryEntry = {
      id: createHash('sha256').update(content).digest('hex').slice(0, 12),
      content: `${role}: ${content}`,
      timestamp: Date.now(),
      priority: role === 'system' ? 'critical' : 'normal',
      tokens: countTokens(content),
    };

    // System prompts are critical and always retained
    if (entry.priority === 'critical') {
      this.workingMemory.unshift(entry);
      return;
    }

    // Check if compression is needed
    const currentTokenCount = this.workingMemory.reduce((sum, e) => sum + (e.tokens || 0), 0);
    
    if (currentTokenCount + (entry.tokens || 0) > this.config.maxTokens) {
      await this.compress();
    }

    this.workingMemory.push(entry);
    
    // Asynchronous long-term storage (non-blocking)
    this.config.vectorStore.insert(entry).catch(err => console.error('Vector insert failed:', err));
  }

  /**
   * Retrieves context for a new query.
   * Combines working memory with relevant long-term memories.
   */
  async retrieve(query: string): Promise<string> {
    // 1. Fetch relevant long-term memories
    const retrieved = await this.config.vectorStore.query(query, 5);
    
    // 2. Deduplicate against working memory
    const workingIds = new Set(this.workingMemory.map(e => e.id));
    const uniqueRetrieved = retrieved.filter(e => !workingIds.has(e.id));

    // 3. Assemble context
    // Priority: Critical/Recent > Retrieved > Compressed Summary
    const contextParts: string[] = [];

    // Add critical system instructions first
    contextParts.push(...this.workingMemory
      .filter(e => e.priority === 'critical')
      .map(e => e.content)
    );

    // Add retrieved context
    contextParts.push(...uniqueRetrieved.map(e => e.content));

    // Add recent working memory
    contextParts.push(...this.workingMemory
      .filter(e => e.priority !== 'critical')
      .slice(-this.config.windowSize)
      .map(e => e.content)
    );

    return contextParts.join('\n');
  }

  /**
   * Compresses lower-priority or older memories into a summary.
   * Preserves critical tokens while reducing footprint.
   */
  private async compress(): Promise<void> {
    const lowPriorityEntries = this.workingMemory
      .filter(e => e.priority === 'low' || e.priority === 'normal')
      .sort((a, b) => a.timestamp - b.timestamp);

    if (lowPriorityEntries.length < 3) return;

    // Summarize oldest entries
    const toCompress = lowPriorityEntries.slice(0, 5);
    const summary = await this.config.summarizer.summarize(toCompress);
    
    const summaryEntry: MemoryEntry = {
      id: `summary_${Date.now()}`,
      content: `[SUMMARY] ${summary}`,
      timestamp: Date.now(),
      priority: 'normal',
      tokens: countTokens(summary),
    };

    // Remove compressed entries and add summary
    const idsToRemove = new Set(toCompress.map(e => e.id));
    this.workingMemory = [
      ...this.workingMemory.filter(e => !idsToRemove.has(e.id)),
      summaryEntry
    ];
  }
}

Rationale

• Token-Aware Logic: The manager calculates token counts dynamically, preventing context overflow before the API call.
• Priority Injection: Critical system prompts are immune to eviction, ensuring behavioral constraints are never lost.
• Async Persistence: Vector storage runs asynchronously to avoid blocking the critical path of user interactions.
• Deduplication: Prevents redundant information in the context, which wastes tokens and confuses the model.

Pitfall Guide

1. Ignoring KV Cache Fragmentation

   • Mistake: Assuming GPU memory is homogeneous. Long variable-length sequences cause fragmentation, leading to OOM errors even when total memory usage is below capacity.
   • Fix: Use serving engines with PagedAttention (e.g., vLLM) or implement static padding strategies. Monitor KV cache hit rates and fragmentation ratios.

2. The "Lost in the Middle" Trap

   • Mistake: Injecting retrieved documents in the middle of the context window.
   • Fix: Place retrieved content at the beginning or end of the context. If using long contexts, use re-ranking to ensure the most relevant documents are positioned where the model's attention is strongest.

3. Over-Vectorization of Exact Data

   • Mistake: Storing IDs, codes, or exact phrases in vector stores. Semantic search fails to retrieve exact matches reliably.
   • Fix: Use hybrid search. Store exact-match fields in a keyword index (e.g., Elasticsearch or BM25) and combine with vector results. Use rerankers to fuse results.

4. Prompt Injection via Memory

   • Mistake: Treating retrieved memory as trusted input. Malicious content in long-term memory can be injected into the context, bypassing input filters.
   • Fix: Implement memory sanitization. Sanitize content upon retrieval. Use separate instruction tokens to demarcate memory content from user input.

5. Stale Memory Accumulation

   • Mistake: Never invalidating or updating memories. The system retains outdated user preferences or deprecated code snippets.
   • Fix: Implement TTLs (Time-To-Live) for memory entries. Use update strategies where new information overwrites old facts. Periodic re-embedding of long-term storage to reflect model improvements.

6. Compression Artifacts

   • Mistake: Summarizing too aggressively, losing numerical data or specific constraints.
   • Fix: Use structured summarization that preserves key-value pairs. Maintain a "facts" layer separate from narrative summaries. Validate summaries against original content using a verification step.

7. Cost Blindness in Retrieval

   • Mistake: Retrieving large chunks without considering embedding and inference costs.
   • Fix: Optimize chunk sizes. Use smaller embedding models for retrieval and larger models only for generation. Cache embedding results for repeated queries.

Production Bundle

Action Checklist

• Token Audit: Instrument all LLM calls to log input/output token counts and track context window utilization over time.
• KV Cache Monitoring: Set up alerts for KV cache memory usage; trigger auto-scaling or request throttling when usage exceeds 85%.
• Eviction Policy: Define and implement LRU or LFU eviction policies for agent sessions; test recovery after eviction.
• Hybrid Search: Deploy a hybrid retrieval pipeline combining vector search and keyword/BM25 search for high-precision requirements.
• Memory Sanitization: Add a sanitization layer to all retrieved memory content to prevent prompt injection attacks.
• Compression Testing: Benchmark summarization quality; ensure critical constraints and numerical data survive compression.
• Cache Strategy: Identify high-frequency prompt prefixes and enable prompt caching (e.g., Anthropic cache control, vLLM caching).
• Drift Detection: Implement periodic checks to detect memory drift where retrieved context no longer aligns with current user state.

Decision Matrix

Configuration Template

memory:
  working_window:
    max_tokens: 32000
    retention_policy: "sliding" # sliding | fixed | hierarchical
    compression_trigger: 0.85 # Compress when 85% full
    summary_strategy: "recursive" # recursive | single_pass
  
  long_term:
    vector_store: "pgvector"
    embedding_model: "text-embedding-3-small"
    chunk_size: 512
    chunk_overlap: 50
    retrieval_top_k: 5
    hybrid_search:
      enabled: true
      bm25_weight: 0.3
      vector_weight: 0.7
  
  caching:
    prompt_cache:
      enabled: true
      ttl: 3600 # seconds
      min_prefix_length: 500
    
    kv_cache:
      quantization: "int8" # none | int8 | fp4
      paged_attention: true
  
  security:
    sanitization: "strict"
    injection_protection: true
    memory_ttls:
      user_data: 7776000 # 90 days
      system_data: null # indefinite

Quick Start Guide

1. Install Dependencies:

   npm install @codcompass/llm-memory tiktoken
   

2. Initialize Manager:

   import { LLMMemoryManager } from '@codcompass/llm-memory';
   
   const memory = new LLMMemoryManager({
     maxTokens: 32000,
     windowSize: 10,
     vectorStore: new PgVectorStore({ connectionString: process.env.DB_URL }),
     summarizer: new LLMSummarizer({ model: 'gpt-4o-mini' })
   });
   

3. Integrate into Pipeline:

   // Before LLM call
   const context = await memory.retrieve(userQuery);
   const prompt = `Context:\n${context}\n\nQuery: ${userQuery}`;
   
   // After LLM response
   await memory.add('assistant', llmResponse);
   

4. Enable Caching: Configure your LLM client to use cache tokens. For Anthropic, set cache_control: { type: "ephemeral" } on system messages.

5. Monitor: Deploy the provided Prometheus metrics exporter to track memory_hit_rate, compression_count, and kv_cache_usage. Adjust compression_trigger based on latency requirements.