From 047e2a4070df98d1689d7373cefcf8bc7902f10c Mon Sep 17 00:00:00 2001 From: kami Date: Wed, 1 Jul 2026 14:29:56 +0400 Subject: [PATCH] =?UTF-8?q?feat(context,infra):=20compression=20pipeline?= =?UTF-8?q?=20stages=204-5=20=E2=80=94=20token=20pruning=20+=20relevance?= =?UTF-8?q?=20+=20ToMe?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- apps/server/build.gradle | 1 + .../kotlin/com/correx/apps/server/Main.kt | 8 +- .../com/correx/core/config/ConfigLoader.kt | 3 + .../com/correx/core/config/CorrexConfig.kt | 7 + .../context/builder/ContextPackBuilder.kt | 2 +- .../builder/DefaultContextPackBuilder.kt | 139 ++++++++++++------ .../context/compression/RelevanceScorer.kt | 21 +++ .../core/context/compression/ToMeMerger.kt | 39 +++++ .../core/context/compression/TokenPruner.kt | 20 +++ .../kotlin/CompressionPipelineStagesTest.kt | 41 +++++- .../test/kotlin/DecisionJournalPinningTest.kt | 3 +- .../inference/EmbeddingRelevanceScorer.kt | 34 +++++ infrastructure/compression/build.gradle | 24 +++ .../compression/HttpTokenPruner.kt | 58 ++++++++ .../src/test/kotlin/HttpTokenPrunerTest.kt | 41 ++++++ settings.gradle | 1 + sidecars/llmlingua/README.md | 37 +++++ sidecars/llmlingua/requirements.txt | 3 + sidecars/llmlingua/server.py | 65 ++++++++ .../kotlin/DefaultContextPackBuilderTest.kt | 16 +- .../test/kotlin/PromptRendererOrderingTest.kt | 3 +- .../src/test/kotlin/GraphRerouteTest.kt | 2 +- .../test/kotlin/NeedsArtifactInjectionTest.kt | 2 +- .../src/test/kotlin/SteeringPreemptionTest.kt | 2 +- .../kotlin/ReplayWorkspaceDerivationTest.kt | 2 +- 25 files changed, 509 insertions(+), 65 deletions(-) create mode 100644 core/context/src/main/kotlin/com/correx/core/context/compression/RelevanceScorer.kt create mode 100644 core/context/src/main/kotlin/com/correx/core/context/compression/ToMeMerger.kt create mode 100644 core/context/src/main/kotlin/com/correx/core/context/compression/TokenPruner.kt create mode 100644 core/inference/src/main/kotlin/com/correx/core/inference/EmbeddingRelevanceScorer.kt create mode 100644 infrastructure/compression/build.gradle create mode 100644 infrastructure/compression/src/main/kotlin/com/correx/infrastructure/compression/HttpTokenPruner.kt create mode 100644 infrastructure/compression/src/test/kotlin/HttpTokenPrunerTest.kt create mode 100644 sidecars/llmlingua/README.md create mode 100644 sidecars/llmlingua/requirements.txt create mode 100644 sidecars/llmlingua/server.py diff --git a/apps/server/build.gradle b/apps/server/build.gradle index 57fd4701..db59d600 100644 --- a/apps/server/build.gradle +++ b/apps/server/build.gradle @@ -43,6 +43,7 @@ dependencies { implementation project(':core:toolintent') implementation project(':infrastructure:tools') implementation project(':infrastructure:tools:filesystem') + implementation project(':infrastructure:compression') implementation "com.github.ajalt.clikt:clikt:5.0.1" diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt index 9b6c7738..2027ac83 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt @@ -17,6 +17,8 @@ import com.correx.core.config.OperatorProfile import com.correx.core.config.ProfileLoader import com.correx.core.config.ProviderConfig import com.correx.core.context.builder.DefaultContextPackBuilder +import com.correx.core.context.compression.CompressionPolicy +import com.correx.infrastructure.compression.HttpTokenPruner import com.correx.core.context.compression.DefaultContextCompressor import com.correx.core.context.model.TokenBudget import com.correx.core.events.EventDispatcher @@ -291,7 +293,11 @@ fun main() { val engines = OrchestratorEngines( transitionResolver = DefaultTransitionResolver { condition, ctx -> condition.evaluate(ctx) }, - contextPackBuilder = DefaultContextPackBuilder(DefaultContextCompressor()), + contextPackBuilder = DefaultContextPackBuilder( + DefaultContextCompressor(), + policy = CompressionPolicy(correxConfig.orchestration.compressionLevel), + tokenPruner = HttpTokenPruner(baseUrl = correxConfig.orchestration.tokenPrunerUrl), + ), inferenceRouter = inferenceRouter, validationPipeline = ValidationPipeline( validators = listOf( diff --git a/core/config/src/main/kotlin/com/correx/core/config/ConfigLoader.kt b/core/config/src/main/kotlin/com/correx/core/config/ConfigLoader.kt index 69e4fc0a..066945a1 100644 --- a/core/config/src/main/kotlin/com/correx/core/config/ConfigLoader.kt +++ b/core/config/src/main/kotlin/com/correx/core/config/ConfigLoader.kt @@ -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( diff --git a/core/config/src/main/kotlin/com/correx/core/config/CorrexConfig.kt b/core/config/src/main/kotlin/com/correx/core/config/CorrexConfig.kt index 2c944050..d5104faf 100644 --- a/core/config/src/main/kotlin/com/correx/core/config/CorrexConfig.kt +++ b/core/config/src/main/kotlin/com/correx/core/config/CorrexConfig.kt @@ -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 diff --git a/core/context/src/main/kotlin/com/correx/core/context/builder/ContextPackBuilder.kt b/core/context/src/main/kotlin/com/correx/core/context/builder/ContextPackBuilder.kt index ea4a06b7..a68bd7d0 100644 --- a/core/context/src/main/kotlin/com/correx/core/context/builder/ContextPackBuilder.kt +++ b/core/context/src/main/kotlin/com/correx/core/context/builder/ContextPackBuilder.kt @@ -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, diff --git a/core/context/src/main/kotlin/com/correx/core/context/builder/DefaultContextPackBuilder.kt b/core/context/src/main/kotlin/com/correx/core/context/builder/DefaultContextPackBuilder.kt index e0311a01..5ab6b8a6 100644 --- a/core/context/src/main/kotlin/com/correx/core/context/builder/DefaultContextPackBuilder.kt +++ b/core/context/src/main/kotlin/com/correx/core/context/builder/DefaultContextPackBuilder.kt @@ -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, 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() - 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): List { + 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): List { + 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): String? = + entries.filter { it.role == EntryRole.USER }.maxByOrNull { it.ordinal }?.content + + private suspend fun rankByRelevance(entries: List, query: String?): List { + 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, diff --git a/core/context/src/main/kotlin/com/correx/core/context/compression/RelevanceScorer.kt b/core/context/src/main/kotlin/com/correx/core/context/compression/RelevanceScorer.kt new file mode 100644 index 00000000..4af78895 --- /dev/null +++ b/core/context/src/main/kotlin/com/correx/core/context/compression/RelevanceScorer.kt @@ -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): List +} + +/** Default (scorer unconfigured): everything equally relevant, so ordering falls back to recency. */ +object UniformRelevanceScorer : RelevanceScorer { + override suspend fun score(query: String, candidates: List): List = + candidates.map { 1.0 } +} diff --git a/core/context/src/main/kotlin/com/correx/core/context/compression/ToMeMerger.kt b/core/context/src/main/kotlin/com/correx/core/context/compression/ToMeMerger.kt new file mode 100644 index 00000000..3f27f906 --- /dev/null +++ b/core/context/src/main/kotlin/com/correx/core/context/compression/ToMeMerger.kt @@ -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): List { + if (entries.size < 2) return entries + val kept = mutableListOf() + // 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 = + text.lowercase().split(WHITESPACE).filter { it.isNotBlank() }.toSet() + + private fun jaccard(a: Set, b: Set): 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+") + } +} diff --git a/core/context/src/main/kotlin/com/correx/core/context/compression/TokenPruner.kt b/core/context/src/main/kotlin/com/correx/core/context/compression/TokenPruner.kt new file mode 100644 index 00000000..93508391 --- /dev/null +++ b/core/context/src/main/kotlin/com/correx/core/context/compression/TokenPruner.kt @@ -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, 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, targetRatio: Double): String = content +} diff --git a/core/context/src/test/kotlin/CompressionPipelineStagesTest.kt b/core/context/src/test/kotlin/CompressionPipelineStagesTest.kt index b47fba64..288b582b 100644 --- a/core/context/src/test/kotlin/CompressionPipelineStagesTest.kt +++ b/core/context/src/test/kotlin/CompressionPipelineStagesTest.kt @@ -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, 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 diff --git a/core/context/src/test/kotlin/DecisionJournalPinningTest.kt b/core/context/src/test/kotlin/DecisionJournalPinningTest.kt index 9fb40717..3cb0b9b8 100644 --- a/core/context/src/test/kotlin/DecisionJournalPinningTest.kt +++ b/core/context/src/test/kotlin/DecisionJournalPinningTest.kt @@ -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 } } diff --git a/core/inference/src/main/kotlin/com/correx/core/inference/EmbeddingRelevanceScorer.kt b/core/inference/src/main/kotlin/com/correx/core/inference/EmbeddingRelevanceScorer.kt new file mode 100644 index 00000000..7101032b --- /dev/null +++ b/core/inference/src/main/kotlin/com/correx/core/inference/EmbeddingRelevanceScorer.kt @@ -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): List { + 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 + } +} diff --git a/infrastructure/compression/build.gradle b/infrastructure/compression/build.gradle new file mode 100644 index 00000000..547d2a0f --- /dev/null +++ b/infrastructure/compression/build.gradle @@ -0,0 +1,24 @@ +plugins { + id 'java-library' + id 'org.jetbrains.kotlin.jvm' + id 'org.jetbrains.kotlin.plugin.serialization' +} + +ext { + ktor_version = '3.0.3' +} + +dependencies { + implementation project(':core:context') + implementation "org.slf4j:slf4j-api" + + implementation "io.ktor:ktor-client-core:$ktor_version" + implementation "io.ktor:ktor-client-cio:$ktor_version" + implementation "io.ktor:ktor-client-content-negotiation:$ktor_version" + implementation "io.ktor:ktor-serialization-kotlinx-json:$ktor_version" + testImplementation "org.junit.jupiter:junit-jupiter" + testImplementation "io.ktor:ktor-client-mock:$ktor_version" + testImplementation "org.jetbrains.kotlin:kotlin-test" +} + +tasks.named("koverVerify").configure { enabled = false } diff --git a/infrastructure/compression/src/main/kotlin/com/correx/infrastructure/compression/HttpTokenPruner.kt b/infrastructure/compression/src/main/kotlin/com/correx/infrastructure/compression/HttpTokenPruner.kt new file mode 100644 index 00000000..11d055c9 --- /dev/null +++ b/infrastructure/compression/src/main/kotlin/com/correx/infrastructure/compression/HttpTokenPruner.kt @@ -0,0 +1,58 @@ +package com.correx.infrastructure.compression + +import com.correx.core.context.compression.TokenPruner +import io.ktor.client.HttpClient +import io.ktor.client.call.body +import io.ktor.client.plugins.contentnegotiation.ContentNegotiation +import io.ktor.client.request.post +import io.ktor.client.request.setBody +import io.ktor.client.engine.cio.CIO +import io.ktor.http.ContentType +import io.ktor.http.contentType +import io.ktor.serialization.kotlinx.json.json +import kotlinx.coroutines.withTimeout +import kotlinx.serialization.Serializable +import kotlinx.serialization.json.Json +import org.slf4j.LoggerFactory + +/** + * [TokenPruner] backed by the LLMLingua-2 sidecar (sidecars/llmlingua/server.py) over localhost + * HTTP. Protected spans are forwarded so the classifier keeps them verbatim. Fails *open*: any + * error (sidecar down, timeout, malformed) returns the original content — compression is an + * optimization, never a correctness dependency, so a missing sidecar must not break inference. + */ +class HttpTokenPruner( + private val baseUrl: String = "http://127.0.0.1:8199", + private val timeoutMs: Long = DEFAULT_TIMEOUT_MS, + private val client: HttpClient = defaultClient(), +) : TokenPruner { + + private val log = LoggerFactory.getLogger(HttpTokenPruner::class.java) + + override suspend fun prune(content: String, protectedSpans: List, targetRatio: Double): String = + runCatching { + withTimeout(timeoutMs) { + val resp: PruneResponse = client.post("$baseUrl/prune") { + contentType(ContentType.Application.Json) + setBody(PruneRequest(content, protectedSpans, targetRatio)) + }.body() + resp.compressed + } + }.getOrElse { e -> + log.debug("token pruning sidecar unavailable ({}), passing content through", e.message) + content + } + + @Serializable + private data class PruneRequest(val text: String, val protected: List, val rate: Double) + + @Serializable + private data class PruneResponse(val compressed: String) + + private companion object { + const val DEFAULT_TIMEOUT_MS = 5000L + fun defaultClient() = HttpClient(CIO) { + install(ContentNegotiation) { json(Json { ignoreUnknownKeys = true }) } + } + } +} diff --git a/infrastructure/compression/src/test/kotlin/HttpTokenPrunerTest.kt b/infrastructure/compression/src/test/kotlin/HttpTokenPrunerTest.kt new file mode 100644 index 00000000..17165e72 --- /dev/null +++ b/infrastructure/compression/src/test/kotlin/HttpTokenPrunerTest.kt @@ -0,0 +1,41 @@ +import com.correx.infrastructure.compression.HttpTokenPruner +import io.ktor.client.HttpClient +import io.ktor.client.engine.mock.MockEngine +import io.ktor.client.engine.mock.respond +import io.ktor.client.plugins.contentnegotiation.ContentNegotiation +import io.ktor.http.HttpHeaders +import io.ktor.http.HttpStatusCode +import io.ktor.http.headersOf +import io.ktor.serialization.kotlinx.json.json +import kotlinx.coroutines.runBlocking +import kotlinx.serialization.json.Json +import kotlin.test.Test +import kotlin.test.assertEquals + +class HttpTokenPrunerTest { + + private fun client(engine: MockEngine) = HttpClient(engine) { + install(ContentNegotiation) { json(Json { ignoreUnknownKeys = true }) } + } + + @Test + fun `returns sidecar compressed output`() = runBlocking { + val engine = MockEngine { + respond( + content = """{"compressed":"short text 10.0.0.1"}""", + status = HttpStatusCode.OK, + headers = headersOf(HttpHeaders.ContentType, "application/json"), + ) + } + val pruner = HttpTokenPruner(client = client(engine)) + assertEquals("short text 10.0.0.1", pruner.prune("a very long text 10.0.0.1", listOf("10.0.0.1"), 0.5)) + } + + @Test + fun `fails open when sidecar errors`() = runBlocking { + val engine = MockEngine { respond("boom", HttpStatusCode.InternalServerError) } + val pruner = HttpTokenPruner(client = client(engine)) + val original = "keep this content untouched" + assertEquals(original, pruner.prune(original, emptyList(), 0.5)) + } +} diff --git a/settings.gradle b/settings.gradle index 9e2c929d..73b6735f 100644 --- a/settings.gradle +++ b/settings.gradle @@ -36,6 +36,7 @@ include ':infrastructure:inference:llama_cpp' include ':infrastructure:inference:openai_compat' include ':infrastructure:router' include ':infrastructure:router:turbovec' +include ':infrastructure:compression' include ':infrastructure:tools' include ':infrastructure:tools:filesystem' include ':infrastructure:workflow' diff --git a/sidecars/llmlingua/README.md b/sidecars/llmlingua/README.md new file mode 100644 index 00000000..d3f04833 --- /dev/null +++ b/sidecars/llmlingua/README.md @@ -0,0 +1,37 @@ +# LLMLingua-2 token-pruning sidecar + +Prunes low-perplexity tokens from freeform prose before it hits the local LLM, so more usable +context fits a bounded window. Implements pipeline stage 3 (`TOKEN_PRUNE`, level 3+) — see +`docs/plans/correx-compression-pipeline.md` §4. + +Python-only because LLMLingua-2 is a torch/BERT classifier with no JVM equivalent. correx calls +it over localhost HTTP via `HttpTokenPruner`, which **fails open**: if this sidecar is down, the +kernel passes context through uncompressed. Nothing breaks; you just don't get token pruning. + +## Run + +```bash +cd sidecars/llmlingua +python -m venv .venv && source .venv/bin/activate +pip install -r requirements.txt +uvicorn server:app --host 127.0.0.1 --port 8199 +``` + +First `/prune` call downloads the model (~1-2 GB) and loads torch; `/health` responds immediately. + +## Wire into correx + +Set compression level ≥ 3 for the workflow and point the kernel at the sidecar: + +```toml +[compression] +level = 4 +token_pruner_url = "http://127.0.0.1:8199" +``` + +## API + +- `GET /health` → `{"status":"ok"}` +- `POST /prune` `{"text": str, "protected": [str], "rate": 0.55}` → `{"compressed": str}` + - `rate` = fraction of tokens to **keep** (0.55 ≈ 45% compression) + - `protected` substrings (IDs, numbers, paths, code) are kept verbatim diff --git a/sidecars/llmlingua/requirements.txt b/sidecars/llmlingua/requirements.txt new file mode 100644 index 00000000..c10a1eb5 --- /dev/null +++ b/sidecars/llmlingua/requirements.txt @@ -0,0 +1,3 @@ +llmlingua>=0.2.2 +fastapi>=0.110 +uvicorn[standard]>=0.29 diff --git a/sidecars/llmlingua/server.py b/sidecars/llmlingua/server.py new file mode 100644 index 00000000..2a3d2bd6 --- /dev/null +++ b/sidecars/llmlingua/server.py @@ -0,0 +1,65 @@ +"""LLMLingua-2 token-pruning sidecar for correx. + +A tiny HTTP service correx calls (via HttpTokenPruner) to prune low-perplexity tokens from +freeform prose before it's sent to the local LLM. Kept in Python because LLMLingua-2 is a +torch/BERT classifier with no JVM equivalent (pipeline §4). + +Contract: + POST /prune {"text": str, "protected": [str], "rate": float} -> {"compressed": str} + rate = fraction of tokens to KEEP (0.55 keeps ~55%, i.e. ~45% compression). + `protected` substrings are force-kept verbatim (IDs, numbers, paths, code). + GET /health -> {"status": "ok"} + +Run: + pip install -r requirements.txt + uvicorn server:app --host 127.0.0.1 --port 8199 +""" +from __future__ import annotations + +import os +from fastapi import FastAPI +from pydantic import BaseModel + +app = FastAPI(title="correx-llmlingua") + +_MODEL = os.environ.get("LLMLINGUA_MODEL", "microsoft/llmlingua-2-xlm-roberta-large-meetingbank") +_compressor = None + + +def _get_compressor(): + # Lazy-load so /health works (and the process starts fast) before torch spins up. + global _compressor + if _compressor is None: + from llmlingua import PromptCompressor + _compressor = PromptCompressor(model_name=_MODEL, use_llmlingua2=True) + return _compressor + + +class PruneRequest(BaseModel): + text: str + protected: list[str] = [] + rate: float = 0.55 + + +class PruneResponse(BaseModel): + compressed: str + + +@app.get("/health") +def health(): + return {"status": "ok"} + + +@app.post("/prune", response_model=PruneResponse) +def prune(req: PruneRequest): + text = req.text.strip() + if not text: + return PruneResponse(compressed=req.text) + # rate is fraction to keep; LLMLingua-2 force_tokens keeps the protected spans verbatim. + result = _get_compressor().compress_prompt( + text, + rate=max(0.1, min(1.0, req.rate)), + force_tokens=req.protected or None, + drop_consecutive=True, + ) + return PruneResponse(compressed=result["compressed_prompt"]) diff --git a/testing/deterministic/src/test/kotlin/DefaultContextPackBuilderTest.kt b/testing/deterministic/src/test/kotlin/DefaultContextPackBuilderTest.kt index 5d7aac25..51f4006e 100644 --- a/testing/deterministic/src/test/kotlin/DefaultContextPackBuilderTest.kt +++ b/testing/deterministic/src/test/kotlin/DefaultContextPackBuilderTest.kt @@ -41,7 +41,7 @@ class DefaultContextPackBuilderTest { ) @Test - fun `oversized tool results from a stage tool loop are trimmed to budget`() { + fun `oversized tool results from a stage tool loop are trimmed to budget`() = kotlinx.coroutines.runBlocking { // Live repro (2026-06-11): analyst stage with three file_read results of ~5.6k/11.9k/10.7k // tokens sailed through a 16384 budget untrimmed and produced a 34k-token prompt. val entries = listOf( @@ -62,7 +62,7 @@ class DefaultContextPackBuilderTest { } @Test - fun `build groups entries by layer`() { + fun `build groups entries by layer`() = kotlinx.coroutines.runBlocking { val entries = listOf( entry("a", ContextLayer.L0, 50), entry("b", ContextLayer.L1, 60), @@ -75,7 +75,7 @@ class DefaultContextPackBuilderTest { } @Test - fun `budgetUsed reflects total token estimate of retained entries`() { + fun `budgetUsed reflects total token estimate of retained entries`() = kotlinx.coroutines.runBlocking { val entries = listOf( entry("a", ContextLayer.L0, 100), entry("b", ContextLayer.L1, 200) @@ -85,14 +85,14 @@ class DefaultContextPackBuilderTest { } @Test - fun `pack budgetUsed never exceeds budgetLimit`() { + fun `pack budgetUsed never exceeds budgetLimit`() = kotlinx.coroutines.runBlocking { val entries = (1..20).map { entry("e$it", ContextLayer.L2, 100) } val pack = builder.build(packId, sessionId, stageId, entries, TokenBudget(limit = 500)) assertTrue(pack.budgetUsed <= 500) } @Test - fun `DecisionPointBuilder returns only L0 and L1 entries`() { + fun `DecisionPointBuilder returns only L0 and L1 entries`() = kotlinx.coroutines.runBlocking { val entries = listOf( entry("a", ContextLayer.L0, 50), entry("b", ContextLayer.L1, 60), @@ -105,7 +105,7 @@ class DefaultContextPackBuilderTest { } @Test - fun `build handles empty entries list`() { + fun `build handles empty entries list`() = kotlinx.coroutines.runBlocking { val pack = builder.build(packId, sessionId, stageId, emptyList(), TokenBudget(limit = 1000)) assertEquals(0, pack.budgetUsed) assertEquals(0, pack.compressionMetadata.entriesDropped) @@ -113,7 +113,7 @@ class DefaultContextPackBuilderTest { } @Test - fun `steeringNote and eventHistory entries are retained even when exceeding budget`() { + fun `steeringNote and eventHistory entries are retained even when exceeding budget`() = kotlinx.coroutines.runBlocking { val entries = listOf( entry("steering", ContextLayer.L0, 500).copy(sourceType = "steeringNote"), entry("history", ContextLayer.L1, 600).copy(sourceType = "eventHistory"), @@ -126,7 +126,7 @@ class DefaultContextPackBuilderTest { } @Test - fun `tool dialogue keeps chronological order despite sourceType grouping`() { + fun `tool dialogue keeps chronological order despite sourceType grouping`() = kotlinx.coroutines.runBlocking { // A multi-round tool loop: assistant call then tool result, interleaved. Compression // groups by sourceType internally; the builder must restore input order so the L2 // layer reads call, result, call, result — not all calls then all results. diff --git a/testing/deterministic/src/test/kotlin/PromptRendererOrderingTest.kt b/testing/deterministic/src/test/kotlin/PromptRendererOrderingTest.kt index d0059afe..e0cbf1a0 100644 --- a/testing/deterministic/src/test/kotlin/PromptRendererOrderingTest.kt +++ b/testing/deterministic/src/test/kotlin/PromptRendererOrderingTest.kt @@ -29,7 +29,7 @@ class PromptRendererOrderingTest { ) @Test - fun `stage tool loop renders chronologically with the task before tool turns`() { + fun `stage tool loop renders chronologically with the task before tool turns`() = kotlinx.coroutines.runBlocking { // The builder stamps ordinals from input order; the renderer must honour them so the // model sees: system, task, assistant call, tool result, assistant call, tool result. val builder = DefaultContextPackBuilder(DefaultContextCompressor()) @@ -49,6 +49,7 @@ class PromptRendererOrderingTest { "tool" to "result1", "assistant" to "call2", "tool" to "result2"), messages.map { it.role to it.content }, ) + Unit } @Test diff --git a/testing/integration/src/test/kotlin/GraphRerouteTest.kt b/testing/integration/src/test/kotlin/GraphRerouteTest.kt index c2def0dd..7791ce51 100644 --- a/testing/integration/src/test/kotlin/GraphRerouteTest.kt +++ b/testing/integration/src/test/kotlin/GraphRerouteTest.kt @@ -68,7 +68,7 @@ import java.util.UUID class GraphRerouteTest { private val passthroughBuilder = object : ContextPackBuilder { - override fun build( + override suspend fun build( id: ContextPackId, sessionId: SessionId, stageId: StageId, diff --git a/testing/integration/src/test/kotlin/NeedsArtifactInjectionTest.kt b/testing/integration/src/test/kotlin/NeedsArtifactInjectionTest.kt index 5da0f197..3183c1da 100644 --- a/testing/integration/src/test/kotlin/NeedsArtifactInjectionTest.kt +++ b/testing/integration/src/test/kotlin/NeedsArtifactInjectionTest.kt @@ -68,7 +68,7 @@ class NeedsArtifactInjectionTest { private val capturedEntries = mutableListOf() private val capturingBuilder = object : ContextPackBuilder { - override fun build( + override suspend fun build( id: ContextPackId, sessionId: SessionId, stageId: StageId, diff --git a/testing/integration/src/test/kotlin/SteeringPreemptionTest.kt b/testing/integration/src/test/kotlin/SteeringPreemptionTest.kt index 4bb3195a..cdef5226 100644 --- a/testing/integration/src/test/kotlin/SteeringPreemptionTest.kt +++ b/testing/integration/src/test/kotlin/SteeringPreemptionTest.kt @@ -63,7 +63,7 @@ class SteeringPreemptionTest { private val capturedEntries = mutableListOf() private val capturingBuilder = object : ContextPackBuilder { - override fun build( + override suspend fun build( id: ContextPackId, sessionId: SessionId, stageId: StageId, diff --git a/testing/kernel/src/test/kotlin/ReplayWorkspaceDerivationTest.kt b/testing/kernel/src/test/kotlin/ReplayWorkspaceDerivationTest.kt index a82f6cfa..b389c1df 100644 --- a/testing/kernel/src/test/kotlin/ReplayWorkspaceDerivationTest.kt +++ b/testing/kernel/src/test/kotlin/ReplayWorkspaceDerivationTest.kt @@ -198,7 +198,7 @@ class ReplayWorkspaceDerivationTest { ) }, contextPackBuilder = object : ContextPackBuilder { - override fun build( + override suspend fun build( id: ContextPackId, sessionId: SessionId, stageId: StageId,