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(':core:toolintent')
|
||||||
implementation project(':infrastructure:tools')
|
implementation project(':infrastructure:tools')
|
||||||
implementation project(':infrastructure:tools:filesystem')
|
implementation project(':infrastructure:tools:filesystem')
|
||||||
|
implementation project(':infrastructure:compression')
|
||||||
|
|
||||||
implementation "com.github.ajalt.clikt:clikt:5.0.1"
|
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.ProfileLoader
|
||||||
import com.correx.core.config.ProviderConfig
|
import com.correx.core.config.ProviderConfig
|
||||||
import com.correx.core.context.builder.DefaultContextPackBuilder
|
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.compression.DefaultContextCompressor
|
||||||
import com.correx.core.context.model.TokenBudget
|
import com.correx.core.context.model.TokenBudget
|
||||||
import com.correx.core.events.EventDispatcher
|
import com.correx.core.events.EventDispatcher
|
||||||
@@ -291,7 +293,11 @@ fun main() {
|
|||||||
|
|
||||||
val engines = OrchestratorEngines(
|
val engines = OrchestratorEngines(
|
||||||
transitionResolver = DefaultTransitionResolver { condition, ctx -> condition.evaluate(ctx) },
|
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,
|
inferenceRouter = inferenceRouter,
|
||||||
validationPipeline = ValidationPipeline(
|
validationPipeline = ValidationPipeline(
|
||||||
validators = listOf(
|
validators = listOf(
|
||||||
|
|||||||
@@ -624,6 +624,9 @@ object ConfigLoader {
|
|||||||
asInt(orchestrationSection["journal_compaction_token_threshold"], 2_000),
|
asInt(orchestrationSection["journal_compaction_token_threshold"], 2_000),
|
||||||
resumeAbandonedMaxAgeMinutes =
|
resumeAbandonedMaxAgeMinutes =
|
||||||
asLong(orchestrationSection["resume_abandoned_max_age_minutes"], 1_440),
|
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(
|
val modelsSettings = ModelsSettings(
|
||||||
|
|||||||
@@ -61,6 +61,13 @@ data class OrchestrationKnobs(
|
|||||||
* 0 disables auto-resume entirely.
|
* 0 disables auto-resume entirely.
|
||||||
*/
|
*/
|
||||||
val resumeAbandonedMaxAgeMinutes: Long = 1_440,
|
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
|
@Serializable
|
||||||
|
|||||||
@@ -8,7 +8,7 @@ import com.correx.core.events.types.SessionId
|
|||||||
import com.correx.core.events.types.StageId
|
import com.correx.core.events.types.StageId
|
||||||
|
|
||||||
interface ContextPackBuilder {
|
interface ContextPackBuilder {
|
||||||
fun build(
|
suspend fun build(
|
||||||
id: ContextPackId,
|
id: ContextPackId,
|
||||||
sessionId: SessionId,
|
sessionId: SessionId,
|
||||||
stageId: StageId,
|
stageId: StageId,
|
||||||
|
|||||||
+87
-42
@@ -8,7 +8,12 @@ import com.correx.core.context.compression.ContextClassifier
|
|||||||
import com.correx.core.context.compression.ContextCompressor
|
import com.correx.core.context.compression.ContextCompressor
|
||||||
import com.correx.core.context.compression.FactSheet
|
import com.correx.core.context.compression.FactSheet
|
||||||
import com.correx.core.context.compression.FormatCompressor
|
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.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.CompressionMetadata
|
||||||
import com.correx.core.context.model.ContextEntry
|
import com.correx.core.context.model.ContextEntry
|
||||||
import com.correx.core.context.model.ContextLayer
|
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.SessionId
|
||||||
import com.correx.core.events.types.StageId
|
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(
|
class DefaultContextPackBuilder(
|
||||||
private val compressor: ContextCompressor,
|
private val compressor: ContextCompressor,
|
||||||
private val policy: CompressionPolicy = CompressionPolicy(),
|
private val policy: CompressionPolicy = CompressionPolicy(),
|
||||||
private val classifier: ContextClassifier = ContextClassifier(),
|
private val classifier: ContextClassifier = ContextClassifier(),
|
||||||
private val formatCompressor: FormatCompressor = FormatCompressor(),
|
private val formatCompressor: FormatCompressor = FormatCompressor(),
|
||||||
|
private val tagger: ProtectedSpanTagger = ProtectedSpanTagger(),
|
||||||
private val factSheet: FactSheet = FactSheet(ProtectedSpanTagger()),
|
private val factSheet: FactSheet = FactSheet(ProtectedSpanTagger()),
|
||||||
|
private val tokenPruner: TokenPruner = NoOpTokenPruner,
|
||||||
|
private val toMeMerger: ToMeMerger = ToMeMerger(),
|
||||||
|
private val relevanceScorer: RelevanceScorer = UniformRelevanceScorer,
|
||||||
) : ContextPackBuilder {
|
) : ContextPackBuilder {
|
||||||
|
|
||||||
private val neverDropSourceTypes = setOf("steeringNote", "eventHistory", "decisionJournal", "factSheet")
|
private val neverDropSourceTypes = setOf("steeringNote", "eventHistory", "decisionJournal", "factSheet")
|
||||||
@@ -33,72 +50,71 @@ class DefaultContextPackBuilder(
|
|||||||
private companion object {
|
private companion object {
|
||||||
const val CHARS_PER_TOKEN = 4
|
const val CHARS_PER_TOKEN = 4
|
||||||
const val FACT_SHEET_ORDINAL = -1
|
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,
|
id: ContextPackId,
|
||||||
sessionId: SessionId,
|
sessionId: SessionId,
|
||||||
stageId: StageId,
|
stageId: StageId,
|
||||||
entries: List<ContextEntry>,
|
entries: List<ContextEntry>,
|
||||||
budget: TokenBudget
|
budget: TokenBudget
|
||||||
): ContextPack {
|
): ContextPack {
|
||||||
// Stamp chronological order from the caller's input sequence. Grouping by
|
// Stamp chronological order from the caller's input sequence so grouping/compression can
|
||||||
// sourceType (for compression) and by layer reorders entries; the ordinal lets
|
// reorder freely and PromptRenderer can restore true turn order from the ordinal.
|
||||||
// us restore true turn order afterwards so a tool loop reads assistant→tool→…
|
|
||||||
// instead of all-assistants-then-all-tools.
|
|
||||||
val stamped = entries.mapIndexed { index, entry -> entry.copy(ordinal = index) }
|
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
|
// Stage 1 (FORMAT_COMPRESS): lossless json→dotted-line compaction of structured entries.
|
||||||
// lossy stages that touch freeform live further down the pipeline (prune/summarize).
|
val formatted = if (policy.enabled(CompressionStage.FORMAT_COMPRESS)) {
|
||||||
val ordered = if (policy.enabled(CompressionStage.FORMAT_COMPRESS)) {
|
|
||||||
stamped.map { entry ->
|
stamped.map { entry ->
|
||||||
if (classifier.classify(entry) == ContextClass.STRUCTURED) {
|
if (classifier.classify(entry) == ContextClass.STRUCTURED) reencode(entry, formatCompressor.compress(entry.content))
|
||||||
val compacted = formatCompressor.compress(entry.content)
|
else entry
|
||||||
if (compacted == entry.content) entry
|
|
||||||
else entry.copy(content = compacted, tokenEstimate = estimateTokens(compacted))
|
|
||||||
} else entry
|
|
||||||
}
|
}
|
||||||
} else stamped
|
} 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
|
// Stage 3 (TOKEN_PRUNE): prune freeform prose, preserving protected spans. When TIER_SPLIT
|
||||||
// insurance — the most compressed tier still can't lose an ID/number/path.
|
// 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 withFactSheet = if (policy.enabled(CompressionStage.FACT_SHEET)) {
|
||||||
val facts = factSheet.extract(ordered)
|
val facts = factSheet.extract(merged)
|
||||||
if (facts.isEmpty()) ordered else ordered + factSheetEntry(factSheet.render(facts))
|
if (facts.isEmpty()) merged else merged + factSheetEntry(factSheet.render(facts))
|
||||||
} else ordered
|
} else merged
|
||||||
|
|
||||||
val (pinned, compressible) = withFactSheet.partition { it.sourceType in neverDropSourceTypes }
|
val (pinned, compressible) = withFactSheet.partition { it.sourceType in neverDropSourceTypes }
|
||||||
val pinnedTokens = pinned.sumOf { it.tokenEstimate }
|
val pinnedTokens = pinned.sumOf { it.tokenEstimate }
|
||||||
var remainingTokens = (budget.limit - pinnedTokens).coerceAtLeast(0)
|
var remainingTokens = (budget.limit - pinnedTokens).coerceAtLeast(0)
|
||||||
|
|
||||||
// Dispatch compression strategy by entry sourceType — applying Conversation
|
// Query-conditioned selection: when a real relevance scorer is wired, reorder each freeform
|
||||||
// strategy uniformly mangles tool logs and artifacts (their shape isn't conversational).
|
// group so the least-relevant turns sit at the front (dropped first) and the most-relevant
|
||||||
// Stage 9 (level 9+): strict allocation reserves budget for structured content by
|
// land at the tail the Conversation strategy keeps. Final layering re-sorts by ordinal, so
|
||||||
// compressing structured-class groups before freeform, so a long conversation can't
|
// chronological render order is unaffected — only *what survives* the budget changes.
|
||||||
// starve tool logs / artifacts out of the window. Below level 9 keep encounter order.
|
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 strict = policy.enabled(CompressionStage.BUDGET_ALLOCATOR_STRICT)
|
||||||
val strategiesUsed = linkedSetOf<String>()
|
val strategiesUsed = linkedSetOf<String>()
|
||||||
val grouped = compressible.groupBy { it.sourceType }.entries
|
val grouped = relevanceRanked.groupBy { it.sourceType }.entries
|
||||||
val orderedGroups = if (strict) {
|
val orderedGroups = if (strict) {
|
||||||
grouped.sortedBy { (_, group) ->
|
grouped.sortedBy { (_, group) -> if (classifier.classify(group.first()) == ContextClass.FREEFORM) 1 else 0 }
|
||||||
if (classifier.classify(group.first()) == ContextClass.FREEFORM) 1 else 0
|
|
||||||
}
|
|
||||||
} else grouped.toList()
|
} else grouped.toList()
|
||||||
val compressed = orderedGroups
|
val compressed = orderedGroups.flatMap { (sourceType, group) ->
|
||||||
.flatMap { (sourceType, group) ->
|
|
||||||
val strategy = strategyFor(sourceType)
|
val strategy = strategyFor(sourceType)
|
||||||
strategiesUsed += strategy::class.simpleName ?: "Unknown"
|
strategiesUsed += strategy::class.simpleName ?: "Unknown"
|
||||||
val groupBudget = TokenBudget(limit = remainingTokens.coerceAtLeast(0))
|
val result = compressor.compress(group, TokenBudget(limit = remainingTokens.coerceAtLeast(0)), strategy)
|
||||||
val result = compressor.compress(group, groupBudget, strategy)
|
|
||||||
remainingTokens -= result.sumOf { it.tokenEstimate }
|
remainingTokens -= result.sumOf { it.tokenEstimate }
|
||||||
result
|
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 retained = (pinned + compressed).sortedBy { it.ordinal }
|
||||||
val layers = retained.groupBy { it.layer }
|
val layers = retained.groupBy { it.layer }
|
||||||
val budgetUsed = retained.sumOf { it.tokenEstimate }
|
val droppedCount = entries.size - retained.count { it.sourceType != "factSheet" }
|
||||||
val droppedCount = entries.size - retained.size
|
|
||||||
val truncatedLayers = if (droppedCount > 0) {
|
val truncatedLayers = if (droppedCount > 0) {
|
||||||
entries.map { it.layer }.distinct().filter { layer ->
|
entries.map { it.layer }.distinct().filter { layer ->
|
||||||
entries.count { it.layer == layer } > retained.count { it.layer == layer }
|
entries.count { it.layer == layer } > retained.count { it.layer == layer }
|
||||||
@@ -110,7 +126,7 @@ class DefaultContextPackBuilder(
|
|||||||
sessionId = sessionId,
|
sessionId = sessionId,
|
||||||
stageId = stageId,
|
stageId = stageId,
|
||||||
layers = layers,
|
layers = layers,
|
||||||
budgetUsed = budgetUsed,
|
budgetUsed = retained.sumOf { it.tokenEstimate },
|
||||||
budgetLimit = budget.limit,
|
budgetLimit = budget.limit,
|
||||||
compressionMetadata = CompressionMetadata(
|
compressionMetadata = CompressionMetadata(
|
||||||
appliedStrategies = strategiesUsed.toList().ifEmpty { listOf("Conversation") },
|
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
|
private suspend fun prune(entries: List<ContextEntry>): List<ContextEntry> {
|
||||||
// math stays consistent after a format-compressed entry shrinks.
|
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
|
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(
|
private fun factSheetEntry(content: String) = ContextEntry(
|
||||||
id = ContextEntryId("factsheet"),
|
id = ContextEntryId("factsheet"),
|
||||||
layer = ContextLayer.L0,
|
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
|
@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 entries = listOf(entry("chat", ContextLayer.L1, EntryRole.USER, "deploy to 10.0.0.1"))
|
||||||
val compressor = com.correx.core.context.compression.DefaultContextCompressor()
|
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)
|
compressor, CompressionPolicy(level)
|
||||||
).build(
|
).build(
|
||||||
com.correx.core.events.types.ContextPackId("p"),
|
com.correx.core.events.types.ContextPackId("p"),
|
||||||
@@ -98,6 +98,43 @@ class CompressionPipelineStagesTest {
|
|||||||
val atSeven = buildAt(7).layers.values.flatten()
|
val atSeven = buildAt(7).layers.values.flatten()
|
||||||
assertTrue(atSeven.any { it.sourceType == "factSheet" && it.content.contains("10.0.0.1") })
|
assertTrue(atSeven.any { it.sourceType == "factSheet" && it.content.contains("10.0.0.1") })
|
||||||
assertFalse(buildAt(6).layers.values.flatten().any { it.sourceType == "factSheet" })
|
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
|
@Test
|
||||||
|
|||||||
@@ -20,7 +20,7 @@ class DecisionJournalPinningTest {
|
|||||||
private val packId = ContextPackId("pack-1")
|
private val packId = ContextPackId("pack-1")
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
fun `decisionJournal entry survives an under-budget build`() {
|
fun `decisionJournal entry survives an under-budget build`() = kotlinx.coroutines.runBlocking {
|
||||||
val entries = listOf(
|
val entries = listOf(
|
||||||
ContextEntry(
|
ContextEntry(
|
||||||
id = ContextEntryId("journal-1"),
|
id = ContextEntryId("journal-1"),
|
||||||
@@ -37,5 +37,6 @@ class DecisionJournalPinningTest {
|
|||||||
allRetained.any { it.sourceType == "decisionJournal" },
|
allRetained.any { it.sourceType == "decisionJournal" },
|
||||||
"decisionJournal entry must be retained even when budget is exceeded"
|
"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:inference:openai_compat'
|
||||||
include ':infrastructure:router'
|
include ':infrastructure:router'
|
||||||
include ':infrastructure:router:turbovec'
|
include ':infrastructure:router:turbovec'
|
||||||
|
include ':infrastructure:compression'
|
||||||
include ':infrastructure:tools'
|
include ':infrastructure:tools'
|
||||||
include ':infrastructure:tools:filesystem'
|
include ':infrastructure:tools:filesystem'
|
||||||
include ':infrastructure:workflow'
|
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
|
@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
|
// 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.
|
// tokens sailed through a 16384 budget untrimmed and produced a 34k-token prompt.
|
||||||
val entries = listOf(
|
val entries = listOf(
|
||||||
@@ -62,7 +62,7 @@ class DefaultContextPackBuilderTest {
|
|||||||
}
|
}
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
fun `build groups entries by layer`() {
|
fun `build groups entries by layer`() = kotlinx.coroutines.runBlocking {
|
||||||
val entries = listOf(
|
val entries = listOf(
|
||||||
entry("a", ContextLayer.L0, 50),
|
entry("a", ContextLayer.L0, 50),
|
||||||
entry("b", ContextLayer.L1, 60),
|
entry("b", ContextLayer.L1, 60),
|
||||||
@@ -75,7 +75,7 @@ class DefaultContextPackBuilderTest {
|
|||||||
}
|
}
|
||||||
|
|
||||||
@Test
|
@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(
|
val entries = listOf(
|
||||||
entry("a", ContextLayer.L0, 100),
|
entry("a", ContextLayer.L0, 100),
|
||||||
entry("b", ContextLayer.L1, 200)
|
entry("b", ContextLayer.L1, 200)
|
||||||
@@ -85,14 +85,14 @@ class DefaultContextPackBuilderTest {
|
|||||||
}
|
}
|
||||||
|
|
||||||
@Test
|
@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 entries = (1..20).map { entry("e$it", ContextLayer.L2, 100) }
|
||||||
val pack = builder.build(packId, sessionId, stageId, entries, TokenBudget(limit = 500))
|
val pack = builder.build(packId, sessionId, stageId, entries, TokenBudget(limit = 500))
|
||||||
assertTrue(pack.budgetUsed <= 500)
|
assertTrue(pack.budgetUsed <= 500)
|
||||||
}
|
}
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
fun `DecisionPointBuilder returns only L0 and L1 entries`() {
|
fun `DecisionPointBuilder returns only L0 and L1 entries`() = kotlinx.coroutines.runBlocking {
|
||||||
val entries = listOf(
|
val entries = listOf(
|
||||||
entry("a", ContextLayer.L0, 50),
|
entry("a", ContextLayer.L0, 50),
|
||||||
entry("b", ContextLayer.L1, 60),
|
entry("b", ContextLayer.L1, 60),
|
||||||
@@ -105,7 +105,7 @@ class DefaultContextPackBuilderTest {
|
|||||||
}
|
}
|
||||||
|
|
||||||
@Test
|
@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))
|
val pack = builder.build(packId, sessionId, stageId, emptyList(), TokenBudget(limit = 1000))
|
||||||
assertEquals(0, pack.budgetUsed)
|
assertEquals(0, pack.budgetUsed)
|
||||||
assertEquals(0, pack.compressionMetadata.entriesDropped)
|
assertEquals(0, pack.compressionMetadata.entriesDropped)
|
||||||
@@ -113,7 +113,7 @@ class DefaultContextPackBuilderTest {
|
|||||||
}
|
}
|
||||||
|
|
||||||
@Test
|
@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(
|
val entries = listOf(
|
||||||
entry("steering", ContextLayer.L0, 500).copy(sourceType = "steeringNote"),
|
entry("steering", ContextLayer.L0, 500).copy(sourceType = "steeringNote"),
|
||||||
entry("history", ContextLayer.L1, 600).copy(sourceType = "eventHistory"),
|
entry("history", ContextLayer.L1, 600).copy(sourceType = "eventHistory"),
|
||||||
@@ -126,7 +126,7 @@ class DefaultContextPackBuilderTest {
|
|||||||
}
|
}
|
||||||
|
|
||||||
@Test
|
@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
|
// 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
|
// 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.
|
// layer reads call, result, call, result — not all calls then all results.
|
||||||
|
|||||||
@@ -29,7 +29,7 @@ class PromptRendererOrderingTest {
|
|||||||
)
|
)
|
||||||
|
|
||||||
@Test
|
@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
|
// 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.
|
// model sees: system, task, assistant call, tool result, assistant call, tool result.
|
||||||
val builder = DefaultContextPackBuilder(DefaultContextCompressor())
|
val builder = DefaultContextPackBuilder(DefaultContextCompressor())
|
||||||
@@ -49,6 +49,7 @@ class PromptRendererOrderingTest {
|
|||||||
"tool" to "result1", "assistant" to "call2", "tool" to "result2"),
|
"tool" to "result1", "assistant" to "call2", "tool" to "result2"),
|
||||||
messages.map { it.role to it.content },
|
messages.map { it.role to it.content },
|
||||||
)
|
)
|
||||||
|
Unit
|
||||||
}
|
}
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
|
|||||||
@@ -68,7 +68,7 @@ import java.util.UUID
|
|||||||
class GraphRerouteTest {
|
class GraphRerouteTest {
|
||||||
|
|
||||||
private val passthroughBuilder = object : ContextPackBuilder {
|
private val passthroughBuilder = object : ContextPackBuilder {
|
||||||
override fun build(
|
override suspend fun build(
|
||||||
id: ContextPackId,
|
id: ContextPackId,
|
||||||
sessionId: SessionId,
|
sessionId: SessionId,
|
||||||
stageId: StageId,
|
stageId: StageId,
|
||||||
|
|||||||
@@ -68,7 +68,7 @@ class NeedsArtifactInjectionTest {
|
|||||||
private val capturedEntries = mutableListOf<ContextEntry>()
|
private val capturedEntries = mutableListOf<ContextEntry>()
|
||||||
|
|
||||||
private val capturingBuilder = object : ContextPackBuilder {
|
private val capturingBuilder = object : ContextPackBuilder {
|
||||||
override fun build(
|
override suspend fun build(
|
||||||
id: ContextPackId,
|
id: ContextPackId,
|
||||||
sessionId: SessionId,
|
sessionId: SessionId,
|
||||||
stageId: StageId,
|
stageId: StageId,
|
||||||
|
|||||||
@@ -63,7 +63,7 @@ class SteeringPreemptionTest {
|
|||||||
private val capturedEntries = mutableListOf<ContextEntry>()
|
private val capturedEntries = mutableListOf<ContextEntry>()
|
||||||
|
|
||||||
private val capturingBuilder = object : ContextPackBuilder {
|
private val capturingBuilder = object : ContextPackBuilder {
|
||||||
override fun build(
|
override suspend fun build(
|
||||||
id: ContextPackId,
|
id: ContextPackId,
|
||||||
sessionId: SessionId,
|
sessionId: SessionId,
|
||||||
stageId: StageId,
|
stageId: StageId,
|
||||||
|
|||||||
@@ -198,7 +198,7 @@ class ReplayWorkspaceDerivationTest {
|
|||||||
)
|
)
|
||||||
},
|
},
|
||||||
contextPackBuilder = object : ContextPackBuilder {
|
contextPackBuilder = object : ContextPackBuilder {
|
||||||
override fun build(
|
override suspend fun build(
|
||||||
id: ContextPackId,
|
id: ContextPackId,
|
||||||
sessionId: SessionId,
|
sessionId: SessionId,
|
||||||
stageId: StageId,
|
stageId: StageId,
|
||||||
|
|||||||
Reference in New Issue
Block a user