feat(context): compression pipeline stage 1 — level policy + deterministic stages
Additive CompressionPolicy (levels 1-9), ProtectedSpanTagger (regex load-bearing spans), FormatCompressor (json→dotted lines), ContextClassifier (static/structured/freeform). Stage 1 (FORMAT_COMPRESS) wired into DefaultContextPackBuilder behind the policy; structured entries only, lossless. PromptRenderer repetition-anchoring of steering directives.
This commit is contained in:
+30
-2
@@ -1,7 +1,12 @@
|
||||
package com.correx.core.context.builder
|
||||
|
||||
import com.correx.core.context.compression.CompressionPolicy
|
||||
import com.correx.core.context.compression.CompressionStage
|
||||
import com.correx.core.context.compression.CompressionStrategy
|
||||
import com.correx.core.context.compression.ContextClass
|
||||
import com.correx.core.context.compression.ContextClassifier
|
||||
import com.correx.core.context.compression.ContextCompressor
|
||||
import com.correx.core.context.compression.FormatCompressor
|
||||
import com.correx.core.context.model.CompressionMetadata
|
||||
import com.correx.core.context.model.ContextEntry
|
||||
import com.correx.core.context.model.ContextPack
|
||||
@@ -11,11 +16,18 @@ import com.correx.core.events.types.SessionId
|
||||
import com.correx.core.events.types.StageId
|
||||
|
||||
class DefaultContextPackBuilder(
|
||||
private val compressor: ContextCompressor
|
||||
private val compressor: ContextCompressor,
|
||||
private val policy: CompressionPolicy = CompressionPolicy(),
|
||||
private val classifier: ContextClassifier = ContextClassifier(),
|
||||
private val formatCompressor: FormatCompressor = FormatCompressor(),
|
||||
) : ContextPackBuilder {
|
||||
|
||||
private val neverDropSourceTypes = setOf("steeringNote", "eventHistory", "decisionJournal")
|
||||
|
||||
private companion object {
|
||||
const val CHARS_PER_TOKEN = 4
|
||||
}
|
||||
|
||||
override fun build(
|
||||
id: ContextPackId,
|
||||
sessionId: SessionId,
|
||||
@@ -27,7 +39,19 @@ class DefaultContextPackBuilder(
|
||||
// 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.
|
||||
val ordered = 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
|
||||
// lossy stages that touch freeform live further down the pipeline (prune/summarize).
|
||||
val ordered = 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
|
||||
}
|
||||
} else stamped
|
||||
val (pinned, compressible) = ordered.partition { it.sourceType in neverDropSourceTypes }
|
||||
val pinnedTokens = pinned.sumOf { it.tokenEstimate }
|
||||
var remainingTokens = (budget.limit - pinnedTokens).coerceAtLeast(0)
|
||||
@@ -73,6 +97,10 @@ 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 fun estimateTokens(text: String): Int = (text.length + (CHARS_PER_TOKEN - 1)) / CHARS_PER_TOKEN
|
||||
|
||||
private fun strategyFor(sourceType: String): CompressionStrategy = when (sourceType) {
|
||||
"toolLog" -> CompressionStrategy.ToolLog
|
||||
"artifact" -> CompressionStrategy.Artifact
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
package com.correx.core.context.compression
|
||||
|
||||
import kotlinx.serialization.Serializable
|
||||
|
||||
/**
|
||||
* The pipeline stages, additively enabled by compression level (see docs/plans/
|
||||
* correx-compression-pipeline.md §5). Level N includes every stage of levels 1..N-1.
|
||||
*/
|
||||
enum class CompressionStage {
|
||||
FORMAT_COMPRESS, // 1: json→toml, dedupe, flatten. free, deterministic
|
||||
CACHE_STATIC, // 2: kv-cache system prompt / tool defs
|
||||
TOKEN_PRUNE, // 3: LLMLingua-2 on freeform (~45%)
|
||||
TIER_SPLIT, // 4: age-tiering (recent full, older pruned)
|
||||
RECURSIVE_SUMMARIZE, // 5: tier2/3 summarization
|
||||
PROTECTED_SPAN_TAG, // 6: entity/regex protection before any lossy stage
|
||||
FACT_SHEET, // 7: standing state block, survives all tiers
|
||||
TOME_MERGE, // 8: merge near-duplicate turns
|
||||
BUDGET_ALLOCATOR_STRICT, // 9: hard per-tier budget enforcement
|
||||
}
|
||||
|
||||
/**
|
||||
* Declarative selection of which pipeline stages run. Additive: raising the level turns
|
||||
* on more stages without removing earlier ones. Level 2 (format + static cache) is the
|
||||
* always-on default — free and deterministic.
|
||||
*/
|
||||
@Serializable
|
||||
data class CompressionPolicy(val level: Int = DEFAULT_LEVEL) {
|
||||
|
||||
val stages: Set<CompressionStage>
|
||||
get() = LEVEL_STAGES.filterKeys { it <= level }.values.flatten().toSet()
|
||||
|
||||
fun enabled(stage: CompressionStage): Boolean = stage in stages
|
||||
|
||||
companion object {
|
||||
const val DEFAULT_LEVEL = 2
|
||||
|
||||
// Stage introduced *at* each level; the getter folds in all lower levels.
|
||||
private val LEVEL_STAGES: Map<Int, List<CompressionStage>> = mapOf(
|
||||
1 to listOf(CompressionStage.FORMAT_COMPRESS),
|
||||
2 to listOf(CompressionStage.CACHE_STATIC),
|
||||
3 to listOf(CompressionStage.TOKEN_PRUNE),
|
||||
4 to listOf(CompressionStage.TIER_SPLIT),
|
||||
5 to listOf(CompressionStage.RECURSIVE_SUMMARIZE),
|
||||
6 to listOf(CompressionStage.PROTECTED_SPAN_TAG),
|
||||
7 to listOf(CompressionStage.FACT_SHEET),
|
||||
8 to listOf(CompressionStage.TOME_MERGE),
|
||||
9 to listOf(CompressionStage.BUDGET_ALLOCATOR_STRICT),
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
package com.correx.core.context.compression
|
||||
|
||||
import com.correx.core.context.model.ContextEntry
|
||||
import com.correx.core.context.model.ContextLayer
|
||||
import com.correx.core.context.model.EntryRole
|
||||
|
||||
/**
|
||||
* Stage 0 of the pipeline: tag each entry as static, structured, or freeform. The class
|
||||
* decides which downstream stages are eligible — static content is cached and never pruned,
|
||||
* structured content is format-compressed but its exact values are never altered, freeform
|
||||
* content is the only class the lossy stages (prune/summarize/merge) touch.
|
||||
*/
|
||||
enum class ContextClass { STATIC, STRUCTURED, FREEFORM }
|
||||
|
||||
class ContextClassifier {
|
||||
|
||||
fun classify(entry: ContextEntry): ContextClass = when {
|
||||
entry.sourceType in STATIC_SOURCES -> ContextClass.STATIC
|
||||
entry.sourceType in STRUCTURED_SOURCES -> ContextClass.STRUCTURED
|
||||
// A pinned system directive that isn't one of the known static prompts is still
|
||||
// exact-value content — treat as structured (format-compress ok, never prune).
|
||||
entry.layer == ContextLayer.L0 && entry.role == EntryRole.SYSTEM -> ContextClass.STATIC
|
||||
entry.role == EntryRole.TOOL -> ContextClass.STRUCTURED
|
||||
else -> ContextClass.FREEFORM
|
||||
}
|
||||
|
||||
private companion object {
|
||||
val STATIC_SOURCES = setOf("systemPrompt", "toolSchema", "fewShot")
|
||||
val STRUCTURED_SOURCES = setOf("toolLog", "artifact", "config", "structured", "steeringNote")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package com.correx.core.context.compression
|
||||
|
||||
import kotlinx.serialization.json.Json
|
||||
import kotlinx.serialization.json.JsonArray
|
||||
import kotlinx.serialization.json.JsonElement
|
||||
import kotlinx.serialization.json.JsonObject
|
||||
import kotlinx.serialization.json.JsonPrimitive
|
||||
|
||||
/**
|
||||
* Stage 1: free, deterministic format-level compression for structured content. JSON is
|
||||
* flattened to dotted `key = value` lines (drops braces, quotes and repeated structural
|
||||
* punctuation the model doesn't need); non-JSON content is returned untouched (line-level
|
||||
* dedupe is already the ToolLog strategy's job). Exact values are preserved verbatim — this
|
||||
* stage is lossless w.r.t. information, it only strips format overhead (pipeline §1).
|
||||
*/
|
||||
class FormatCompressor {
|
||||
|
||||
fun compress(content: String): String {
|
||||
val trimmed = content.trim()
|
||||
if (!trimmed.startsWith("{") && !trimmed.startsWith("[")) return content
|
||||
return runCatching { json.parseToJsonElement(trimmed) }
|
||||
.getOrNull()
|
||||
?.let { flatten(it).takeIf { lines -> lines.isNotEmpty() }?.joinToString("\n") }
|
||||
?: content
|
||||
}
|
||||
|
||||
private fun flatten(element: JsonElement, prefix: String = ""): List<String> = when (element) {
|
||||
is JsonObject -> element.entries.flatMap { (k, v) -> flatten(v, join(prefix, k)) }
|
||||
is JsonArray -> element.flatMapIndexed { i, v -> flatten(v, join(prefix, i.toString())) }
|
||||
is JsonPrimitive -> listOf("${prefix.ifEmpty { "value" }} = ${element.content}")
|
||||
}
|
||||
|
||||
private fun join(prefix: String, key: String) = if (prefix.isEmpty()) key else "$prefix.$key"
|
||||
|
||||
private companion object {
|
||||
val json = Json { ignoreUnknownKeys = true }
|
||||
}
|
||||
}
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
package com.correx.core.context.compression
|
||||
|
||||
/**
|
||||
* Deterministic, regex-only extraction of load-bearing spans that no lossy stage (token
|
||||
* pruning, summarization, ToMe merge) may alter: numbers, IDs, hashes, IPs, file paths,
|
||||
* config keys and anything inside code fences or backticks. Runs FIRST, before any lossy
|
||||
* stage — this is what keeps compression from becoming a hallucination risk (pipeline §2).
|
||||
*
|
||||
* [extraEntities] are caller-supplied domain terms (service/model/domain names) protected
|
||||
* verbatim in addition to the structural patterns.
|
||||
*/
|
||||
class ProtectedSpanTagger(private val extraEntities: List<String> = emptyList()) {
|
||||
|
||||
fun protectedSpans(content: String): List<String> {
|
||||
val spans = LinkedHashSet<String>()
|
||||
// Code fences first (greedy, multi-line) so their inner content is protected whole,
|
||||
// then inline backticks, then the structural token patterns.
|
||||
for (pattern in STRUCTURAL_PATTERNS) {
|
||||
pattern.findAll(content).forEach { spans += it.value }
|
||||
}
|
||||
for (entity in extraEntities) {
|
||||
if (entity.isNotBlank() && content.contains(entity)) spans += entity
|
||||
}
|
||||
return spans.toList()
|
||||
}
|
||||
|
||||
private companion object {
|
||||
val STRUCTURAL_PATTERNS: List<Regex> = listOf(
|
||||
Regex("```.*?```", setOf(RegexOption.DOT_MATCHES_ALL)), // fenced code
|
||||
Regex("`[^`]+`"), // inline code
|
||||
Regex("""\b\d{1,3}(?:\.\d{1,3}){3}\b"""), // IPv4
|
||||
Regex("""\b0x[0-9a-fA-F]+\b"""), // hex literal
|
||||
Regex("""\b[0-9a-fA-F]{7,40}\b"""), // hashes / long hex
|
||||
Regex("""(?:[\w.-]+)?/[\w./-]+"""), // file paths
|
||||
Regex("""\b\w+(?:[._-]\w+)*\s*[:=]"""), // config keys (key: / key =)
|
||||
Regex("""-?\d+(?:\.\d+)?"""), // bare numbers
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
import com.correx.core.context.compression.CompressionPolicy
|
||||
import com.correx.core.context.compression.CompressionStage
|
||||
import com.correx.core.context.compression.ContextClass
|
||||
import com.correx.core.context.compression.ContextClassifier
|
||||
import com.correx.core.context.compression.FormatCompressor
|
||||
import com.correx.core.context.compression.ProtectedSpanTagger
|
||||
import com.correx.core.context.model.ContextEntry
|
||||
import com.correx.core.context.model.ContextLayer
|
||||
import com.correx.core.context.model.EntryRole
|
||||
import com.correx.core.events.types.ContextEntryId
|
||||
import org.junit.jupiter.api.Assertions.assertEquals
|
||||
import org.junit.jupiter.api.Assertions.assertFalse
|
||||
import org.junit.jupiter.api.Assertions.assertTrue
|
||||
import org.junit.jupiter.api.Test
|
||||
|
||||
class CompressionPipelineStagesTest {
|
||||
|
||||
private fun entry(sourceType: String, layer: ContextLayer, role: EntryRole, content: String = "x") =
|
||||
ContextEntry(
|
||||
id = ContextEntryId("e"),
|
||||
layer = layer,
|
||||
content = content,
|
||||
sourceType = sourceType,
|
||||
sourceId = "s",
|
||||
tokenEstimate = 1,
|
||||
role = role,
|
||||
)
|
||||
|
||||
@Test
|
||||
fun `levels are additive`() {
|
||||
val l1 = CompressionPolicy(1)
|
||||
assertTrue(l1.enabled(CompressionStage.FORMAT_COMPRESS))
|
||||
assertFalse(l1.enabled(CompressionStage.CACHE_STATIC))
|
||||
|
||||
val l7 = CompressionPolicy(7)
|
||||
// includes everything from 1..7
|
||||
assertTrue(l7.enabled(CompressionStage.FORMAT_COMPRESS))
|
||||
assertTrue(l7.enabled(CompressionStage.FACT_SHEET))
|
||||
assertTrue(l7.enabled(CompressionStage.PROTECTED_SPAN_TAG))
|
||||
assertFalse(l7.enabled(CompressionStage.TOME_MERGE)) // level 8
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `protected spans capture ids numbers paths and code`() {
|
||||
val tagger = ProtectedSpanTagger(extraEntities = listOf("GatedDeltaNet"))
|
||||
val spans = tagger.protectedSpans(
|
||||
"connect to 10.0.0.1 at path /etc/nginx.conf value 0xFF port = 8080 " +
|
||||
"`inline` on GatedDeltaNet"
|
||||
)
|
||||
assertTrue(spans.any { it == "10.0.0.1" }, spans.toString())
|
||||
assertTrue(spans.any { it.contains("/etc/nginx.conf") }, spans.toString())
|
||||
assertTrue(spans.any { it == "0xFF" }, spans.toString())
|
||||
assertTrue(spans.any { it == "`inline`" }, spans.toString())
|
||||
assertTrue(spans.contains("GatedDeltaNet"), spans.toString())
|
||||
assertTrue(spans.any { it == "8080" }, spans.toString())
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `format compressor flattens json and preserves exact values`() {
|
||||
val out = FormatCompressor().compress("""{"a":{"b":1},"c":"10.0.0.1"}""")
|
||||
assertEquals("a.b = 1\nc = 10.0.0.1", out)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `format compressor leaves non-json content untouched`() {
|
||||
val out = FormatCompressor().compress("log line\nlog line\nother")
|
||||
assertEquals("log line\nlog line\nother", out)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `classifier separates static structured freeform`() {
|
||||
val c = ContextClassifier()
|
||||
assertEquals(ContextClass.STATIC, c.classify(entry("systemPrompt", ContextLayer.L0, EntryRole.SYSTEM)))
|
||||
assertEquals(ContextClass.STRUCTURED, c.classify(entry("toolLog", ContextLayer.L2, EntryRole.TOOL)))
|
||||
assertEquals(ContextClass.STRUCTURED, c.classify(entry("steeringNote", ContextLayer.L2, EntryRole.SYSTEM)))
|
||||
assertEquals(ContextClass.FREEFORM, c.classify(entry("chat", ContextLayer.L1, EntryRole.USER)))
|
||||
}
|
||||
}
|
||||
@@ -37,9 +37,18 @@ object PromptRenderer {
|
||||
val conversationMessages = conversationEntries
|
||||
.sortedWith(compareBy({ it.second.ordinal }, { layerPriority(it.first) }))
|
||||
.map { (_, entry) -> entry.toChatMessage() }
|
||||
// Repetition anchoring: steering directives fold into the leading system message, far
|
||||
// from the final query — weak local models forget them (lost-in-the-middle). Restate
|
||||
// them once as a trailing user turn, where models attend strongest. Template-safe: a
|
||||
// user message at the end never trips strict system-must-be-first templates.
|
||||
val anchor = systemEntries
|
||||
.filter { it.second.sourceType == "steeringNote" }
|
||||
.joinToString("\n") { it.second.content }
|
||||
.takeIf { it.isNotBlank() }
|
||||
val messages = buildList {
|
||||
systemContent?.let { add(ChatMessage("system", it)) }
|
||||
addAll(conversationMessages)
|
||||
anchor?.let { add(ChatMessage("user", "Reminder — active steering directive(s):\n$it")) }
|
||||
}
|
||||
return messages.ifEmpty { listOf(ChatMessage("user", "")) }
|
||||
}
|
||||
|
||||
@@ -70,6 +70,33 @@ class PromptRendererOrderingTest {
|
||||
assertEquals(listOf("system" to "memory", "user" to "question"), messages.map { it.role to it.content })
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `steering directives are re-anchored as a trailing user reminder`() {
|
||||
// A steering note folds into the leading system block AND is restated at the tail so a
|
||||
// weak local model still sees the active constraint next to the final query.
|
||||
val pack = ContextPack(
|
||||
id = ContextPackId("p"),
|
||||
sessionId = sessionId,
|
||||
stageId = stageId,
|
||||
layers = mapOf(
|
||||
ContextLayer.L0 to listOf(entry("sys", ContextLayer.L0, EntryRole.SYSTEM, "systemPrompt")),
|
||||
ContextLayer.L2 to listOf(entry("directive", ContextLayer.L2, EntryRole.SYSTEM, "steeringNote")),
|
||||
ContextLayer.L1 to listOf(entry("question", ContextLayer.L1, EntryRole.USER, "chat")),
|
||||
),
|
||||
budgetUsed = 30,
|
||||
budgetLimit = 4000,
|
||||
)
|
||||
val messages = PromptRenderer.render(pack)
|
||||
assertEquals(
|
||||
listOf(
|
||||
"system" to "sys\n\ndirective",
|
||||
"user" to "question",
|
||||
"user" to "Reminder — active steering directive(s):\ndirective",
|
||||
),
|
||||
messages.map { it.role to it.content },
|
||||
)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `non-L0 SYSTEM entries fold into the single leading system message`() {
|
||||
// Strict chat templates (Qwen) 400 on any system message after index 0. Recalled L3
|
||||
|
||||
Reference in New Issue
Block a user