Files
poimen-memory/docs/CONTEXT_OPTIMIZER.md
T
Story Crater Bot 262478f7f2
Build and Push / Test (push) Failing after 1m55s
Build and Push / Build and push image (push) Skipped
plan: add Magika ML classifier to content router
2026-08-28 09:12:09 -07:00

12 KiB

Context Optimizer — Pre-LLM Compression Layer

Motivation

Agent transcripts and tool outputs are noisy. A 50-chunk ingestion run might feed the GRU-Mem gate evidence that's 43% tool results, full of timestamps, temp paths, ANSI codes, and verbose JSON. The model wastes tokens parsing noise, risks hallucinating on irrelevant details, and we pay full price for bloated input. LLM provider KV caches miss because dynamic content (timestamps, session IDs) pollutes the prefix.

This layer sits between retrieval and the LLM call. Search indexes (pgvector + OpenSearch) stay untouched at full fidelity. Only the evidence chunks entering the prompt get optimized.

Architecture

Query → Hybrid Search (pgvector 60% + OpenSearch 40%)
            │
            │  full-fidelity chunks (untouched)
            ▼
    ┌───────────────────────┐
    │   CONTEXT OPTIMIZER   │  ← THIS MODULE
    │                       │
    │  1. CacheAligner      │  Move dynamic content (timestamps, UUIDs)
    │                       │  to end of context. Keep static prefix
    │                       │  stable for KV cache hits.
    │                       │
    │  2. ContentRouter     │  Auto-detect content type per chunk:
    │                       │  JSON, code, logs, diffs, plain text.
    │                       │  Route each to best compressor.
    │                       │
    │  3. Compressors       │  Per-type compression:
    │     - JsonCrusher     │    Statistical field analysis, keep keys
    │     - LogCompressor   │    Keep errors/stack traces, drop noise
    │     - CodeCompressor  │    AST-aware: keep signatures, drop bodies
    │     - TextCompressor  │    Token importance scoring
    │     - DiffCompressor  │    Keep change hunks, drop context
    │                       │
    │  4. CCR Store         │  Cache full originals with hash reference.
    │                       │  Inject retrieval hint so model can fetch
    │                       │  full content if needed.
    │                       │
    └───────────┬───────────┘
                │
                │  optimized chunks (smaller, cleaner)
                ▼
    ┌───────────────────────┐
    │  Cache-Aligned Prompt │  (already built)
    │  system | query | turn│
    └───────────┬───────────┘
                │
                ▼
           LLM Gateway

What Each Stage Does

Stage 1: CacheAligner

Goal: Maximize LLM provider KV cache hits by stabilizing the prompt prefix.

LLM providers (Anthropic, OpenAI) cache based on exact prefix match. A single changing timestamp or session ID early in the prompt invalidates the entire cache. CacheAligner:

  1. Scans for dynamic patterns in the prompt prefix:

    • ISO timestamps (2026-08-28T...)
    • UUIDs (550e8400-e29b-...)
    • Session tokens, run IDs
    • Temp paths (/tmp/abc123)
  2. Moves detected dynamic content to the end of the context (after static instructions and query), preserving the stable prefix for cache hits.

  3. Reports drift metrics: how much of the prefix changed vs. previous call.

Implementation: Regex-based pattern detection + reordering. No ML needed. Reuses existing normalisation patterns from lesson.rs (M3.7.7).

pub struct CacheAligner;

impl CacheAligner {
    /// Stabilize prompt prefix by moving dynamic content to tail.
    /// Returns (stable_prefix, dynamic_tail).
    pub fn align(content: &str) -> AlignedContent {
        // Detect and extract dynamic patterns
        // Reorder so static content comes first
    }
}

Stage 2: ContentRouter (Magika ML + regex fallback)

Goal: Auto-detect content type and route to the best compressor.

Primary classifier: Google Magika (magika crate v1.1.0) — fast encoder-only ONNX model that classifies content into 100+ types. <1ms per classification. No LLM calls, no network — runs locally with embedded ONNX model.

Fallback: Regex heuristics for content types Magika doesn't distinguish well (e.g., build logs vs. plain text) or when confidence is below threshold.

use magika::Session;

pub struct ContentRouter {
    magika: Session,
    confidence_threshold: f32, // default 0.7
}

pub enum ContentType {
    Json,
    Code { language: String },
    Log,
    Diff,
    Config,
    Text,
}

impl ContentRouter {
    pub fn detect(&self, content: &str) -> ContentType {
        // 1. Try Magika ML classification
        if let Ok(result) = self.magika.identify_content_sync(content.as_bytes()) {
            let label = result.info().label;
            let score = result.score();
            
            if score >= self.confidence_threshold {
                return match label {
                    "json" | "jsonl" => ContentType::Json,
                    "python" | "javascript" | "typescript" | "rust" | "go" | "shell"
                        => ContentType::Code { language: label.to_string() },
                    "diff" => ContentType::Diff,
                    "yaml" | "toml" | "ini" | "xml" => ContentType::Config,
                    _ => self.regex_fallback(content),
                };
            }
        }
        
        // 2. Fallback to regex heuristics
        self.regex_fallback(content)
    }
    
    fn regex_fallback(&self, content: &str) -> ContentType {
        if is_json(content) { return ContentType::Json; }
        if is_log(content)  { return ContentType::Log; }
        if is_diff(content) { return ContentType::Diff; }
        if is_code(content) { return ContentType::Code { language: "unknown".into() }; }
        ContentType::Text
    }
}

Magika label → compressor mapping:

Magika Label ContentType Compressor
json, jsonl Json JsonCrusher
python, javascript, rust, go, typescript, shell Code CodeCompressor
diff Diff DiffCompressor
yaml, toml, ini, xml Config (passthrough, already compact)
txt + log heuristics Log LogCompressor
everything else Text TextCompressor

Stage 3: Compressors

Goal: Reduce token count per content type while preserving signal.

JsonCrusher (70-90% savings)

  • Analyse field-level variance across JSON array elements
  • Keep: keys, structure, error markers, boundary items (first/last)
  • Drop: redundant mid-array elements, long string values, whitespace
  • Allocation: 30% start (schema), 15% end (recency), 55% importance

LogCompressor (85-95% savings)

  • Keep: error lines, stack traces, exit codes, FAIL markers
  • Drop: passing test output, INFO-level noise, repeated patterns
  • Reuses M3.7.7 signature extraction patterns (markers, cascade detection)

CodeCompressor (40-70% savings, opt-in)

  • Keep: imports, function/method signatures, type annotations
  • Drop: function bodies, inline comments, blank lines
  • Uses simple AST heuristics (brace counting), not full parser

DiffCompressor (60-80% savings)

  • Keep: change hunks (+/- lines), hunk headers
  • Drop: unchanged context lines (the @@ surrounding context)

TextCompressor (30-50% savings)

  • Token importance scoring: keep high-entropy tokens (IDs, hashes, error codes)
  • Drop: low-information prose, repeated phrases, filler words
pub trait Compressor {
    fn compress(&self, content: &str, budget: usize) -> CompressResult;
}

pub struct CompressResult {
    pub compressed: String,
    pub original_tokens: usize,
    pub compressed_tokens: usize,
    pub content_type: ContentType,
    /// Hash of original content for CCR retrieval
    pub ccr_hash: Option<String>,
}

Stage 4: CCR Store (Compress-Cache-Retrieve)

Goal: Lossless compression — model can retrieve full originals if needed.

When a chunk is compressed, the full original is cached with a SHA256 hash. A retrieval hint is injected into the compressed output:

[compressed content...]
<!-- CCR:abc123 — full content available via retrieval -->

If the model determines it needs more detail, it can request the full content. This makes compression aggressive but reversible.

pub struct CcrStore {
    cache: HashMap<String, String>,  // hash → original content
    max_entries: usize,
    ttl: Duration,
}

impl CcrStore {
    pub fn store(&mut self, content: &str) -> String; // returns hash
    pub fn retrieve(&self, hash: &str) -> Option<&str>;
}

Integration with Existing Code

Where It Plugs In

The context optimizer sits in mem-core as a new module, called by PromptBuilder::build_cache_aligned() before assembling the final prompt:

// In PromptBuilder::build_cache_aligned()
let chunk_text = Self::render_chunk(chunk)?;

// NEW: Optimize before prompt assembly
let optimized = ContextOptimizer::optimize(&chunk_text, &OptimizeConfig {
    cache_align: true,
    compress: true,
    ccr_enabled: true,
    token_budget: BUDGET_CHUNK_MAX,
})?;

let turn_msg = CACHE_TURN
    .replace("{memory}", memory_text)
    .replace("{chunk}", &optimized.compressed);

What Already Exists (Reuse)

Existing Code Reuse For
lesson.rs normalise() CacheAligner pattern detection (timestamps, paths, SHAs)
lesson.rs markers() LogCompressor error line detection
lesson.rs is_cascade() LogCompressor cascade suppression
lesson.rs strip_ansi() Pre-processing for all compressors
symptom_projection.rs stop words TextCompressor low-value token detection

What's New

New Code Location Est. LOC
context_optimizer.rs crates/mem-core/src/ 150
content_router.rs crates/mem-core/src/ 150
compressors/json.rs crates/mem-core/src/ 200
compressors/log.rs crates/mem-core/src/ 150
compressors/code.rs crates/mem-core/src/ 150
compressors/diff.rs crates/mem-core/src/ 100
compressors/text.rs crates/mem-core/src/ 100
ccr_store.rs crates/mem-core/src/ 80
Total ~1030

Performance Targets

Metric Target
Detection + compression < 10ms per chunk
JSON compression ratio 70-90%
Log compression ratio 85-95%
Code compression ratio 40-70%
Cache hit rate improvement > 50% on repeated queries
Zero false negatives Model should never miss real errors

Phasing

Phase 1: Foundation (1-2 days)

  • ContextOptimizer orchestrator
  • ContentRouter with detection heuristics
  • LogCompressor (reuses M3.7.7 lesson.rs patterns)
  • 10 unit tests

Phase 2: JSON + Diff (1-2 days)

  • JsonCrusher with field variance analysis
  • DiffCompressor with hunk preservation
  • 10 unit tests

Phase 3: CCR + CacheAligner (1 day)

  • CcrStore with LRU cache
  • CacheAligner with dynamic pattern extraction
  • Integration with PromptBuilder::build_cache_aligned()
  • 10 unit tests

Phase 4: Code + Text (1 day)

  • CodeCompressor (opt-in, brace-counting heuristics)
  • TextCompressor (token importance scoring)
  • 10 unit tests

Key Design Decisions

  1. Magika ML for detection, rules for compression. Content type detection uses Google's Magika ONNX model (<1ms, local, no network). Compression itself uses deterministic algorithms (no LLM calls in hot path).

  2. Search indexes untouched. Compression happens AFTER retrieval. pgvector and OpenSearch see full-fidelity text. Only the LLM prompt is optimized.

  3. Reversible via CCR. Every compression is reversible. The model can request full originals via hash lookup. Aggressive compression is safe because nothing is permanently lost.

  4. Reuse existing code. M3.7.7's normalisation patterns, marker detection, and cascade suppression are directly reusable for log compression. M3.7.8's stop words help text compression.

  5. Budget-aware. Each compressor respects a token budget. If the chunk is already under budget, no compression is applied (zero overhead).


Inspired by Headroom. Adapted for Rust, integrated with poimen-memory's hybrid search architecture.