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
This commit is contained in:
@@ -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"
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
+92
-47
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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 }
|
||||
+58
@@ -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<String>, 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<String>, 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 }) }
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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'
|
||||
|
||||
@@ -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
|
||||
@@ -0,0 +1,3 @@
|
||||
llmlingua>=0.2.2
|
||||
fastapi>=0.110
|
||||
uvicorn[standard]>=0.29
|
||||
@@ -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"])
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -68,7 +68,7 @@ class NeedsArtifactInjectionTest {
|
||||
private val capturedEntries = mutableListOf<ContextEntry>()
|
||||
|
||||
private val capturingBuilder = object : ContextPackBuilder {
|
||||
override fun build(
|
||||
override suspend fun build(
|
||||
id: ContextPackId,
|
||||
sessionId: SessionId,
|
||||
stageId: StageId,
|
||||
|
||||
@@ -63,7 +63,7 @@ class SteeringPreemptionTest {
|
||||
private val capturedEntries = mutableListOf<ContextEntry>()
|
||||
|
||||
private val capturingBuilder = object : ContextPackBuilder {
|
||||
override fun build(
|
||||
override suspend fun build(
|
||||
id: ContextPackId,
|
||||
sessionId: SessionId,
|
||||
stageId: StageId,
|
||||
|
||||
@@ -198,7 +198,7 @@ class ReplayWorkspaceDerivationTest {
|
||||
)
|
||||
},
|
||||
contextPackBuilder = object : ContextPackBuilder {
|
||||
override fun build(
|
||||
override suspend fun build(
|
||||
id: ContextPackId,
|
||||
sessionId: SessionId,
|
||||
stageId: StageId,
|
||||
|
||||
Reference in New Issue
Block a user