mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
feat(transforms): dispatch kompress/text via the compressor registry + forward question (#2411)
## What
Completes the if/elif → registry migration in the content router:
**KOMPRESS and TEXT** now dispatch through the `kompress` built-in
adapter (`_registry_compress`), like every other strategy. Also **fixes
a latent bug** in `_invoke_kompress` that dropped the QA-aware
`question` argument (hardcoded `None`) — `question` now rides
`CompressInput.config['question']` and is forwarded into
`_try_ml_compressor`, so QA-aware compression content is preserved.
## Intentionally NOT byte-identical (one approved change)
The sole behavior change is the KOMPRESS/TEXT **token metric**: reported
`compressed_tokens` is now `_estimate_tokens(output.content)` — the
router's calibrated estimate, consistent with `original_tokens` and
every other registry-dispatched strategy — instead of the Kompress
model's own tuple count. **Compressed content is preserved byte-for-byte
in all paths.**
## Decision-impact analysis (traced every reader of `compressed_tokens`)
No content, routing, keep/drop, fallback, or lossless-then-lossy
decision reads this metric for KOMPRESS/TEXT: they're not in
`fallback_eligible_strategy` nor `{SEARCH,LOG,HTML}`, and the
STAGE-0/general layering calls `_try_ml_compressor` directly (unchanged,
already forwards `question`). The only downstream value-reader is
`_record_to_toin`'s skip gate (`original_tokens <= compressed_tokens`) —
**telemetry/learning only**, never affects returned content or routing,
and arguably more correct now (both sides on the same `_estimate_tokens`
scale). Consciously accepted.
## Tests
Rewrote the PR-C2 deferral-pinning tests →
registry-dispatch-matches-direct (content matches the direct
`_try_ml_compressor(..., question)` call; token assertion switched
`==<model count>` → `==_estimate_tokens(output)`, the only assertion
change, solely due to the approved metric switch). Added a
QA-differential test (question changes content) + an adapter-level
`question`-forwarding test. Offline suite: 96 passed; ruff 0.15.17 +
mypy clean.
**Note:** the full content-router CI suite may require further test
updates for any test that exercises the real KOMPRESS/TEXT branch and
asserts the returned count equals the model's tuple `compressed_tokens`
— those should switch to `_estimate_tokens(output)`. (The broad
content_router/compression selection wasn't run locally — it needs
ONNX/HF.)
After this, the router's per-strategy dispatch is fully
registry-resolved.
This commit is contained in:
parent
f9cbdd6e39
commit
446ec26003
4 changed files with 185 additions and 48 deletions
|
|
@ -298,7 +298,15 @@ def _invoke_kompress(router: ContentRouter, inp: CompressInput) -> str | None:
|
|||
# The router dispatches KOMPRESS through ``_try_ml_compressor`` (size gate,
|
||||
# tag protection, background load, marker policy), so the adapter delegates
|
||||
# to the SAME method to stay byte-identical to the router's kompress path.
|
||||
compressed, _tokens = router._try_ml_compressor(inp.content, inp.query, None)
|
||||
# ``question`` (QA-aware compression) rides the pure-data contract via
|
||||
# ``config['question']`` — the router sets it in ``_registry_compress`` — and
|
||||
# is forwarded here so the compressed CONTENT matches the router's direct
|
||||
# ``_try_ml_compressor(content, context, question)`` call. A missing/None
|
||||
# ``question`` forwards None, exactly the no-question path. (Previously this
|
||||
# hardcoded ``None``, silently dropping the QA-aware ``question`` — that bug is
|
||||
# fixed here so the flip preserves content.)
|
||||
question = inp.config.get("question")
|
||||
compressed, _tokens = router._try_ml_compressor(inp.content, inp.query, question)
|
||||
return compressed
|
||||
|
||||
|
||||
|
|
@ -2686,6 +2694,7 @@ class ContentRouter(Transform):
|
|||
context: str,
|
||||
bias: float,
|
||||
config: dict[str, Any] | None = None,
|
||||
question: str | None = None,
|
||||
) -> CompressOutput | None:
|
||||
"""Compress ``content`` with a built-in via the registry, full output.
|
||||
|
||||
|
|
@ -2706,7 +2715,18 @@ class ContentRouter(Transform):
|
|||
real compression the returned content is byte-identical to the direct
|
||||
call. Callers read ``.compressed`` to reproduce the historical
|
||||
``compressed is None`` fallback/passthrough branches exactly.
|
||||
|
||||
``question`` (QA-aware compression) is carried on the pure-data contract
|
||||
via ``CompressInput.config['question']`` — a free ``dict`` — so the
|
||||
contract shape is unchanged. The ``kompress`` adapter reads it back with
|
||||
``inp.config.get('question')`` and forwards it into
|
||||
``_try_ml_compressor(content, context, question)``, matching the router's
|
||||
historical direct call. When ``None`` it is not injected, leaving other
|
||||
built-ins' config untouched.
|
||||
"""
|
||||
merged_config = dict(config or {})
|
||||
if question is not None:
|
||||
merged_config["question"] = question
|
||||
entry = self.compressor_registry.get(name)
|
||||
if entry is None:
|
||||
return None
|
||||
|
|
@ -2717,7 +2737,7 @@ class ContentRouter(Transform):
|
|||
self._content_type_from_strategy(strategy), "text/plain"
|
||||
),
|
||||
query=context,
|
||||
config=config or {},
|
||||
config=merged_config,
|
||||
budget={"bias": bias},
|
||||
)
|
||||
)
|
||||
|
|
@ -3076,14 +3096,46 @@ class ContentRouter(Transform):
|
|||
decision_reason = "html_extractor"
|
||||
|
||||
elif strategy == CompressionStrategy.KOMPRESS:
|
||||
compressed, compressed_tokens = self._try_ml_compressor(content, context, question)
|
||||
# Registry-resolved dispatch: the built-in "kompress" adapter
|
||||
# delegates to the SAME ``_try_ml_compressor(content, context,
|
||||
# question)`` the router historically called here — with
|
||||
# ``question`` forwarded via the CompressInput config — so the
|
||||
# compressed CONTENT is byte-identical to the direct call.
|
||||
# ``_try_ml_compressor`` always returns a str (passthrough on a
|
||||
# no-op / unavailable model), so the adapter always reports
|
||||
# ``compressed=True``; ``output`` is ``None`` only in the
|
||||
# defensive not-registered case, which falls through to the
|
||||
# bottom passthrough exactly as before. The token count is now
|
||||
# ``_estimate_tokens(output.content)`` — the router's calibrated
|
||||
# estimate — replacing the Kompress model's own tuple
|
||||
# ``compressed_tokens``. This is the ONE approved,
|
||||
# non-byte-identical change (a reported metric only; see the
|
||||
# decision-impact note: no keep/drop, fallback, or lossless-
|
||||
# then-lossy gate reads the KOMPRESS/TEXT ``compressed_tokens``).
|
||||
output = self._registry_compress(
|
||||
"kompress", strategy, content, context, bias, question=question
|
||||
)
|
||||
if output is not None:
|
||||
compressed = output.content
|
||||
compressed_tokens = _estimate_tokens(output.content)
|
||||
compressor_name = "KompressCompressor"
|
||||
decision_reason = "kompress"
|
||||
|
||||
elif strategy == CompressionStrategy.TEXT:
|
||||
# Prefer Kompress ML compressor for text
|
||||
# Passes through unchanged if Kompress not available
|
||||
compressed, compressed_tokens = self._try_ml_compressor(content, context, question)
|
||||
# Prefer Kompress ML compressor for text; passes through unchanged
|
||||
# if Kompress is not available. Registry-resolved dispatch via the
|
||||
# SAME built-in "kompress" adapter (TEXT and KOMPRESS share the ML
|
||||
# compressor) with ``question`` forwarded via config, so the
|
||||
# compressed CONTENT is byte-identical to the historical direct
|
||||
# ``_try_ml_compressor(content, context, question)`` call. The
|
||||
# token count is now ``_estimate_tokens(output.content)`` (the same
|
||||
# approved metric change as the KOMPRESS branch above).
|
||||
output = self._registry_compress(
|
||||
"kompress", strategy, content, context, bias, question=question
|
||||
)
|
||||
if output is not None:
|
||||
compressed = output.content
|
||||
compressed_tokens = _estimate_tokens(output.content)
|
||||
compressor_name = "KompressCompressor"
|
||||
decision_reason = "text_uses_kompress"
|
||||
|
||||
|
|
|
|||
|
|
@ -191,6 +191,39 @@ def test_kompress_adapter_maps_mocked_result(monkeypatch: pytest.MonkeyPatch) ->
|
|||
_assert_output_contract(out, inp, entry)
|
||||
|
||||
|
||||
def test_kompress_adapter_forwards_question_from_config(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# The kompress adapter reads ``question`` from ``CompressInput.config`` and
|
||||
# forwards it into _try_ml_compressor (fixing the latent bug where it hardcoded
|
||||
# None). A question-aware fake model embeds the question, so the adapter output
|
||||
# differs when the config question differs — proving end-to-end forwarding.
|
||||
router = _router()
|
||||
seen: dict[str, object] = {}
|
||||
|
||||
def _compress(text: str, *, question: object = None, **kwargs: object) -> SimpleNamespace:
|
||||
seen["question"] = question
|
||||
return SimpleNamespace(compressed=f"Q[{question}]::{text}", compressed_tokens=5)
|
||||
|
||||
monkeypatch.setattr(
|
||||
router,
|
||||
"_get_kompress",
|
||||
lambda: SimpleNamespace(
|
||||
is_ready=lambda: True, ensure_background_load=lambda: None, compress=_compress
|
||||
),
|
||||
)
|
||||
content = "plain text conditioned on the question " * 4
|
||||
entry = _entry(router, "kompress")
|
||||
out = entry.compress(
|
||||
CompressInput(content=content, content_type="text/plain", config={"question": "why"})
|
||||
)
|
||||
assert seen["question"] == "why"
|
||||
assert out.content == f"Q[why]::{content}"
|
||||
# No question in config forwards None (the historical no-question path).
|
||||
seen.clear()
|
||||
out_none = entry.compress(CompressInput(content=content, content_type="text/plain"))
|
||||
assert seen["question"] is None
|
||||
assert out_none.content == f"Q[None]::{content}"
|
||||
|
||||
|
||||
def test_kompress_adapter_passthrough_when_ml_disabled() -> None:
|
||||
# With ML disabled the router's kompress path is a passthrough (no model load
|
||||
# / network), so the adapter returns the content unchanged, never raising.
|
||||
|
|
|
|||
|
|
@ -12,13 +12,15 @@ Each FLIPPED strategy has a differential test comparing the router's dispatch
|
|||
output to the built-in's direct output obtained via its ``_get_*`` getter — i.e.
|
||||
"registry dispatch == old dispatch". SMART_CRUSHER is now also flipped (its
|
||||
``.crush`` primary invocation goes through the registry while the shared
|
||||
Kompress→Log fallback block stays direct); the KOMPRESS/TEXT ML boundary
|
||||
(``_try_ml_compressor``) remains DEFERRED and is asserted unchanged. See
|
||||
Kompress→Log fallback block stays direct); KOMPRESS/TEXT now dispatch via the
|
||||
registry too — the ``kompress`` adapter delegates to the SAME ``_try_ml_compressor``
|
||||
call with ``question`` forwarded via ``config['question']``, so CONTENT is
|
||||
preserved and only the reported token metric changed. See
|
||||
``test_router_registry_smartcrusher.py`` for the full SMART_CRUSHER fallback-chain
|
||||
and KOMPRESS/TEXT differential coverage.
|
||||
|
||||
Offline guardrails:
|
||||
* No real ML/ONNX/HF inference — the deferred KOMPRESS path is mocked.
|
||||
* No real ML/ONNX/HF inference — the KOMPRESS ML boundary is mocked.
|
||||
* The flipped strategies shrink their representative content, so no zero-savings
|
||||
Kompress fallback fires (that would touch the ML boundary and append KOMPRESS
|
||||
to the chain).
|
||||
|
|
@ -327,11 +329,11 @@ def test_diff_deferred_no_registry_entry() -> None:
|
|||
assert "diff" not in {d.name for d in _BUILTIN_COMPRESSOR_DESCRIPTORS}
|
||||
|
||||
|
||||
def test_kompress_deferred_unchanged(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# KOMPRESS is DEFERRED (it is the ML boundary — dispatched through
|
||||
# _try_ml_compressor, not a built-in adapter). Mock the underlying model so no
|
||||
# real ONNX/HF inference runs, and assert the router still routes through
|
||||
# _try_ml_compressor rather than the registry.
|
||||
def test_kompress_registry_dispatch_matches_direct(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# KOMPRESS now dispatches via the registry "kompress" adapter, which still
|
||||
# delegates to _try_ml_compressor (the ML boundary). Mock the underlying model
|
||||
# so no real ONNX/HF inference runs, and assert the compressed CONTENT is
|
||||
# byte-identical to the direct _try_ml_compressor call (question is None here).
|
||||
router = _router()
|
||||
_isolate_branch(monkeypatch, router)
|
||||
fake = SimpleNamespace(
|
||||
|
|
@ -348,5 +350,7 @@ def test_kompress_deferred_unchanged(monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
)
|
||||
assert out == "KOMPRESSED::" + content
|
||||
assert chain == [CompressionStrategy.KOMPRESS.value]
|
||||
# Still the bespoke ML path (unchanged), not a registry round-trip.
|
||||
# CONTENT is byte-identical to the direct ML call (registry round-trip preserves
|
||||
# content); the approved token-metric change is covered in
|
||||
# test_router_registry_smartcrusher.py.
|
||||
assert out == router._try_ml_compressor(content, "", None)[0]
|
||||
|
|
|
|||
|
|
@ -1,19 +1,24 @@
|
|||
"""Byte-identical differential tests for the SMART_CRUSHER / KOMPRESS / TEXT flip.
|
||||
"""Differential tests for the SMART_CRUSHER / KOMPRESS / TEXT registry flip.
|
||||
|
||||
PR-C2 flips the PRIMARY compressor invocation of SMART_CRUSHER to the compressor
|
||||
registry (``_registry_compress("smart_crusher", ...)``), while the shared
|
||||
post-strategy Kompress -> Log fallback block stays a direct dispatch. KOMPRESS and
|
||||
TEXT are DEFERRED: they are the ``_try_ml_compressor`` ML boundary, and the
|
||||
``kompress`` built-in adapter (a) hardcodes ``question=None`` (dropping the real
|
||||
QA-aware ``question`` argument) and (b) recomputes the token count via
|
||||
``_estimate_tokens`` instead of returning ``_try_ml_compressor``'s tuple token
|
||||
count (Kompress's own word-count ``compressed_tokens``, computed pre-CCR-marker),
|
||||
so a registry round-trip cannot reproduce either the content (when a question is
|
||||
supplied) or the token metric byte-for-byte. These tests pin both facts.
|
||||
SMART_CRUSHER flips the PRIMARY compressor invocation to the compressor registry
|
||||
(``_registry_compress("smart_crusher", ...)``), while the shared post-strategy
|
||||
Kompress -> Log fallback block stays a direct dispatch. KOMPRESS and TEXT now ALSO
|
||||
dispatch via the registry (``_registry_compress("kompress", ...)``): the
|
||||
``kompress`` built-in adapter delegates to the SAME ``_try_ml_compressor(content,
|
||||
context, question)`` the router historically called, with ``question`` forwarded
|
||||
via ``CompressInput.config['question']`` (previously the adapter dropped it,
|
||||
passing ``None`` — that latent bug is fixed here). Content is therefore PRESERVED
|
||||
byte-for-byte against the direct ``_try_ml_compressor`` call.
|
||||
|
||||
Every path asserts registry-dispatch output == the historical direct-dispatch
|
||||
output (content, token metric, ``strategy_chain``, and — via the recorded call
|
||||
args — the query/bias/question that flow through).
|
||||
The ONE approved, non-byte-identical change is the token metric: KOMPRESS/TEXT now
|
||||
report ``_estimate_tokens(output.content)`` (the router's calibrated estimate)
|
||||
instead of ``_try_ml_compressor``'s tuple token count (Kompress's own
|
||||
``compressed_tokens``). No content/routing/fallback/lossless-then-lossy decision
|
||||
reads that metric for these two branches, so only the reported number changes.
|
||||
|
||||
Every path asserts registry-dispatch CONTENT == the historical direct-dispatch
|
||||
content (and, via recorded call args, that query/bias/question flow through), and
|
||||
that the token metric equals ``_estimate_tokens(output.content)``.
|
||||
|
||||
Offline guardrails:
|
||||
* No real ML/ONNX/HF inference — the ML boundary is mocked at
|
||||
|
|
@ -187,7 +192,7 @@ def test_smart_crusher_log_fallback_matches_direct(monkeypatch: pytest.MonkeyPat
|
|||
]
|
||||
|
||||
|
||||
# ───────────────────── KOMPRESS / TEXT: deferred (ML boundary) ─────────────────
|
||||
# ─────────────── KOMPRESS / TEXT: flipped (registry, ML boundary) ──────────────
|
||||
|
||||
|
||||
def _fake_kompress(seen: dict[str, object]) -> SimpleNamespace:
|
||||
|
|
@ -195,10 +200,11 @@ def _fake_kompress(seen: dict[str, object]) -> SimpleNamespace:
|
|||
|
||||
def _compress(text: str, **kwargs: object) -> SimpleNamespace:
|
||||
seen.update(kwargs)
|
||||
# ``compressed_tokens`` is the model's OWN word-count (7), deliberately
|
||||
# unequal to ``_estimate_tokens`` of the output — this is exactly the
|
||||
# value the registry adapter would discard, which is why the branch is
|
||||
# deferred.
|
||||
# ``compressed_tokens`` is the model's OWN count (7), deliberately unequal
|
||||
# to ``_estimate_tokens`` of the output. After the flip the branch reports
|
||||
# ``_estimate_tokens(output.content)`` and DISCARDS this tuple value — the
|
||||
# single approved metric change — so the tests below assert the metric is
|
||||
# the estimate, not 7.
|
||||
return SimpleNamespace(compressed="KOMPRESSED::" + text, compressed_tokens=7)
|
||||
|
||||
return SimpleNamespace(
|
||||
|
|
@ -208,7 +214,9 @@ def _fake_kompress(seen: dict[str, object]) -> SimpleNamespace:
|
|||
)
|
||||
|
||||
|
||||
def test_kompress_deferred_ml_path_unchanged(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def test_kompress_registry_dispatch_matches_direct(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# KOMPRESS now dispatches through the registry "kompress" adapter, which
|
||||
# delegates to _try_ml_compressor with ``question`` forwarded via config.
|
||||
router = _router()
|
||||
_isolate_branch(monkeypatch, router)
|
||||
seen: dict[str, object] = {}
|
||||
|
|
@ -218,21 +226,25 @@ def test_kompress_deferred_ml_path_unchanged(monkeypatch: pytest.MonkeyPatch) ->
|
|||
out, tokens, chain = router._apply_strategy_to_content(
|
||||
content, CompressionStrategy.KOMPRESS, "ctx", question="my question", bias=1.0
|
||||
)
|
||||
# (a) CONTENT preserved — byte-identical to the direct _try_ml_compressor call
|
||||
# with the SAME question (so question is forwarded through the adapter).
|
||||
assert out == "KOMPRESSED::" + content
|
||||
assert chain == [CompressionStrategy.KOMPRESS.value]
|
||||
# Token count is the model's tuple value (7), NOT _estimate_tokens(out) — the
|
||||
# exact divergence that makes the registry round-trip non-byte-identical.
|
||||
assert tokens == 7
|
||||
assert tokens != _estimate_tokens(out)
|
||||
# The real ``question`` is forwarded (the kompress adapter would pass None).
|
||||
assert seen["question"] == "my question"
|
||||
# Still the bespoke ML path, byte-identical to a direct _try_ml_compressor call.
|
||||
direct, direct_tokens = router._try_ml_compressor(content, "ctx", "my question")
|
||||
assert out == direct
|
||||
assert tokens == direct_tokens
|
||||
# The real ``question`` reached the model (config['question'] -> adapter ->
|
||||
# _try_ml_compressor -> compressor.compress), fixing the latent adapter bug
|
||||
# that dropped it (passed None).
|
||||
assert seen["question"] == "my question"
|
||||
# (b) APPROVED metric change: the token count is now _estimate_tokens(output),
|
||||
# NOT the model's own tuple value (7) that the direct call returns.
|
||||
assert tokens == _estimate_tokens(out)
|
||||
assert tokens != direct_tokens
|
||||
assert direct_tokens == 7
|
||||
|
||||
|
||||
def test_text_deferred_ml_path_unchanged(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def test_text_registry_dispatch_matches_direct(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# TEXT shares the "kompress" adapter — same guarantees as the KOMPRESS branch.
|
||||
router = _router()
|
||||
_isolate_branch(monkeypatch, router)
|
||||
seen: dict[str, object] = {}
|
||||
|
|
@ -244,9 +256,45 @@ def test_text_deferred_ml_path_unchanged(monkeypatch: pytest.MonkeyPatch) -> Non
|
|||
)
|
||||
assert out == "KOMPRESSED::" + content
|
||||
assert chain == [CompressionStrategy.TEXT.value]
|
||||
assert tokens == 7
|
||||
assert tokens != _estimate_tokens(out)
|
||||
assert seen["question"] == "q2"
|
||||
direct, direct_tokens = router._try_ml_compressor(content, "ctx", "q2")
|
||||
assert out == direct
|
||||
assert tokens == direct_tokens
|
||||
assert seen["question"] == "q2"
|
||||
assert tokens == _estimate_tokens(out)
|
||||
assert tokens != direct_tokens
|
||||
assert direct_tokens == 7
|
||||
|
||||
|
||||
def test_kompress_question_forwarding_changes_content(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# QA-differential: a question-aware fake model embeds the question in its
|
||||
# output, so a DIFFERENT ``question`` yields DIFFERENT content. Proves the
|
||||
# router forwards ``question`` end-to-end via the registry adapter
|
||||
# (config['question'] -> _invoke_kompress -> _try_ml_compressor ->
|
||||
# compressor.compress) — the fix for the adapter that previously dropped it.
|
||||
router = _router()
|
||||
_isolate_branch(monkeypatch, router)
|
||||
|
||||
def _qa_model() -> SimpleNamespace:
|
||||
def _compress(text: str, *, question: object = None, **kwargs: object) -> SimpleNamespace:
|
||||
return SimpleNamespace(compressed=f"Q[{question}]::{text}", compressed_tokens=5)
|
||||
|
||||
return SimpleNamespace(
|
||||
is_ready=lambda: True,
|
||||
ensure_background_load=lambda: None,
|
||||
compress=_compress,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(router, "_get_kompress", _qa_model)
|
||||
content = "the body the model compresses conditioned on the question. " * 3
|
||||
|
||||
out_a, _tok_a, _ = router._apply_strategy_to_content(
|
||||
content, CompressionStrategy.KOMPRESS, "ctx", question="alpha", bias=1.0
|
||||
)
|
||||
out_b, _tok_b, _ = router._apply_strategy_to_content(
|
||||
content, CompressionStrategy.KOMPRESS, "ctx", question="beta", bias=1.0
|
||||
)
|
||||
assert out_a != out_b
|
||||
assert out_a.startswith("Q[alpha]::")
|
||||
assert out_b.startswith("Q[beta]::")
|
||||
# Each matches the direct _try_ml_compressor call with the SAME question.
|
||||
assert out_a == router._try_ml_compressor(content, "ctx", "alpha")[0]
|
||||
assert out_b == router._try_ml_compressor(content, "ctx", "beta")[0]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue