feat(context,infra): compression pipeline stages 4-5 — token pruning + relevance + ToMe

- build() is now suspend: pipeline runs all stages in fixed order, each gated by level - TOKEN_PRUNE (level 3): TokenPruner interface + LLMLingua-2 sidecar (sidecars/llmlingua) + HttpTokenPruner adapter (fails open if sidecar down); prunes freeform, preserves protected spans, skips tier-0 turns when TIER_SPLIT on - TOME_MERGE (level 8): ToMeMerger collapses near-duplicate freeform turns (Jaccard) - Stage 5 selection: RelevanceScorer + EmbeddingRelevanceScorer (cosine over Embedder); query-conditioned reorder so least-relevant freeform drops first under budget - [orchestration] compression_level + token_pruner_url config, wired in Main - suspend ripple fixed across builder callers/stubs
2026-07-01 14:29:56 +04:00
parent e0c222392c
commit 047e2a4070
25 changed files with 509 additions and 65 deletions
@@ -624,6 +624,9 @@ object ConfigLoader {
                asInt(orchestrationSection["journal_compaction_token_threshold"], 2_000),
            resumeAbandonedMaxAgeMinutes =
                asLong(orchestrationSection["resume_abandoned_max_age_minutes"], 1_440),
+            compressionLevel = asInt(orchestrationSection["compression_level"], 2),
+            tokenPrunerUrl =
+                asString(orchestrationSection["token_pruner_url"], "http://127.0.0.1:8199"),
        )

        val modelsSettings = ModelsSettings(
@@ -61,6 +61,13 @@ data class OrchestrationKnobs(
     * 0 disables auto-resume entirely.
     */
    val resumeAbandonedMaxAgeMinutes: Long = 1_440,
+    /**
+     * Context compression pipeline level (docs/plans/correx-compression-pipeline.md §5), additive
+     * 1..9. Default 2 = free format-compress + static cache. Raise toward the 16k wall; ≥3 needs
+     * the LLMLingua-2 sidecar at [tokenPrunerUrl] (fails open if absent).
+     */
+    val compressionLevel: Int = 2,
+    val tokenPrunerUrl: String = "http://127.0.0.1:8199",
 )

@Serializable
@@ -8,7 +8,7 @@ import com.correx.core.events.types.SessionId
 import com.correx.core.events.types.StageId

 interface ContextPackBuilder {
-    fun build(
+    suspend fun build(
        id: ContextPackId,
        sessionId: SessionId,
        stageId: StageId,
@@ -8,7 +8,12 @@ import com.correx.core.context.compression.ContextClassifier
 import com.correx.core.context.compression.ContextCompressor
 import com.correx.core.context.compression.FactSheet
 import com.correx.core.context.compression.FormatCompressor
+import com.correx.core.context.compression.NoOpTokenPruner
 import com.correx.core.context.compression.ProtectedSpanTagger
+import com.correx.core.context.compression.RelevanceScorer
+import com.correx.core.context.compression.ToMeMerger
+import com.correx.core.context.compression.TokenPruner
+import com.correx.core.context.compression.UniformRelevanceScorer
 import com.correx.core.context.model.CompressionMetadata
 import com.correx.core.context.model.ContextEntry
 import com.correx.core.context.model.ContextLayer
@@ -20,12 +25,24 @@ import com.correx.core.events.types.ContextPackId
 import com.correx.core.events.types.SessionId
 import com.correx.core.events.types.StageId

+/**
+ * Runs the compression pipeline (docs/plans/correx-compression-pipeline.md) in fixed pipeline
+ * order; each stage is skipped unless its [CompressionStage] is enabled by the [policy] level.
+ * Lossy stages (token pruning, relevance reordering) are suspend because they may hit an
+ * out-of-process sidecar; the deterministic stages run in-line. Recursive summarization
+ * (level 5) is handled by the orchestrator's TierContextSummarizer, not here, because it emits
+ * events — the builder stays free of side effects.
+ */
 class DefaultContextPackBuilder(
    private val compressor: ContextCompressor,
    private val policy: CompressionPolicy = CompressionPolicy(),
    private val classifier: ContextClassifier = ContextClassifier(),
    private val formatCompressor: FormatCompressor = FormatCompressor(),
+    private val tagger: ProtectedSpanTagger = ProtectedSpanTagger(),
    private val factSheet: FactSheet = FactSheet(ProtectedSpanTagger()),
+    private val tokenPruner: TokenPruner = NoOpTokenPruner,
+    private val toMeMerger: ToMeMerger = ToMeMerger(),
+    private val relevanceScorer: RelevanceScorer = UniformRelevanceScorer,
 ) : ContextPackBuilder {

    private val neverDropSourceTypes = setOf("steeringNote", "eventHistory", "decisionJournal", "factSheet")
@@ -33,72 +50,71 @@ class DefaultContextPackBuilder(
    private companion object {
        const val CHARS_PER_TOKEN = 4
        const val FACT_SHEET_ORDINAL = -1
+        const val DEFAULT_PRUNE_RATIO = 0.45
+        const val TIER0_TURNS = 3
    }

-    override fun build(
+    override suspend fun build(
        id: ContextPackId,
        sessionId: SessionId,
        stageId: StageId,
        entries: List<ContextEntry>,
        budget: TokenBudget
    ): ContextPack {
-        // Stamp chronological order from the caller's input sequence. Grouping by
-        // sourceType (for compression) and by layer reorders entries; the ordinal lets
-        // us restore true turn order afterwards so a tool loop reads assistant→tool→…
-        // instead of all-assistants-then-all-tools.
+        // Stamp chronological order from the caller's input sequence so grouping/compression can
+        // reorder freely and PromptRenderer can restore true turn order from the ordinal.
        val stamped = entries.mapIndexed { index, entry -> entry.copy(ordinal = index) }
-        // Stage 1 (level 1+): free, lossless format compression of structured entries only —
-        // json→dotted lines / dedupe. Static and freeform content is left byte-identical; the
-        // lossy stages that touch freeform live further down the pipeline (prune/summarize).
-        val ordered = if (policy.enabled(CompressionStage.FORMAT_COMPRESS)) {
+
+        // Stage 1 (FORMAT_COMPRESS): lossless json→dotted-line compaction of structured entries.
+        val formatted = if (policy.enabled(CompressionStage.FORMAT_COMPRESS)) {
            stamped.map { entry ->
-                if (classifier.classify(entry) == ContextClass.STRUCTURED) {
-                    val compacted = formatCompressor.compress(entry.content)
-                    if (compacted == entry.content) entry
-                    else entry.copy(content = compacted, tokenEstimate = estimateTokens(compacted))
-                } else entry
+                if (classifier.classify(entry) == ContextClass.STRUCTURED) reencode(entry, formatCompressor.compress(entry.content))
+                else entry
            }
        } else stamped
-        // Stage 7 (level 7+): derive the fact sheet from load-bearing spans and pin it as an
-        // L0 SYSTEM entry (folds into the leading system message, never dropped). Accuracy
-        // insurance — the most compressed tier still can't lose an ID/number/path.
+
+        // Stage 3 (TOKEN_PRUNE): prune freeform prose, preserving protected spans. When TIER_SPLIT
+        // is on, the newest TIER0_TURNS freeform turns are left full-fidelity (tier 0).
+        val pruned = if (policy.enabled(CompressionStage.TOKEN_PRUNE)) prune(formatted) else formatted
+
+        // Stage 8 (TOME_MERGE): collapse near-duplicate freeform turns into the newest survivor.
+        val merged = if (policy.enabled(CompressionStage.TOME_MERGE)) applyToMe(pruned) else pruned
+
+        // Stage 7 (FACT_SHEET): pin load-bearing spans as an L0 SYSTEM entry that never drops.
        val withFactSheet = if (policy.enabled(CompressionStage.FACT_SHEET)) {
-            val facts = factSheet.extract(ordered)
-            if (facts.isEmpty()) ordered else ordered + factSheetEntry(factSheet.render(facts))
-        } else ordered
+            val facts = factSheet.extract(merged)
+            if (facts.isEmpty()) merged else merged + factSheetEntry(factSheet.render(facts))
+        } else merged
+
        val (pinned, compressible) = withFactSheet.partition { it.sourceType in neverDropSourceTypes }
        val pinnedTokens = pinned.sumOf { it.tokenEstimate }
        var remainingTokens = (budget.limit - pinnedTokens).coerceAtLeast(0)

-        // Dispatch compression strategy by entry sourceType — applying Conversation
-        // strategy uniformly mangles tool logs and artifacts (their shape isn't conversational).
-        // Stage 9 (level 9+): strict allocation reserves budget for structured content by
-        // compressing structured-class groups before freeform, so a long conversation can't
-        // starve tool logs / artifacts out of the window. Below level 9 keep encounter order.
+        // Query-conditioned selection: when a real relevance scorer is wired, reorder each freeform
+        // group so the least-relevant turns sit at the front (dropped first) and the most-relevant
+        // land at the tail the Conversation strategy keeps. Final layering re-sorts by ordinal, so
+        // chronological render order is unaffected — only *what survives* the budget changes.
+        val relevanceRanked = rankByRelevance(compressible, currentQuery(withFactSheet))
+
+        // Stage 9 (BUDGET_ALLOCATOR_STRICT): compress structured groups before freeform so a long
+        // conversation can't starve tool logs / artifacts out of the window.
        val strict = policy.enabled(CompressionStage.BUDGET_ALLOCATOR_STRICT)
        val strategiesUsed = linkedSetOf<String>()
-        val grouped = compressible.groupBy { it.sourceType }.entries
+        val grouped = relevanceRanked.groupBy { it.sourceType }.entries
        val orderedGroups = if (strict) {
-            grouped.sortedBy { (_, group) ->
-                if (classifier.classify(group.first()) == ContextClass.FREEFORM) 1 else 0
-            }
+            grouped.sortedBy { (_, group) -> if (classifier.classify(group.first()) == ContextClass.FREEFORM) 1 else 0 }
        } else grouped.toList()
-        val compressed = orderedGroups
-            .flatMap { (sourceType, group) ->
-                val strategy = strategyFor(sourceType)
-                strategiesUsed += strategy::class.simpleName ?: "Unknown"
-                val groupBudget = TokenBudget(limit = remainingTokens.coerceAtLeast(0))
-                val result = compressor.compress(group, groupBudget, strategy)
-                remainingTokens -= result.sumOf { it.tokenEstimate }
-                result
-            }
+        val compressed = orderedGroups.flatMap { (sourceType, group) ->
+            val strategy = strategyFor(sourceType)
+            strategiesUsed += strategy::class.simpleName ?: "Unknown"
+            val result = compressor.compress(group, TokenBudget(limit = remainingTokens.coerceAtLeast(0)), strategy)
+            remainingTokens -= result.sumOf { it.tokenEstimate }
+            result
+        }

-        // Restore chronological order before layering — groupBy preserves encounter
-        // order, so each layer's list comes out in true turn order.
        val retained = (pinned + compressed).sortedBy { it.ordinal }
        val layers = retained.groupBy { it.layer }
-        val budgetUsed = retained.sumOf { it.tokenEstimate }
-        val droppedCount = entries.size - retained.size
+        val droppedCount = entries.size - retained.count { it.sourceType != "factSheet" }
        val truncatedLayers = if (droppedCount > 0) {
            entries.map { it.layer }.distinct().filter { layer ->
                entries.count { it.layer == layer } > retained.count { it.layer == layer }
@@ -110,7 +126,7 @@ class DefaultContextPackBuilder(
            sessionId = sessionId,
            stageId = stageId,
            layers = layers,
-            budgetUsed = budgetUsed,
+            budgetUsed = retained.sumOf { it.tokenEstimate },
            budgetLimit = budget.limit,
            compressionMetadata = CompressionMetadata(
                appliedStrategies = strategiesUsed.toList().ifEmpty { listOf("Conversation") },
@@ -120,12 +136,41 @@ class DefaultContextPackBuilder(
        )
    }

-    // ~4 chars/token heuristic; matches the estimate used across the orchestrator so budget
-    // math stays consistent after a format-compressed entry shrinks.
+    private suspend fun prune(entries: List<ContextEntry>): List<ContextEntry> {
+        val tierSplit = policy.enabled(CompressionStage.TIER_SPLIT)
+        val freeformOrdinals = entries.filter { classifier.classify(it) == ContextClass.FREEFORM }
+            .map { it.ordinal }.sortedDescending()
+        val tier0 = if (tierSplit) freeformOrdinals.take(TIER0_TURNS).toSet() else emptySet()
+        return entries.map { entry ->
+            if (classifier.classify(entry) != ContextClass.FREEFORM || entry.ordinal in tier0) entry
+            else reencode(entry, tokenPruner.prune(entry.content, tagger.protectedSpans(entry.content), DEFAULT_PRUNE_RATIO))
+        }
+    }
+
+    private fun applyToMe(entries: List<ContextEntry>): List<ContextEntry> {
+        val (freeform, rest) = entries.partition { classifier.classify(it) == ContextClass.FREEFORM }
+        return (rest + toMeMerger.merge(freeform)).sortedBy { it.ordinal }
+    }
+
+    // The current query is the most-recent user turn — the anchor relevance is scored against.
+    private fun currentQuery(entries: List<ContextEntry>): String? =
+        entries.filter { it.role == EntryRole.USER }.maxByOrNull { it.ordinal }?.content
+
+    private suspend fun rankByRelevance(entries: List<ContextEntry>, query: String?): List<ContextEntry> {
+        if (relevanceScorer is UniformRelevanceScorer || query == null) return entries
+        val (freeform, rest) = entries.partition { classifier.classify(it) == ContextClass.FREEFORM }
+        if (freeform.isEmpty()) return entries
+        val scores = relevanceScorer.score(query, freeform.map { it.content })
+        // Least-relevant first so the Conversation strategy's takeLast keeps the most-relevant.
+        val ranked = freeform.zip(scores).sortedBy { it.second }.map { it.first }
+        return rest + ranked
+    }
+
+    private fun reencode(entry: ContextEntry, content: String): ContextEntry =
+        if (content == entry.content) entry else entry.copy(content = content, tokenEstimate = estimateTokens(content))
+
    private fun estimateTokens(text: String): Int = (text.length + (CHARS_PER_TOKEN - 1)) / CHARS_PER_TOKEN

-    // Negative ordinal so it sorts ahead of every turn; L0 SYSTEM so PromptRenderer folds it
-    // into the leading system message rather than emitting an illegal mid-stream system turn.
    private fun factSheetEntry(content: String) = ContextEntry(
        id = ContextEntryId("factsheet"),
        layer = ContextLayer.L0,
@@ -0,0 +1,21 @@
+package com.correx.core.context.compression
+
+/**
+ * Stage 5 ordering/selection: score a candidate text's relevance to the current query (cheap
+ * embedding cosine similarity). Recency picks the compression *level*; relevance picks what to
+ * *prioritize* when the budget forces a drop — turn 3 may matter more than turn 40 (pipeline §6,
+ * "relevance-reranked history" + "query-conditioned compression").
+ *
+ * Returns a score in [0,1]; higher = keep. Suspend because embedding may hit an out-of-process
+ * embedder. Environment observations (the embeddings) are the caller's to record as events if
+ * they cross into replay-relevant state (invariant #9) — for ephemeral prompt assembly they don't.
+ */
+interface RelevanceScorer {
+    suspend fun score(query: String, candidates: List<String>): List<Double>
+}
+
+/** Default (scorer unconfigured): everything equally relevant, so ordering falls back to recency. */
+object UniformRelevanceScorer : RelevanceScorer {
+    override suspend fun score(query: String, candidates: List<String>): List<Double> =
+        candidates.map { 1.0 }
+}
@@ -0,0 +1,39 @@
+package com.correx.core.context.compression
+
+import com.correx.core.context.model.ContextEntry
+
+/**
+ * Stage 8 (TOME_MERGE): merge near-duplicate freeform turns into one representative instead of
+ * dropping them outright, preserving partial signal (pipeline §1, "Token merging"). Deterministic:
+ * similarity is Jaccard overlap of word shingles, so no model call. Adjacent-in-relevance entries
+ * whose overlap exceeds [threshold] collapse to the newest (highest ordinal), which carries the
+ * merged content forward; the older duplicates are dropped.
+ */
+class ToMeMerger(private val threshold: Double = DEFAULT_THRESHOLD) {
+
+    fun merge(entries: List<ContextEntry>): List<ContextEntry> {
+        if (entries.size < 2) return entries
+        val kept = mutableListOf<ContextEntry>()
+        // Process newest-first so the survivor of a near-duplicate pair is the most recent turn.
+        for (entry in entries.sortedByDescending { it.ordinal }) {
+            val dup = kept.any { jaccard(shingles(it.content), shingles(entry.content)) >= threshold }
+            if (!dup) kept += entry
+        }
+        return kept.sortedBy { it.ordinal }
+    }
+
+    private fun shingles(text: String): Set<String> =
+        text.lowercase().split(WHITESPACE).filter { it.isNotBlank() }.toSet()
+
+    private fun jaccard(a: Set<String>, b: Set<String>): Double {
+        if (a.isEmpty() && b.isEmpty()) return 1.0
+        val inter = a.intersect(b).size.toDouble()
+        val union = a.union(b).size.toDouble()
+        return if (union == 0.0) 0.0 else inter / union
+    }
+
+    private companion object {
+        const val DEFAULT_THRESHOLD = 0.9
+        val WHITESPACE = Regex("\\s+")
+    }
+}
@@ -0,0 +1,20 @@
+package com.correx.core.context.compression
+
+/**
+ * Stage 3 (TOKEN_PRUNE): LLMLingua-2-style token pruning of freeform prose. A small classifier
+ * scores tokens by predictability and drops the low-perplexity connective tissue, keeping the
+ * load-bearing tokens — ~40-50% compression before quality drops (pipeline §1).
+ *
+ * [protectedSpans] (from [ProtectedSpanTagger]) must survive verbatim; the implementation may
+ * not alter any substring in that set. This is what keeps pruning from becoming a hallucination
+ * risk (pipeline §2). Pure prose in, pruned prose out; suspend because the reference
+ * implementation is an out-of-process sidecar.
+ */
+interface TokenPruner {
+    suspend fun prune(content: String, protectedSpans: List<String>, targetRatio: Double): String
+}
+
+/** Default when no pruner is configured (levels < 3, or sidecar unavailable): identity. */
+object NoOpTokenPruner : TokenPruner {
+    override suspend fun prune(content: String, protectedSpans: List<String>, targetRatio: Double): String = content
+}
@@ -81,11 +81,11 @@ class CompressionPipelineStagesTest {
    }

    @Test
-    fun `builder pins a fact sheet at level 7 and omits it below`() {
+    fun `builder pins a fact sheet at level 7 and omits it below`() = kotlinx.coroutines.runBlocking {
        val entries = listOf(entry("chat", ContextLayer.L1, EntryRole.USER, "deploy to 10.0.0.1"))
        val compressor = com.correx.core.context.compression.DefaultContextCompressor()

-        fun buildAt(level: Int) = com.correx.core.context.builder.DefaultContextPackBuilder(
+        suspend fun buildAt(level: Int) = com.correx.core.context.builder.DefaultContextPackBuilder(
            compressor, CompressionPolicy(level)
        ).build(
            com.correx.core.events.types.ContextPackId("p"),
@@ -98,6 +98,43 @@ class CompressionPipelineStagesTest {
        val atSeven = buildAt(7).layers.values.flatten()
        assertTrue(atSeven.any { it.sourceType == "factSheet" && it.content.contains("10.0.0.1") })
        assertFalse(buildAt(6).layers.values.flatten().any { it.sourceType == "factSheet" })
+        Unit
+    }
+
+    @Test
+    fun `tome merges near-duplicate freeform turns keeping the newest`() {
+        val merger = com.correx.core.context.compression.ToMeMerger(threshold = 0.8)
+        val a = entry("chat", ContextLayer.L2, EntryRole.USER, "the deploy failed on host web").copy(ordinal = 0)
+        val b = entry("chat", ContextLayer.L2, EntryRole.USER, "the deploy failed on host web again").copy(ordinal = 1)
+        val c = entry("chat", ContextLayer.L2, EntryRole.USER, "totally unrelated content here").copy(ordinal = 2)
+        val out = merger.merge(listOf(a, b, c))
+        // a and b are near-duplicates -> only the newest (b) survives; c is distinct
+        assertEquals(listOf(1, 2), out.map { it.ordinal })
+    }
+
+    @Test
+    fun `token pruning at level 3 rewrites freeform preserving protected spans`() = kotlinx.coroutines.runBlocking {
+        val pruner = object : com.correx.core.context.compression.TokenPruner {
+            override suspend fun prune(content: String, protectedSpans: List<String>, targetRatio: Double): String =
+                "PRUNED " + protectedSpans.joinToString(" ")
+        }
+        val builder = com.correx.core.context.builder.DefaultContextPackBuilder(
+            com.correx.core.context.compression.DefaultContextCompressor(),
+            CompressionPolicy(3),
+            tokenPruner = pruner,
+        )
+        val entries = listOf(entry("chat", ContextLayer.L1, EntryRole.USER, "please redeploy 10.0.0.1 now"))
+        val pack = builder.build(
+            com.correx.core.events.types.ContextPackId("p"),
+            com.correx.core.events.types.SessionId("s"),
+            com.correx.core.events.types.StageId("st"),
+            entries,
+            com.correx.core.context.model.TokenBudget(limit = 4000),
+        )
+        val content = pack.layers.values.flatten().first { it.sourceType == "chat" }.content
+        assertTrue(content.startsWith("PRUNED"), content)
+        assertTrue(content.contains("10.0.0.1"), content)
+        Unit
    }

    @Test
@@ -20,7 +20,7 @@ class DecisionJournalPinningTest {
    private val packId = ContextPackId("pack-1")

    @Test
-    fun `decisionJournal entry survives an under-budget build`() {
+    fun `decisionJournal entry survives an under-budget build`() = kotlinx.coroutines.runBlocking {
        val entries = listOf(
            ContextEntry(
                id = ContextEntryId("journal-1"),
@@ -37,5 +37,6 @@ class DecisionJournalPinningTest {
            allRetained.any { it.sourceType == "decisionJournal" },
            "decisionJournal entry must be retained even when budget is exceeded"
        )
+        Unit
    }
 }
@@ -0,0 +1,34 @@
+package com.correx.core.inference
+
+import com.correx.core.context.compression.RelevanceScorer
+import kotlin.math.sqrt
+
+/**
+ * [RelevanceScorer] backed by an [Embedder]: cosine similarity of each candidate to the current
+ * query, rescaled from [-1,1] to [0,1]. Feeds the pipeline's query-conditioned selection — under
+ * budget pressure the least query-relevant turns are dropped first, regardless of age (pipeline §6).
+ * Embeds the query once, then each candidate; a failed embedding scores 0.5 (neutral) so one bad
+ * call doesn't distort the ranking.
+ */
+class EmbeddingRelevanceScorer(private val embedder: Embedder) : RelevanceScorer {
+
+    override suspend fun score(query: String, candidates: List<String>): List<Double> {
+        val q = runCatching { embedder.embed(query) }.getOrNull() ?: return candidates.map { NEUTRAL }
+        return candidates.map { c ->
+            runCatching { cosine(q, embedder.embed(c)) }.getOrDefault(NEUTRAL)
+        }
+    }
+
+    private fun cosine(a: FloatArray, b: FloatArray): Double {
+        if (a.isEmpty() || b.isEmpty() || a.size != b.size) return NEUTRAL
+        var dot = 0.0; var na = 0.0; var nb = 0.0
+        for (i in a.indices) { dot += a[i] * b[i]; na += a[i] * a[i]; nb += b[i] * b[i] }
+        if (na == 0.0 || nb == 0.0) return NEUTRAL
+        val cos = dot / (sqrt(na) * sqrt(nb))
+        return ((cos + 1.0) / 2.0).coerceIn(0.0, 1.0)
+    }
+
+    private companion object {
+        const val NEUTRAL = 0.5
+    }
+}