akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(transforms): add pluggable compressor registry + headroom.compressor entry point (#2370)

Browse source 
## What
Adds a pluggable compressor registry and a `headroom.compressor`
entry-point group so compressors can be registered, discovered, and
selected by name.

- Pure-data contract (`CompressorDescriptor`, `CompressInput`,
`CompressOutput`, `Compressor` Protocol). Only plain types
(`str`/`int`/`bool`/`list`/`dict`) cross the boundary — no tokenizer,
store, or config objects — so the same contract can be implemented
outside Python.
- `CompressorRegistry`: starts empty and accepts explicit registrations
by name; discovers external compressors from the `headroom.compressor`
group fail-open (mirrors the existing pipeline-extension discovery);
resolves an opt-in selection (`select`/`active`) — nothing active by
default, `"*"` for all, otherwise a name allowlist with unknown names
logged and skipped. Discovery loads compressors but never invokes
`compress`.

## Why
Compressors are currently constructed and dispatched via a hardcoded
chain in the content router; there is no way to add or select one
without editing the router. This lands a name-addressable seam so that
becomes possible.

## Behavior change
None. Purely additive — not wired into `content_router`, the proxy
server, or config, and constructing the registry has no global side
effects. Router integration is a deliberate follow-up.

## Testing
- `pytest tests/test_compressor_registry.py -q` → 11 passed (contract
round-trip, registration, opt-in selection semantics, wildcard,
unknown-name skip, monkeypatched entry-point discovery,
discovery-never-runs-compress).
- `ruff check` / `ruff format` → clean; `mypy` → no issues.
This commit is contained in:
Tejas Chopra 2026-07-17 21:55:16 -07:00 • committed by GitHub
parent 44136ed042
commit a02073e332
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 527 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

282
headroom/transforms/compressor_registry.py Normal file
View file

@ -0,0 +1,282 @@
"""Pluggable compressor registry and ``headroom.compressor`` entry-point seam.
This module lands a *name-addressable* seam for compressors: built-in
compressors can be registered explicitly, third-party compressors can be
discovered from the ``headroom.compressor`` entry-point group, and a caller can
opt in to a specific set of them by name.
It is deliberately **additive**. Nothing here is wired into the content router,
proxy server, or config in this change — registering or discovering a compressor
has no effect on request handling until an integration point is added in a
follow-up. Constructing a ``CompressorRegistry`` has no global side effects.
Rust-portable contract
----------------------
The compressor boundary is **pure data in / data out**. No Python-only objects
(tokenizer instances, live store handles, rich config classes) cross it:
* :class:`CompressorDescriptor` — static, declarative capability metadata.
* :class:`CompressInput` — the content plus plain ``dict`` config/budget.
* :class:`CompressOutput` — the compressed content plus plain counts, string
markers, a ``hash -> original`` recovery map, and string warnings.
Every field is a ``str``, ``int``, ``bool``, ``list``, or ``dict`` of those, so
an equivalent contract can be implemented in another language (e.g. a Rust
compressor invoked over the same shapes) without carrying Python objects.
Discovery vs. selection (opt-in)
--------------------------------
Discovery (:meth:`CompressorRegistry.discover`) enumerates and *loads* every
registered entry point — it may import a module and construct the compressor
object — but it never invokes ``compress``. Selection
(:meth:`CompressorRegistry.select` / :meth:`CompressorRegistry.active`) is
opt-in: with no names (or an empty set) nothing is active; the literal ``"*"``
selects everything registered; otherwise only the named-and-registered
compressors are active. This mirrors the opt-in model used for proxy extensions
so that merely installing a third-party package cannot silently change behavior.
External packages register a compressor like this::
[project.entry-points."headroom.compressor"]
my_compressor = "my_pkg.compressor:MyCompressor"
The value may be a :class:`Compressor` instance or a zero-arg class implementing
the protocol; a class is instantiated during discovery.
"""
from __future__ import annotations
import importlib.metadata
import logging
from dataclasses import dataclass, field
from typing import Any, Protocol, runtime_checkable
log = logging.getLogger(__name__)
ENTRY_POINT_GROUP = "headroom.compressor"
#: Recognized values for :attr:`CompressorDescriptor.cost_tier`.
COST_TIERS: tuple[str, ...] = ("fast", "ml", "remote")
@dataclass(frozen=True)
class CompressorDescriptor:
"""Static, declarative metadata describing a compressor's capabilities.
Attributes:
name: Canonical, unique compressor name used for registration/selection.
content_types: Content types this compressor handles (e.g. ``["text/plain"]``).
lossless: Whether compression is losslessly reversible.
cost_tier: One of :data:`COST_TIERS` — ``"fast"`` (local/cheap),
``"ml"`` (local model inference), or ``"remote"`` (network call).
recoverable: Whether the compressor can emit a ``hash -> original``
recovery map in :attr:`CompressOutput.recoverable`.
"""
name: str
content_types: list[str]
lossless: bool
cost_tier: str
recoverable: bool
@dataclass
class CompressInput:
"""Pure-data input to :meth:`Compressor.compress`.
Attributes:
content: The raw content to compress.
content_type: The content type of ``content``.
query: Optional task/query hint for relevance-aware compressors.
config: Plain compressor-specific configuration.
budget: Plain budget hints, e.g. ``target_ratio`` (float),
``time_ms`` (int), ``max_items`` (int).
"""
content: str
content_type: str
query: str = ""
config: dict[str, Any] = field(default_factory=dict)
budget: dict[str, Any] = field(default_factory=dict)
@dataclass
class CompressOutput:
"""Pure-data output from :meth:`Compressor.compress`.
Attributes:
content: The compressed content.
tokens_before: Token count of the input content.
tokens_after: Token count of the compressed content.
lossless: Whether this particular result is losslessly reversible.
markers: Marker strings describing what was applied (e.g. for routing).
recoverable: ``hash -> original`` map for recovering dropped content.
warnings: Non-fatal warning strings emitted during compression.
"""
content: str
tokens_before: int
tokens_after: int
lossless: bool
markers: list[str] = field(default_factory=list)
recoverable: dict[str, str] = field(default_factory=dict)
warnings: list[str] = field(default_factory=list)
@runtime_checkable
class Compressor(Protocol):
"""Name-addressable compressor contract (pure data in / data out)."""
@property
def descriptor(self) -> CompressorDescriptor:
"""Return this compressor's static capability metadata."""
...
def compress(self, inp: CompressInput) -> CompressOutput:
"""Compress ``inp`` and return a :class:`CompressOutput`."""
...
class CompressorRegistry:
"""Registry of compressors addressable by :attr:`CompressorDescriptor.name`.
Starts empty. Built-in compressors are added via :meth:`register`; external
compressors are added via :meth:`discover`. Selection is opt-in.
"""
def __init__(self) -> None:
self._compressors: dict[str, Compressor] = {}
def register(self, compressor: Compressor, *, replace: bool = False) -> str:
"""Register ``compressor`` under its ``descriptor.name``.
Args:
compressor: The compressor to register.
replace: If ``True``, replace an existing registration of the same
name instead of raising.
Returns:
The registered name.
Raises:
ValueError: If the name is empty, or already registered and
``replace`` is ``False``.
"""
name = compressor.descriptor.name
if not name:
raise ValueError("compressor descriptor.name must be non-empty")
if name in self._compressors and not replace:
raise ValueError(f"compressor {name!r} is already registered")
self._compressors[name] = compressor
return name
def get(self, name: str) -> Compressor | None:
"""Return the registered compressor named ``name``, or ``None``."""
return self._compressors.get(name)
def names(self) -> list[str]:
"""Return all registered compressor names, sorted."""
return sorted(self._compressors)
def descriptors(self) -> list[CompressorDescriptor]:
"""Return the descriptors of all registered compressors, sorted by name."""
return [self._compressors[n].descriptor for n in sorted(self._compressors)]
def discover(self, *, replace: bool = False) -> list[str]:
"""Load and register compressors from the ``headroom.compressor`` group.
Mirrors the pipeline extension discovery helper: entry points are
enumerated fail-open, each is loaded fail-open, and a class value is
instantiated fail-open. A broken third-party package is logged and
skipped rather than aborting discovery. ``compress`` is never invoked
here — discovery only loads and registers.
Args:
replace: Passed through to :meth:`register` for name collisions.
Returns:
The list of newly registered compressor names.
"""
discovered: list[str] = []
try:
entries = importlib.metadata.entry_points(group=ENTRY_POINT_GROUP)
except Exception as exc: # noqa: BLE001 - importlib metadata varies by runtime
log.debug("compressor registry: entry-point enumeration failed: %s", exc)
return discovered
for entry in entries:
try:
obj = entry.load()
except Exception as exc: # noqa: BLE001 - third-party load failures are isolated
log.warning("compressor %r failed to load: %s", entry.name, exc)
continue
if isinstance(obj, type):
try:
obj = obj()
except Exception as exc: # noqa: BLE001
log.warning("compressor %r failed to initialize: %s", entry.name, exc)
continue
try:
name = self.register(obj, replace=replace)
except Exception as exc: # noqa: BLE001 - bad descriptor / duplicate name
log.warning("compressor %r failed to register: %s", entry.name, exc)
continue
discovered.append(name)
return discovered
@staticmethod
def _resolve_selection(names: set[str] | None) -> set[str]:
"""Normalize a requested selection: strip whitespace, drop empties."""
if not names:
return set()
return {stripped for n in names if (stripped := n.strip())}
def select(self, names: set[str] | None) -> set[str]:
"""Resolve an opt-in selection to the set of *active* registered names.
Mirrors the proxy-extension opt-in model:
* ``None`` or an empty set selects nothing (opt-in default).
* ``"*"`` selects every registered compressor.
* Otherwise only names that are both requested and registered are
active; requested-but-unregistered names are logged and skipped.
Discovery is never triggered here and ``compress`` is never invoked;
this only decides which already-registered compressors are active.
"""
requested = self._resolve_selection(names)
registered = set(self._compressors)
if not requested:
if registered:
log.info(
"compressors registered but none selected (opt-in): %s. "
"Select by name or use '*' for all.",
",".join(sorted(registered)),
)
return set()
if "*" in requested:
return registered
active = requested & registered
missing = requested - registered
if missing:
log.warning(
"compressors requested but not registered: %s (available: %s)",
",".join(sorted(missing)),
",".join(sorted(registered)) or "<none>",
)
return active
def active(self, selection: set[str] | None) -> list[Compressor]:
"""Return the active compressor objects for ``selection``, sorted by name.
``selection`` is the raw opt-in request; it is resolved via
:meth:`select`.
"""
return [self._compressors[name] for name in sorted(self.select(selection))]

245
tests/test_compressor_registry.py Normal file
View file

@ -0,0 +1,245 @@
"""Tests for the pluggable compressor registry and entry-point discovery."""
from __future__ import annotations
import importlib.metadata
import logging
from dataclasses import asdict, dataclass
import pytest
from headroom.transforms import compressor_registry
from headroom.transforms.compressor_registry import (
ENTRY_POINT_GROUP,
CompressInput,
CompressorDescriptor,
CompressorRegistry,
CompressOutput,
)
class FakeCompressor:
"""Minimal in-memory compressor implementing the Compressor protocol."""
def __init__(
self,
name: str = "fake",
*,
content_types: list[str] | None = None,
lossless: bool = True,
cost_tier: str = "fast",
recoverable: bool = True,
raise_on_compress: bool = False,
) -> None:
self._descriptor = CompressorDescriptor(
name=name,
content_types=content_types or ["text/plain"],
lossless=lossless,
cost_tier=cost_tier,
recoverable=recoverable,
)
self._raise_on_compress = raise_on_compress
@property
def descriptor(self) -> CompressorDescriptor:
return self._descriptor
def compress(self, inp: CompressInput) -> CompressOutput:
if self._raise_on_compress:
raise RuntimeError("compress must not run during discovery/selection")
half = inp.content[: max(1, len(inp.content) // 2)]
return CompressOutput(
content=half,
tokens_before=len(inp.content),
tokens_after=len(half),
lossless=self._descriptor.lossless,
markers=[f"fake:{inp.content_type}"],
recoverable={"deadbeef": inp.content} if self._descriptor.recoverable else {},
warnings=[],
)
@dataclass
class FakeEntryPoint:
"""Stand-in for ``importlib.metadata.EntryPoint`` used in discovery tests."""
name: str
value: object
def load(self) -> object:
if isinstance(self.value, Exception):
raise self.value
return self.value
def test_registry_starts_empty() -> None:
reg = CompressorRegistry()
assert reg.names() == []
assert reg.get("fake") is None
assert reg.active(None) == []
assert reg.active({"*"}) == []
def test_descriptor_round_trips() -> None:
desc = CompressorDescriptor(
name="fake",
content_types=["text/plain", "text/markdown"],
lossless=True,
cost_tier="fast",
recoverable=True,
)
assert desc.name == "fake"
assert desc.content_types == ["text/plain", "text/markdown"]
assert desc.lossless is True
assert desc.cost_tier in compressor_registry.COST_TIERS
assert desc.recoverable is True
# Pure-data: fully representable as a plain dict for a cross-language boundary.
assert asdict(desc) == {
"name": "fake",
"content_types": ["text/plain", "text/markdown"],
"lossless": True,
"cost_tier": "fast",
"recoverable": True,
}
def test_register_and_get_by_name() -> None:
reg = CompressorRegistry()
comp = FakeCompressor("fake")
assert reg.register(comp) == "fake"
assert reg.get("fake") is comp
assert reg.names() == ["fake"]
assert [d.name for d in reg.descriptors()] == ["fake"]
def test_register_rejects_empty_name_and_duplicates() -> None:
reg = CompressorRegistry()
with pytest.raises(ValueError):
reg.register(FakeCompressor(""))
reg.register(FakeCompressor("fake"))
with pytest.raises(ValueError):
reg.register(FakeCompressor("fake"))
replacement = FakeCompressor("fake")
assert reg.register(replacement, replace=True) == "fake"
assert reg.get("fake") is replacement
def test_selection_is_opt_in_none_and_empty_select_nothing(caplog) -> None:
reg = CompressorRegistry()
reg.register(FakeCompressor("a"))
reg.register(FakeCompressor("b"))
with caplog.at_level(logging.INFO, logger=compressor_registry.log.name):
assert reg.select(None) == set()
assert reg.select(set()) == set()
assert reg.active(None) == []
# Opt-in default is surfaced.
assert "none selected (opt-in)" in caplog.text
def test_selection_wildcard_selects_all() -> None:
reg = CompressorRegistry()
reg.register(FakeCompressor("a"))
reg.register(FakeCompressor("b"))
assert reg.select({"*"}) == {"a", "b"}
active = reg.active({"*"})
assert [c.descriptor.name for c in active] == ["a", "b"]
def test_selection_by_name_and_unknown_is_skipped(caplog) -> None:
reg = CompressorRegistry()
reg.register(FakeCompressor("a"))
reg.register(FakeCompressor("b"))
assert reg.select({"a"}) == {"a"}
assert [c.descriptor.name for c in reg.active({"a", "b"})] == ["a", "b"]
with caplog.at_level(logging.WARNING, logger=compressor_registry.log.name):
assert reg.select({"a", "nope"}) == {"a"}
assert "requested but not registered: nope" in caplog.text
def test_selection_normalizes_whitespace_and_empties() -> None:
reg = CompressorRegistry()
reg.register(FakeCompressor("a"))
assert reg.select({" a ", "", " "}) == {"a"}
def test_compress_round_trips_input_to_output() -> None:
reg = CompressorRegistry()
reg.register(FakeCompressor("fake"))
comp = reg.active({"fake"})[0]
inp = CompressInput(
content="hello world, this is content",
content_type="text/plain",
query="summarize",
config={"aggressive": True},
budget={"target_ratio": 0.5, "time_ms": 10, "max_items": 3},
)
out = comp.compress(inp)
assert isinstance(out, CompressOutput)
assert out.tokens_before == len(inp.content)
assert out.tokens_after == len(out.content)
assert out.tokens_after < out.tokens_before
assert out.lossless is True
assert out.markers == ["fake:text/plain"]
assert out.recoverable == {"deadbeef": inp.content}
assert out.warnings == []
def test_discovery_registers_selects_but_does_not_run_compress(monkeypatch, caplog) -> None:
# A compressor whose compress() would raise proves discovery never calls it.
guarded = FakeCompressor("guarded", raise_on_compress=True)
class NeedsInit:
def __init__(self) -> None:
raise RuntimeError("bad init")
monkeypatch.setattr(
importlib.metadata,
"entry_points",
lambda group=None: (
[
FakeEntryPoint("guarded-instance", guarded),
FakeEntryPoint("working-class", FakeCompressor),
FakeEntryPoint("bad-load", RuntimeError("bad load")),
FakeEntryPoint("bad-init", NeedsInit),
]
if group == ENTRY_POINT_GROUP
else []
),
)
reg = CompressorRegistry()
with caplog.at_level(logging.WARNING, logger=compressor_registry.log.name):
discovered = reg.discover()
# Only the loadable/constructible compressors are registered.
assert sorted(discovered) == ["fake", "guarded"]
assert reg.get("guarded") is guarded
# Failures are logged and skipped, not raised.
assert "bad-load" in caplog.text
assert "bad-init" in caplog.text
# Discovery did not run compress; selection alone does not either.
active = reg.active({"guarded"})
assert active == [guarded]
# compress only runs when the caller invokes it.
with pytest.raises(RuntimeError):
active[0].compress(CompressInput(content="x", content_type="text/plain"))
def test_discovery_handles_enumeration_failure(monkeypatch) -> None:
monkeypatch.setattr(
importlib.metadata,
"entry_points",
lambda group=None: (_ for _ in ()).throw(RuntimeError("boom")),
)
reg = CompressorRegistry()
assert reg.discover() == []
assert reg.names() == []