mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
## Description On a cold-start large context, kompress (ModernBERT ONNX) runs **synchronously on the request thread** — ~200–300s for ~1M tokens. It blows the 30s compression budget, leaks a non-preemptible worker, and cascades (executor saturation → queue timeouts on healthy requests); on timeout the request is forwarded **uncompressed** after eating 30s. This adds four layered, **default-off, fail-open** mitigations so the request path is never blocked on ML compression. Closes #1171 ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [x] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - **Phase 0 — size gate** (`HEADROOM_KOMPRESS_MAX_TOKENS`, default 50000): route oversized text away from ModernBERT (→ LogCompressor / TextCrusher / passthrough) at the single `_try_ml_compressor` boundary. - **Phase 1 — cooperative deadline** (`HEADROOM_COMPRESSION_DEADLINE_MS`, default 20000): any kompress run self-terminates at the next chunk boundary past the budget, keeping the unprocessed tail verbatim. - **Phase 2 — TextCrusher** (`HEADROOM_TEXT_CRUSHER`): a new **native Rust** extractive prose compressor in `crates/headroom-core/src/transforms/text_crusher/`, exposed via PyO3 as `headroom._core.TextCrusher` with a thin Python wrapper. It **reuses the shared `crate::relevance::BM25Scorer`** rather than reimplementing BM25, and ships record/replay parity fixtures (mirroring the SmartCrusher Rust-core + Python-shim pattern). - **Phase 3 — off-path compression** (`HEADROOM_BACKGROUND_COMPRESSION`): forward uncompressed immediately and compress in a per-process background drain; a byte-identical cache hit on a later turn means the request never blocks on ML. - Benchmark (`benchmarks/text_crusher_quality_eval.py`), CHANGELOG entry, and docstrings documenting the fail-open limits. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy`, new modules) - [x] New tests added for new functionality - [ ] Manual testing performed (Phase 0/1 gate-fire + deadline observed on real traffic in earlier iterations; Phase 3 off-path is unit- + byte-identity-tested, not yet live-validated) ### Test Output ```text $ pytest tests/test_transforms/ tests/test_cache/ \ tests/test_proxy/test_background_compression.py tests/test_proxy/test_phase3_byte_identity.py -q 501 passed, 37 skipped in 40.33s $ cargo test -p headroom-core --lib text_crusher test result: ok. 3 passed; 0 failed; 0 ignored; 0 measured; 834 filtered out $ ruff check <changed files> All checks passed! $ mypy headroom/proxy/background_compression.py headroom/transforms/text_crusher.py Success: no issues found in 2 source files ``` New coverage: size-gate incl. the strategy-dispatch funnel (KOMPRESS + TEXT); partial-run deadline (chunk-0 compressed + chunk-1 verbatim tail); BackgroundCompressor (dedup / queue-full / fail-open); Phase 3 byte-identity round-trip; TextCrusher unit + parity. ## Real Behavior Proof - Environment: macOS, local dev — `uv` venv, Rust `_core` built via `uv pip install -e .`. - Exact command / steps: the `pytest` / `cargo test` / `ruff` / `mypy` commands shown under Test Output; quality eval `python benchmarks/text_crusher_quality_eval.py /tmp/squad_dev.json`. - Observed result: 501 Python + 3 Rust tests pass; ruff + mypy clean on changed/new modules. Quality eval: TextCrusher keeps ~94% of buried SQuAD answers at 30% size vs ~36% truncate/random; self-contained speed run ~333k words in ~76ms (one O(n) pass) — sub-second where ModernBERT takes minutes (fast-vs-slow contrast, not a same-input run). - Not tested: Phase 3 off-path on live traffic; multi-worker (per-process by design — see Additional Notes). ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes - **All four features are off by default and fail-open** — with the env flags unset the paths are no-ops for realistic inputs; on any error the request is forwarded (compressed if possible, else verbatim), never dropped. A full background queue / duplicate key surfaces as `deferred:dropped`. - **Known limits (documented in `background_compression.py`):** Phase 3 is per-process, in-memory, and token-mode-only — these are **lost-savings, never lost-correctness**, and consistent with the project's existing per-process compression cache + sticky-session multi-worker model. The startup multi-worker warning now names off-path background compression. - Phase 2 reuses the existing BM25 scorer; reuse did not improve answer-retention over a Python prototype (query-awareness dominates) — its value is the Rust speed + repo-conventional Rust-core/Python-shim shape. --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
c7f75b27e9
commit
6c68ff4e9f
28 changed files with 1855 additions and 14 deletions
16
.codegraph/.gitignore
vendored
Normal file
16
.codegraph/.gitignore
vendored
Normal file
|
|
@ -0,0 +1,16 @@
|
||||||
|
# CodeGraph data files
|
||||||
|
# These are local to each machine and should not be committed
|
||||||
|
|
||||||
|
# Database
|
||||||
|
*.db
|
||||||
|
*.db-wal
|
||||||
|
*.db-shm
|
||||||
|
|
||||||
|
# Cache
|
||||||
|
cache/
|
||||||
|
|
||||||
|
# Logs
|
||||||
|
*.log
|
||||||
|
|
||||||
|
# Hook markers
|
||||||
|
.dirty
|
||||||
|
|
@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||||
|
|
||||||
* **learn:** weight loops in `headroom learn`. A new loop detector (`headroom/learn/loops.py`) recognizes repeated tool-call patterns — including RTK re-fetch loops, where RTK's output truncation makes the agent re-run larger-limit variants of a *successful* command — collapses output-limit variants to one signature, measures the wasted tokens, surfaces loops as a highest-priority digest section, and weights loop guardrails above one-off rules by their measured waste. Previously loops had no special weight and a no-failure re-fetch loop was skipped entirely. Adds an RTK-loop eval (`benchmarks/rtk_loop_learn_eval.py`) that reproduces a loop, runs it through Learn, and asserts the generated guardrail ranks first and prevents re-triggering.
|
* **learn:** weight loops in `headroom learn`. A new loop detector (`headroom/learn/loops.py`) recognizes repeated tool-call patterns — including RTK re-fetch loops, where RTK's output truncation makes the agent re-run larger-limit variants of a *successful* command — collapses output-limit variants to one signature, measures the wasted tokens, surfaces loops as a highest-priority digest section, and weights loop guardrails above one-off rules by their measured waste. Previously loops had no special weight and a no-failure re-fetch loop was skipped entirely. Adds an RTK-loop eval (`benchmarks/rtk_loop_learn_eval.py`) that reproduces a loop, runs it through Learn, and asserts the generated guardrail ranks first and prevents re-triggering.
|
||||||
* **learn:** write per-project learnings to the personal, gitignored `CLAUDE.local.md` by default instead of the team-shared `CLAUDE.md`, matching Claude Code's memory convention so machine-specific paths and tool-discovery byproducts no longer pollute the shared file. Adds a `--target` flag to override the destination (e.g. `--target CLAUDE.md` to opt back into the shared file, or any custom path), and auto-migrates a stale learned-patterns block out of an existing `CLAUDE.md` into `CLAUDE.local.md` with a warning ([#1072](https://github.com/chopratejas/headroom/issues/1072)).
|
* **learn:** write per-project learnings to the personal, gitignored `CLAUDE.local.md` by default instead of the team-shared `CLAUDE.md`, matching Claude Code's memory convention so machine-specific paths and tool-discovery byproducts no longer pollute the shared file. Adds a `--target` flag to override the destination (e.g. `--target CLAUDE.md` to opt back into the shared file, or any custom path), and auto-migrates a stale learned-patterns block out of an existing `CLAUDE.md` into `CLAUDE.local.md` with a warning ([#1072](https://github.com/chopratejas/headroom/issues/1072)).
|
||||||
|
* **proxy/transforms:** take large cold-start contexts off the synchronous kompress path — the root cause behind the `compression_first_stage` 30s-timeout + leaked-thread → executor-saturation cascade ([#1171](https://github.com/chopratejas/headroom/issues/1171)). A token size-gate inside the ML boundary routes oversized text away from ModernBERT (`HEADROOM_KOMPRESS_MAX_TOKENS`); a cooperative chunk-deadline bounds any kompress run that does proceed (`HEADROOM_COMPRESSION_DEADLINE_MS`); an opt-in off-path mode forwards uncompressed immediately and compresses in a single per-process background drain so the request never blocks on ML (`HEADROOM_BACKGROUND_COMPRESSION`); and a new native `TextCrusher` — a fast deterministic extractive prose compressor in `headroom._core` that reuses the shared BM25 relevance scorer — is the fast alternative to ModernBERT for large plain text (`HEADROOM_TEXT_CRUSHER`). All default off and fail-open. On a SQuAD answer-retention eval (requires the SQuAD dev set) TextCrusher keeps ~94% of buried answers at 30% size vs ~36% for truncate/random, and runs in one O(n) pass -- sub-second where ModernBERT takes minutes (self-contained speed benchmark in `benchmarks/text_crusher_quality_eval.py`).
|
||||||
* **proxy:** measure and surface rolling and current token throughput metrics (active/wall-clock input, compression, effective forward, and streamed generation) in `headroom perf` CLI and the dashboard ([#959](https://github.com/chopratejas/headroom/issues/959)).
|
* **proxy:** measure and surface rolling and current token throughput metrics (active/wall-clock input, compression, effective forward, and streamed generation) in `headroom perf` CLI and the dashboard ([#959](https://github.com/chopratejas/headroom/issues/959)).
|
||||||
* **vibe:** add Mistral Vibe CLI support with `headroom wrap vibe`.
|
* **vibe:** add Mistral Vibe CLI support with `headroom wrap vibe`.
|
||||||
* **proxy:** per-project savings breakdown on the dashboard for all wrapped agents — Claude Code, Codex, aider, Copilot, and Cursor ([#802](https://github.com/chopratejas/headroom/issues/802)). `headroom wrap claude`/`codex` tag requests with an `X-Headroom-Project` header (launch-directory name); `wrap aider`/`copilot`/`cursor` — whose clients cannot send custom headers — use a `/p/<name>` base-URL prefix the proxy strips. Savings are aggregated per project (persisted, schema v3 with transparent v2 migration), exposed as `savings.per_project` in `/stats` and `projects` in `/stats-history`, and shown in a Per-Project Savings dashboard table.
|
* **proxy:** per-project savings breakdown on the dashboard for all wrapped agents — Claude Code, Codex, aider, Copilot, and Cursor ([#802](https://github.com/chopratejas/headroom/issues/802)). `headroom wrap claude`/`codex` tag requests with an `X-Headroom-Project` header (launch-directory name); `wrap aider`/`copilot`/`cursor` — whose clients cannot send custom headers — use a `/p/<name>` base-URL prefix the proxy strips. Savings are aggregated per project (persisted, schema v3 with transparent v2 migration), exposed as `savings.per_project` in `/stats` and `projects` in `/stats-history`, and shown in a Per-Project Savings dashboard table.
|
||||||
|
|
|
||||||
221
benchmarks/text_crusher_quality_eval.py
Normal file
221
benchmarks/text_crusher_quality_eval.py
Normal file
|
|
@ -0,0 +1,221 @@
|
||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Quality eval for TextCrusher (Phase 2, #1171): does extractive compression
|
||||||
|
preserve the answer-bearing content? No LLM/API calls -- fully local.
|
||||||
|
|
||||||
|
Part A -- SQuAD answer-retention (the strong, labeled metric): bury a real QA
|
||||||
|
answer in a haystack of distractor paragraphs, compress to a target ratio, and
|
||||||
|
measure whether the gold answer SURVIVES. TextCrusher (query-aware) vs truncate
|
||||||
|
(keep-recent) vs random baselines. Mirrors kompress's published
|
||||||
|
must_keep_recall (0.977 on its own labeled set).
|
||||||
|
|
||||||
|
Part B -- real-transcript fidelity: compress large text blocks from a real
|
||||||
|
Claude Code transcript (ANONYMIZED), measuring ratio, speed, and salient-token
|
||||||
|
retention (identifiers/numbers/errors -- the must-keep info in coding contexts).
|
||||||
|
Only aggregate metrics are printed; raw content is never echoed.
|
||||||
|
|
||||||
|
Usage: python benchmarks/text_crusher_quality_eval.py [squad_dev.json] [transcript.jsonl]
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import glob
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import random
|
||||||
|
import re
|
||||||
|
import sys
|
||||||
|
import time
|
||||||
|
|
||||||
|
from headroom.transforms.text_crusher import TextCrusher
|
||||||
|
|
||||||
|
_SEG = re.compile(r"(?<=[.!?])\s+|\n+")
|
||||||
|
_SALIENT = re.compile(
|
||||||
|
r"\b(?:error|exception|fail(?:ed|ure)?|warning|traceback|assert|todo|fixme)\b"
|
||||||
|
r"|\b[A-Z]{2,}\b|\b[A-Za-z_][A-Za-z0-9_]*\.[A-Za-z_][A-Za-z0-9_]*\b|\b\d+\b"
|
||||||
|
)
|
||||||
|
|
||||||
|
# --- anonymization (脱敏): scrub before any processing; never echo raw content ---
|
||||||
|
_REDACT = [
|
||||||
|
(re.compile(r"/Users/[^/\s]+"), "/Users/USER"),
|
||||||
|
(re.compile(r"\b[\w.+-]+@[\w-]+\.[\w.-]+\b"), "EMAIL"),
|
||||||
|
(re.compile(r"\b(?:sk|pk|ghp|gho|xox[baprs])-[A-Za-z0-9_-]{10,}\b"), "TOKEN"),
|
||||||
|
(re.compile(r"\bBearer\s+[A-Za-z0-9._-]{10,}"), "Bearer TOKEN"),
|
||||||
|
(re.compile(r"\b[A-Fa-f0-9]{40,}\b"), "HEX"),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def anon(t: str) -> str:
|
||||||
|
for rx, rep in _REDACT:
|
||||||
|
t = rx.sub(rep, t)
|
||||||
|
return t
|
||||||
|
|
||||||
|
|
||||||
|
def norm(s: str) -> str:
|
||||||
|
return re.sub(r"\s+", " ", s.lower()).strip()
|
||||||
|
|
||||||
|
|
||||||
|
def _segs(text: str) -> list[str]:
|
||||||
|
return [s for s in _SEG.split(text) if s.strip()]
|
||||||
|
|
||||||
|
|
||||||
|
def truncate_keep_last(text: str, ratio: float) -> str:
|
||||||
|
segs = _segs(text)
|
||||||
|
budget = int(sum(len(s) for s in segs) * ratio)
|
||||||
|
kept: list[str] = []
|
||||||
|
c = 0
|
||||||
|
for s in reversed(segs):
|
||||||
|
if c >= budget:
|
||||||
|
break
|
||||||
|
kept.append(s)
|
||||||
|
c += len(s)
|
||||||
|
return "\n".join(reversed(kept))
|
||||||
|
|
||||||
|
|
||||||
|
def random_keep(text: str, ratio: float, seed: int) -> str:
|
||||||
|
segs = _segs(text)
|
||||||
|
idx = list(range(len(segs)))
|
||||||
|
random.Random(seed).shuffle(idx)
|
||||||
|
budget = int(sum(len(s) for s in segs) * ratio)
|
||||||
|
kept: set[int] = set()
|
||||||
|
c = 0
|
||||||
|
for i in idx:
|
||||||
|
if c >= budget:
|
||||||
|
break
|
||||||
|
kept.add(i)
|
||||||
|
c += len(segs[i])
|
||||||
|
return "\n".join(segs[i] for i in sorted(kept))
|
||||||
|
|
||||||
|
|
||||||
|
def eval_squad(path: str, n: int = 200, n_distract: int = 40, ratio: float = 0.3, seed: int = 0):
|
||||||
|
data = json.load(open(path))
|
||||||
|
paras = [(p["context"], p["qas"]) for a in data["data"] for p in a["paragraphs"]]
|
||||||
|
all_ctx = [c for c, _ in paras]
|
||||||
|
examples = [
|
||||||
|
(ctx, qas[0]["question"], qas[0]["answers"][0]["text"])
|
||||||
|
for ctx, qas in paras
|
||||||
|
if qas and qas[0]["answers"]
|
||||||
|
]
|
||||||
|
rnd = random.Random(seed)
|
||||||
|
rnd.shuffle(examples)
|
||||||
|
examples = examples[:n]
|
||||||
|
tc = TextCrusher()
|
||||||
|
hit = {"text_crusher": 0, "truncate": 0, "random": 0}
|
||||||
|
tc_ratios: list[float] = []
|
||||||
|
for gold_ctx, q, ans in examples:
|
||||||
|
docs = rnd.sample(all_ctx, n_distract) + [gold_ctx]
|
||||||
|
rnd.shuffle(docs)
|
||||||
|
haystack = "\n\n".join(docs)
|
||||||
|
a = norm(ans)
|
||||||
|
out_tc = tc.compress(haystack, context=q, target_ratio=ratio).compressed
|
||||||
|
hit["text_crusher"] += a in norm(out_tc)
|
||||||
|
hit["truncate"] += a in norm(truncate_keep_last(haystack, ratio))
|
||||||
|
hit["random"] += a in norm(random_keep(haystack, ratio, seed))
|
||||||
|
tc_ratios.append(len(out_tc) / max(1, len(haystack)))
|
||||||
|
nn = len(examples)
|
||||||
|
print(
|
||||||
|
f"\n=== Part A: SQuAD answer-retention (n={nn}, distractors={n_distract}, target_ratio={ratio}) ==="
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
f" TextCrusher (query-aware): {hit['text_crusher'] / nn:6.1%} answer survives compression"
|
||||||
|
)
|
||||||
|
print(f" Truncate (keep recent): {hit['truncate'] / nn:6.1%}")
|
||||||
|
print(f" Random keep: {hit['random'] / nn:6.1%}")
|
||||||
|
print(
|
||||||
|
f" TextCrusher mean char-ratio: {sum(tc_ratios) / nn:.2f} (kept ~{sum(tc_ratios) / nn:.0%} of bytes)"
|
||||||
|
)
|
||||||
|
print(" reference: kompress published must_keep_recall = 0.977 (its own labeled set)")
|
||||||
|
|
||||||
|
|
||||||
|
def _block_texts(jsonl_path: str, min_words: int, limit: int) -> list[str]:
|
||||||
|
out: list[str] = []
|
||||||
|
with open(jsonl_path) as fh:
|
||||||
|
for line in fh:
|
||||||
|
try:
|
||||||
|
o = json.loads(line)
|
||||||
|
except json.JSONDecodeError:
|
||||||
|
continue
|
||||||
|
m = o.get("message") or {}
|
||||||
|
c = m.get("content")
|
||||||
|
parts = (
|
||||||
|
[c]
|
||||||
|
if isinstance(c, str)
|
||||||
|
else [
|
||||||
|
p["text"] for p in c if isinstance(p, dict) and isinstance(p.get("text"), str)
|
||||||
|
]
|
||||||
|
if isinstance(c, list)
|
||||||
|
else []
|
||||||
|
)
|
||||||
|
for t in parts:
|
||||||
|
if len(t.split()) >= min_words:
|
||||||
|
out.append(anon(t))
|
||||||
|
if len(out) >= limit:
|
||||||
|
break
|
||||||
|
return out[:limit]
|
||||||
|
|
||||||
|
|
||||||
|
def eval_transcript(jsonl_path: str, ratio: float = 0.4, min_words: int = 1500, limit: int = 40):
|
||||||
|
blocks = _block_texts(jsonl_path, min_words, limit)
|
||||||
|
if not blocks:
|
||||||
|
print(
|
||||||
|
f"\n=== Part B: no text blocks >= {min_words} words in {os.path.basename(jsonl_path)} ==="
|
||||||
|
)
|
||||||
|
return
|
||||||
|
tc = TextCrusher()
|
||||||
|
ratios: list[float] = []
|
||||||
|
times: list[float] = []
|
||||||
|
retentions: list[float] = []
|
||||||
|
for b in blocks:
|
||||||
|
sal_before = set(_SALIENT.findall(b))
|
||||||
|
t0 = time.perf_counter()
|
||||||
|
out = tc.compress(b, target_ratio=ratio).compressed
|
||||||
|
times.append((time.perf_counter() - t0) * 1000)
|
||||||
|
sal_after = set(_SALIENT.findall(out))
|
||||||
|
retentions.append(len(sal_before & sal_after) / max(1, len(sal_before)))
|
||||||
|
ratios.append(len(out.split()) / max(1, len(b.split())))
|
||||||
|
n = len(blocks)
|
||||||
|
print(
|
||||||
|
f"\n=== Part B: real transcript fidelity (n={n} large blocks, anonymized, target_ratio={ratio}) ==="
|
||||||
|
)
|
||||||
|
print(f" mean token-ratio kept: {sum(ratios) / n:.2f}")
|
||||||
|
print(f" mean speed: {sum(times) / n:.1f} ms/block")
|
||||||
|
print(
|
||||||
|
f" salient-token retention: {sum(retentions) / n:6.1%} (identifiers/numbers/errors kept)"
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
f" -> keeps salient info at {sum(retentions) / n:.0%} while dropping to {sum(ratios) / n:.0%} of tokens"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def eval_speed(scale_words: int = 250_000):
|
||||||
|
# Reproducible throughput on a large synthetic prose block (no external data).
|
||||||
|
text = " ".join(
|
||||||
|
f"Sentence {i} discusses subsystem {i} and its failure mode {i % 7} in detail."
|
||||||
|
for i in range(scale_words // 9)
|
||||||
|
)
|
||||||
|
nwords = len(text.split())
|
||||||
|
tc = TextCrusher()
|
||||||
|
t0 = time.perf_counter()
|
||||||
|
out = tc.compress(text, target_ratio=0.3)
|
||||||
|
ms = (time.perf_counter() - t0) * 1000
|
||||||
|
print(f"\n=== Part C: speed (synthetic, {nwords:,} words, fully reproducible) ===")
|
||||||
|
print(f" TextCrusher compress: {ms:.0f} ms ({nwords / max(ms / 1000, 1e-6):,.0f} words/sec)")
|
||||||
|
print(f" kept ratio: {out.compressed_tokens / max(1, out.original_tokens):.2f}")
|
||||||
|
print(" reference: kompress (ModernBERT ONNX) ~272s for ~1M tokens (measured, query-blind)")
|
||||||
|
print(" -> fast-vs-slow CONTRAST, not a same-input side-by-side run")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
eval_speed()
|
||||||
|
squad = sys.argv[1] if len(sys.argv) > 1 else "/tmp/squad_dev.json"
|
||||||
|
tx = sys.argv[2] if len(sys.argv) > 2 else None
|
||||||
|
if os.path.exists(squad):
|
||||||
|
eval_squad(squad)
|
||||||
|
else:
|
||||||
|
print(f"SQuAD not found at {squad}; skipping Part A")
|
||||||
|
if tx is None:
|
||||||
|
found = glob.glob(os.path.expanduser("~/.claude/projects/*headroom*/*.jsonl"))
|
||||||
|
tx = max(found, key=os.path.getsize) if found else None
|
||||||
|
if tx and os.path.exists(tx):
|
||||||
|
eval_transcript(tx)
|
||||||
|
else:
|
||||||
|
print("no transcript jsonl found; skipping Part B")
|
||||||
|
|
@ -29,6 +29,7 @@ pub mod safety;
|
||||||
pub mod search_compressor;
|
pub mod search_compressor;
|
||||||
pub mod smart_crusher;
|
pub mod smart_crusher;
|
||||||
pub mod tag_protector;
|
pub mod tag_protector;
|
||||||
|
pub mod text_crusher;
|
||||||
pub mod unidiff_detector;
|
pub mod unidiff_detector;
|
||||||
|
|
||||||
pub use content_detector::{
|
pub use content_detector::{
|
||||||
|
|
@ -61,4 +62,5 @@ pub use search_compressor::{
|
||||||
SearchCompressorStats, SearchMatch,
|
SearchCompressorStats, SearchMatch,
|
||||||
};
|
};
|
||||||
pub use tag_protector::{is_known_html_tag, protect_tags, restore_tags, ProtectStats};
|
pub use tag_protector::{is_known_html_tag, protect_tags, restore_tags, ProtectStats};
|
||||||
|
pub use text_crusher::{TextCrusher, TextCrusherConfig, TextCrusherResult};
|
||||||
pub use unidiff_detector::{detect_diff, is_diff};
|
pub use unidiff_detector::{detect_diff, is_diff};
|
||||||
|
|
|
||||||
34
crates/headroom-core/src/transforms/text_crusher/config.rs
Normal file
34
crates/headroom-core/src/transforms/text_crusher/config.rs
Normal file
|
|
@ -0,0 +1,34 @@
|
||||||
|
//! TextCrusher configuration (Phase 2, #1171).
|
||||||
|
//!
|
||||||
|
//! Mirrors the Python `TextCrusherConfig`. Weights and thresholds are tuning
|
||||||
|
//! knobs for the recency + relevance + salience scoring.
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct TextCrusherConfig {
|
||||||
|
/// Keep roughly this fraction of characters.
|
||||||
|
pub target_ratio: f64,
|
||||||
|
pub w_recency: f64,
|
||||||
|
pub w_relevance: f64,
|
||||||
|
pub w_salience: f64,
|
||||||
|
/// Segments shorter than this are de-prioritized (× 0.25).
|
||||||
|
pub min_segment_chars: usize,
|
||||||
|
/// Skip a candidate when this fraction of its word-shingles is already
|
||||||
|
/// covered by kept segments (near-duplicate suppression).
|
||||||
|
pub near_dup_threshold: f64,
|
||||||
|
/// Below this many segments, pass through unchanged (nothing to gain).
|
||||||
|
pub min_segments_for_crush: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for TextCrusherConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
TextCrusherConfig {
|
||||||
|
target_ratio: 0.5,
|
||||||
|
w_recency: 1.0,
|
||||||
|
w_relevance: 2.0,
|
||||||
|
w_salience: 1.5,
|
||||||
|
min_segment_chars: 12,
|
||||||
|
near_dup_threshold: 0.85,
|
||||||
|
min_segments_for_crush: 6,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
317
crates/headroom-core/src/transforms/text_crusher/crusher.rs
Normal file
317
crates/headroom-core/src/transforms/text_crusher/crusher.rs
Normal file
|
|
@ -0,0 +1,317 @@
|
||||||
|
//! TextCrusher: fast deterministic extractive prose compressor (Phase 2, #1171).
|
||||||
|
//!
|
||||||
|
//! Splits prose into sentence segments, scores each by recency + query
|
||||||
|
//! relevance + structural salience, suppresses near-duplicates via a global
|
||||||
|
//! word-shingle index, and keeps the top segments (in original order) up to a
|
||||||
|
//! target ratio. Output is extractive: the kept sentences are verbatim words
|
||||||
|
//! (each segment trimmed, re-joined with `\n`) -- no invented words, no rewrite.
|
||||||
|
//!
|
||||||
|
//! The relevance term REUSES the shared [`BM25Scorer`](crate::relevance) rather
|
||||||
|
//! than reimplementing BM25 -- only the prose-specific splitting + selection
|
||||||
|
//! lives here.
|
||||||
|
|
||||||
|
use std::cmp::Ordering;
|
||||||
|
use std::collections::HashSet;
|
||||||
|
|
||||||
|
use super::config::TextCrusherConfig;
|
||||||
|
use crate::relevance::{BM25Scorer, RelevanceScorer};
|
||||||
|
|
||||||
|
const KEYWORDS: [&str; 10] = [
|
||||||
|
"error",
|
||||||
|
"exception",
|
||||||
|
"failed",
|
||||||
|
"failure",
|
||||||
|
"fail",
|
||||||
|
"warning",
|
||||||
|
"traceback",
|
||||||
|
"assert",
|
||||||
|
"todo",
|
||||||
|
"fixme",
|
||||||
|
];
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct TextCrusherResult {
|
||||||
|
pub compressed: String,
|
||||||
|
pub original_tokens: usize,
|
||||||
|
pub compressed_tokens: usize,
|
||||||
|
pub compression_ratio: f64,
|
||||||
|
pub kept_segments: usize,
|
||||||
|
pub total_segments: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct TextCrusher {
|
||||||
|
config: TextCrusherConfig,
|
||||||
|
scorer: BM25Scorer,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for TextCrusher {
|
||||||
|
fn default() -> Self {
|
||||||
|
TextCrusher::new(TextCrusherConfig::default())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TextCrusher {
|
||||||
|
pub fn new(config: TextCrusherConfig) -> Self {
|
||||||
|
TextCrusher {
|
||||||
|
config,
|
||||||
|
scorer: BM25Scorer::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn passthrough(content: &str, n_segments: usize) -> TextCrusherResult {
|
||||||
|
let toks = content.split_whitespace().count();
|
||||||
|
TextCrusherResult {
|
||||||
|
compressed: content.to_string(),
|
||||||
|
original_tokens: toks,
|
||||||
|
compressed_tokens: toks,
|
||||||
|
compression_ratio: 1.0,
|
||||||
|
kept_segments: n_segments,
|
||||||
|
total_segments: n_segments,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn compress(
|
||||||
|
&self,
|
||||||
|
content: &str,
|
||||||
|
context: &str,
|
||||||
|
target_ratio: Option<f64>,
|
||||||
|
) -> TextCrusherResult {
|
||||||
|
let cfg = &self.config;
|
||||||
|
let ratio = target_ratio.unwrap_or(cfg.target_ratio).clamp(0.05, 1.0);
|
||||||
|
|
||||||
|
let segments = split_segments(content);
|
||||||
|
if segments.len() < cfg.min_segments_for_crush {
|
||||||
|
return Self::passthrough(content, segments.len());
|
||||||
|
}
|
||||||
|
|
||||||
|
let n = segments.len();
|
||||||
|
let total_chars: usize = segments.iter().map(|s| s.len()).sum();
|
||||||
|
// .max(1) so a tiny input never truncates the budget to 0 (which would
|
||||||
|
// admit nothing and silently fall back to a 100% passthrough).
|
||||||
|
let target_chars = ((total_chars as f64 * ratio) as usize).max(1);
|
||||||
|
|
||||||
|
// Relevance via the shared BM25 scorer (already [0, 1]).
|
||||||
|
let seg_refs: Vec<&str> = segments.iter().map(|s| s.as_str()).collect();
|
||||||
|
let relevance = self.scorer.score_batch(&seg_refs, context);
|
||||||
|
|
||||||
|
let seg_tokens: Vec<Vec<String>> = segments.iter().map(|s| tokens(s)).collect();
|
||||||
|
|
||||||
|
let mut scores = vec![0.0f64; n];
|
||||||
|
for i in 0..n {
|
||||||
|
let recency = (i as f64 + 1.0) / n as f64;
|
||||||
|
let rel = relevance.get(i).map(|r| r.score).unwrap_or(0.0);
|
||||||
|
let words: Vec<&str> = segments[i].split_whitespace().collect();
|
||||||
|
let salient = words.iter().filter(|w| is_salient(w)).count();
|
||||||
|
let salience = salient as f64 / (words.len() as f64 + 1.0);
|
||||||
|
let mut score =
|
||||||
|
cfg.w_recency * recency + cfg.w_relevance * rel + cfg.w_salience * salience;
|
||||||
|
if segments[i].len() < cfg.min_segment_chars {
|
||||||
|
score *= 0.25;
|
||||||
|
}
|
||||||
|
scores[i] = score;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Highest score first; stable tiebreak by index for determinism.
|
||||||
|
let mut order: Vec<usize> = (0..n).collect();
|
||||||
|
order.sort_by(|&a, &b| {
|
||||||
|
scores[b]
|
||||||
|
.partial_cmp(&scores[a])
|
||||||
|
.unwrap_or(Ordering::Equal)
|
||||||
|
.then(a.cmp(&b))
|
||||||
|
});
|
||||||
|
|
||||||
|
let mut kept = vec![false; n];
|
||||||
|
let mut seen: HashSet<String> = HashSet::new();
|
||||||
|
let mut kept_chars = 0usize;
|
||||||
|
let mut kept_count = 0usize;
|
||||||
|
for &i in &order {
|
||||||
|
if kept_chars >= target_chars {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
let sh = shingles(&seg_tokens[i], 3);
|
||||||
|
if !sh.is_empty() {
|
||||||
|
let covered =
|
||||||
|
sh.iter().filter(|s| seen.contains(*s)).count() as f64 / sh.len() as f64;
|
||||||
|
if covered >= cfg.near_dup_threshold {
|
||||||
|
continue; // near-duplicate: most shingles already kept
|
||||||
|
}
|
||||||
|
}
|
||||||
|
kept[i] = true;
|
||||||
|
kept_count += 1;
|
||||||
|
for s in sh {
|
||||||
|
seen.insert(s);
|
||||||
|
}
|
||||||
|
kept_chars += segments[i].len();
|
||||||
|
}
|
||||||
|
|
||||||
|
if kept_count == 0 {
|
||||||
|
return Self::passthrough(content, n);
|
||||||
|
}
|
||||||
|
|
||||||
|
let compressed = (0..n)
|
||||||
|
.filter(|&i| kept[i])
|
||||||
|
.map(|i| segments[i].as_str())
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.join("\n");
|
||||||
|
let orig_tok = content.split_whitespace().count();
|
||||||
|
let comp_tok = compressed.split_whitespace().count();
|
||||||
|
TextCrusherResult {
|
||||||
|
compression_ratio: if orig_tok > 0 {
|
||||||
|
comp_tok as f64 / orig_tok as f64
|
||||||
|
} else {
|
||||||
|
1.0
|
||||||
|
},
|
||||||
|
compressed,
|
||||||
|
original_tokens: orig_tok,
|
||||||
|
compressed_tokens: comp_tok,
|
||||||
|
kept_segments: kept_count,
|
||||||
|
total_segments: n,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Split into sentence/line segments: on newlines, and after `.`/`!`/`?`
|
||||||
|
/// followed by whitespace. Byte-faithful (kept segments are joined verbatim).
|
||||||
|
fn split_segments(text: &str) -> Vec<String> {
|
||||||
|
let mut segs = Vec::new();
|
||||||
|
for line in text.split('\n') {
|
||||||
|
let trimmed = line.trim();
|
||||||
|
if trimmed.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let mut cur = String::new();
|
||||||
|
let mut prev_term = false;
|
||||||
|
for c in trimmed.chars() {
|
||||||
|
if prev_term && c.is_whitespace() {
|
||||||
|
let s = cur.trim();
|
||||||
|
if !s.is_empty() {
|
||||||
|
segs.push(s.to_string());
|
||||||
|
}
|
||||||
|
cur.clear();
|
||||||
|
prev_term = false;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
cur.push(c);
|
||||||
|
prev_term = matches!(c, '.' | '!' | '?');
|
||||||
|
}
|
||||||
|
let s = cur.trim();
|
||||||
|
if !s.is_empty() {
|
||||||
|
segs.push(s.to_string());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
segs
|
||||||
|
}
|
||||||
|
|
||||||
|
fn tokens(text: &str) -> Vec<String> {
|
||||||
|
let mut out = Vec::new();
|
||||||
|
let mut cur = String::new();
|
||||||
|
for c in text.chars() {
|
||||||
|
if c.is_alphanumeric() || c == '_' {
|
||||||
|
for lc in c.to_lowercase() {
|
||||||
|
cur.push(lc);
|
||||||
|
}
|
||||||
|
} else if !cur.is_empty() {
|
||||||
|
out.push(std::mem::take(&mut cur));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !cur.is_empty() {
|
||||||
|
out.push(cur);
|
||||||
|
}
|
||||||
|
out
|
||||||
|
}
|
||||||
|
|
||||||
|
fn shingles(words: &[String], k: usize) -> HashSet<String> {
|
||||||
|
let mut set = HashSet::new();
|
||||||
|
if words.is_empty() {
|
||||||
|
return set;
|
||||||
|
}
|
||||||
|
if words.len() < k {
|
||||||
|
// Short segment: emit every sub-window (1..=len) so identical/overlapping
|
||||||
|
// short segments still near-dup-match each other. (They can't match a
|
||||||
|
// longer segment's k-grams, but short segments are score-penalized and
|
||||||
|
// rarely survive selection anyway.)
|
||||||
|
for size in 1..=words.len() {
|
||||||
|
for w in words.windows(size) {
|
||||||
|
set.insert(w.join("\u{1}"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return set;
|
||||||
|
}
|
||||||
|
for w in words.windows(k) {
|
||||||
|
set.insert(w.join("\u{1}"));
|
||||||
|
}
|
||||||
|
set
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A word carries specific, hard-to-reconstruct information if it has a digit,
|
||||||
|
/// is an error/status keyword, is ALLCAPS (2+ letters), or is a dotted
|
||||||
|
/// identifier (`foo.bar`).
|
||||||
|
fn is_salient(word: &str) -> bool {
|
||||||
|
if word.chars().any(|c| c.is_ascii_digit()) {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
let lower = word
|
||||||
|
.trim_matches(|c: char| !c.is_alphanumeric())
|
||||||
|
.to_lowercase();
|
||||||
|
if KEYWORDS.contains(&lower.as_str()) {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
let alpha: Vec<char> = word.chars().filter(|c| c.is_alphabetic()).collect();
|
||||||
|
if alpha.len() >= 2 && alpha.iter().all(|c| c.is_uppercase()) {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
if let Some(dot) = word.find('.') {
|
||||||
|
let a = &word[..dot];
|
||||||
|
let b = &word[dot + 1..];
|
||||||
|
if !a.is_empty()
|
||||||
|
&& !b.is_empty()
|
||||||
|
&& a.chars()
|
||||||
|
.next()
|
||||||
|
.is_some_and(|c| c.is_alphabetic() || c == '_')
|
||||||
|
&& b.chars()
|
||||||
|
.next()
|
||||||
|
.is_some_and(|c| c.is_alphabetic() || c == '_')
|
||||||
|
{
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
false
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn doc(n: usize) -> String {
|
||||||
|
(0..n)
|
||||||
|
.map(|i| format!("Sentence number {i} describes a distinct topic {i} in some detail."))
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.join(" ")
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn extractive_and_compresses() {
|
||||||
|
let content = doc(40);
|
||||||
|
let r = TextCrusher::default().compress(&content, "", Some(0.3));
|
||||||
|
assert!(r.compressed_tokens < r.original_tokens);
|
||||||
|
// extractive: every output word appears in the input
|
||||||
|
let orig: HashSet<&str> = content.split_whitespace().collect();
|
||||||
|
assert!(r.compressed.split_whitespace().all(|w| orig.contains(w)));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn deterministic() {
|
||||||
|
let content = doc(40);
|
||||||
|
let tc = TextCrusher::default();
|
||||||
|
assert_eq!(
|
||||||
|
tc.compress(&content, "", Some(0.4)).compressed,
|
||||||
|
tc.compress(&content, "", Some(0.4)).compressed
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn passthrough_when_small() {
|
||||||
|
let r = TextCrusher::default().compress("one. two. three.", "", None);
|
||||||
|
assert_eq!(r.compression_ratio, 1.0);
|
||||||
|
}
|
||||||
|
}
|
||||||
11
crates/headroom-core/src/transforms/text_crusher/mod.rs
Normal file
11
crates/headroom-core/src/transforms/text_crusher/mod.rs
Normal file
|
|
@ -0,0 +1,11 @@
|
||||||
|
//! TextCrusher — fast deterministic extractive prose compressor (Phase 2, #1171).
|
||||||
|
//!
|
||||||
|
//! The request-path-safe alternative to ModernBERT (kompress) for large plain
|
||||||
|
//! text: heuristic sentence scoring (recency + reused BM25 relevance +
|
||||||
|
//! salience) with near-duplicate suppression, in one O(n) pass.
|
||||||
|
|
||||||
|
mod config;
|
||||||
|
mod crusher;
|
||||||
|
|
||||||
|
pub use config::TextCrusherConfig;
|
||||||
|
pub use crusher::{TextCrusher, TextCrusherResult};
|
||||||
|
|
@ -40,6 +40,10 @@ use headroom_core::transforms::{
|
||||||
SearchCompressionResult as RustSearchResult, SearchCompressor as RustSearchCompressor,
|
SearchCompressionResult as RustSearchResult, SearchCompressor as RustSearchCompressor,
|
||||||
SearchCompressorConfig as RustSearchConfig, SearchCompressorStats as RustSearchStats,
|
SearchCompressorConfig as RustSearchConfig, SearchCompressorStats as RustSearchStats,
|
||||||
};
|
};
|
||||||
|
use headroom_core::transforms::{
|
||||||
|
TextCrusher as RustTextCrusher, TextCrusherConfig as RustTextCrusherConfig,
|
||||||
|
TextCrusherResult as RustTextCrusherResult,
|
||||||
|
};
|
||||||
use pyo3::prelude::*;
|
use pyo3::prelude::*;
|
||||||
use pyo3::types::{PyBytes, PyDict};
|
use pyo3::types::{PyBytes, PyDict};
|
||||||
|
|
||||||
|
|
@ -1622,6 +1626,145 @@ fn compress_openai_responses_live_zone(
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// --- TextCrusher (Phase 2, #1171): fast extractive prose compressor ---
|
||||||
|
|
||||||
|
#[pyclass(name = "TextCrusherConfig", module = "headroom._core")]
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct PyTextCrusherConfig {
|
||||||
|
inner: RustTextCrusherConfig,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[pymethods]
|
||||||
|
impl PyTextCrusherConfig {
|
||||||
|
#[new]
|
||||||
|
#[pyo3(signature = (
|
||||||
|
target_ratio = 0.5,
|
||||||
|
w_recency = 1.0,
|
||||||
|
w_relevance = 2.0,
|
||||||
|
w_salience = 1.5,
|
||||||
|
min_segment_chars = 12,
|
||||||
|
near_dup_threshold = 0.85,
|
||||||
|
min_segments_for_crush = 6,
|
||||||
|
))]
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
fn new(
|
||||||
|
target_ratio: f64,
|
||||||
|
w_recency: f64,
|
||||||
|
w_relevance: f64,
|
||||||
|
w_salience: f64,
|
||||||
|
min_segment_chars: usize,
|
||||||
|
near_dup_threshold: f64,
|
||||||
|
min_segments_for_crush: usize,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
inner: RustTextCrusherConfig {
|
||||||
|
target_ratio,
|
||||||
|
w_recency,
|
||||||
|
w_relevance,
|
||||||
|
w_salience,
|
||||||
|
min_segment_chars,
|
||||||
|
near_dup_threshold,
|
||||||
|
min_segments_for_crush,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[getter]
|
||||||
|
fn target_ratio(&self) -> f64 {
|
||||||
|
self.inner.target_ratio
|
||||||
|
}
|
||||||
|
#[getter]
|
||||||
|
fn near_dup_threshold(&self) -> f64 {
|
||||||
|
self.inner.near_dup_threshold
|
||||||
|
}
|
||||||
|
#[getter]
|
||||||
|
fn min_segments_for_crush(&self) -> usize {
|
||||||
|
self.inner.min_segments_for_crush
|
||||||
|
}
|
||||||
|
#[getter]
|
||||||
|
fn w_recency(&self) -> f64 {
|
||||||
|
self.inner.w_recency
|
||||||
|
}
|
||||||
|
#[getter]
|
||||||
|
fn w_relevance(&self) -> f64 {
|
||||||
|
self.inner.w_relevance
|
||||||
|
}
|
||||||
|
#[getter]
|
||||||
|
fn w_salience(&self) -> f64 {
|
||||||
|
self.inner.w_salience
|
||||||
|
}
|
||||||
|
#[getter]
|
||||||
|
fn min_segment_chars(&self) -> usize {
|
||||||
|
self.inner.min_segment_chars
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[pyclass(name = "TextCrusherResult", module = "headroom._core")]
|
||||||
|
struct PyTextCrusherResult {
|
||||||
|
inner: RustTextCrusherResult,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[pymethods]
|
||||||
|
impl PyTextCrusherResult {
|
||||||
|
#[getter]
|
||||||
|
fn compressed(&self) -> String {
|
||||||
|
self.inner.compressed.clone()
|
||||||
|
}
|
||||||
|
#[getter]
|
||||||
|
fn original_tokens(&self) -> usize {
|
||||||
|
self.inner.original_tokens
|
||||||
|
}
|
||||||
|
#[getter]
|
||||||
|
fn compressed_tokens(&self) -> usize {
|
||||||
|
self.inner.compressed_tokens
|
||||||
|
}
|
||||||
|
#[getter]
|
||||||
|
fn compression_ratio(&self) -> f64 {
|
||||||
|
self.inner.compression_ratio
|
||||||
|
}
|
||||||
|
#[getter]
|
||||||
|
fn kept_segments(&self) -> usize {
|
||||||
|
self.inner.kept_segments
|
||||||
|
}
|
||||||
|
#[getter]
|
||||||
|
fn total_segments(&self) -> usize {
|
||||||
|
self.inner.total_segments
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[pyclass(name = "TextCrusher", module = "headroom._core")]
|
||||||
|
struct PyTextCrusher {
|
||||||
|
inner: RustTextCrusher,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[pymethods]
|
||||||
|
impl PyTextCrusher {
|
||||||
|
#[new]
|
||||||
|
#[pyo3(signature = (config = None))]
|
||||||
|
fn new(config: Option<&PyTextCrusherConfig>) -> Self {
|
||||||
|
let cfg = config.map(|c| c.inner.clone()).unwrap_or_default();
|
||||||
|
Self {
|
||||||
|
inner: RustTextCrusher::new(cfg),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `compress(content, context="", target_ratio=None) -> TextCrusherResult`.
|
||||||
|
/// Releases the GIL across the Rust compress call.
|
||||||
|
#[pyo3(signature = (content, context = "", target_ratio = None))]
|
||||||
|
fn compress(
|
||||||
|
&self,
|
||||||
|
py: Python<'_>,
|
||||||
|
content: &str,
|
||||||
|
context: &str,
|
||||||
|
target_ratio: Option<f64>,
|
||||||
|
) -> PyTextCrusherResult {
|
||||||
|
let content = content.to_string();
|
||||||
|
let context = context.to_string();
|
||||||
|
let inner = py.allow_threads(|| self.inner.compress(&content, &context, target_ratio));
|
||||||
|
PyTextCrusherResult { inner }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[pymodule]
|
#[pymodule]
|
||||||
fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||||
// Bridge Rust diagnostics into Python's `logging`. headroom-core emits
|
// Bridge Rust diagnostics into Python's `logging`. headroom-core emits
|
||||||
|
|
@ -1647,6 +1790,9 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||||
m.add_class::<PySmartCrusherConfig>()?;
|
m.add_class::<PySmartCrusherConfig>()?;
|
||||||
m.add_class::<PyCrushResult>()?;
|
m.add_class::<PyCrushResult>()?;
|
||||||
m.add_class::<PySmartCrusher>()?;
|
m.add_class::<PySmartCrusher>()?;
|
||||||
|
m.add_class::<PyTextCrusherConfig>()?;
|
||||||
|
m.add_class::<PyTextCrusherResult>()?;
|
||||||
|
m.add_class::<PyTextCrusher>()?;
|
||||||
m.add_class::<PyDetectionResult>()?;
|
m.add_class::<PyDetectionResult>()?;
|
||||||
m.add_class::<PyLogCompressorConfig>()?;
|
m.add_class::<PyLogCompressorConfig>()?;
|
||||||
m.add_class::<PyLogCompressionResult>()?;
|
m.add_class::<PyLogCompressionResult>()?;
|
||||||
|
|
|
||||||
144
headroom/proxy/background_compression.py
Normal file
144
headroom/proxy/background_compression.py
Normal file
|
|
@ -0,0 +1,144 @@
|
||||||
|
"""Off-path background compression (Phase 3, #1171).
|
||||||
|
|
||||||
|
The request path must never block on ML compression. When a cold-start-large
|
||||||
|
request would otherwise run kompress synchronously under the 30s budget (and
|
||||||
|
leak a non-preemptible worker on timeout -> executor saturation -> cascade),
|
||||||
|
it instead forwards the already-cached/uncompressed messages immediately and
|
||||||
|
enqueues the compression here. A single per-process drain runs it with NO
|
||||||
|
request-coupled deadline and stores the result in the session
|
||||||
|
``CompressionCache``, so the next turn is a cache hit and the forwarded bytes
|
||||||
|
become (and stay) the compressed form.
|
||||||
|
|
||||||
|
This is per-process by design: ``CompressionCache`` is already per-process
|
||||||
|
(``HeadroomProxy._compression_caches``), and multi-worker deployments are
|
||||||
|
already warned to use ``--workers 1`` or sticky sessions, so a per-process
|
||||||
|
drain matches the existing cache semantics without any new cross-process lock.
|
||||||
|
|
||||||
|
Limitations, all fail-open (lost savings, never lost correctness): only the
|
||||||
|
token-mode cold-start path defers here -- other modes compress synchronously;
|
||||||
|
the queue is in-memory, so a restart mid-drain drops queued jobs (they re-defer
|
||||||
|
on a later turn); and a full queue or a duplicate in-flight key drops the job,
|
||||||
|
surfaced to telemetry as ``deferred:dropped``. Background work is bounded by the
|
||||||
|
Phase 1 kompress deadline (a non-terminating compressor would pin the single
|
||||||
|
drain thread, but the compressors terminate).
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
from collections.abc import Awaitable, Callable
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class _Job:
|
||||||
|
key: str
|
||||||
|
compress: Callable[[], Any] # sync callable, runs in the executor (no timeout)
|
||||||
|
store: Callable[[Any], None] # sync callable, stores the result into the cache
|
||||||
|
|
||||||
|
|
||||||
|
class BackgroundCompressor:
|
||||||
|
"""Single per-process async drain that compresses enqueued work off the
|
||||||
|
request path, with no request-coupled deadline.
|
||||||
|
|
||||||
|
``run_in_executor`` is injected so the drain reuses the proxy's compression
|
||||||
|
ThreadPoolExecutor (without the request-path ``asyncio.wait_for`` timeout),
|
||||||
|
and so tests can supply a trivial runner.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
run_in_executor: Callable[[Callable[[], Any]], Awaitable[Any]],
|
||||||
|
*,
|
||||||
|
max_queue: int = 256,
|
||||||
|
) -> None:
|
||||||
|
self._run_in_executor = run_in_executor
|
||||||
|
self._queue: asyncio.Queue[_Job] = asyncio.Queue(maxsize=max_queue)
|
||||||
|
self._pending: set[str] = set()
|
||||||
|
self._task: asyncio.Task[None] | None = None
|
||||||
|
self._processed = 0
|
||||||
|
self._dropped = 0
|
||||||
|
self._errors = 0
|
||||||
|
|
||||||
|
def enqueue(
|
||||||
|
self,
|
||||||
|
key: str,
|
||||||
|
compress: Callable[[], Any],
|
||||||
|
store: Callable[[Any], None],
|
||||||
|
) -> bool:
|
||||||
|
"""Queue a compression job. Returns False (and drops) if the key is
|
||||||
|
already in flight or the queue is full -- both are safe: the request
|
||||||
|
has already been forwarded uncompressed, so a drop just defers the
|
||||||
|
savings to a later turn."""
|
||||||
|
if key in self._pending:
|
||||||
|
return False # already queued / in flight -- dedup
|
||||||
|
# Claim the slot BEFORE the job is observable in the queue so dedup is
|
||||||
|
# atomic against another enqueue of the same key.
|
||||||
|
self._pending.add(key)
|
||||||
|
try:
|
||||||
|
self._queue.put_nowait(_Job(key, compress, store))
|
||||||
|
except asyncio.QueueFull:
|
||||||
|
self._pending.discard(key)
|
||||||
|
self._dropped += 1
|
||||||
|
logger.warning(
|
||||||
|
"background compression queue full (%d); dropping %s "
|
||||||
|
"(request already forwarded uncompressed)",
|
||||||
|
self._queue.maxsize,
|
||||||
|
key,
|
||||||
|
)
|
||||||
|
return False
|
||||||
|
return True
|
||||||
|
|
||||||
|
async def _process_one(self, job: _Job) -> None:
|
||||||
|
try:
|
||||||
|
result = await self._run_in_executor(job.compress)
|
||||||
|
job.store(result)
|
||||||
|
self._processed += 1
|
||||||
|
except Exception as e: # noqa: BLE001 -- fail-open: request already went out uncompressed
|
||||||
|
self._errors += 1
|
||||||
|
logger.warning("background compression failed for %s: %s", job.key, e)
|
||||||
|
finally:
|
||||||
|
self._pending.discard(job.key)
|
||||||
|
|
||||||
|
async def _drain(self) -> None:
|
||||||
|
while True:
|
||||||
|
job = await self._queue.get()
|
||||||
|
try:
|
||||||
|
await self._process_one(job)
|
||||||
|
finally:
|
||||||
|
self._queue.task_done()
|
||||||
|
|
||||||
|
async def start(self) -> None:
|
||||||
|
if self._task is None or self._task.done():
|
||||||
|
self._task = asyncio.create_task(self._drain(), name="headroom-bg-compress")
|
||||||
|
|
||||||
|
async def stop(self, *, drain: bool = True, timeout: float = 5.0) -> None:
|
||||||
|
if self._task is None:
|
||||||
|
return
|
||||||
|
if drain:
|
||||||
|
try:
|
||||||
|
await asyncio.wait_for(self._queue.join(), timeout=timeout)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
logger.warning(
|
||||||
|
"background compression drain timed out with %d queued",
|
||||||
|
self._queue.qsize(),
|
||||||
|
)
|
||||||
|
self._task.cancel()
|
||||||
|
try:
|
||||||
|
await self._task
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
pass
|
||||||
|
self._task = None
|
||||||
|
|
||||||
|
def stats(self) -> dict[str, int]:
|
||||||
|
return {
|
||||||
|
"queued": self._queue.qsize(),
|
||||||
|
"pending": len(self._pending),
|
||||||
|
"processed": self._processed,
|
||||||
|
"dropped": self._dropped,
|
||||||
|
"errors": self._errors,
|
||||||
|
}
|
||||||
|
|
@ -1064,30 +1064,86 @@ class AnthropicHandlerMixin:
|
||||||
# Record all tool_results in the verified frozen prefix as stable
|
# Record all tool_results in the verified frozen prefix as stable
|
||||||
comp_cache.mark_stable_from_messages(messages, frozen_message_count)
|
comp_cache.mark_stable_from_messages(messages, frozen_message_count)
|
||||||
|
|
||||||
async with stage_timer.measure("compression_first_stage"):
|
# Phase 3 (#1171): off-path deferral gate. On a cold-
|
||||||
result = await self._run_compression_in_executor(
|
# start-large request (frozen=0 + large live zone) the
|
||||||
|
# synchronous kompress run would blow the 30s budget and
|
||||||
|
# leak a non-preemptible worker. Forward the cache-
|
||||||
|
# swapped messages uncompressed NOW and compress off the
|
||||||
|
# request path; the result lands in the SAME
|
||||||
|
# CompressionCache and apply_cached swaps it in (live
|
||||||
|
# zone) on a later turn. Byte-identity holds — see the
|
||||||
|
# frozen/live invariant; one-time upstream cache miss
|
||||||
|
# when the compressed form first lands, then stable.
|
||||||
|
if (
|
||||||
|
getattr(self, "_background_compression_enabled", False)
|
||||||
|
and frozen_message_count == 0
|
||||||
|
and original_tokens >= self._background_compression_min_tokens
|
||||||
|
):
|
||||||
|
# Snapshot refs for the async job. The handler must
|
||||||
|
# NOT mutate these lists/dicts in-place after this
|
||||||
|
# point -- the background job reads them on a later
|
||||||
|
# turn. It doesn't today; keep it that way.
|
||||||
|
_bg_messages = messages
|
||||||
|
_bg_working = working_messages
|
||||||
|
_bg_frozen = frozen_message_count
|
||||||
|
# Dedup key: the gate only fires at frozen==0 (the
|
||||||
|
# first in-flight deferral of a session episode), so
|
||||||
|
# session_id alone is the right granularity -- one
|
||||||
|
# background job per session in flight -- and avoids
|
||||||
|
# JSON-serializing the large message list for a key.
|
||||||
|
accepted = self._background_compressor.enqueue(
|
||||||
|
session_id,
|
||||||
lambda: self.anthropic_pipeline.apply(
|
lambda: self.anthropic_pipeline.apply(
|
||||||
messages=working_messages,
|
messages=_bg_working,
|
||||||
model=model,
|
model=model,
|
||||||
model_limit=context_limit,
|
model_limit=context_limit,
|
||||||
context=extract_user_query(working_messages),
|
context=extract_user_query(_bg_working),
|
||||||
frozen_message_count=frozen_message_count,
|
frozen_message_count=_bg_frozen,
|
||||||
biases=biases,
|
biases=biases,
|
||||||
request_id=request_id,
|
request_id=request_id,
|
||||||
compression_policy=compression_policy,
|
compression_policy=compression_policy,
|
||||||
**proxy_pipeline_kwargs(self.config),
|
**proxy_pipeline_kwargs(self.config),
|
||||||
),
|
),
|
||||||
timeout=COMPRESSION_TIMEOUT_SECONDS,
|
lambda result: comp_cache.update_from_result(
|
||||||
|
_bg_messages, result.messages
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
# Forward uncompressed either way (the request can't
|
||||||
|
# wait); only CLAIM deferral when the job was actually
|
||||||
|
# queued. A full-queue drop is visible in telemetry as
|
||||||
|
# "deferred:dropped" and self-heals on a later turn.
|
||||||
|
optimized_messages = working_messages
|
||||||
|
transforms_applied = [
|
||||||
|
"deferred:background_compression"
|
||||||
|
if accepted
|
||||||
|
else "deferred:dropped"
|
||||||
|
]
|
||||||
|
pipeline_timing = {}
|
||||||
|
else:
|
||||||
|
async with stage_timer.measure("compression_first_stage"):
|
||||||
|
result = await self._run_compression_in_executor(
|
||||||
|
lambda: self.anthropic_pipeline.apply(
|
||||||
|
messages=working_messages,
|
||||||
|
model=model,
|
||||||
|
model_limit=context_limit,
|
||||||
|
context=extract_user_query(working_messages),
|
||||||
|
frozen_message_count=frozen_message_count,
|
||||||
|
biases=biases,
|
||||||
|
request_id=request_id,
|
||||||
|
compression_policy=compression_policy,
|
||||||
|
**proxy_pipeline_kwargs(self.config),
|
||||||
|
),
|
||||||
|
timeout=COMPRESSION_TIMEOUT_SECONDS,
|
||||||
|
)
|
||||||
|
|
||||||
# Cache newly compressed messages (index-aligned diff)
|
# Cache newly compressed messages (index-aligned diff)
|
||||||
if result.messages != working_messages:
|
if result.messages != working_messages:
|
||||||
comp_cache.update_from_result(messages, result.messages)
|
comp_cache.update_from_result(messages, result.messages)
|
||||||
|
|
||||||
# Always use pipeline result — Zone 1 swaps are already applied
|
# Always use pipeline result — Zone 1 swaps applied
|
||||||
optimized_messages = result.messages
|
optimized_messages = result.messages
|
||||||
transforms_applied = result.transforms_applied
|
transforms_applied = result.transforms_applied
|
||||||
pipeline_timing = result.timing
|
pipeline_timing = result.timing
|
||||||
# Issue #327 / Bug 3: pipeline.apply uses the provider-
|
# Issue #327 / Bug 3: pipeline.apply uses the provider-
|
||||||
# side tokenizer (AnthropicProvider tiktoken estimator),
|
# side tokenizer (AnthropicProvider tiktoken estimator),
|
||||||
# which counts ~25% higher than the proxy-side
|
# which counts ~25% higher than the proxy-side
|
||||||
|
|
|
||||||
|
|
@ -105,6 +105,7 @@ from headroom.providers.registry import (
|
||||||
)
|
)
|
||||||
from headroom.proxy import runtime_env
|
from headroom.proxy import runtime_env
|
||||||
from headroom.proxy.auth_mode import should_stamp_codex_client
|
from headroom.proxy.auth_mode import should_stamp_codex_client
|
||||||
|
from headroom.proxy.background_compression import BackgroundCompressor
|
||||||
|
|
||||||
# =============================================================================
|
# =============================================================================
|
||||||
# Extracted modules (re-exported for backward compatibility)
|
# Extracted modules (re-exported for backward compatibility)
|
||||||
|
|
@ -812,6 +813,28 @@ class HeadroomProxy(
|
||||||
max_workers=_compression_max,
|
max_workers=_compression_max,
|
||||||
thread_name_prefix="headroom-compress",
|
thread_name_prefix="headroom-compress",
|
||||||
)
|
)
|
||||||
|
# Phase 3 (#1171): off-path background compression. When enabled, a
|
||||||
|
# cold-start-large request (frozen=0 + large live zone) forwards
|
||||||
|
# uncompressed immediately and enqueues the compression here instead of
|
||||||
|
# blocking the request thread under the 30s budget (which leaks a
|
||||||
|
# non-preemptible worker -> executor saturation -> cascade). Default
|
||||||
|
# off (opt-in), fail-open. Per-process, matching _compression_caches.
|
||||||
|
self._background_compression_enabled: bool = os.environ.get(
|
||||||
|
"HEADROOM_BACKGROUND_COMPRESSION", ""
|
||||||
|
).strip().lower() in ("1", "true", "yes", "on")
|
||||||
|
try:
|
||||||
|
self._background_compression_min_tokens: int = int(
|
||||||
|
os.environ.get("HEADROOM_BACKGROUND_COMPRESSION_MIN_TOKENS", "50000")
|
||||||
|
)
|
||||||
|
except ValueError:
|
||||||
|
self._background_compression_min_tokens = 50000
|
||||||
|
# Dedicated single thread: no-timeout background jobs never contend with
|
||||||
|
# the request-path executor (Phase 3, #1171). Lazy -- no thread spawns
|
||||||
|
# until the first off-path job is submitted.
|
||||||
|
self._background_compression_executor = concurrent.futures.ThreadPoolExecutor(
|
||||||
|
max_workers=1, thread_name_prefix="headroom-bg-compress"
|
||||||
|
)
|
||||||
|
self._background_compressor = BackgroundCompressor(self._run_compression_background)
|
||||||
# Gauge: currently-running compression tasks. Mutated under
|
# Gauge: currently-running compression tasks. Mutated under
|
||||||
# ``_compression_metrics_lock`` from worker threads + the asyncio
|
# ``_compression_metrics_lock`` from worker threads + the asyncio
|
||||||
# event loop.
|
# event loop.
|
||||||
|
|
@ -1112,6 +1135,18 @@ class HeadroomProxy(
|
||||||
self._compression_queue_timeouts += 1
|
self._compression_queue_timeouts += 1
|
||||||
raise
|
raise
|
||||||
|
|
||||||
|
async def _run_compression_background(self, fn): # noqa: ANN001, ANN201
|
||||||
|
"""Run a compression callable on the shared executor with NO request-
|
||||||
|
coupled deadline (Phase 3 off-path, #1171).
|
||||||
|
|
||||||
|
Unlike ``_run_compression_in_executor`` there is no ``asyncio.wait_for``
|
||||||
|
and no leaked-thread accounting: no caller is waiting, so a slow run
|
||||||
|
backs up the background queue rather than starving the request executor.
|
||||||
|
Runs on the dedicated single-thread background executor.
|
||||||
|
"""
|
||||||
|
loop = asyncio.get_running_loop()
|
||||||
|
return await loop.run_in_executor(self._background_compression_executor, fn)
|
||||||
|
|
||||||
def _get_compression_cache(self, session_id: str) -> CompressionCache:
|
def _get_compression_cache(self, session_id: str) -> CompressionCache:
|
||||||
"""Get or create a CompressionCache for a session.
|
"""Get or create a CompressionCache for a session.
|
||||||
|
|
||||||
|
|
@ -1851,6 +1886,8 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
||||||
await proxy.usage_reporter.start(proxy)
|
await proxy.usage_reporter.start(proxy)
|
||||||
if proxy.traffic_learner:
|
if proxy.traffic_learner:
|
||||||
await proxy.traffic_learner.start()
|
await proxy.traffic_learner.start()
|
||||||
|
if proxy._background_compression_enabled:
|
||||||
|
await proxy._background_compressor.start()
|
||||||
|
|
||||||
# Only start beacon if we acquire the lock (first worker wins)
|
# Only start beacon if we acquire the lock (first worker wins)
|
||||||
_beacon_is_owner[0] = _try_acquire_beacon_lock()
|
_beacon_is_owner[0] = _try_acquire_beacon_lock()
|
||||||
|
|
@ -1885,6 +1922,9 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
||||||
await proxy.usage_reporter.stop()
|
await proxy.usage_reporter.stop()
|
||||||
if proxy.traffic_learner:
|
if proxy.traffic_learner:
|
||||||
await proxy.traffic_learner.stop()
|
await proxy.traffic_learner.stop()
|
||||||
|
if proxy._background_compression_enabled:
|
||||||
|
await proxy._background_compressor.stop()
|
||||||
|
proxy._background_compression_executor.shutdown(wait=False)
|
||||||
if proxy.code_graph_watcher:
|
if proxy.code_graph_watcher:
|
||||||
proxy.code_graph_watcher.stop()
|
proxy.code_graph_watcher.stop()
|
||||||
await proxy.shutdown()
|
await proxy.shutdown()
|
||||||
|
|
@ -3924,7 +3964,8 @@ def run_server(
|
||||||
else:
|
else:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"Headroom is running with workers=%d. The in-memory CCR store, "
|
"Headroom is running with workers=%d. The in-memory CCR store, "
|
||||||
"compression cache, prefix tracker, TOIN state, and CostTracker are all "
|
"compression cache (incl. off-path background compression), prefix "
|
||||||
|
"tracker, TOIN state, and CostTracker are all "
|
||||||
"per-process; multi-worker deployments produce silent CCR retrieval "
|
"per-process; multi-worker deployments produce silent CCR retrieval "
|
||||||
"failures, avoidable cache busts, and an unstable dashboard 'Proxy $ Saved' "
|
"failures, avoidable cache busts, and an unstable dashboard 'Proxy $ Saved' "
|
||||||
"hero tile (each /stats poll hits a different worker's partial total) when "
|
"hero tile (each /stats poll hits a different worker's partial total) when "
|
||||||
|
|
|
||||||
|
|
@ -970,6 +970,27 @@ class ContentRouter(Transform):
|
||||||
self._tabular_compressor: Any = None
|
self._tabular_compressor: Any = None
|
||||||
self._kompress: Any = None
|
self._kompress: Any = None
|
||||||
|
|
||||||
|
# Phase 0 (#1171): cap the input size handed to kompress (ModernBERT
|
||||||
|
# ONNX). Its inference scales O(tokens) and runs synchronously on the
|
||||||
|
# request thread under the 30s compression budget; above this ceiling we
|
||||||
|
# route to the fast LogCompressor instead so the request path stays
|
||||||
|
# bounded. ~4 chars/token is a cheap proxy (no tokenizer needed; counts
|
||||||
|
# dense JSON/code correctly, unlike word count). 0 disables the gate.
|
||||||
|
try:
|
||||||
|
self._kompress_max_tokens: int = int(
|
||||||
|
os.environ.get("HEADROOM_KOMPRESS_MAX_TOKENS", "50000")
|
||||||
|
)
|
||||||
|
except ValueError:
|
||||||
|
self._kompress_max_tokens = 50000
|
||||||
|
self._kompress_gate_fires: int = 0
|
||||||
|
# Phase 2 (#1171): when enabled, the size-gate routes oversized text to
|
||||||
|
# the fast extractive TextCrusher (real prose savings) instead of the
|
||||||
|
# LogCompressor (~0 savings on prose). Opt-in, default off.
|
||||||
|
self._text_crusher_enabled: bool = os.environ.get(
|
||||||
|
"HEADROOM_TEXT_CRUSHER", ""
|
||||||
|
).strip().lower() in ("1", "true", "yes", "on")
|
||||||
|
self._text_crusher: Any = None
|
||||||
|
|
||||||
# TOIN integration for cross-strategy learning
|
# TOIN integration for cross-strategy learning
|
||||||
self._toin: Any = None
|
self._toin: Any = None
|
||||||
|
|
||||||
|
|
@ -1713,6 +1734,46 @@ class ContentRouter(Transform):
|
||||||
compressed: str | None = None
|
compressed: str | None = None
|
||||||
compressed_tokens: int | None = None
|
compressed_tokens: int | None = None
|
||||||
|
|
||||||
|
# Phase 0 (#1171): size gate. This is the single ML boundary, so gating
|
||||||
|
# here covers EVERY kompress entry point -- TEXT, KOMPRESS-direct,
|
||||||
|
# CODE_AWARE->KOMPRESS, and the strategy-fallback path all route through
|
||||||
|
# _try_ml_compressor. Kompress ONNX inference is O(tokens) and runs
|
||||||
|
# synchronously on the request thread; on a large/cold context it
|
||||||
|
# exceeds the 30s budget and leaks a non-preemptible worker (#1171).
|
||||||
|
# Above the ceiling, route to the fast LogCompressor (or pass through)
|
||||||
|
# rather than ModernBERT, keeping the request path bounded.
|
||||||
|
if self._kompress_max_tokens > 0 and len(text_to_compress) > self._kompress_max_tokens * 4:
|
||||||
|
self._kompress_gate_fires += 1
|
||||||
|
logger.info(
|
||||||
|
"kompress size-gate fired: ~%d tok (>%d) routed off ML (fire #%d)",
|
||||||
|
len(text_to_compress) // 4,
|
||||||
|
self._kompress_max_tokens,
|
||||||
|
self._kompress_gate_fires,
|
||||||
|
)
|
||||||
|
out = text_to_compress
|
||||||
|
crusher = self._get_text_crusher()
|
||||||
|
if crusher is not None:
|
||||||
|
try:
|
||||||
|
out = crusher.compress(text_to_compress, context=context or "").compressed
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(
|
||||||
|
"Kompress size-gate -> TextCrusher failed (%s); passing through", e
|
||||||
|
)
|
||||||
|
out = text_to_compress
|
||||||
|
elif self.config.enable_log_compressor:
|
||||||
|
lc = self._get_log_compressor()
|
||||||
|
if lc:
|
||||||
|
try:
|
||||||
|
out = lc.compress(text_to_compress).compressed
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(
|
||||||
|
"Kompress size-gate -> LogCompressor failed (%s); passing through", e
|
||||||
|
)
|
||||||
|
out = text_to_compress
|
||||||
|
if protected:
|
||||||
|
out = restore_tags(out, protected)
|
||||||
|
return out, len(out.split())
|
||||||
|
|
||||||
# Primary: Kompress. On a cold cache the model is fetched once in the
|
# Primary: Kompress. On a cold cache the model is fetched once in the
|
||||||
# background (ensure_background_load) instead of blocking this request
|
# background (ensure_background_load) instead of blocking this request
|
||||||
# thread on a 274MB download that races the compression timeout and
|
# thread on a 274MB download that races the compression timeout and
|
||||||
|
|
@ -1844,6 +1905,22 @@ class ContentRouter(Transform):
|
||||||
logger.debug("LogCompressor not available")
|
logger.debug("LogCompressor not available")
|
||||||
return self._log_compressor
|
return self._log_compressor
|
||||||
|
|
||||||
|
def _get_text_crusher(self) -> Any:
|
||||||
|
"""Get TextCrusher (Phase 2, lazy load). Returns None when disabled, or
|
||||||
|
when the native ``headroom._core`` extension is not built (mirrors the
|
||||||
|
ImportError handling of the other ``_get_*`` compressor getters)."""
|
||||||
|
if not getattr(self, "_text_crusher_enabled", False):
|
||||||
|
return None
|
||||||
|
if self._text_crusher is None:
|
||||||
|
try:
|
||||||
|
from .text_crusher import TextCrusher
|
||||||
|
|
||||||
|
self._text_crusher = TextCrusher()
|
||||||
|
except ImportError:
|
||||||
|
logger.debug("TextCrusher (headroom._core) unavailable; disabling gate route")
|
||||||
|
self._text_crusher_enabled = False
|
||||||
|
return self._text_crusher
|
||||||
|
|
||||||
def _get_tabular_compressor(self) -> Any:
|
def _get_tabular_compressor(self) -> Any:
|
||||||
"""Get TabularCompressor (lazy load)."""
|
"""Get TabularCompressor (lazy load)."""
|
||||||
if self._tabular_compressor is None:
|
if self._tabular_compressor is None:
|
||||||
|
|
|
||||||
|
|
@ -856,6 +856,24 @@ class KompressCompressor(Transform):
|
||||||
if n_words < 10:
|
if n_words < 10:
|
||||||
return self._passthrough(content, n_words)
|
return self._passthrough(content, n_words)
|
||||||
|
|
||||||
|
# Cooperative wall-clock budget (#1171): kompress ONNX inference is
|
||||||
|
# O(tokens) and non-preemptible once the request's asyncio timeout fires,
|
||||||
|
# so one large block can run for minutes holding a worker (the leak ->
|
||||||
|
# executor-saturation -> queue-timeout cascade). Bail at the next chunk
|
||||||
|
# boundary past this budget, keeping the unprocessed tail verbatim. 0
|
||||||
|
# disables. Env HEADROOM_COMPRESSION_DEADLINE_MS overrides (default 20s).
|
||||||
|
# Cached per instance: operator config, read once -- not per compress() call.
|
||||||
|
deadline_s = getattr(self, "_deadline_s", None)
|
||||||
|
if deadline_s is None:
|
||||||
|
try:
|
||||||
|
deadline_s = max(
|
||||||
|
0.0,
|
||||||
|
float(os.environ.get("HEADROOM_COMPRESSION_DEADLINE_MS", "20000")) / 1000.0,
|
||||||
|
)
|
||||||
|
except ValueError:
|
||||||
|
deadline_s = 20.0
|
||||||
|
self._deadline_s = deadline_s
|
||||||
|
|
||||||
try:
|
try:
|
||||||
model, tokenizer, backend = _load_kompress(
|
model, tokenizer, backend = _load_kompress(
|
||||||
self.config.model_id, self.config.device, allow_download=allow_download
|
self.config.model_id, self.config.device, allow_download=allow_download
|
||||||
|
|
@ -879,8 +897,23 @@ class KompressCompressor(Transform):
|
||||||
kept_ids: set[int] = set()
|
kept_ids: set[int] = set()
|
||||||
inference_ms = 0.0
|
inference_ms = 0.0
|
||||||
chunk_count = 0
|
chunk_count = 0
|
||||||
|
t_deadline = time.perf_counter()
|
||||||
|
|
||||||
for chunk_start in range(0, n_words, max_chunk_words):
|
for chunk_start in range(0, n_words, max_chunk_words):
|
||||||
|
if deadline_s and (time.perf_counter() - t_deadline) > deadline_s:
|
||||||
|
# Keep everything from here on verbatim and stop: a partial
|
||||||
|
# compression that returns NOW beats a full one that leaks a
|
||||||
|
# non-preemptible worker for minutes (#1171).
|
||||||
|
kept_ids.update(range(chunk_start, n_words))
|
||||||
|
logger.warning(
|
||||||
|
"Kompress hit %.1fs deadline after %d/%d words (%d chunks done); "
|
||||||
|
"kept remainder verbatim to free the request thread (#1171)",
|
||||||
|
deadline_s,
|
||||||
|
chunk_start,
|
||||||
|
n_words,
|
||||||
|
chunk_count,
|
||||||
|
)
|
||||||
|
break
|
||||||
chunk_count += 1
|
chunk_count += 1
|
||||||
chunk_words = words[chunk_start : chunk_start + max_chunk_words]
|
chunk_words = words[chunk_start : chunk_start + max_chunk_words]
|
||||||
|
|
||||||
|
|
|
||||||
59
headroom/transforms/text_crusher.py
Normal file
59
headroom/transforms/text_crusher.py
Normal file
|
|
@ -0,0 +1,59 @@
|
||||||
|
"""TextCrusher — fast deterministic extractive prose compressor (Phase 2, #1171).
|
||||||
|
|
||||||
|
Thin Python wrapper over the native ``headroom._core.TextCrusher``. The
|
||||||
|
algorithm (sentence scoring with the SHARED BM25 relevance scorer + global
|
||||||
|
word-shingle near-dup suppression) lives in Rust (``crates/headroom-core``),
|
||||||
|
reusing the same scorer SmartCrusher uses rather than reimplementing it. This
|
||||||
|
wrapper only keeps the Python-facing interface stable for ContentRouter + tests.
|
||||||
|
|
||||||
|
TextCrusher is the request-path-safe alternative to ModernBERT (kompress) for
|
||||||
|
large plain text: ~milliseconds instead of minutes. Extractive -- the kept
|
||||||
|
sentences are verbatim words (each segment trimmed, re-joined with newlines),
|
||||||
|
never paraphrased; it selects, it does not rewrite.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from headroom._core import TextCrusher as _RustTextCrusher
|
||||||
|
from headroom._core import TextCrusherConfig as _RustTextCrusherConfig
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class TextCrusherConfig:
|
||||||
|
target_ratio: float = 0.5
|
||||||
|
w_recency: float = 1.0
|
||||||
|
w_relevance: float = 2.0
|
||||||
|
w_salience: float = 1.5
|
||||||
|
min_segment_chars: int = 12
|
||||||
|
near_dup_threshold: float = 0.85
|
||||||
|
min_segments_for_crush: int = 6
|
||||||
|
|
||||||
|
|
||||||
|
class TextCrusher:
|
||||||
|
"""Extractive prose compressor. ``compress`` returns a result whose
|
||||||
|
``compressed`` text is the kept input sentences verbatim (each trimmed,
|
||||||
|
re-joined with newlines) in original order -- selection, not rewriting.
|
||||||
|
Backed by the Rust core."""
|
||||||
|
|
||||||
|
def __init__(self, config: TextCrusherConfig | None = None) -> None:
|
||||||
|
cfg = config or TextCrusherConfig()
|
||||||
|
self._rust = _RustTextCrusher(
|
||||||
|
_RustTextCrusherConfig(
|
||||||
|
target_ratio=cfg.target_ratio,
|
||||||
|
w_recency=cfg.w_recency,
|
||||||
|
w_relevance=cfg.w_relevance,
|
||||||
|
w_salience=cfg.w_salience,
|
||||||
|
min_segment_chars=cfg.min_segment_chars,
|
||||||
|
near_dup_threshold=cfg.near_dup_threshold,
|
||||||
|
min_segments_for_crush=cfg.min_segments_for_crush,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
def compress(self, content: str, context: str = "", target_ratio: float | None = None) -> Any:
|
||||||
|
"""Returns a ``TextCrusherResult`` (Rust pyclass) with ``.compressed``,
|
||||||
|
``.original_tokens``, ``.compressed_tokens``, ``.compression_ratio``,
|
||||||
|
``.kept_segments``, ``.total_segments``."""
|
||||||
|
return self._rust.compress(content, context, target_ratio)
|
||||||
|
|
@ -0,0 +1,18 @@
|
||||||
|
{
|
||||||
|
"transform": "text_crusher",
|
||||||
|
"label": "plain_prose",
|
||||||
|
"input": {
|
||||||
|
"content": "Sentence number 0 explains how distributed systems reconcile state across topic 0. Sentence number 1 explains how distributed systems reconcile state across topic 1. Sentence number 2 explains how distributed systems reconcile state across topic 2. Sentence number 3 explains how distributed systems reconcile state across topic 3. Sentence number 4 explains how distributed systems reconcile state across topic 4. Sentence number 5 explains how distributed systems reconcile state across topic 5. Sentence number 6 explains how distributed systems reconcile state across topic 6. Sentence number 7 explains how distributed systems reconcile state across topic 7. Sentence number 8 explains how distributed systems reconcile state across topic 8. Sentence number 9 explains how distributed systems reconcile state across topic 9. Sentence number 10 explains how distributed systems reconcile state across topic 10. Sentence number 11 explains how distributed systems reconcile state across topic 11. Sentence number 12 explains how distributed systems reconcile state across topic 12. Sentence number 13 explains how distributed systems reconcile state across topic 13. Sentence number 14 explains how distributed systems reconcile state across topic 14. Sentence number 15 explains how distributed systems reconcile state across topic 15. Sentence number 16 explains how distributed systems reconcile state across topic 16. Sentence number 17 explains how distributed systems reconcile state across topic 17. Sentence number 18 explains how distributed systems reconcile state across topic 18. Sentence number 19 explains how distributed systems reconcile state across topic 19. Sentence number 20 explains how distributed systems reconcile state across topic 20. Sentence number 21 explains how distributed systems reconcile state across topic 21. Sentence number 22 explains how distributed systems reconcile state across topic 22. Sentence number 23 explains how distributed systems reconcile state across topic 23. Sentence number 24 explains how distributed systems reconcile state across topic 24. Sentence number 25 explains how distributed systems reconcile state across topic 25. Sentence number 26 explains how distributed systems reconcile state across topic 26. Sentence number 27 explains how distributed systems reconcile state across topic 27. Sentence number 28 explains how distributed systems reconcile state across topic 28. Sentence number 29 explains how distributed systems reconcile state across topic 29.",
|
||||||
|
"context": "how do distributed systems reconcile state",
|
||||||
|
"target_ratio": 0.3
|
||||||
|
},
|
||||||
|
"output": {
|
||||||
|
"compressed": "Sentence number 21 explains how distributed systems reconcile state across topic 21.\nSentence number 22 explains how distributed systems reconcile state across topic 22.\nSentence number 23 explains how distributed systems reconcile state across topic 23.\nSentence number 24 explains how distributed systems reconcile state across topic 24.\nSentence number 25 explains how distributed systems reconcile state across topic 25.\nSentence number 26 explains how distributed systems reconcile state across topic 26.\nSentence number 27 explains how distributed systems reconcile state across topic 27.\nSentence number 28 explains how distributed systems reconcile state across topic 28.\nSentence number 29 explains how distributed systems reconcile state across topic 29.",
|
||||||
|
"original_tokens": 360,
|
||||||
|
"compressed_tokens": 108,
|
||||||
|
"compression_ratio": 0.3,
|
||||||
|
"kept_segments": 9,
|
||||||
|
"total_segments": 30
|
||||||
|
},
|
||||||
|
"input_sha256": "197195471bb9b8be5fc4d5f555cdd0d4857ec2ad3cd23135d8df3dc81ff50c35"
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,18 @@
|
||||||
|
{
|
||||||
|
"transform": "text_crusher",
|
||||||
|
"label": "plain_prose_no_query",
|
||||||
|
"input": {
|
||||||
|
"content": "Sentence number 0 explains how distributed systems reconcile state across topic 0. Sentence number 1 explains how distributed systems reconcile state across topic 1. Sentence number 2 explains how distributed systems reconcile state across topic 2. Sentence number 3 explains how distributed systems reconcile state across topic 3. Sentence number 4 explains how distributed systems reconcile state across topic 4. Sentence number 5 explains how distributed systems reconcile state across topic 5. Sentence number 6 explains how distributed systems reconcile state across topic 6. Sentence number 7 explains how distributed systems reconcile state across topic 7. Sentence number 8 explains how distributed systems reconcile state across topic 8. Sentence number 9 explains how distributed systems reconcile state across topic 9. Sentence number 10 explains how distributed systems reconcile state across topic 10. Sentence number 11 explains how distributed systems reconcile state across topic 11. Sentence number 12 explains how distributed systems reconcile state across topic 12. Sentence number 13 explains how distributed systems reconcile state across topic 13. Sentence number 14 explains how distributed systems reconcile state across topic 14. Sentence number 15 explains how distributed systems reconcile state across topic 15. Sentence number 16 explains how distributed systems reconcile state across topic 16. Sentence number 17 explains how distributed systems reconcile state across topic 17. Sentence number 18 explains how distributed systems reconcile state across topic 18. Sentence number 19 explains how distributed systems reconcile state across topic 19. Sentence number 20 explains how distributed systems reconcile state across topic 20. Sentence number 21 explains how distributed systems reconcile state across topic 21. Sentence number 22 explains how distributed systems reconcile state across topic 22. Sentence number 23 explains how distributed systems reconcile state across topic 23. Sentence number 24 explains how distributed systems reconcile state across topic 24. Sentence number 25 explains how distributed systems reconcile state across topic 25. Sentence number 26 explains how distributed systems reconcile state across topic 26. Sentence number 27 explains how distributed systems reconcile state across topic 27. Sentence number 28 explains how distributed systems reconcile state across topic 28. Sentence number 29 explains how distributed systems reconcile state across topic 29.",
|
||||||
|
"context": "",
|
||||||
|
"target_ratio": 0.5
|
||||||
|
},
|
||||||
|
"output": {
|
||||||
|
"compressed": "Sentence number 15 explains how distributed systems reconcile state across topic 15.\nSentence number 16 explains how distributed systems reconcile state across topic 16.\nSentence number 17 explains how distributed systems reconcile state across topic 17.\nSentence number 18 explains how distributed systems reconcile state across topic 18.\nSentence number 19 explains how distributed systems reconcile state across topic 19.\nSentence number 20 explains how distributed systems reconcile state across topic 20.\nSentence number 21 explains how distributed systems reconcile state across topic 21.\nSentence number 22 explains how distributed systems reconcile state across topic 22.\nSentence number 23 explains how distributed systems reconcile state across topic 23.\nSentence number 24 explains how distributed systems reconcile state across topic 24.\nSentence number 25 explains how distributed systems reconcile state across topic 25.\nSentence number 26 explains how distributed systems reconcile state across topic 26.\nSentence number 27 explains how distributed systems reconcile state across topic 27.\nSentence number 28 explains how distributed systems reconcile state across topic 28.\nSentence number 29 explains how distributed systems reconcile state across topic 29.",
|
||||||
|
"original_tokens": 360,
|
||||||
|
"compressed_tokens": 180,
|
||||||
|
"compression_ratio": 0.5,
|
||||||
|
"kept_segments": 15,
|
||||||
|
"total_segments": 30
|
||||||
|
},
|
||||||
|
"input_sha256": "197195471bb9b8be5fc4d5f555cdd0d4857ec2ad3cd23135d8df3dc81ff50c35"
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,18 @@
|
||||||
|
{
|
||||||
|
"transform": "text_crusher",
|
||||||
|
"label": "redundant",
|
||||||
|
"input": {
|
||||||
|
"content": "The quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nA distinct fact about subsystem 0 is recorded plainly here.\nA distinct fact about subsystem 1 is recorded plainly here.\nA distinct fact about subsystem 2 is recorded plainly here.\nA distinct fact about subsystem 3 is recorded plainly here.\nA distinct fact about subsystem 4 is recorded plainly here.\nA distinct fact about subsystem 5 is recorded plainly here.\nA distinct fact about subsystem 6 is recorded plainly here.\nA distinct fact about subsystem 7 is recorded plainly here.",
|
||||||
|
"context": "",
|
||||||
|
"target_ratio": 0.9
|
||||||
|
},
|
||||||
|
"output": {
|
||||||
|
"compressed": "The quick brown fox jumps over the very lazy dog every single morning.\nA distinct fact about subsystem 0 is recorded plainly here.\nA distinct fact about subsystem 1 is recorded plainly here.\nA distinct fact about subsystem 2 is recorded plainly here.\nA distinct fact about subsystem 3 is recorded plainly here.\nA distinct fact about subsystem 4 is recorded plainly here.\nA distinct fact about subsystem 5 is recorded plainly here.\nA distinct fact about subsystem 6 is recorded plainly here.\nA distinct fact about subsystem 7 is recorded plainly here.",
|
||||||
|
"original_tokens": 210,
|
||||||
|
"compressed_tokens": 93,
|
||||||
|
"compression_ratio": 0.44285714285714284,
|
||||||
|
"kept_segments": 9,
|
||||||
|
"total_segments": 18
|
||||||
|
},
|
||||||
|
"input_sha256": "2eaa453c47191d4687669467586834d66ab14aedbbe6ee12d01487dbd0922960"
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,18 @@
|
||||||
|
{
|
||||||
|
"transform": "text_crusher",
|
||||||
|
"label": "salient_heavy",
|
||||||
|
"input": {
|
||||||
|
"content": "ERROR connection refused at host 10.0.0.42 after 3 retries.\nThe authentication module validated tokens against auth.registry before forwarding.\nA traceback was logged with code 500 and request_id req-9182.\nJust some generic filler text without any specific identifiers here.\nMore plain filler describing the overall behavior in vague terms.\nWarning: cache hit ratio dropped to 71 percent during the spike.\nAnother unremarkable sentence with no salient tokens at all today.\nThe pipeline.apply call returned 42 kept rows out of 1000 total.",
|
||||||
|
"context": "authentication tokens errors",
|
||||||
|
"target_ratio": 0.4
|
||||||
|
},
|
||||||
|
"output": {
|
||||||
|
"compressed": "The authentication module validated tokens against auth.registry before forwarding.\nAnother unremarkable sentence with no salient tokens at all today.\nThe pipeline.apply call returned 42 kept rows out of 1000 total.",
|
||||||
|
"original_tokens": 80,
|
||||||
|
"compressed_tokens": 30,
|
||||||
|
"compression_ratio": 0.375,
|
||||||
|
"kept_segments": 3,
|
||||||
|
"total_segments": 8
|
||||||
|
},
|
||||||
|
"input_sha256": "57cb8ef778b5c69b199541ebd017a34883e0084343d92ebb54e91773041b4217"
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,18 @@
|
||||||
|
{
|
||||||
|
"transform": "text_crusher",
|
||||||
|
"label": "short_passthrough",
|
||||||
|
"input": {
|
||||||
|
"content": "one thing. two thing. three thing.",
|
||||||
|
"context": "",
|
||||||
|
"target_ratio": null
|
||||||
|
},
|
||||||
|
"output": {
|
||||||
|
"compressed": "one thing. two thing. three thing.",
|
||||||
|
"original_tokens": 6,
|
||||||
|
"compressed_tokens": 6,
|
||||||
|
"compression_ratio": 1.0,
|
||||||
|
"kept_segments": 3,
|
||||||
|
"total_segments": 3
|
||||||
|
},
|
||||||
|
"input_sha256": "d30100c6f67fdae02d1c9312e9572a85951191111f39233b7aa8a86ea782959b"
|
||||||
|
}
|
||||||
18
tests/parity/fixtures/text_crusher/unicode_d46b3fb761b1.json
Normal file
18
tests/parity/fixtures/text_crusher/unicode_d46b3fb761b1.json
Normal file
|
|
@ -0,0 +1,18 @@
|
||||||
|
{
|
||||||
|
"transform": "text_crusher",
|
||||||
|
"label": "unicode",
|
||||||
|
"input": {
|
||||||
|
"content": "句子 0 描述了系统在主题 0 上的行为细节。 句子 1 描述了系统在主题 1 上的行为细节。 句子 2 描述了系统在主题 2 上的行为细节。 句子 3 描述了系统在主题 3 上的行为细节。 句子 4 描述了系统在主题 4 上的行为细节。 句子 5 描述了系统在主题 5 上的行为细节。 句子 6 描述了系统在主题 6 上的行为细节。 句子 7 描述了系统在主题 7 上的行为细节。 句子 8 描述了系统在主题 8 上的行为细节。 句子 9 描述了系统在主题 9 上的行为细节。 句子 10 描述了系统在主题 10 上的行为细节。 句子 11 描述了系统在主题 11 上的行为细节。",
|
||||||
|
"context": "系统",
|
||||||
|
"target_ratio": 0.4
|
||||||
|
},
|
||||||
|
"output": {
|
||||||
|
"compressed": "句子 0 描述了系统在主题 0 上的行为细节。 句子 1 描述了系统在主题 1 上的行为细节。 句子 2 描述了系统在主题 2 上的行为细节。 句子 3 描述了系统在主题 3 上的行为细节。 句子 4 描述了系统在主题 4 上的行为细节。 句子 5 描述了系统在主题 5 上的行为细节。 句子 6 描述了系统在主题 6 上的行为细节。 句子 7 描述了系统在主题 7 上的行为细节。 句子 8 描述了系统在主题 8 上的行为细节。 句子 9 描述了系统在主题 9 上的行为细节。 句子 10 描述了系统在主题 10 上的行为细节。 句子 11 描述了系统在主题 11 上的行为细节。",
|
||||||
|
"original_tokens": 60,
|
||||||
|
"compressed_tokens": 60,
|
||||||
|
"compression_ratio": 1.0,
|
||||||
|
"kept_segments": 1,
|
||||||
|
"total_segments": 1
|
||||||
|
},
|
||||||
|
"input_sha256": "d46b3fb761b16a7771144b388721042e322655450fb9f10e80459d118aabf600"
|
||||||
|
}
|
||||||
94
tests/parity/record_text_crusher.py
Normal file
94
tests/parity/record_text_crusher.py
Normal file
|
|
@ -0,0 +1,94 @@
|
||||||
|
"""Record TextCrusher parity fixtures (Phase 2, #1171).
|
||||||
|
|
||||||
|
Locks the Rust core's ``compress`` output for a fixed set of deterministic
|
||||||
|
scenarios so a future change to the Rust algorithm is caught as a regression.
|
||||||
|
The Python wrapper delegates to ``headroom._core.TextCrusher``, so these
|
||||||
|
fixtures are recorded from (and verified against) the native implementation.
|
||||||
|
|
||||||
|
Re-record after an intentional algorithm change:
|
||||||
|
python tests/parity/record_text_crusher.py
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import hashlib
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
|
||||||
|
from headroom.transforms.text_crusher import TextCrusher
|
||||||
|
|
||||||
|
FIXTURE_DIR = os.path.join(os.path.dirname(__file__), "fixtures", "text_crusher")
|
||||||
|
|
||||||
|
|
||||||
|
def _prose(n: int) -> str:
|
||||||
|
return " ".join(
|
||||||
|
f"Sentence number {i} explains how distributed systems reconcile state across topic {i}."
|
||||||
|
for i in range(n)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _redundant() -> str:
|
||||||
|
dup = "The quick brown fox jumps over the very lazy dog every single morning."
|
||||||
|
uniques = [f"A distinct fact about subsystem {i} is recorded plainly here." for i in range(8)]
|
||||||
|
return "\n".join([dup] * 10 + uniques)
|
||||||
|
|
||||||
|
|
||||||
|
def _salient() -> str:
|
||||||
|
return "\n".join(
|
||||||
|
[
|
||||||
|
"ERROR connection refused at host 10.0.0.42 after 3 retries.",
|
||||||
|
"The authentication module validated tokens against auth.registry before forwarding.",
|
||||||
|
"A traceback was logged with code 500 and request_id req-9182.",
|
||||||
|
"Just some generic filler text without any specific identifiers here.",
|
||||||
|
"More plain filler describing the overall behavior in vague terms.",
|
||||||
|
"Warning: cache hit ratio dropped to 71 percent during the spike.",
|
||||||
|
"Another unremarkable sentence with no salient tokens at all today.",
|
||||||
|
"The pipeline.apply call returned 42 kept rows out of 1000 total.",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# (label, content, context, target_ratio)
|
||||||
|
SCENARIOS: list[tuple[str, str, str, float | None]] = [
|
||||||
|
("plain_prose", _prose(30), "how do distributed systems reconcile state", 0.3),
|
||||||
|
("plain_prose_no_query", _prose(30), "", 0.5),
|
||||||
|
("redundant", _redundant(), "", 0.9),
|
||||||
|
("salient_heavy", _salient(), "authentication tokens errors", 0.4),
|
||||||
|
("short_passthrough", "one thing. two thing. three thing.", "", None),
|
||||||
|
(
|
||||||
|
"unicode",
|
||||||
|
" ".join(f"句子 {i} 描述了系统在主题 {i} 上的行为细节。" for i in range(12)),
|
||||||
|
"系统",
|
||||||
|
0.4,
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def record() -> None:
|
||||||
|
os.makedirs(FIXTURE_DIR, exist_ok=True)
|
||||||
|
tc = TextCrusher()
|
||||||
|
for label, content, context, ratio in SCENARIOS:
|
||||||
|
r = tc.compress(content, context, ratio)
|
||||||
|
digest = hashlib.sha256(content.encode("utf-8")).hexdigest()
|
||||||
|
fixture = {
|
||||||
|
"transform": "text_crusher",
|
||||||
|
"label": label,
|
||||||
|
"input": {"content": content, "context": context, "target_ratio": ratio},
|
||||||
|
"output": {
|
||||||
|
"compressed": r.compressed,
|
||||||
|
"original_tokens": r.original_tokens,
|
||||||
|
"compressed_tokens": r.compressed_tokens,
|
||||||
|
"compression_ratio": r.compression_ratio,
|
||||||
|
"kept_segments": r.kept_segments,
|
||||||
|
"total_segments": r.total_segments,
|
||||||
|
},
|
||||||
|
"input_sha256": digest,
|
||||||
|
}
|
||||||
|
path = os.path.join(FIXTURE_DIR, f"{label}_{digest[:12]}.json")
|
||||||
|
with open(path, "w", encoding="utf-8") as fh:
|
||||||
|
json.dump(fixture, fh, indent=2, ensure_ascii=False)
|
||||||
|
print(f"wrote {path}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
record()
|
||||||
109
tests/test_proxy/test_background_compression.py
Normal file
109
tests/test_proxy/test_background_compression.py
Normal file
|
|
@ -0,0 +1,109 @@
|
||||||
|
"""Phase 3 (#1171): off-path BackgroundCompressor.
|
||||||
|
|
||||||
|
A single per-process drain compresses enqueued work with NO request-coupled
|
||||||
|
deadline and stores the result -- so the request path never blocks on ML.
|
||||||
|
Tests use asyncio.run so they do not depend on a pytest-asyncio config.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
from headroom.proxy.background_compression import BackgroundCompressor
|
||||||
|
|
||||||
|
|
||||||
|
async def _passthrough_executor(fn):
|
||||||
|
return fn()
|
||||||
|
|
||||||
|
|
||||||
|
def test_enqueue_compresses_and_stores():
|
||||||
|
async def main():
|
||||||
|
stored: list[str] = []
|
||||||
|
bc = BackgroundCompressor(_passthrough_executor)
|
||||||
|
await bc.start()
|
||||||
|
ok = bc.enqueue("k1", lambda: "COMPRESSED", lambda r: stored.append(r))
|
||||||
|
await asyncio.wait_for(bc._queue.join(), timeout=2)
|
||||||
|
await bc.stop()
|
||||||
|
return ok, stored, bc.stats()
|
||||||
|
|
||||||
|
ok, stored, stats = asyncio.run(main())
|
||||||
|
assert ok is True
|
||||||
|
assert stored == ["COMPRESSED"]
|
||||||
|
assert stats["processed"] == 1
|
||||||
|
assert stats["pending"] == 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_dedup_skips_inflight_key():
|
||||||
|
async def main():
|
||||||
|
calls: list[int] = []
|
||||||
|
|
||||||
|
async def slow_executor(fn):
|
||||||
|
await asyncio.sleep(0.05)
|
||||||
|
return fn()
|
||||||
|
|
||||||
|
bc = BackgroundCompressor(slow_executor)
|
||||||
|
await bc.start()
|
||||||
|
first = bc.enqueue("same", lambda: calls.append(1), lambda r: None)
|
||||||
|
second = bc.enqueue("same", lambda: calls.append(1), lambda r: None)
|
||||||
|
await asyncio.wait_for(bc._queue.join(), timeout=2)
|
||||||
|
await bc.stop()
|
||||||
|
return first, second, len(calls)
|
||||||
|
|
||||||
|
first, second, n = asyncio.run(main())
|
||||||
|
assert first is True
|
||||||
|
assert second is False # duplicate key skipped
|
||||||
|
assert n == 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_fail_open_on_compress_error():
|
||||||
|
async def main():
|
||||||
|
stored: list[str] = []
|
||||||
|
bc = BackgroundCompressor(_passthrough_executor)
|
||||||
|
await bc.start()
|
||||||
|
|
||||||
|
def boom():
|
||||||
|
raise RuntimeError("kompress exploded")
|
||||||
|
|
||||||
|
bc.enqueue("bad", boom, lambda r: stored.append(r))
|
||||||
|
bc.enqueue("good", lambda: "ok", lambda r: stored.append(r))
|
||||||
|
await asyncio.wait_for(bc._queue.join(), timeout=2)
|
||||||
|
await bc.stop()
|
||||||
|
return stored, bc.stats()
|
||||||
|
|
||||||
|
stored, stats = asyncio.run(main())
|
||||||
|
assert stored == ["ok"] # failing job did not store; later job still processed
|
||||||
|
assert stats["errors"] == 1
|
||||||
|
assert stats["processed"] == 1
|
||||||
|
assert stats["pending"] == 0 # key released even on error
|
||||||
|
|
||||||
|
|
||||||
|
def test_queue_full_drops_without_raising():
|
||||||
|
async def main():
|
||||||
|
# Never start the drain, so the queue fills and stays full.
|
||||||
|
bc = BackgroundCompressor(_passthrough_executor, max_queue=2)
|
||||||
|
results = [bc.enqueue(f"k{i}", lambda: None, lambda r: None) for i in range(5)]
|
||||||
|
return results, bc.stats()
|
||||||
|
|
||||||
|
results, stats = asyncio.run(main())
|
||||||
|
assert results[:2] == [True, True]
|
||||||
|
assert results[2:] == [False, False, False] # overflow dropped, no exception
|
||||||
|
assert stats["dropped"] == 3
|
||||||
|
|
||||||
|
|
||||||
|
def test_stop_drains_remaining_jobs():
|
||||||
|
async def main():
|
||||||
|
stored: list[str] = []
|
||||||
|
|
||||||
|
async def slow_executor(fn):
|
||||||
|
await asyncio.sleep(0.02)
|
||||||
|
return fn()
|
||||||
|
|
||||||
|
bc = BackgroundCompressor(slow_executor)
|
||||||
|
await bc.start()
|
||||||
|
for i in range(4):
|
||||||
|
bc.enqueue(f"k{i}", (lambda i=i: f"c{i}"), lambda r: stored.append(r))
|
||||||
|
await bc.stop() # drain=True by default -> waits for queue.join()
|
||||||
|
return sorted(stored)
|
||||||
|
|
||||||
|
stored = asyncio.run(main())
|
||||||
|
assert stored == ["c0", "c1", "c2", "c3"]
|
||||||
73
tests/test_proxy/test_phase3_byte_identity.py
Normal file
73
tests/test_proxy/test_phase3_byte_identity.py
Normal file
|
|
@ -0,0 +1,73 @@
|
||||||
|
"""Phase 3 (#1171) byte-identity + off-path data flow.
|
||||||
|
|
||||||
|
The off-path design's correctness argument is: forwarding the uncompressed
|
||||||
|
messages on turn N and the compressed form on turn N+1 does NOT corrupt the
|
||||||
|
upstream prefix cache, because ``apply_cached`` swaps in the stored compressed
|
||||||
|
bytes verbatim (one-time miss, then stable). These tests pin that claim and the
|
||||||
|
end-to-end enqueue -> drain -> store -> cache-hit flow the handler gate relies
|
||||||
|
on, without standing up the full Anthropic handler.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
from headroom.cache.compression_cache import CompressionCache
|
||||||
|
from headroom.proxy.background_compression import BackgroundCompressor
|
||||||
|
|
||||||
|
|
||||||
|
def _tool(content: str) -> dict:
|
||||||
|
return {"role": "tool", "content": content}
|
||||||
|
|
||||||
|
|
||||||
|
def test_apply_cached_is_byte_identical_and_stable():
|
||||||
|
cache = CompressionCache()
|
||||||
|
originals = [_tool("x " * 10000)] # a large tool result
|
||||||
|
compressed = [_tool("COMPRESSED")]
|
||||||
|
|
||||||
|
cache.update_from_result(originals, compressed)
|
||||||
|
|
||||||
|
# One-time miss already paid; from here the swap is verbatim AND stable
|
||||||
|
# across repeated turns (no every-turn thrash).
|
||||||
|
out1 = cache.apply_cached(originals)
|
||||||
|
out2 = cache.apply_cached(originals)
|
||||||
|
assert out1[0]["content"] == "COMPRESSED"
|
||||||
|
assert out1 == out2
|
||||||
|
# apply_cached never mutates its input.
|
||||||
|
assert originals[0]["content"] == "x " * 10000
|
||||||
|
|
||||||
|
|
||||||
|
def test_unchanged_content_is_not_swapped():
|
||||||
|
cache = CompressionCache()
|
||||||
|
originals = [_tool("same bytes")]
|
||||||
|
# Pipeline returned identical content (nothing to compress) -> no mapping.
|
||||||
|
cache.update_from_result(originals, [_tool("same bytes")])
|
||||||
|
assert cache.apply_cached(originals)[0]["content"] == "same bytes"
|
||||||
|
|
||||||
|
|
||||||
|
def test_offpath_enqueue_drain_store_then_cache_hit():
|
||||||
|
async def main():
|
||||||
|
cache = CompressionCache()
|
||||||
|
originals = [_tool("x " * 10000)]
|
||||||
|
|
||||||
|
async def run(fn): # trivial in-loop executor stand-in
|
||||||
|
return fn()
|
||||||
|
|
||||||
|
bc = BackgroundCompressor(run)
|
||||||
|
await bc.start()
|
||||||
|
# Exactly the two lambdas the deferral gate passes: compress -> produce
|
||||||
|
# the compressed messages; store -> fold them into the session cache.
|
||||||
|
bc.enqueue(
|
||||||
|
"sess:42",
|
||||||
|
lambda: [_tool("COMPRESSED")],
|
||||||
|
lambda result: cache.update_from_result(originals, result),
|
||||||
|
)
|
||||||
|
await bc._queue.join()
|
||||||
|
await bc.stop()
|
||||||
|
return cache.apply_cached(originals), bc.stats()
|
||||||
|
|
||||||
|
out, stats = asyncio.run(main())
|
||||||
|
# After the background job ran, the next turn is a byte-identical cache hit.
|
||||||
|
assert out[0]["content"] == "COMPRESSED"
|
||||||
|
assert stats["processed"] == 1
|
||||||
|
assert stats["errors"] == 0
|
||||||
82
tests/test_transforms/test_kompress_deadline.py
Normal file
82
tests/test_transforms/test_kompress_deadline.py
Normal file
|
|
@ -0,0 +1,82 @@
|
||||||
|
"""Phase 1 (#1171): kompress cooperative chunk-boundary deadline.
|
||||||
|
|
||||||
|
Kompress ONNX inference is O(tokens) and non-preemptible once the request's
|
||||||
|
asyncio timeout fires, so one large block can run for minutes holding a worker
|
||||||
|
(the leak -> executor-saturation -> queue-timeout cascade). compress() checks a
|
||||||
|
wall-clock budget at each chunk boundary and, when over, keeps the unprocessed
|
||||||
|
tail verbatim and returns -- a partial compression that returns now beats a full
|
||||||
|
one that leaks.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from headroom.transforms import kompress_compressor as kc
|
||||||
|
|
||||||
|
|
||||||
|
def test_compress_bails_at_deadline_keeping_tail_verbatim(monkeypatch):
|
||||||
|
# Fake clock: the pre-loop stamp reads 0s, the first loop-top check reads
|
||||||
|
# 999s elapsed -> deadline trips on chunk 0 before any model/tokenizer use.
|
||||||
|
clock = iter([0.0] + [999.0] * 50)
|
||||||
|
monkeypatch.setattr(kc.time, "perf_counter", lambda: next(clock))
|
||||||
|
monkeypatch.setattr(kc, "_load_kompress", lambda *a, **k: (object(), object(), "onnx"))
|
||||||
|
monkeypatch.setenv("HEADROOM_COMPRESSION_DEADLINE_MS", "20000")
|
||||||
|
|
||||||
|
comp = kc.KompressCompressor()
|
||||||
|
monkeypatch.setattr(comp, "_should_batch_single_content", lambda *a, **k: False)
|
||||||
|
|
||||||
|
content = " ".join(f"w{i}" for i in range(1000))
|
||||||
|
result = comp.compress(content)
|
||||||
|
|
||||||
|
# Deadline tripped on the first chunk -> nothing dropped, tail kept verbatim.
|
||||||
|
assert result.compressed_tokens == 1000
|
||||||
|
assert result.compressed.split() == content.split()
|
||||||
|
|
||||||
|
|
||||||
|
def test_compress_partial_run_keeps_processed_head_plus_verbatim_tail(monkeypatch):
|
||||||
|
# The real partial case: chunk 0 processes (gets compressed), chunk 1 trips
|
||||||
|
# the deadline (kept verbatim). Output must be compressed-head + verbatim-tail.
|
||||||
|
# Clock: call1=t_deadline(0); calls 2-4 are chunk-0's check+inference reads
|
||||||
|
# (under budget); call 5+ is chunk-1's check -> trips.
|
||||||
|
# Robust clock: jump past the deadline only AFTER chunk 0 is processed
|
||||||
|
# (tracked via the model mock), so adding perf_counter calls inside the chunk
|
||||||
|
# body -- e.g. sub-stage timing -- can't shift when the deadline trips.
|
||||||
|
state = {"chunks_done": 0}
|
||||||
|
|
||||||
|
def fake_clock():
|
||||||
|
return 999.0 if state["chunks_done"] >= 1 else 0.0
|
||||||
|
|
||||||
|
monkeypatch.setattr(kc.time, "perf_counter", fake_clock)
|
||||||
|
|
||||||
|
class _Enc(dict):
|
||||||
|
def word_ids(self, batch_index=0):
|
||||||
|
return self["_word_ids"]
|
||||||
|
|
||||||
|
class _Tok:
|
||||||
|
def __call__(self, chunk_words, **kw):
|
||||||
|
n = len(chunk_words)
|
||||||
|
return _Enc(input_ids=[[0] * n], attention_mask=[[1] * n], _word_ids=list(range(n)))
|
||||||
|
|
||||||
|
class _Model:
|
||||||
|
def get_keep_mask(self, input_ids, attention_mask):
|
||||||
|
n = len(input_ids[0])
|
||||||
|
mask = [[i < n // 2 for i in range(n)]] # keep first half of the chunk
|
||||||
|
state["chunks_done"] += 1 # after chunk 0, the clock trips the deadline
|
||||||
|
return mask
|
||||||
|
|
||||||
|
monkeypatch.setattr(kc, "_load_kompress", lambda *a, **k: (_Model(), _Tok(), "onnx"))
|
||||||
|
monkeypatch.setattr(kc, "_model_device_type", lambda *a, **k: "cpu")
|
||||||
|
monkeypatch.setenv("HEADROOM_COMPRESSION_DEADLINE_MS", "20000")
|
||||||
|
|
||||||
|
comp = kc.KompressCompressor()
|
||||||
|
comp.config.chunk_words = 10 # 20 words -> 2 chunks
|
||||||
|
monkeypatch.setattr(comp, "_should_batch_single_content", lambda *a, **k: False)
|
||||||
|
|
||||||
|
words = [f"w{i}" for i in range(20)]
|
||||||
|
out = comp.compress(" ".join(words)).compressed.split()
|
||||||
|
|
||||||
|
# chunk 0 processed: first half kept (w0..w4), second half dropped (w5..w9)
|
||||||
|
assert "w0" in out and "w4" in out
|
||||||
|
assert "w5" not in out and "w9" not in out
|
||||||
|
# chunk 1 tripped the deadline -> its words kept verbatim (w10..w19 all present)
|
||||||
|
for i in range(10, 20):
|
||||||
|
assert f"w{i}" in out
|
||||||
76
tests/test_transforms/test_kompress_size_gate.py
Normal file
76
tests/test_transforms/test_kompress_size_gate.py
Normal file
|
|
@ -0,0 +1,76 @@
|
||||||
|
"""Phase 0 (#1171): kompress input-size gate.
|
||||||
|
|
||||||
|
Kompress (ModernBERT ONNX) inference scales O(tokens) and runs synchronously on
|
||||||
|
the request thread under the 30s compression budget. Above a size ceiling the
|
||||||
|
router must route around the ML path (to the fast LogCompressor, else
|
||||||
|
passthrough) so a large/cold context can't blow the timeout and leak a
|
||||||
|
non-preemptible worker. The gate lives inside ``_try_ml_compressor`` so it
|
||||||
|
covers EVERY kompress entry point (TEXT, KOMPRESS-direct, CODE_AWARE, and the
|
||||||
|
strategy-fallback path), all of which funnel through that single boundary.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from headroom.transforms import ContentRouter
|
||||||
|
from headroom.transforms.content_router import CompressionStrategy
|
||||||
|
|
||||||
|
|
||||||
|
def test_oversized_input_is_gated_away_from_kompress(monkeypatch):
|
||||||
|
router = ContentRouter()
|
||||||
|
router._kompress_max_tokens = 50 # tiny ceiling (×4 = 200 chars)
|
||||||
|
|
||||||
|
def _boom(): # kompress must never be fetched for a gated input
|
||||||
|
raise AssertionError("kompress must not be invoked for oversized (gated) input")
|
||||||
|
|
||||||
|
monkeypatch.setattr(router, "_get_kompress", _boom)
|
||||||
|
|
||||||
|
big = "the quick brown fox jumps over the lazy dog. " * 100 # >200 chars
|
||||||
|
out, ntok = router._try_ml_compressor(big, "")
|
||||||
|
|
||||||
|
assert router._kompress_gate_fires == 1
|
||||||
|
assert isinstance(out, str) and ntok > 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_small_input_still_routes_to_kompress(monkeypatch):
|
||||||
|
router = ContentRouter()
|
||||||
|
router._kompress_max_tokens = 50000 # default ceiling; small input is under it
|
||||||
|
|
||||||
|
# Return no kompressor so the ML block is a fast no-op (no model load in test).
|
||||||
|
monkeypatch.setattr(router, "_get_kompress", lambda: None)
|
||||||
|
|
||||||
|
out, ntok = router._try_ml_compressor("short text", "")
|
||||||
|
|
||||||
|
assert router._kompress_gate_fires == 0 # gate did NOT fire for small input
|
||||||
|
assert isinstance(out, str)
|
||||||
|
|
||||||
|
|
||||||
|
def test_gate_disabled_when_threshold_zero(monkeypatch):
|
||||||
|
router = ContentRouter()
|
||||||
|
router._kompress_max_tokens = 0 # disabled
|
||||||
|
|
||||||
|
monkeypatch.setattr(router, "_get_kompress", lambda: None)
|
||||||
|
|
||||||
|
big = "x " * 100000
|
||||||
|
out, ntok = router._try_ml_compressor(big, "")
|
||||||
|
|
||||||
|
assert router._kompress_gate_fires == 0 # disabled → never gates
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("strategy", [CompressionStrategy.KOMPRESS, CompressionStrategy.TEXT])
|
||||||
|
def test_gate_fires_through_strategy_dispatch(monkeypatch, strategy):
|
||||||
|
# Funnel check: drive the strategy dispatch (not _try_ml_compressor directly)
|
||||||
|
# and confirm both ML strategies reach the single gate boundary.
|
||||||
|
router = ContentRouter()
|
||||||
|
router._kompress_max_tokens = 50
|
||||||
|
|
||||||
|
def _boom():
|
||||||
|
raise AssertionError("kompress must not run for gated input")
|
||||||
|
|
||||||
|
monkeypatch.setattr(router, "_get_kompress", _boom)
|
||||||
|
|
||||||
|
big = "the quick brown fox jumps over the lazy dog. " * 100
|
||||||
|
router._apply_strategy_to_content(big, strategy, "")
|
||||||
|
|
||||||
|
assert router._kompress_gate_fires == 1
|
||||||
65
tests/test_transforms/test_text_crusher.py
Normal file
65
tests/test_transforms/test_text_crusher.py
Normal file
|
|
@ -0,0 +1,65 @@
|
||||||
|
"""Phase 2 (#1171): TextCrusher fast extractive compressor.
|
||||||
|
|
||||||
|
Validates the core contract: extractive (no invented words), deterministic,
|
||||||
|
actually compresses, suppresses near-duplicates, and preferentially keeps
|
||||||
|
query-relevant segments. End-to-end answer-quality vs kompress is validated
|
||||||
|
separately via headroom/evals before defaulting it on.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from headroom.transforms.text_crusher import TextCrusher, TextCrusherConfig
|
||||||
|
|
||||||
|
|
||||||
|
def _doc(n: int = 40) -> str:
|
||||||
|
return " ".join(
|
||||||
|
f"Sentence number {i} describes a distinct topic {i} in some detail." for i in range(n)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_extractive_invents_no_new_words():
|
||||||
|
content = _doc()
|
||||||
|
r = TextCrusher().compress(content, target_ratio=0.5)
|
||||||
|
orig_words = set(content.split())
|
||||||
|
assert set(r.compressed.split()) <= orig_words
|
||||||
|
|
||||||
|
|
||||||
|
def test_deterministic():
|
||||||
|
content = _doc()
|
||||||
|
a = TextCrusher().compress(content, target_ratio=0.4).compressed
|
||||||
|
b = TextCrusher().compress(content, target_ratio=0.4).compressed
|
||||||
|
assert a == b
|
||||||
|
|
||||||
|
|
||||||
|
def test_actually_compresses_large_text():
|
||||||
|
r = TextCrusher().compress(_doc(60), target_ratio=0.3)
|
||||||
|
assert r.compressed_tokens < r.original_tokens
|
||||||
|
assert r.compression_ratio < 0.6
|
||||||
|
|
||||||
|
|
||||||
|
def test_passthrough_when_too_few_segments():
|
||||||
|
content = "one thing. two thing. three thing." # < min_segments_for_crush (6)
|
||||||
|
r = TextCrusher().compress(content)
|
||||||
|
assert r.compressed == content
|
||||||
|
assert r.compression_ratio == 1.0
|
||||||
|
|
||||||
|
|
||||||
|
def test_near_duplicates_suppressed():
|
||||||
|
dup = "The quick brown fox jumps over the very lazy dog today."
|
||||||
|
uniques = [f"A unique fact about item {i} stated plainly here." for i in range(8)]
|
||||||
|
content = "\n".join([dup] * 10 + uniques)
|
||||||
|
r = TextCrusher(TextCrusherConfig(near_dup_threshold=0.8)).compress(content, target_ratio=0.9)
|
||||||
|
# The duplicated sentence must not be kept 10 times.
|
||||||
|
assert r.compressed.count("quick brown fox") <= 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_relevance_keeps_query_relevant_segment():
|
||||||
|
filler = [f"Filler line number {i} with generic words and padding here." for i in range(30)]
|
||||||
|
needle = "The authentication token expires after thirty minutes of inactivity."
|
||||||
|
content = "\n".join(filler[:15] + [needle] + filler[15:])
|
||||||
|
r = TextCrusher().compress(
|
||||||
|
content,
|
||||||
|
context="how long until the authentication token expires",
|
||||||
|
target_ratio=0.2,
|
||||||
|
)
|
||||||
|
assert "authentication token expires" in r.compressed
|
||||||
38
tests/test_transforms/test_text_crusher_parity.py
Normal file
38
tests/test_transforms/test_text_crusher_parity.py
Normal file
|
|
@ -0,0 +1,38 @@
|
||||||
|
"""TextCrusher parity (Phase 2, #1171): the native Rust core must keep
|
||||||
|
reproducing the recorded ``compress`` output for every fixture. Catches drift
|
||||||
|
in the Rust algorithm. Re-record intentional changes with:
|
||||||
|
python tests/parity/record_text_crusher.py
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import glob
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from headroom.transforms.text_crusher import TextCrusher
|
||||||
|
|
||||||
|
_FIXTURE_GLOB = os.path.join(
|
||||||
|
os.path.dirname(__file__), "..", "parity", "fixtures", "text_crusher", "*.json"
|
||||||
|
)
|
||||||
|
FIXTURES = sorted(glob.glob(_FIXTURE_GLOB))
|
||||||
|
|
||||||
|
|
||||||
|
def test_fixtures_exist():
|
||||||
|
assert FIXTURES, "no text_crusher parity fixtures; run tests/parity/record_text_crusher.py"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("path", FIXTURES, ids=lambda p: os.path.basename(p))
|
||||||
|
def test_text_crusher_matches_recorded(path):
|
||||||
|
with open(path, encoding="utf-8") as fh:
|
||||||
|
fx = json.load(fh)
|
||||||
|
inp = fx["input"]
|
||||||
|
r = TextCrusher().compress(inp["content"], inp["context"], inp["target_ratio"])
|
||||||
|
exp = fx["output"]
|
||||||
|
assert r.compressed == exp["compressed"]
|
||||||
|
assert r.compressed_tokens == exp["compressed_tokens"]
|
||||||
|
assert r.kept_segments == exp["kept_segments"]
|
||||||
|
assert r.total_segments == exp["total_segments"]
|
||||||
|
assert abs(r.compression_ratio - exp["compression_ratio"]) < 1e-9
|
||||||
38
tests/test_transforms/test_text_crusher_routing.py
Normal file
38
tests/test_transforms/test_text_crusher_routing.py
Normal file
|
|
@ -0,0 +1,38 @@
|
||||||
|
"""Phase 2 (#1171): the kompress size-gate routes oversized text to TextCrusher
|
||||||
|
when HEADROOM_TEXT_CRUSHER is enabled (real prose savings), instead of the
|
||||||
|
LogCompressor (which yields ~0 on prose) or ModernBERT (slow)."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from headroom.transforms.content_router import ContentRouter
|
||||||
|
|
||||||
|
|
||||||
|
def _prose() -> str:
|
||||||
|
return " ".join(
|
||||||
|
f"Sentence {i} about distributed systems and authentication tokens expiring soon."
|
||||||
|
for i in range(300)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_gate_routes_to_text_crusher_when_enabled(monkeypatch):
|
||||||
|
monkeypatch.setenv("HEADROOM_TEXT_CRUSHER", "1")
|
||||||
|
router = ContentRouter()
|
||||||
|
router._kompress_max_tokens = 50 # tiny ceiling so the gate fires
|
||||||
|
|
||||||
|
def _boom():
|
||||||
|
raise AssertionError("kompress must not run for gated input")
|
||||||
|
|
||||||
|
monkeypatch.setattr(router, "_get_kompress", _boom)
|
||||||
|
|
||||||
|
prose = _prose()
|
||||||
|
out, ntok = router._try_ml_compressor(prose, "authentication tokens")
|
||||||
|
|
||||||
|
assert router._kompress_gate_fires == 1
|
||||||
|
assert ntok < len(prose.split()) # TextCrusher actually compressed (LogCompressor ~0)
|
||||||
|
assert set(out.split()) <= set(prose.split()) # extractive: no invented words
|
||||||
|
|
||||||
|
|
||||||
|
def test_text_crusher_disabled_by_default(monkeypatch):
|
||||||
|
monkeypatch.delenv("HEADROOM_TEXT_CRUSHER", raising=False)
|
||||||
|
router = ContentRouter()
|
||||||
|
assert router._get_text_crusher() is None
|
||||||
Loading…
Add table
Add a link
Reference in a new issue