mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat: gated Markdown-KV compaction formatter (serialization-aware output) (#859)
Closes #858. ## What Adds an opt-in **Markdown-KV** renderer to the lossless-first compaction stage, plus the plumbing to pick a compaction formatter by name. Default behavior is unchanged (`csv-schema`). Format-comprehension benchmarks show models retrieve values from Markdown-KV substantially more reliably than from CSV (~60.7% vs ~44.3%) — token-cheapest is not the same as most comprehensible. This makes the trade-off selectable per workload. ## How - **`MarkdownKvFormatter`** (`compaction/formatter.rs`): keeps the `[N]{cols}` declaration line, renders each row as a Markdown list item with `key: value` lines. - Missing cells omitted entirely (the KV advantage over positional CSV). - Strings ambiguous on a line (newlines, leading/trailing whitespace, empty) render JSON-quoted; everything else raw — commas and quotes need no escaping. - Nested cells inline compact JSON; opaque cells keep the fixed `<<ccr:HASH,KIND,SIZE>>` marker contract shared by all formatters. - **`CompactionStage::from_format_name`** maps `"csv-schema" | "json" | "markdown-kv"` to presets. - **Core**: `SmartCrusher::with_compaction_format(config, name)` — standard OSS composition with the named formatter. - **PyO3 bridge**: `SmartCrusher.with_compaction_format(config, format_name)` staticmethod; `ValueError` on unknown names (loud, no silent fallback). - **Python**: `SmartCrusher(compaction_format=...)` kwarg, falling back to the `HEADROOM_COMPACTION_FORMAT` env var, default `"csv-schema"`. ## Safety - **Default-off**: the default constructor path still calls the Rust `new()` constructor, so byte-parity coverage stays on the exact production codepath. A test asserts default output is byte-identical to an explicit `csv-schema` opt-in. - The existing `lossless_min_savings_ratio` gate (0.30) still applies. Markdown-KV repeats field names per row, so it clears the gate less often than CSV and falls through to the lossy path — we never inline a "lossless" rendering that isn't actually smaller. - CCR marker format unchanged across formatters; downstream retrieval pattern-matching keeps working. - No user/assistant content dropped — the formatter is a pure rendering of the same Compaction IR. ## Tests - Rust: 10 new unit tests in `compaction/formatter.rs` (table/buckets rendering, missing-cell omission, string quoting, CCR markers, drop summary, byte-size sanity vs raw JSON). `cargo test -p headroom-core`: 894 passed. Clippy + fmt clean. - Python: `tests/test_compaction_markdown_kv.py` (10 tests) — bridge rendering end-to-end, name→preset parity with the default constructor, kwarg/env knob precedence, loud failure on unknown names, default-output-unchanged guarantee. Existing smart_crusher suite: 38 passed. --------- Co-authored-by: Ash Rhodes <ashley.rhodes@king.com>
This commit is contained in:
parent
c2106cbdab
commit
06b2625b17
7 changed files with 620 additions and 6 deletions
|
|
@ -1,4 +1,4 @@
|
|||
//! Formatter trait + two built-in implementations.
|
||||
//! Formatter trait + the built-in implementations.
|
||||
//!
|
||||
//! [`Formatter`] walks a [`Compaction`] tree and renders bytes. It's the
|
||||
//! pluggable seam where users (or Enterprise plugins) choose how the
|
||||
|
|
@ -14,10 +14,14 @@
|
|||
//! TOON's most useful idea (the `[N]{cols}` declaration) without
|
||||
//! adopting TOON's bespoke escaping rules — every model has seen
|
||||
//! millions of CSV examples in training.
|
||||
//! - [`MarkdownKvFormatter`] — the same `[N]{cols}` declaration +
|
||||
//! one Markdown list item per row with `key: value` lines.
|
||||
//! Token-heavier than CSV (field names repeat per row) but
|
||||
//! format-comprehension benchmarks favor KV for read-back accuracy.
|
||||
//!
|
||||
//! # Nested cells
|
||||
//!
|
||||
//! Both formatters handle [`CellValue::Nested`] by recursively
|
||||
//! The formatters handle [`CellValue::Nested`] by recursively
|
||||
//! formatting the sub-compaction and embedding the result. The CSV
|
||||
//! formatter wraps nested output in CSV-quoted form; the JSON
|
||||
//! formatter embeds it as a structured JSON object.
|
||||
|
|
@ -26,7 +30,7 @@
|
|||
//!
|
||||
//! [`CellValue::OpaqueRef`] renders as a structured marker the model
|
||||
//! can recognize: `<<ccr:HASH,KIND,SIZE>>`. This format is fixed across
|
||||
//! both built-in formatters so downstream consumers can pattern-match
|
||||
//! all built-in formatters so downstream consumers can pattern-match
|
||||
//! markers regardless of which formatter produced them.
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
|
@ -368,6 +372,200 @@ fn csv_quote(s: &str) -> String {
|
|||
out
|
||||
}
|
||||
|
||||
// ─────────────────────────── Markdown-KV formatter ───────────────────────────
|
||||
|
||||
/// Renders a `Compaction` as a `[N]{cols}` declaration followed by one
|
||||
/// Markdown list item per row, each cell on its own `key: value` line.
|
||||
///
|
||||
/// Token-heavier than [`CsvSchemaFormatter`] (field names repeat per
|
||||
/// row), but format-comprehension benchmarks show models retrieve
|
||||
/// values from Markdown-KV substantially more reliably than from CSV.
|
||||
/// Offered as an opt-in trade of tokens for read accuracy.
|
||||
///
|
||||
/// Rendering rules:
|
||||
/// - Missing cells are omitted entirely (no `key:` line) — sparse rows
|
||||
/// cost nothing, unlike CSV's positional empty cells.
|
||||
/// - Strings that would be ambiguous on a line (contain newlines,
|
||||
/// leading/trailing whitespace, or are empty) render JSON-quoted;
|
||||
/// everything else renders raw.
|
||||
/// - Nested cells render as compact inline JSON, matching
|
||||
/// [`CsvSchemaFormatter`].
|
||||
/// - Opaque cells keep the fixed `<<ccr:HASH,KIND,SIZE>>` marker
|
||||
/// contract shared by all formatters.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct MarkdownKvFormatter {
|
||||
/// If true, emit a `__dropped:N` note on the declaration line when
|
||||
/// rows were dropped under budget. Mirrors
|
||||
/// [`CsvSchemaFormatter::include_drop_summary`].
|
||||
pub include_drop_summary: bool,
|
||||
}
|
||||
|
||||
impl MarkdownKvFormatter {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
pub fn with_drop_summary(mut self) -> Self {
|
||||
self.include_drop_summary = true;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl Formatter for MarkdownKvFormatter {
|
||||
fn name(&self) -> &str {
|
||||
"markdown-kv"
|
||||
}
|
||||
|
||||
fn format(&self, c: &Compaction) -> String {
|
||||
let mut out = String::new();
|
||||
write_compaction_kv(&mut out, c, self);
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
fn write_compaction_kv(out: &mut String, c: &Compaction, fmt: &MarkdownKvFormatter) {
|
||||
match c {
|
||||
Compaction::Table {
|
||||
schema,
|
||||
rows,
|
||||
original_count,
|
||||
} => {
|
||||
write_kv_table(out, schema, rows, *original_count, fmt);
|
||||
}
|
||||
Compaction::Buckets {
|
||||
discriminator,
|
||||
buckets,
|
||||
original_count,
|
||||
} => {
|
||||
out.push_str("__buckets:");
|
||||
out.push_str(discriminator);
|
||||
if fmt.include_drop_summary {
|
||||
let kept: usize = buckets.iter().map(|b| b.rows.len()).sum();
|
||||
if kept < *original_count {
|
||||
out.push_str(&format!(" __dropped:{}", original_count - kept));
|
||||
}
|
||||
}
|
||||
out.push('\n');
|
||||
for b in buckets {
|
||||
out.push_str(&format!("__key:{}\n", kv_scalar(&b.key)));
|
||||
write_kv_table(out, &b.schema, &b.rows, b.rows.len(), fmt);
|
||||
}
|
||||
}
|
||||
Compaction::OpaqueRef {
|
||||
ccr_hash,
|
||||
byte_size,
|
||||
kind,
|
||||
} => {
|
||||
out.push_str(&format_ccr_marker(ccr_hash, *byte_size, kind));
|
||||
}
|
||||
Compaction::Untouched(v) => {
|
||||
out.push_str(&serde_json::to_string(v).unwrap_or_default());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn write_kv_table(
|
||||
out: &mut String,
|
||||
schema: &Schema,
|
||||
rows: &[Row],
|
||||
original_count: usize,
|
||||
fmt: &MarkdownKvFormatter,
|
||||
) {
|
||||
// Same declaration line as the CSV formatter: keeps row count and
|
||||
// typed shape up front where the model (and telemetry) expect it.
|
||||
// Unlike CSV (pre-existing exposure, kept byte-identical), KV quotes
|
||||
// pathological field names here so the declaration parses the same
|
||||
// way as the row lines below.
|
||||
out.push('[');
|
||||
out.push_str(&rows.len().to_string());
|
||||
out.push_str("]{");
|
||||
let col_decl: Vec<String> = schema
|
||||
.fields
|
||||
.iter()
|
||||
.map(|f| {
|
||||
let name = kv_field_name(&f.name);
|
||||
if f.nullable {
|
||||
format!("{}:{}?", name, f.type_tag)
|
||||
} else {
|
||||
format!("{}:{}", name, f.type_tag)
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
out.push_str(&col_decl.join(","));
|
||||
out.push('}');
|
||||
if fmt.include_drop_summary && rows.len() < original_count {
|
||||
out.push_str(&format!(" __dropped:{}", original_count - rows.len()));
|
||||
}
|
||||
out.push('\n');
|
||||
|
||||
for row in rows {
|
||||
// Compactor invariant: one cell per schema field. zip() would
|
||||
// silently drop extras — fail loudly in debug builds instead.
|
||||
debug_assert_eq!(row.0.len(), schema.fields.len());
|
||||
let mut wrote_first = false;
|
||||
for (field, cell) in schema.fields.iter().zip(row.0.iter()) {
|
||||
let rendered = match cell {
|
||||
CellValue::Missing => continue,
|
||||
CellValue::Scalar(v) => kv_scalar(v),
|
||||
CellValue::Nested(sub) => JsonFormatter::new().format(sub),
|
||||
CellValue::OpaqueRef {
|
||||
ccr_hash,
|
||||
byte_size,
|
||||
kind,
|
||||
} => format_ccr_marker(ccr_hash, *byte_size, kind),
|
||||
};
|
||||
out.push_str(if wrote_first { " " } else { "- " });
|
||||
out.push_str(&kv_field_name(&field.name));
|
||||
out.push_str(": ");
|
||||
out.push_str(&rendered);
|
||||
out.push('\n');
|
||||
wrote_first = true;
|
||||
}
|
||||
// All-missing row: keep a bare list item so the rendered row
|
||||
// count still matches the declaration.
|
||||
if !wrote_first {
|
||||
out.push_str("-\n");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn kv_scalar(v: &Value) -> String {
|
||||
match v {
|
||||
Value::Null => "null".to_string(),
|
||||
Value::Bool(b) => if *b { "true" } else { "false" }.to_string(),
|
||||
Value::Number(n) => n.to_string(),
|
||||
Value::String(s) => {
|
||||
if needs_kv_quote(s) {
|
||||
serde_json::to_string(s).unwrap_or_default()
|
||||
} else {
|
||||
s.clone()
|
||||
}
|
||||
}
|
||||
// Object/array fall back to compact JSON (rare — usually
|
||||
// already promoted to Nested by the compactor).
|
||||
_ => serde_json::to_string(v).unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn needs_kv_quote(s: &str) -> bool {
|
||||
s.is_empty()
|
||||
|| s.contains('\n')
|
||||
|| s.contains('\r')
|
||||
|| s.starts_with(char::is_whitespace)
|
||||
|| s.ends_with(char::is_whitespace)
|
||||
}
|
||||
|
||||
/// Field names are normally bare identifiers, but nothing upstream
|
||||
/// enforces that. Quote the pathological ones the same way as values:
|
||||
/// an embedded newline would inject fake row lines, and `": "` inside
|
||||
/// a key would split the line at the wrong colon on read-back.
|
||||
fn kv_field_name(name: &str) -> String {
|
||||
if needs_kv_quote(name) || name.contains(": ") {
|
||||
serde_json::to_string(name).unwrap_or_default()
|
||||
} else {
|
||||
name.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
|
@ -591,4 +789,173 @@ mod tests {
|
|||
raw_json.len()
|
||||
);
|
||||
}
|
||||
|
||||
// ── MarkdownKvFormatter ──
|
||||
|
||||
#[test]
|
||||
fn markdown_kv_renders_table() {
|
||||
let items = vec![
|
||||
json!({"id": 1, "name": "alice"}),
|
||||
json!({"id": 2, "name": "bob"}),
|
||||
];
|
||||
let c = compact(&items, &cfg());
|
||||
let out = MarkdownKvFormatter::new().format(&c);
|
||||
let lines: Vec<&str> = out.trim_end().lines().collect();
|
||||
assert!(lines[0].starts_with("[2]{"), "got line[0]: {}", lines[0]);
|
||||
assert!(lines[0].contains("id:int"));
|
||||
assert!(out.contains("- id: 1\n name: alice\n"), "got: {out}");
|
||||
assert!(out.contains("- id: 2\n name: bob\n"), "got: {out}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn markdown_kv_omits_missing_cells() {
|
||||
let items = vec![json!({"id": 1, "note": "has note"}), json!({"id": 2})];
|
||||
let c = compact(&items, &cfg());
|
||||
let out = MarkdownKvFormatter::new().format(&c);
|
||||
assert!(out.contains("note: has note"), "got: {out}");
|
||||
// Row 2 has no `note:` line at all.
|
||||
let row2 = out.split("- id: 2").nth(1).expect("row 2 present");
|
||||
assert!(!row2.contains("note:"), "got row2 tail: {row2}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn markdown_kv_quotes_ambiguous_strings() {
|
||||
let items = vec![
|
||||
json!({"id": 1, "msg": "line one\nline two"}),
|
||||
json!({"id": 2, "msg": "plain"}),
|
||||
];
|
||||
let c = compact(&items, &cfg());
|
||||
let out = MarkdownKvFormatter::new().format(&c);
|
||||
assert!(
|
||||
out.contains(r#"msg: "line one\nline two""#),
|
||||
"multiline must be JSON-quoted, got: {out}"
|
||||
);
|
||||
assert!(out.contains("msg: plain\n"), "got: {out}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn markdown_kv_quotes_pathological_field_names() {
|
||||
// A newline in a key would inject fake row lines; ": " in a key
|
||||
// would split read-back at the wrong colon. Both get JSON-quoted
|
||||
// in the declaration and in every row line.
|
||||
let items = vec![
|
||||
json!({"bad\nkey": 1, "note: extra": "x"}),
|
||||
json!({"bad\nkey": 2, "note: extra": "y"}),
|
||||
];
|
||||
let c = compact(&items, &cfg());
|
||||
let out = MarkdownKvFormatter::new().format(&c);
|
||||
assert!(!out.contains("bad\nkey"), "raw newline key leaked: {out}");
|
||||
assert!(out.contains(r#""bad\nkey""#), "got: {out}");
|
||||
assert!(out.contains(r#""note: extra": x"#), "got: {out}");
|
||||
let decl = out.lines().next().unwrap();
|
||||
assert!(decl.contains(r#""bad\nkey":int"#), "got decl: {decl}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn markdown_kv_plain_strings_unquoted() {
|
||||
let items = vec![
|
||||
json!({"id": 1, "name": "alice, the \"great\""}),
|
||||
json!({"id": 2, "name": "bob"}),
|
||||
];
|
||||
let c = compact(&items, &cfg());
|
||||
let out = MarkdownKvFormatter::new().format(&c);
|
||||
// Commas and quotes are fine on a KV line — no CSV-style quoting.
|
||||
assert!(out.contains(r#"name: alice, the "great""#), "got: {out}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn markdown_kv_emits_ccr_marker() {
|
||||
let big = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=".repeat(8);
|
||||
let items = vec![
|
||||
json!({"id": 1, "blob": big.clone()}),
|
||||
json!({"id": 2, "blob": big.clone()}),
|
||||
];
|
||||
let c = compact(&items, &cfg());
|
||||
let out = MarkdownKvFormatter::new().format(&c);
|
||||
assert!(out.contains("<<ccr:"), "got: {out}");
|
||||
assert!(out.contains(",base64,"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn markdown_kv_renders_buckets() {
|
||||
let items = vec![
|
||||
json!({"type": "user", "id": 1, "name": "alice"}),
|
||||
json!({"type": "user", "id": 2, "name": "bob"}),
|
||||
json!({"type": "order", "id": 99, "total": 50}),
|
||||
json!({"type": "order", "id": 100, "total": 75}),
|
||||
];
|
||||
let c = compact(&items, &cfg());
|
||||
let out = MarkdownKvFormatter::new().format(&c);
|
||||
assert!(out.starts_with("__buckets:type"), "got: {out}");
|
||||
assert!(out.contains("__key:order"));
|
||||
assert!(out.contains("__key:user"));
|
||||
assert!(out.contains("- id:"), "got: {out}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn markdown_kv_drop_summary_opt_in() {
|
||||
let rows = vec![
|
||||
Row::new(vec![CellValue::Scalar(json!(1))]),
|
||||
Row::new(vec![CellValue::Scalar(json!(2))]),
|
||||
];
|
||||
let c = Compaction::Table {
|
||||
schema: Schema {
|
||||
fields: vec![super::super::ir::FieldSpec {
|
||||
name: "x".into(),
|
||||
type_tag: "int".into(),
|
||||
nullable: false,
|
||||
}],
|
||||
},
|
||||
rows,
|
||||
original_count: 5, // 3 dropped
|
||||
};
|
||||
let with_summary = MarkdownKvFormatter::new().with_drop_summary().format(&c);
|
||||
assert!(with_summary.contains("__dropped:3"));
|
||||
let without = MarkdownKvFormatter::new().format(&c);
|
||||
assert!(!without.contains("__dropped"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn markdown_kv_nested_cell_inline_json() {
|
||||
let items = vec![
|
||||
json!({"event": "batch", "payload": r#"[{"x":1},{"x":2},{"x":3}]"#}),
|
||||
json!({"event": "batch", "payload": r#"[{"x":4},{"x":5}]"#}),
|
||||
];
|
||||
let c = compact(&items, &cfg());
|
||||
let out = MarkdownKvFormatter::new().format(&c);
|
||||
assert!(out.contains("_compaction"), "got: {out}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn markdown_kv_estimate_matches_format_len() {
|
||||
let items = vec![json!({"a": 1}), json!({"a": 2})];
|
||||
let c = compact(&items, &cfg());
|
||||
let f = MarkdownKvFormatter::new();
|
||||
assert_eq!(f.estimate_bytes(&c), f.format(&c).len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn markdown_kv_smaller_than_raw_json() {
|
||||
// KV repeats field names per row, so it loses to CSV on bytes —
|
||||
// but it should still beat naïve JSON (quotes + braces + commas).
|
||||
let items: Vec<Value> = (0..50)
|
||||
.map(|i| {
|
||||
json!({
|
||||
"id": i,
|
||||
"name": format!("user_{i}"),
|
||||
"email": format!("user_{i}@example.com"),
|
||||
"status": if i % 3 == 0 { "ok" } else { "pending" },
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
let c = compact(&items, &cfg());
|
||||
let kv_out = MarkdownKvFormatter::new().format(&c);
|
||||
let raw_json = serde_json::to_string(&Value::Array(items.clone())).unwrap();
|
||||
assert!(
|
||||
kv_out.len() < raw_json.len(),
|
||||
"kv {} bytes vs raw json {} bytes",
|
||||
kv_out.len(),
|
||||
raw_json.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ pub mod walker;
|
|||
|
||||
pub use classifier::{classify_cell, CellClass, ClassifyConfig};
|
||||
pub use compactor::{compact, CompactConfig};
|
||||
pub use formatter::{CsvSchemaFormatter, Formatter, JsonFormatter};
|
||||
pub use formatter::{CsvSchemaFormatter, Formatter, JsonFormatter, MarkdownKvFormatter};
|
||||
pub use ir::{Bucket, CellValue, Compaction, FieldSpec, OpaqueKind, Row, Schema};
|
||||
pub use walker::{
|
||||
compact_document, emit_opaque_ccr_marker, try_parse_json_container, DocumentCompactor,
|
||||
|
|
@ -69,6 +69,34 @@ impl CompactionStage {
|
|||
}
|
||||
}
|
||||
|
||||
/// Markdown-KV formatter, default config — opt-in trade of tokens
|
||||
/// for model read accuracy (field names repeat per row, but
|
||||
/// format-comprehension benchmarks favor KV over CSV).
|
||||
pub fn default_markdown_kv() -> Self {
|
||||
Self {
|
||||
config: CompactConfig::default(),
|
||||
formatter: Box::new(MarkdownKvFormatter::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Formatter names accepted by [`Self::from_format_name`]. The
|
||||
/// single source of truth for caller error messages (the PyO3
|
||||
/// bridge renders this list) — keep in sync with the match below.
|
||||
pub const SUPPORTED_FORMAT_NAMES: &'static [&'static str] =
|
||||
&["csv-schema", "json", "markdown-kv"];
|
||||
|
||||
/// Look up a preset by its formatter name (see
|
||||
/// [`Self::SUPPORTED_FORMAT_NAMES`]). `None` for unknown names —
|
||||
/// callers own the fallback/error policy.
|
||||
pub fn from_format_name(name: &str) -> Option<Self> {
|
||||
match name {
|
||||
"csv-schema" => Some(Self::default_csv_schema()),
|
||||
"json" => Some(Self::default_json()),
|
||||
"markdown-kv" => Some(Self::default_markdown_kv()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Run the stage end-to-end: compact + format. Returns the
|
||||
/// [`Compaction`] tree (so callers can inspect kept/total row
|
||||
/// counts) alongside the rendered bytes.
|
||||
|
|
|
|||
|
|
@ -179,6 +179,21 @@ impl SmartCrusher {
|
|||
.build()
|
||||
}
|
||||
|
||||
/// Construct like [`SmartCrusher::new`] but with the compaction
|
||||
/// stage's formatter chosen by name (`"csv-schema"`, `"json"`,
|
||||
/// `"markdown-kv"`). `None` for unknown names — callers own the
|
||||
/// fallback/error policy. `"csv-schema"` is equivalent to `new`.
|
||||
pub fn with_compaction_format(config: SmartCrusherConfig, format_name: &str) -> Option<Self> {
|
||||
let stage = CompactionStage::from_format_name(format_name)?;
|
||||
Some(
|
||||
SmartCrusherBuilder::new(config)
|
||||
.with_default_oss_setup()
|
||||
.with_compaction(stage)
|
||||
.with_default_ccr_store()
|
||||
.build(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Begin a builder chain for custom composition. The Enterprise
|
||||
/// entry point: swap the scorer, add business-rule constraints,
|
||||
/// attach an audit observer.
|
||||
|
|
|
|||
|
|
@ -677,6 +677,26 @@ impl PySmartCrusher {
|
|||
}
|
||||
}
|
||||
|
||||
/// Construct with the lossless-first compaction stage's formatter
|
||||
/// chosen by name: `"csv-schema"` (the `new()` default), `"json"`,
|
||||
/// or `"markdown-kv"`. Raises `ValueError` on unknown names so a
|
||||
/// misconfigured knob is visible instead of silently falling back.
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (config = None, format_name = "csv-schema"))]
|
||||
fn with_compaction_format(
|
||||
config: Option<&PySmartCrusherConfig>,
|
||||
format_name: &str,
|
||||
) -> PyResult<Self> {
|
||||
let cfg = config.map(|c| c.inner.clone()).unwrap_or_default();
|
||||
match RustSmartCrusher::with_compaction_format(cfg, format_name) {
|
||||
Some(inner) => Ok(Self { inner }),
|
||||
None => Err(pyo3::exceptions::PyValueError::new_err(format!(
|
||||
"unknown compaction format {format_name:?}; expected one of: {}",
|
||||
headroom_core::transforms::smart_crusher::compaction::CompactionStage::SUPPORTED_FORMAT_NAMES.join(", ")
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// `crush(content, query="", bias=1.0) -> CrushResult`. Argument
|
||||
/// order and keyword names mirror the Python implementation.
|
||||
///
|
||||
|
|
|
|||
|
|
@ -358,6 +358,11 @@ class SmartCrusherConfig:
|
|||
first_fraction: float = 0.3 # 30% of K from start of array
|
||||
last_fraction: float = 0.15 # 15% of K from end of array
|
||||
|
||||
# Lossless compaction only replaces the original when it saves at
|
||||
# least this byte fraction vs the (minified) input. Mirrors the
|
||||
# Rust default.
|
||||
lossless_min_savings_ratio: float = 0.30
|
||||
|
||||
|
||||
@dataclass
|
||||
class CacheOptimizerConfig:
|
||||
|
|
|
|||
|
|
@ -45,6 +45,7 @@ from __future__ import annotations
|
|||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
|
|
@ -56,6 +57,12 @@ from .base import Transform
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Lossless-compaction renderers known to the Rust core — mirrors
|
||||
# `CompactionStage::SUPPORTED_FORMAT_NAMES` in
|
||||
# `crates/headroom-core/.../compaction/mod.rs`.
|
||||
_SUPPORTED_COMPACTION_FORMATS = ("csv-schema", "json", "markdown-kv")
|
||||
|
||||
|
||||
# ─── CCR sentinel ─────────────────────────────────────────────────────────
|
||||
#
|
||||
# When SmartCrusher's lossy path drops rows, it appends a sentinel object
|
||||
|
|
@ -172,6 +179,11 @@ class SmartCrusherConfig:
|
|||
dedup_identical_items: bool = True
|
||||
first_fraction: float = 0.3
|
||||
last_fraction: float = 0.15
|
||||
# Lossless compaction only replaces the original when it saves at
|
||||
# least this byte fraction vs the (minified) input. Mirrors the Rust
|
||||
# default; mainly lowered in tests and KV experiments — KV repeats
|
||||
# field names per row, so it clears the gate less often than CSV.
|
||||
lossless_min_savings_ratio: float = 0.30
|
||||
|
||||
|
||||
# ─── Rust-backed SmartCrusher ─────────────────────────────────────────────
|
||||
|
|
@ -198,6 +210,7 @@ class SmartCrusher(Transform):
|
|||
ccr_config: CCRConfig | None = None,
|
||||
with_compaction: bool = True,
|
||||
observer: Any = None,
|
||||
compaction_format: str | None = None,
|
||||
):
|
||||
# Hard import — no Python fallback. If the wheel is missing the
|
||||
# caller must build it (scripts/build_rust_extension.sh) or
|
||||
|
|
@ -303,6 +316,7 @@ class SmartCrusher(Transform):
|
|||
dedup_identical_items=cfg.dedup_identical_items,
|
||||
first_fraction=cfg.first_fraction,
|
||||
last_fraction=cfg.last_fraction,
|
||||
lossless_min_savings_ratio=cfg.lossless_min_savings_ratio,
|
||||
relevance_threshold=0.3,
|
||||
enable_ccr_marker=(
|
||||
self._ccr_config.enabled and self._ccr_config.inject_retrieval_marker
|
||||
|
|
@ -314,10 +328,34 @@ class SmartCrusher(Transform):
|
|||
# markers. Pass `with_compaction=False` to opt into the
|
||||
# pre-PR4 lossy-only path (used by retention-property tests
|
||||
# that depend on row-level item preservation).
|
||||
if with_compaction:
|
||||
#
|
||||
# `compaction_format` picks the lossless renderer:
|
||||
# "csv-schema" (default), "json", or "markdown-kv" (opt-in
|
||||
# trade of tokens for model read accuracy). Falls back to the
|
||||
# HEADROOM_COMPACTION_FORMAT env var when the kwarg is None.
|
||||
# Ignored when with_compaction=False.
|
||||
resolved_format = compaction_format or os.environ.get(
|
||||
"HEADROOM_COMPACTION_FORMAT", "csv-schema"
|
||||
)
|
||||
# Validate even when with_compaction=False: an explicit bogus
|
||||
# format (kwarg or env var) is a misconfiguration that should be
|
||||
# visible, not silently accepted because the knob happens to be
|
||||
# ignored on this path.
|
||||
if resolved_format not in _SUPPORTED_COMPACTION_FORMATS:
|
||||
raise ValueError(
|
||||
f"unknown compaction format {resolved_format!r}; "
|
||||
f"expected one of: {', '.join(_SUPPORTED_COMPACTION_FORMATS)}"
|
||||
)
|
||||
self._compaction_format = resolved_format if with_compaction else None
|
||||
if not with_compaction:
|
||||
self._rust = _RustSmartCrusher.without_compaction(rust_cfg)
|
||||
elif resolved_format == "csv-schema":
|
||||
# Keep the `new()` constructor for the default path so its
|
||||
# byte-parity coverage stays on the exact production
|
||||
# codepath.
|
||||
self._rust = _RustSmartCrusher(rust_cfg)
|
||||
else:
|
||||
self._rust = _RustSmartCrusher.without_compaction(rust_cfg)
|
||||
self._rust = _RustSmartCrusher.with_compaction_format(rust_cfg, resolved_format)
|
||||
|
||||
def crush(self, content: str, query: str = "", bias: float = 1.0) -> CrushResult:
|
||||
"""Crush a single JSON content string.
|
||||
|
|
|
|||
141
tests/test_compaction_markdown_kv.py
Normal file
141
tests/test_compaction_markdown_kv.py
Normal file
|
|
@ -0,0 +1,141 @@
|
|||
"""Markdown-KV compaction formatter — opt-in serialization-aware output.
|
||||
|
||||
Covers the plumbing added for issue #858:
|
||||
|
||||
- ``headroom._core.SmartCrusher.with_compaction_format`` renders uniform
|
||||
arrays as Markdown-KV when the lossless gate passes.
|
||||
- The high-level ``SmartCrusher`` exposes the knob via the
|
||||
``compaction_format`` kwarg and the ``HEADROOM_COMPACTION_FORMAT`` env
|
||||
var, defaulting to the unchanged ``csv-schema`` path.
|
||||
- Unknown format names fail loudly (``ValueError``) instead of silently
|
||||
falling back.
|
||||
|
||||
The formatter's rendering rules themselves (missing-cell omission,
|
||||
string quoting, CCR marker contract, buckets) are covered by the Rust
|
||||
unit tests in ``compaction/formatter.rs``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom._core import SmartCrusher as RustSmartCrusher
|
||||
from headroom._core import SmartCrusherConfig as RustSmartCrusherConfig
|
||||
from headroom.transforms.smart_crusher import SmartCrusher
|
||||
from headroom.transforms.smart_crusher import SmartCrusherConfig as PySmartCrusherConfig
|
||||
|
||||
|
||||
def _tabular_json(n: int = 50) -> str:
|
||||
return json.dumps(
|
||||
[
|
||||
{
|
||||
"id": i,
|
||||
"name": f"user_{i}",
|
||||
"email": f"user_{i}@example.com",
|
||||
"status": "ok" if i % 3 == 0 else "pending",
|
||||
}
|
||||
for i in range(n)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
# ── Rust bridge: with_compaction_format ──
|
||||
|
||||
|
||||
def test_markdown_kv_renders_kv_lines() -> None:
|
||||
# Lower the lossless gate: Markdown-KV repeats field names per row,
|
||||
# so its savings vs raw JSON are real but below the 30% default.
|
||||
cfg = RustSmartCrusherConfig(lossless_min_savings_ratio=0.01)
|
||||
crusher = RustSmartCrusher.with_compaction_format(cfg, "markdown-kv")
|
||||
result = crusher.crush(_tabular_json(), "", 1.0)
|
||||
assert result.was_modified
|
||||
assert "lossless" in result.strategy
|
||||
# Columns are schema-sorted, so `email` opens each row and `id`
|
||||
# renders as a continuation line.
|
||||
assert "- email: user_0@example.com" in result.compressed
|
||||
assert "id: 0" in result.compressed
|
||||
assert "name: user_1" in result.compressed
|
||||
# Declaration line survives (shared with the CSV formatter).
|
||||
assert "[50]{" in result.compressed
|
||||
|
||||
|
||||
def test_csv_schema_format_name_matches_default() -> None:
|
||||
cfg = RustSmartCrusherConfig(lossless_min_savings_ratio=0.01)
|
||||
via_name = RustSmartCrusher.with_compaction_format(cfg, "csv-schema")
|
||||
via_default = RustSmartCrusher(cfg)
|
||||
content = _tabular_json()
|
||||
assert (
|
||||
via_name.crush(content, "", 1.0).compressed
|
||||
== via_default.crush(content, "", 1.0).compressed
|
||||
)
|
||||
|
||||
|
||||
def test_unknown_format_name_raises() -> None:
|
||||
with pytest.raises(ValueError, match="markdown-kv"):
|
||||
RustSmartCrusher.with_compaction_format(None, "toon")
|
||||
|
||||
|
||||
# ── High-level SmartCrusher knob ──
|
||||
|
||||
|
||||
def test_default_format_is_csv_schema() -> None:
|
||||
crusher = SmartCrusher()
|
||||
assert crusher._compaction_format == "csv-schema"
|
||||
|
||||
|
||||
def test_kwarg_opts_into_markdown_kv() -> None:
|
||||
crusher = SmartCrusher(compaction_format="markdown-kv")
|
||||
assert crusher._compaction_format == "markdown-kv"
|
||||
|
||||
|
||||
def test_env_var_opts_in(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_COMPACTION_FORMAT", "markdown-kv")
|
||||
crusher = SmartCrusher()
|
||||
assert crusher._compaction_format == "markdown-kv"
|
||||
|
||||
|
||||
def test_kwarg_overrides_env_var(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_COMPACTION_FORMAT", "json")
|
||||
crusher = SmartCrusher(compaction_format="markdown-kv")
|
||||
assert crusher._compaction_format == "markdown-kv"
|
||||
|
||||
|
||||
def test_unknown_format_kwarg_raises() -> None:
|
||||
with pytest.raises(ValueError):
|
||||
SmartCrusher(compaction_format="bogus")
|
||||
|
||||
|
||||
def test_without_compaction_ignores_format() -> None:
|
||||
crusher = SmartCrusher(with_compaction=False, compaction_format="markdown-kv")
|
||||
assert crusher._compaction_format is None
|
||||
|
||||
|
||||
def test_without_compaction_still_validates_format() -> None:
|
||||
# An explicit bogus format is a misconfiguration even when the knob
|
||||
# is ignored on this path — fail loudly, don't silently accept.
|
||||
with pytest.raises(ValueError, match="bogus"):
|
||||
SmartCrusher(with_compaction=False, compaction_format="bogus")
|
||||
|
||||
|
||||
def test_end_to_end_crush_emits_markdown_kv() -> None:
|
||||
# Through the high-level Python SmartCrusher, not the Rust bridge:
|
||||
# proves the kwarg changes crush() output, not just the stored
|
||||
# attribute. Same lowered gate as the bridge test — KV's savings on
|
||||
# minified JSON sit below the 30% default.
|
||||
config = PySmartCrusherConfig(lossless_min_savings_ratio=0.01)
|
||||
crusher = SmartCrusher(config=config, compaction_format="markdown-kv")
|
||||
result = crusher.crush(_tabular_json())
|
||||
assert result.was_modified
|
||||
assert "- email: user_0@example.com" in result.compressed
|
||||
assert "[50]{" in result.compressed
|
||||
|
||||
|
||||
def test_default_output_unchanged_by_feature() -> None:
|
||||
# The default constructor path must stay byte-identical to an
|
||||
# explicit csv-schema opt-in — proves the gate is truly default-off.
|
||||
content = _tabular_json()
|
||||
default_out = SmartCrusher().crush(content)
|
||||
explicit_out = SmartCrusher(compaction_format="csv-schema").crush(content)
|
||||
assert default_out.compressed == explicit_out.compressed
|
||||
Loading…
Add table
Add a link
Reference in a new issue