mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(proxy): route selected external compressors through the content router (#2388)
## What
Scope 3 of the pluggable-compressor system: a **selected external
`headroom.compressor` plugin now compresses real traffic**. Opt-in via
`--compressor` / `HEADROOM_COMPRESSORS` — external (non-built-in) names
flow to `ContentRouterConfig.active_external_compressors`, resolved once
against the registry in `__init__` (built-in inventory entries filtered
out).
## How
A single guarded branch at the top of `_apply_strategy_to_content`,
immediately before the built-in if/elif. When a selected external
compressor declares the block's detected content type (exact MIME,
`text/*`, or `*` wildcard), the block runs through the pure-data
`Compressor` contract; otherwise it falls through to the built-in path
unchanged.
## Cache safety (by construction)
The branch lives **inside the per-block strategy dispatch**, which only
runs on non-frozen, already-compressible blocks — the frozen/cached
prefix is split off upstream in `apply()`. So a selected external
compressor **can never rewrite cached-prefix content and bust the prompt
cache**; it inherits the exact same cache-preservation the built-ins
have.
## Fail-open + fidelity
Raise, malformed/non-`CompressOutput`, empty-from-non-empty, or
expansion all fall back to the built-in path. Tokens are counted with
the router's own estimator (never the compressor's self-report). Any
`recoverable` (hash→original) map is mirrored to the CCR store like
SmartCrusher, so `/v1/retrieve/{hash}` resolves. Reached only in
lossy/CCR mode (lossless-only sessions return earlier), so it can't
inject unrecoverable loss.
## Behavior change
**None by default.** With no external compressor selected, the branch is
a single cheap guard and everything below is byte-identical to today.
## Testing
`tests/test_router_external_dispatch.py` — end-to-end dispatch of a
selected external compressor, recoverable-map retrievability,
non-hex-hash skip, fail-open on raise/malformed/empty/expansion,
not-selected & non-matching-content-type leave the built-in path
unchanged, wildcard selection. Offline suite: 84 passed (this file +
selection + registry + settings_store). ruff + mypy clean.
Note: the broad content-router/compression suite exercises real HF-Hub
model downloads + local ONNX inference and is slow/flaky in some local
envs — deferred to CI.
Stacks on #2370/#2371/#2373 (all merged).
This commit is contained in:
parent
d7a8cdbee1
commit
e3c7964038
3 changed files with 651 additions and 0 deletions
|
|
@ -715,6 +715,32 @@ def _apply_compressor_selection(
|
|||
setattr(router_config, flag, select_all or name in selected)
|
||||
|
||||
|
||||
def _external_compressor_selection(compressors: set[str] | None) -> list[str] | None:
|
||||
"""Return the selected EXTERNAL (non-built-in) compressor names, or ``None``.
|
||||
|
||||
The built-in selection (:func:`_apply_compressor_selection`) consumes only
|
||||
the names in :data:`BUILTIN_COMPRESSOR_FLAGS`; every OTHER selected name is a
|
||||
third-party ``headroom.compressor`` entry point. This threads those to the
|
||||
router (via ``ContentRouterConfig.active_external_compressors``) so it can
|
||||
route matching blocks through them.
|
||||
|
||||
Returns ``None`` — the router's external-dispatch branch stays inert, so the
|
||||
request path is byte-identical to today — when the selection is empty or
|
||||
contains only recognized built-in names. ``"*"`` is preserved so the router
|
||||
activates every discovered external compressor (mirroring the wildcard's
|
||||
"select everything" meaning on the built-in side).
|
||||
"""
|
||||
if not compressors:
|
||||
return None
|
||||
selected = {name.strip() for name in compressors if name.strip()}
|
||||
if not selected:
|
||||
return None
|
||||
if "*" in selected:
|
||||
return ["*"]
|
||||
external = sorted(selected - set(BUILTIN_COMPRESSOR_FLAGS))
|
||||
return external or None
|
||||
|
||||
|
||||
class HeadroomProxy(
|
||||
StreamingMixin,
|
||||
AnthropicHandlerMixin,
|
||||
|
|
@ -820,6 +846,13 @@ class HeadroomProxy(
|
|||
# Runs BEFORE the disable_kompress override below so that flag stays
|
||||
# authoritative for turning Kompress off.
|
||||
_apply_compressor_selection(router_config, config.compressors)
|
||||
# External (non-built-in) `headroom.compressor` selections are ignored by
|
||||
# `_apply_compressor_selection` (they have no enable_* flag). Thread them
|
||||
# to the router here so it can route matching blocks through them; None
|
||||
# (no external selected) keeps the external-dispatch branch inert.
|
||||
router_config.active_external_compressors = _external_compressor_selection(
|
||||
config.compressors
|
||||
)
|
||||
# No-CCR lossless mode: compress tool outputs with format-native
|
||||
# lossless compaction and marker-free SmartCrusher, and suppress every
|
||||
# retrieval marker + the retrieve-tool injection so no MCP round-trip is
|
||||
|
|
|
|||
|
|
@ -244,6 +244,44 @@ def _build_compressor_registry() -> CompressorRegistry:
|
|||
return registry
|
||||
|
||||
|
||||
# Canonical map from the router's internal :class:`ContentType` to the MIME
|
||||
# string an external ``headroom.compressor`` declares in
|
||||
# ``CompressorDescriptor.content_types``. Mirrors the MIME strings used by the
|
||||
# built-in descriptors above (so an external JSON compressor declares the same
|
||||
# ``application/json`` a built-in would), with ``text/x-diff`` added for
|
||||
# ``GIT_DIFF`` (no built-in descriptor covers diffs). This is the ONLY bridge
|
||||
# between the enum the router routes on and the pure-string content type the
|
||||
# registry contract carries; it is read solely by the opt-in external-dispatch
|
||||
# branch and never on the default request path.
|
||||
_CONTENT_TYPE_TO_MIME: dict[ContentType, str] = {
|
||||
ContentType.JSON_ARRAY: "application/json",
|
||||
ContentType.SOURCE_CODE: "text/x-code",
|
||||
ContentType.SEARCH_RESULTS: "text/x-search-results",
|
||||
ContentType.BUILD_OUTPUT: "text/x-log",
|
||||
ContentType.GIT_DIFF: "text/x-diff",
|
||||
ContentType.HTML: "text/html",
|
||||
ContentType.TABULAR: "text/csv",
|
||||
ContentType.STRUCTURED_CONFIG: "text/x-config",
|
||||
ContentType.PLAIN_TEXT: "text/plain",
|
||||
}
|
||||
|
||||
|
||||
def _external_compressor_matches(descriptor: CompressorDescriptor, content_mime: str) -> bool:
|
||||
"""True if ``descriptor`` declares support for ``content_mime``.
|
||||
|
||||
Accepts an exact MIME match, a full wildcard (``"*"`` or ``"*/*"``), or a
|
||||
type wildcard (``"text/*"`` matches ``"text/plain"``). Anything else is a
|
||||
non-match, so a selected external compressor only ever sees content it
|
||||
explicitly declared it can handle.
|
||||
"""
|
||||
declared = descriptor.content_types or []
|
||||
if content_mime in declared:
|
||||
return True
|
||||
top = content_mime.split("/", 1)[0]
|
||||
type_wildcard = f"{top}/*"
|
||||
return any(d in ("*", "*/*") or d == type_wildcard for d in declared)
|
||||
|
||||
|
||||
def _tool_call_args_text(raw: Any) -> str:
|
||||
"""Compact, query-usable text from a tool call's args.
|
||||
|
||||
|
|
@ -1231,6 +1269,18 @@ class ContentRouterConfig:
|
|||
# Tool exclusion (Read/Glob/...) and reversibility gates still apply.
|
||||
force_kompress_all: bool = False
|
||||
|
||||
# Opt-in selection of EXTERNAL (non-built-in) `headroom.compressor` names to
|
||||
# route real traffic through. `None`/empty (the default) means the external-
|
||||
# dispatch branch in `_apply_strategy_to_content` is inert and the request
|
||||
# path is byte-identical to today. Built-in names are NOT put here — they are
|
||||
# selected via the `enable_*` flags above (see the proxy's
|
||||
# `_apply_compressor_selection`). `"*"` activates every discovered external
|
||||
# compressor. The router resolves these names against `compressor_registry`
|
||||
# and, when a block's detected content type matches an active external
|
||||
# compressor's declared `content_types`, runs it via the registry contract
|
||||
# instead of the built-in if/elif — fail-open back to the built-in path.
|
||||
active_external_compressors: list[str] | None = None
|
||||
|
||||
# No-CCR lossless mode. When True the router compresses LOG/SEARCH/DIFF
|
||||
# content with format-native lossless compaction (headroom.transforms.
|
||||
# lossless_compaction) instead of the lossy Rust drop path, and never
|
||||
|
|
@ -1496,6 +1546,15 @@ class ContentRouter(Transform):
|
|||
logger.debug("compressor registry unavailable: %s", exc)
|
||||
self.compressor_registry = CompressorRegistry()
|
||||
|
||||
# Resolve the opt-in EXTERNAL compressor selection ONCE — the registry
|
||||
# and `config.active_external_compressors` are both fixed after
|
||||
# construction. Empty unless the operator selected a non-built-in
|
||||
# compressor (via `--compressor`), so the external-dispatch branch in
|
||||
# `_apply_strategy_to_content` is a single cheap guard and the default
|
||||
# request path stays byte-identical. Built-in registry entries are
|
||||
# filtered out here so they are only ever dispatched by the if/elif.
|
||||
self._active_external_compressors: list[Any] = self._resolve_active_external_compressors()
|
||||
|
||||
# Lazy-loaded compressors
|
||||
self._code_compressor: Any = None
|
||||
self._smart_crusher: Any = None
|
||||
|
|
@ -2259,6 +2318,188 @@ class ContentRouter(Transform):
|
|||
return False
|
||||
return self._lossless_first(content, CompressionStrategy.PASSTHROUGH)[1] is not None
|
||||
|
||||
# ── External compressor dispatch (opt-in; fail-open) ──────────────────────
|
||||
|
||||
def _resolve_active_external_compressors(self) -> list[Any]:
|
||||
"""Resolve the opt-in external compressor selection against the registry.
|
||||
|
||||
Returns the active EXTERNAL compressor objects (built-in inventory
|
||||
entries filtered out — they own the if/elif dispatch, never the
|
||||
registry). Empty when nothing external is selected or resolution fails,
|
||||
so the caller's external-dispatch branch is inert by default.
|
||||
"""
|
||||
selection = self.config.active_external_compressors
|
||||
if not selection:
|
||||
return []
|
||||
try:
|
||||
active = self.compressor_registry.active(set(selection))
|
||||
except Exception as exc: # noqa: BLE001 - selection is non-critical
|
||||
logger.debug("external compressor resolution failed: %s", exc)
|
||||
return []
|
||||
return [c for c in active if not isinstance(c, _BuiltinCompressorEntry)]
|
||||
|
||||
def _try_external_compressor(
|
||||
self,
|
||||
content: str,
|
||||
strategy: CompressionStrategy,
|
||||
context: str,
|
||||
question: str | None,
|
||||
) -> tuple[str, int, list[str]] | None:
|
||||
"""Route a block through a *selected* external compressor, or ``None``.
|
||||
|
||||
Opt-in and fail-open. Returns ``None`` — leaving the built-in if/elif
|
||||
dispatch to run UNCHANGED — whenever:
|
||||
|
||||
* no external compressor was selected (the default: a single cheap
|
||||
guard, so the request path is byte-identical to today);
|
||||
* none of the active external compressors declares this block's
|
||||
detected content type;
|
||||
* the chosen compressor raises, returns malformed/empty output, or
|
||||
would expand the content.
|
||||
|
||||
On success it returns the router's normal ``(content, tokens, chain)``
|
||||
shape: the external output, tokens counted with the router's OWN
|
||||
estimator (never the compressor's self-reported count), and an
|
||||
``["external:<name>"]`` chain. Any ``recoverable`` (hash -> original)
|
||||
map is persisted to the CCR store exactly like SmartCrusher's mirror,
|
||||
so ``/v1/retrieve/{hash}`` resolves.
|
||||
|
||||
Reached only in lossy/CCR mode: ``_apply_strategy_to_content`` returns
|
||||
earlier in lossless-only mode (STAGE 0), so an external compressor can
|
||||
never inject unrecoverable loss into a lossless-only session.
|
||||
"""
|
||||
active = self._active_external_compressors
|
||||
if not active:
|
||||
return None
|
||||
content_mime = _CONTENT_TYPE_TO_MIME.get(self._content_type_from_strategy(strategy))
|
||||
if content_mime is None:
|
||||
return None
|
||||
for compressor in active:
|
||||
try:
|
||||
descriptor = compressor.descriptor
|
||||
except Exception as exc: # noqa: BLE001 - a broken external is isolated
|
||||
logger.debug("external compressor descriptor unavailable: %s", exc)
|
||||
continue
|
||||
if not _external_compressor_matches(descriptor, content_mime):
|
||||
continue
|
||||
result = self._run_external_compressor(
|
||||
compressor, descriptor.name, content, content_mime, context, question
|
||||
)
|
||||
if result is not None:
|
||||
return result
|
||||
return None
|
||||
|
||||
def _run_external_compressor(
|
||||
self,
|
||||
compressor: Any,
|
||||
name: str,
|
||||
content: str,
|
||||
content_mime: str,
|
||||
context: str,
|
||||
question: str | None,
|
||||
) -> tuple[str, int, list[str]] | None:
|
||||
"""Invoke one external compressor via the contract; fail open to ``None``."""
|
||||
inp = CompressInput(
|
||||
content=content,
|
||||
content_type=content_mime,
|
||||
query=question or context or "",
|
||||
config={},
|
||||
budget={},
|
||||
)
|
||||
try:
|
||||
out = compressor.compress(inp)
|
||||
except Exception as exc: # noqa: BLE001 - fail open to the built-in path
|
||||
logger.warning(
|
||||
"external compressor %r raised (%s); falling back to built-in", name, exc
|
||||
)
|
||||
return None
|
||||
if not isinstance(out, CompressOutput) or not isinstance(out.content, str):
|
||||
logger.warning(
|
||||
"external compressor %r returned malformed output (%s); falling back",
|
||||
name,
|
||||
type(out).__name__,
|
||||
)
|
||||
return None
|
||||
compressed = out.content
|
||||
# Never blank out a non-empty block (an empty user/tool block makes
|
||||
# providers reject the request); fall back so the built-in path runs.
|
||||
if content.strip() and not compressed.strip():
|
||||
logger.warning(
|
||||
"external compressor %r produced empty output; falling back to built-in", name
|
||||
)
|
||||
return None
|
||||
# Never let an external compressor expand a block; fall back so the
|
||||
# built-in path (or passthrough) can do better.
|
||||
if len(compressed) > len(content):
|
||||
logger.debug(
|
||||
"external compressor %r expanded content (%d -> %d chars); falling back",
|
||||
name,
|
||||
len(content),
|
||||
len(compressed),
|
||||
)
|
||||
return None
|
||||
# Count with the router's OWN estimator, not the compressor's self-report.
|
||||
compressed_tokens = _estimate_tokens(compressed)
|
||||
# Persist the hash -> original recovery map the SAME way SmartCrusher
|
||||
# mirrors its markers, so a later /v1/retrieve resolves each hash.
|
||||
self._persist_external_recoverable(out.recoverable, name, context)
|
||||
if out.warnings:
|
||||
logger.debug(
|
||||
"external compressor %r warnings: %s", name, "; ".join(map(str, out.warnings))
|
||||
)
|
||||
if out.markers and logger.isEnabledFor(logging.DEBUG):
|
||||
logger.debug(
|
||||
"external compressor %r markers: %s", name, "; ".join(map(str, out.markers))
|
||||
)
|
||||
return compressed, compressed_tokens, [f"external:{name}"]
|
||||
|
||||
def _persist_external_recoverable(
|
||||
self, recoverable: dict[str, str], name: str, context: str
|
||||
) -> None:
|
||||
"""Mirror an external compressor's hash -> original map into the CCR store.
|
||||
|
||||
Mirrors SmartCrusher's ``_mirror_single_hash_to_python_store``: each
|
||||
entry is stored under its own hash via ``explicit_hash`` so a
|
||||
``/v1/retrieve/{hash}`` lookup returns the original. Best-effort — a
|
||||
store failure or a non-hex hash is logged and never breaks the request
|
||||
(the compressed block is still returned; only that entry is unretrievable).
|
||||
"""
|
||||
if not recoverable:
|
||||
return
|
||||
try:
|
||||
from ..cache.compression_store import get_compression_store
|
||||
|
||||
store = get_compression_store()
|
||||
except Exception as exc: # noqa: BLE001 - CCR store optional/stripped builds
|
||||
logger.debug("external compressor %r: CCR store unavailable (%s)", name, exc)
|
||||
return
|
||||
strategy_label = f"external:{name}"
|
||||
for ccr_hash, original in recoverable.items():
|
||||
if not isinstance(ccr_hash, str) or not isinstance(original, str):
|
||||
logger.debug("external compressor %r: skipping non-str recoverable entry", name)
|
||||
continue
|
||||
try:
|
||||
store.store(
|
||||
original=original,
|
||||
# The compressed payload isn't meaningfully addressable per
|
||||
# hash here; use a placeholder marker (as SmartCrusher does
|
||||
# — /v1/retrieve returns original_content, not compressed).
|
||||
compressed=f"<<external:{name}:{ccr_hash}>>",
|
||||
query_context=context or None,
|
||||
compression_strategy=strategy_label,
|
||||
explicit_hash=ccr_hash,
|
||||
)
|
||||
except ValueError:
|
||||
# explicit_hash must be hex; a malformed hash means this entry
|
||||
# won't be retrievable, but the request must not break.
|
||||
logger.warning(
|
||||
"external compressor %r: recoverable hash %r is not hex; not stored",
|
||||
name,
|
||||
ccr_hash,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 - defensive; never break the request
|
||||
logger.debug("external compressor %r: store.store raised (%s)", name, exc)
|
||||
|
||||
def _apply_strategy_to_content(
|
||||
self,
|
||||
content: str,
|
||||
|
|
@ -2381,6 +2622,19 @@ class ContentRouter(Transform):
|
|||
# split → fall through to the lossy compressors below (kompress /
|
||||
# smart_crusher / code), which attach CCR retrieval markers when enabled.
|
||||
|
||||
# ── External compressor dispatch (opt-in) ────────────────────────────
|
||||
# Immediately before the built-in if/elif, give a *selected* external
|
||||
# `headroom.compressor` first crack at this block IFF its declared
|
||||
# content_types match the block's detected content type. This runs only
|
||||
# when the operator selected a non-built-in compressor, so with no such
|
||||
# selection it is a single cheap guard and everything below is
|
||||
# byte-identical to today. Fully fail-open: a non-match, an error,
|
||||
# malformed/empty output, or an expansion all return None and fall
|
||||
# through to the EXISTING built-in dispatch UNCHANGED.
|
||||
external = self._try_external_compressor(content, strategy, context, question)
|
||||
if external is not None:
|
||||
return external
|
||||
|
||||
try:
|
||||
if strategy == CompressionStrategy.CODE_AWARE:
|
||||
if self.config.enable_code_aware:
|
||||
|
|
|
|||
364
tests/test_router_external_dispatch.py
Normal file
364
tests/test_router_external_dispatch.py
Normal file
|
|
@ -0,0 +1,364 @@
|
|||
"""Tests for routing a selected EXTERNAL compressor through the content router.
|
||||
|
||||
Scope 3 of the pluggable-compressor work: an opt-in ``headroom.compressor``
|
||||
entry point, when SELECTED, actually compresses matching real traffic through
|
||||
:meth:`ContentRouter._apply_strategy_to_content`, with fail-open fallback to the
|
||||
built-in dispatch. BACKWARD COMPATIBILITY is the hard requirement — with nothing
|
||||
selected the branch is inert and the default request path is byte-identical.
|
||||
|
||||
Covers:
|
||||
(a) a selected external compressor compresses a matching block end-to-end;
|
||||
(b) its ``recoverable`` (hash -> original) map is retrievable from the CCR store;
|
||||
(c) fail-open: an external that raises / malforms / expands falls back to the
|
||||
built-in path (never breaks the request);
|
||||
(d) NOT selected, or a non-matching content type, leaves the built-in path
|
||||
untouched (the external is never even invoked);
|
||||
plus the proxy seam that threads external names to the router.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.cache.compression_store import (
|
||||
get_compression_store,
|
||||
reset_compression_store,
|
||||
)
|
||||
from headroom.proxy.server import _external_compressor_selection
|
||||
from headroom.transforms.compressor_registry import (
|
||||
CompressInput,
|
||||
CompressorDescriptor,
|
||||
CompressOutput,
|
||||
)
|
||||
from headroom.transforms.content_router import (
|
||||
CompressionStrategy,
|
||||
ContentRouter,
|
||||
ContentRouterConfig,
|
||||
)
|
||||
|
||||
# A JSON array reliably routes to SMART_CRUSHER (content type application/json)
|
||||
# and is not touched by the STAGE-0 lossless fold, so it reaches the external
|
||||
# dispatch branch. Big enough that the reference compressor's output shrinks it.
|
||||
_JSON_ARRAY = (
|
||||
"["
|
||||
+ ",".join(
|
||||
f'{{"id":{i},"name":"item-{i}","status":"active","value":{i * 7},'
|
||||
f'"note":"a fairly long descriptive field number {i} to add bulk"}}'
|
||||
for i in range(40)
|
||||
)
|
||||
+ "]"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _memory_ccr(monkeypatch):
|
||||
"""Isolated in-memory CCR store + offline content detection per test.
|
||||
|
||||
``HEADROOM_DETECT_BACKEND=python`` forces the pure-Python regex detector so
|
||||
``compress()`` never touches the native Magika/ONNX detector (which needs a
|
||||
model download and blocks in this offline environment). The external-dispatch
|
||||
branch under test is independent of the detector backend.
|
||||
"""
|
||||
monkeypatch.setenv("HEADROOM_CCR_BACKEND", "memory")
|
||||
monkeypatch.setenv("HEADROOM_DETECT_BACKEND", "python")
|
||||
reset_compression_store()
|
||||
yield
|
||||
reset_compression_store()
|
||||
|
||||
|
||||
class _WordTokenizer:
|
||||
"""Word-count tokenizer stub — no model, deterministic, offline-safe."""
|
||||
|
||||
def count_text(self, text: object) -> int:
|
||||
return len(str(text).split())
|
||||
|
||||
def count_messages(self, messages: list[dict]) -> int:
|
||||
return sum(self.count_text(m.get("content", "")) for m in messages)
|
||||
|
||||
|
||||
def _json_array(tag: str) -> str:
|
||||
"""A distinct, spaced JSON array (>50 word-tokens, >500 chars) that the
|
||||
pure-Python detector routes to SMART_CRUSHER (application/json)."""
|
||||
return (
|
||||
"["
|
||||
+ ",".join(
|
||||
f'{{"id":{i},"tag":"{tag}","note":"long descriptive field number {i} '
|
||||
f'to add real bulk here"}}'
|
||||
for i in range(40)
|
||||
)
|
||||
+ "]"
|
||||
)
|
||||
|
||||
|
||||
def _tool_msg(call_id: str, content: str) -> dict:
|
||||
# tool_call_id with no matching assistant tool_calls -> not excluded -> the
|
||||
# non-frozen one reaches compression (matches Bash/shell output).
|
||||
return {"role": "tool", "tool_call_id": call_id, "content": content}
|
||||
|
||||
|
||||
class _RecordingExternal:
|
||||
"""Reference in-process external ``Compressor`` for the router tests.
|
||||
|
||||
Deterministically shrinks its input to a short marker string and, unless
|
||||
``recoverable=False``, returns a ``{hash: original}`` recovery map keyed by
|
||||
the same hash it embeds in the output — mirroring how SmartCrusher's markers
|
||||
point back into the CCR store.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str = "ext_json",
|
||||
content_types: tuple[str, ...] = ("application/json",),
|
||||
*,
|
||||
lossless: bool = False,
|
||||
recoverable: bool = True,
|
||||
raises: bool = False,
|
||||
expand: bool = False,
|
||||
malformed: bool = False,
|
||||
empty: bool = False,
|
||||
bad_hash: bool = False,
|
||||
) -> None:
|
||||
self._name = name
|
||||
self._content_types = list(content_types)
|
||||
self._lossless = lossless
|
||||
self._recoverable = recoverable
|
||||
self._raises = raises
|
||||
self._expand = expand
|
||||
self._malformed = malformed
|
||||
self._empty = empty
|
||||
self._bad_hash = bad_hash
|
||||
self.calls: list[CompressInput] = []
|
||||
|
||||
@property
|
||||
def descriptor(self) -> CompressorDescriptor:
|
||||
return CompressorDescriptor(
|
||||
name=self._name,
|
||||
content_types=self._content_types,
|
||||
lossless=self._lossless,
|
||||
cost_tier="fast",
|
||||
recoverable=self._recoverable,
|
||||
)
|
||||
|
||||
def compress(self, inp: CompressInput): # noqa: ANN201 - matches protocol
|
||||
self.calls.append(inp)
|
||||
if self._raises:
|
||||
raise RuntimeError("external boom")
|
||||
if self._malformed:
|
||||
return {"not": "a CompressOutput"}
|
||||
original = inp.content
|
||||
digest = hashlib.sha256(original.encode()).hexdigest()[:24]
|
||||
key = "zznothex" if self._bad_hash else digest
|
||||
if self._empty:
|
||||
content = ""
|
||||
elif self._expand:
|
||||
content = original + (" PADDING" * 200)
|
||||
else:
|
||||
content = f"[external-compressed <<ccr:{digest}>>]"
|
||||
recoverable = {key: original} if self._recoverable else {}
|
||||
return CompressOutput(
|
||||
content=content,
|
||||
tokens_before=len(original.split()),
|
||||
tokens_after=len(content.split()),
|
||||
lossless=self._lossless,
|
||||
markers=[f"external:{self._name}"],
|
||||
recoverable=recoverable,
|
||||
warnings=["reference-warning"],
|
||||
)
|
||||
|
||||
|
||||
def _cfg(**kwargs) -> ContentRouterConfig:
|
||||
"""Router config with Kompress OFF so the built-in fallback path never
|
||||
loads the ModernBERT ML model (keeps these unit tests fast and offline).
|
||||
The external-dispatch branch under test is independent of this flag."""
|
||||
kwargs.setdefault("enable_kompress", False)
|
||||
return ContentRouterConfig(**kwargs)
|
||||
|
||||
|
||||
def _router_with_external(comp: _RecordingExternal, selection):
|
||||
"""Build a router with ``comp`` registered and ``selection`` active."""
|
||||
router = ContentRouter(_cfg(active_external_compressors=selection))
|
||||
router.compressor_registry.register(comp, replace=True)
|
||||
# Re-resolve now that the external compressor is registered (the router
|
||||
# resolves the selection once at construction, before this injection).
|
||||
router._active_external_compressors = router._resolve_active_external_compressors()
|
||||
return router
|
||||
|
||||
|
||||
# ─────────────────────────── (a) end-to-end ──────────────────────────────────
|
||||
|
||||
|
||||
def test_selected_external_compresses_matching_block_end_to_end(_memory_ccr):
|
||||
comp = _RecordingExternal(name="ext_json", content_types=("application/json",))
|
||||
router = _router_with_external(comp, ["ext_json"])
|
||||
|
||||
result = router.compress(_JSON_ARRAY)
|
||||
|
||||
assert comp.calls, "external compressor should have been invoked"
|
||||
assert "external:ext_json" in result.strategy_chain
|
||||
assert "external-compressed" in result.compressed
|
||||
assert len(result.compressed) < len(_JSON_ARRAY)
|
||||
# The CompressInput carried the block + its detected MIME content type.
|
||||
assert comp.calls[0].content == _JSON_ARRAY
|
||||
assert comp.calls[0].content_type == "application/json"
|
||||
|
||||
|
||||
def test_external_dispatch_via_apply_strategy_returns_normal_shape(_memory_ccr):
|
||||
comp = _RecordingExternal(name="ext_json", content_types=("application/json",))
|
||||
router = _router_with_external(comp, ["ext_json"])
|
||||
|
||||
compressed, tokens, chain = router._apply_strategy_to_content(
|
||||
_JSON_ARRAY, CompressionStrategy.SMART_CRUSHER, ""
|
||||
)
|
||||
|
||||
assert chain == ["external:ext_json"]
|
||||
assert "external-compressed" in compressed
|
||||
# Tokens counted with the router's OWN estimator (a positive int), not the
|
||||
# compressor's self-report.
|
||||
assert isinstance(tokens, int) and tokens > 0
|
||||
|
||||
|
||||
# ─────────────────────────── (b) recoverable map ─────────────────────────────
|
||||
|
||||
|
||||
def test_recoverable_map_is_retrievable(_memory_ccr):
|
||||
comp = _RecordingExternal(name="ext_json", content_types=("application/json",))
|
||||
router = _router_with_external(comp, ["ext_json"])
|
||||
|
||||
router.compress(_JSON_ARRAY)
|
||||
|
||||
digest = hashlib.sha256(_JSON_ARRAY.encode()).hexdigest()[:24]
|
||||
entry = get_compression_store().retrieve(digest)
|
||||
assert entry is not None, "recoverable entry should be in the CCR store"
|
||||
assert entry.original_content == _JSON_ARRAY
|
||||
assert entry.compression_strategy == "external:ext_json"
|
||||
|
||||
|
||||
def test_non_hex_recoverable_hash_is_skipped_without_breaking(_memory_ccr):
|
||||
# A malformed (non-hex) recovery hash must not break the request; the block
|
||||
# is still compressed, only that entry is not retrievable.
|
||||
comp = _RecordingExternal(name="ext_json", content_types=("application/json",), bad_hash=True)
|
||||
router = _router_with_external(comp, ["ext_json"])
|
||||
|
||||
result = router.compress(_JSON_ARRAY)
|
||||
|
||||
assert "external:ext_json" in result.strategy_chain
|
||||
assert get_compression_store().retrieve("zznothex") is None
|
||||
|
||||
|
||||
# ─────────────────────────── (c) fail-open ───────────────────────────────────
|
||||
|
||||
|
||||
def test_external_raise_falls_back_to_builtin(_memory_ccr):
|
||||
comp = _RecordingExternal(name="ext_json", content_types=("application/json",), raises=True)
|
||||
router = _router_with_external(comp, ["ext_json"])
|
||||
|
||||
result = router.compress(_JSON_ARRAY)
|
||||
|
||||
assert comp.calls, "external should have been attempted"
|
||||
# Fell back: no external marker in the chain, request not broken.
|
||||
assert "external:ext_json" not in result.strategy_chain
|
||||
assert result.compressed and result.compressed.strip()
|
||||
|
||||
|
||||
def test_external_malformed_output_falls_back_to_builtin(_memory_ccr):
|
||||
comp = _RecordingExternal(name="ext_json", content_types=("application/json",), malformed=True)
|
||||
router = _router_with_external(comp, ["ext_json"])
|
||||
|
||||
result = router.compress(_JSON_ARRAY)
|
||||
|
||||
assert comp.calls
|
||||
assert "external:ext_json" not in result.strategy_chain
|
||||
|
||||
|
||||
def test_external_expansion_falls_back_to_builtin(_memory_ccr):
|
||||
comp = _RecordingExternal(name="ext_json", content_types=("application/json",), expand=True)
|
||||
router = _router_with_external(comp, ["ext_json"])
|
||||
|
||||
result = router.compress(_JSON_ARRAY)
|
||||
|
||||
assert comp.calls
|
||||
assert "external:ext_json" not in result.strategy_chain
|
||||
# Never expands: the returned block is no larger than the input.
|
||||
assert len(result.compressed) <= len(_JSON_ARRAY)
|
||||
|
||||
|
||||
def test_external_empty_output_falls_back_to_builtin(_memory_ccr):
|
||||
comp = _RecordingExternal(name="ext_json", content_types=("application/json",), empty=True)
|
||||
router = _router_with_external(comp, ["ext_json"])
|
||||
|
||||
result = router.compress(_JSON_ARRAY)
|
||||
|
||||
assert comp.calls
|
||||
assert "external:ext_json" not in result.strategy_chain
|
||||
assert result.compressed.strip(), "non-empty input must never blank out"
|
||||
|
||||
|
||||
# ─────────────── (d) not selected / non-matching → built-in ──────────────────
|
||||
|
||||
|
||||
def test_not_selected_leaves_builtin_path_unchanged(_memory_ccr):
|
||||
comp = _RecordingExternal(name="ext_json", content_types=("application/json",))
|
||||
# Registered but NOT selected.
|
||||
router = _router_with_external(comp, None)
|
||||
|
||||
baseline = ContentRouter(_cfg()).compress(_JSON_ARRAY)
|
||||
result = router.compress(_JSON_ARRAY)
|
||||
|
||||
assert comp.calls == [], "unselected external must never be invoked"
|
||||
assert "external:ext_json" not in result.strategy_chain
|
||||
# Byte-identical to a plain router with no external registered at all.
|
||||
assert result.compressed == baseline.compressed
|
||||
assert result.strategy_chain == baseline.strategy_chain
|
||||
|
||||
|
||||
def test_non_matching_content_type_leaves_builtin_path_unchanged(_memory_ccr):
|
||||
# Selected, but declares a content type the JSON block never has.
|
||||
comp = _RecordingExternal(name="ext_diff", content_types=("text/x-diff",))
|
||||
router = _router_with_external(comp, ["ext_diff"])
|
||||
|
||||
baseline = ContentRouter(_cfg()).compress(_JSON_ARRAY)
|
||||
result = router.compress(_JSON_ARRAY)
|
||||
|
||||
assert comp.calls == [], "non-matching external must never be invoked"
|
||||
assert "external:ext_diff" not in result.strategy_chain
|
||||
assert result.compressed == baseline.compressed
|
||||
|
||||
|
||||
def test_default_config_has_no_external_selection():
|
||||
assert ContentRouterConfig().active_external_compressors is None
|
||||
router = ContentRouter(ContentRouterConfig())
|
||||
assert router._active_external_compressors == []
|
||||
|
||||
|
||||
# ─────────────────────────── proxy seam ──────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"selection,expected",
|
||||
[
|
||||
(None, None),
|
||||
(set(), None),
|
||||
({"", " "}, None),
|
||||
({"smart_crusher"}, None), # built-in only → no external
|
||||
({"smart_crusher", "kompress"}, None),
|
||||
({"my_ext"}, ["my_ext"]),
|
||||
({"kompress", "my_ext"}, ["my_ext"]),
|
||||
({"b_ext", "a_ext"}, ["a_ext", "b_ext"]), # sorted
|
||||
({"*"}, ["*"]),
|
||||
({"*", "my_ext"}, ["*"]), # wildcard wins
|
||||
],
|
||||
)
|
||||
def test_external_compressor_selection_helper(selection, expected):
|
||||
assert _external_compressor_selection(selection) == expected
|
||||
|
||||
|
||||
def test_wildcard_selection_activates_registered_external(_memory_ccr):
|
||||
comp = _RecordingExternal(name="ext_json", content_types=("application/json",))
|
||||
router = _router_with_external(comp, ["*"])
|
||||
|
||||
result = router.compress(_JSON_ARRAY)
|
||||
|
||||
assert comp.calls
|
||||
assert "external:ext_json" in result.strategy_chain
|
||||
Loading…
Add table
Add a link
Reference in a new issue