Build a Context Compression Engine

import { Redis } from "ioredis";
import { createHash } from "node:crypto";
const redis = new Redis(process.env.REDIS_URL!);

interface ContextChunk { id: string; content: string; tokens: number; source: string; relevance: number; type: "system" | "conversation" | "document" | "tool_result"; timestamp: number; }
interface CompressedContext { chunks: ContextChunk[]; totalTokens: number; originalTokens: number; compressionRatio: number; droppedChunks: number; }

const PRIORITY_WEIGHTS: Record<string, number> = { system: 1.0, tool_result: 0.8, conversation: 0.6, document: 0.4 };

// Compress context to fit within token budget
export function compressContext(chunks: ContextChunk[], maxTokens: number): CompressedContext {
  const originalTokens = chunks.reduce((s, c) => s + c.tokens, 0);
  if (originalTokens <= maxTokens) return { chunks, totalTokens: originalTokens, originalTokens, compressionRatio: 1, droppedChunks: 0 };

  // Step 1: Deduplicate semantically similar chunks
  const deduped = deduplicateChunks(chunks);

  // Step 2: Score each chunk
  const scored = deduped.map((chunk) => ({
    chunk,
    score: calculateChunkScore(chunk, chunks),
  }));

  // Step 3: Sort by score (highest first)
  scored.sort((a, b) => b.score - a.score);

  // Step 4: Greedily select chunks within budget
  // Always include system prompts
  const selected: ContextChunk[] = [];
  let currentTokens = 0;

  // System chunks always included
  for (const { chunk } of scored.filter((s) => s.chunk.type === "system")) {
    selected.push(chunk);
    currentTokens += chunk.tokens;
  }

  // Recent conversation (last 5 turns) always included
  const recentConversation = scored
    .filter((s) => s.chunk.type === "conversation")
    .sort((a, b) => b.chunk.timestamp - a.chunk.timestamp)
    .slice(0, 10);
  for (const { chunk } of recentConversation) {
    if (currentTokens + chunk.tokens <= maxTokens) {
      selected.push(chunk);
      currentTokens += chunk.tokens;
    }
  }

  // Fill remaining budget with highest-scored chunks
  for (const { chunk } of scored) {
    if (selected.includes(chunk)) continue;
    if (currentTokens + chunk.tokens > maxTokens) {
      // Try to compress the chunk itself
      const compressed = compressChunk(chunk, maxTokens - currentTokens);
      if (compressed && compressed.tokens > 50) {
        selected.push(compressed);
        currentTokens += compressed.tokens;
      }
      continue;
    }
    selected.push(chunk);
    currentTokens += chunk.tokens;
  }

  // Sort selected by original order (maintain coherence)
  selected.sort((a, b) => {
    const typeOrder = { system: 0, conversation: 1, document: 2, tool_result: 3 };
    const typeA = typeOrder[a.type] ?? 4;
    const typeB = typeOrder[b.type] ?? 4;
    if (typeA !== typeB) return typeA - typeB;
    return a.timestamp - b.timestamp;
  });

  return {
    chunks: selected,
    totalTokens: currentTokens,
    originalTokens,
    compressionRatio: Math.round((currentTokens / originalTokens) * 100) / 100,
    droppedChunks: chunks.length - selected.length,
  };
}

// Deduplicate semantically similar chunks
function deduplicateChunks(chunks: ContextChunk[]): ContextChunk[] {
  const seen = new Map<string, ContextChunk>();

  for (const chunk of chunks) {
    // Create fingerprint from content (simplified — production uses embeddings)
    const words = chunk.content.toLowerCase().split(/\s+/).sort().slice(0, 20).join(" ");
    const fingerprint = createHash("md5").update(words).digest("hex").slice(0, 8);

    const existing = seen.get(fingerprint);
    if (existing) {
      // Keep the one with higher relevance
      if (chunk.relevance > existing.relevance) seen.set(fingerprint, chunk);
    } else {
      seen.set(fingerprint, chunk);
    }
  }

  return [...seen.values()];
}

function calculateChunkScore(chunk: ContextChunk, allChunks: ContextChunk[]): number {
  const typeWeight = PRIORITY_WEIGHTS[chunk.type] || 0.3;
  const relevanceScore = chunk.relevance;

  // Recency bonus (newer = higher)
  const maxTime = Math.max(...allChunks.map((c) => c.timestamp));
  const minTime = Math.min(...allChunks.map((c) => c.timestamp));
  const recencyScore = maxTime > minTime ? (chunk.timestamp - minTime) / (maxTime - minTime) : 0.5;

  // Token efficiency (more info per token = better)
  const efficiency = Math.min(1, 500 / chunk.tokens); // prefer chunks under 500 tokens

  return typeWeight * 0.3 + relevanceScore * 0.4 + recencyScore * 0.2 + efficiency * 0.1;
}

// Compress individual chunk (extractive summarization)
function compressChunk(chunk: ContextChunk, maxTokens: number): ContextChunk | null {
  if (maxTokens < 50) return null;

  const sentences = chunk.content.split(/[.!?]+/).filter((s) => s.trim().length > 10);
  if (sentences.length <= 1) return null;

  // Keep sentences with highest information density
  const scored = sentences.map((s) => ({
    text: s.trim(),
    score: s.match(/\d/g)?.length || 0 + // numbers = data
      (s.match(/\b(is|are|means|shows|indicates|requires)\b/gi)?.length || 0) * 2 + // key verbs
      (s.length > 50 ? 1 : 0), // prefer substantial sentences
    tokens: Math.ceil(s.length / 4),
  }));

  scored.sort((a, b) => b.score - a.score);

  let compressed = "";
  let tokens = 0;
  for (const s of scored) {
    if (tokens + s.tokens > maxTokens) break;
    compressed += s.text + ". ";
    tokens += s.tokens;
  }

  return { ...chunk, content: compressed.trim(), tokens };
}

// Estimate tokens
export function estimateTokens(text: string): number {
  return Math.ceil(text.length / 4);
}