From b6eb7a7613f73d4a1bce790a9da217f66d2621ad Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Tue, 14 Jul 2026 01:01:11 -0400 Subject: [PATCH] feat(kompress): optional remote compression endpoint (HEADROOM_KOMPRESS_ENDPOINT) (#2171) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Adds an **opt-in remote Kompress backend** so the proxy can offload Kompress ML inference to a hosted `/compress` endpoint instead of loading the ONNX model in-process. This lets Headroom run as a lean proxy in a sandbox installed with only `[proxy]` deps while the model runs elsewhere. The feature is purely additive: with `HEADROOM_KOMPRESS_ENDPOINT` unset, behavior remains the existing in-process Kompress path. ## Type of Change - [x] New feature (non-breaking change that adds functionality) ## Changes Made - `headroom/transforms/kompress_remote.py`: adds `RemoteKompressCompressor`, a `KompressCompressor`-compatible HTTP client that posts to `/compress`, sends optional bearer auth, skips network for tiny inputs, and fails open on HTTP/network/malformed-response errors. - `headroom/transforms/kompress_compressor.py`: extracts `store_kompress_in_ccr()` so the remote client reuses the same proxy-local CCR marker/storage policy without importing the ML model. - `headroom/transforms/content_router.py`: selects the remote compressor when `HEADROOM_KOMPRESS_ENDPOINT` is set, while `"disabled"` still wins and the unset path remains local Kompress. - `tests/test_transforms/test_kompress_remote.py`: covers mocked remote success, auth/header/request behavior, tiny-input no-call behavior, HTTP fail-open, malformed-success fail-open, and router env selection. ## Testing - [x] Unit tests pass (`uv run --extra dev pytest tests/test_transforms/test_kompress_remote.py -q`) - [x] Linting passes (`uvx ruff@0.15.17 check headroom/transforms/kompress_remote.py headroom/transforms/kompress_compressor.py headroom/transforms/content_router.py tests/test_transforms/test_kompress_remote.py`) - [x] Formatting passes (`uvx ruff@0.15.17 format --check headroom/transforms/kompress_remote.py headroom/transforms/kompress_compressor.py headroom/transforms/content_router.py tests/test_transforms/test_kompress_remote.py`) - [x] Type checking passes (`uv run --extra dev mypy headroom/transforms/kompress_remote.py headroom/transforms/kompress_compressor.py headroom/transforms/content_router.py`) - [x] New tests added for new functionality - [x] Manual testing performed by the author against a live endpoint ## Real Behavior Proof - Environment: Windows 11 review worktree, Python 3.13.3 for mocked tests; author also manually tested against a Modal deployment of `chopratejas/kompress-v2-base`. - Exact command / steps: ran the focused mocked endpoint test file plus lint/format/mypy on the changed modules. - Observed result: remote success maps endpoint response into `KompressResult`; short inputs do not call the network; 503 responses and malformed 200 responses return the original content; router selects the remote compressor only when the env var is set. - Not tested: full `pytest` suite; production concurrency/latency under load; endpoints other than the author's Modal reference deployment. ## 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 - [ ] I have made corresponding changes to the documentation — follow-up README flag section - [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 - [ ] I have updated the CHANGELOG.md if applicable ## Additional Notes - The endpoint/deploy artifact (`modal_serve.py`) lives in the separate `kompress` repo; this PR is only the client-side flag. - The endpoint is intentionally stateless for CCR. Original-content storage and retrieval markers remain proxy-local. - Design note: this capability is intentionally in OSS as an opt-in flag. The same flag serves self-hosted endpoints and, later, a hosted endpoint. --------- Co-authored-by: JerrettDavis --- headroom/transforms/content_router.py | 32 +++++ headroom/transforms/kompress_compressor.py | 72 ++++++---- headroom/transforms/kompress_remote.py | 133 ++++++++++++++++++ pyproject.toml | 19 +++ tests/test_transforms/test_kompress_remote.py | 111 +++++++++++++++ 5 files changed, 337 insertions(+), 30 deletions(-) create mode 100644 headroom/transforms/kompress_remote.py create mode 100644 tests/test_transforms/test_kompress_remote.py diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index e9e07912c..164bdb2d7 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -2833,6 +2833,15 @@ class ContentRouter(Transform): if model_id == "disabled": return None + # Remote Kompress (HEADROOM_KOMPRESS_ENDPOINT): offload inference to a + # hosted /compress endpoint so a sandboxed proxy needs no local ML deps. + # Intercepts BOTH default and custom-model paths (the endpoint's deployed + # model is authoritative) and bypasses is_kompress_available() — there is + # nothing to load locally. The CCR store stays proxy-local. + remote = self._get_remote_kompress() + if remote is not None: + return remote + # Custom model — don't touch self._kompress (that's the default cache) if model_id: try: @@ -2874,6 +2883,29 @@ class ContentRouter(Transform): logger.debug("Kompress dependencies not available") return self._kompress + def _get_remote_kompress(self) -> Any: + """Return a cached RemoteKompressCompressor when HEADROOM_KOMPRESS_ENDPOINT + is set, else None. + + The endpoint runs the model, so this needs no local ML deps and no + is_kompress_available() gate. Cached per ContentRouter instance so the + httpx connection pool is reused across requests. + """ + endpoint = os.environ.get("HEADROOM_KOMPRESS_ENDPOINT", "").strip() + if not endpoint: + return None + if getattr(self, "_kompress_remote", None) is None: + from .kompress_compressor import KompressConfig + from .kompress_remote import RemoteKompressCompressor + + self._kompress_remote = RemoteKompressCompressor( + endpoint=endpoint, + token=os.environ.get("HEADROOM_KOMPRESS_ENDPOINT_TOKEN") or None, + config=KompressConfig(enable_ccr=self.config.ccr_inject_marker), + ) + logger.info("Kompress: using remote endpoint %s", endpoint) + return self._kompress_remote + def _get_image_optimizer(self) -> Any: """Create an ImageCompressor for one optimization pass. diff --git a/headroom/transforms/kompress_compressor.py b/headroom/transforms/kompress_compressor.py index 9ad410399..e31d91258 100644 --- a/headroom/transforms/kompress_compressor.py +++ b/headroom/transforms/kompress_compressor.py @@ -890,6 +890,47 @@ class KompressResult: return (self.tokens_saved / self.original_tokens) * 100 +def store_kompress_in_ccr(original: str, compressed: str, original_tokens: int) -> str | None: + """Store an original->compressed mapping in the proxy-local CCR store and + return its retrieval hash (or None on any failure). + + Module-level so both the in-process compressor and the remote client + (:mod:`headroom.transforms.kompress_remote`) share one CCR policy. Model-free + — touches only the compression store + telemetry, never the ONNX model — so + it works in a sandboxed proxy installed without the ``[ml]`` extra. + """ + try: + from ..cache.compression_store import get_compression_store + + signature = _kompress_content_signature(original) + compressed_tokens = len(compressed.split()) + store = get_compression_store() + cache_key = store.store( + original, + compressed, + original_tokens=original_tokens, + compressed_tokens=compressed_tokens, + original_item_count=original_tokens, + compressed_item_count=compressed_tokens, + tool_signature_hash=signature.structure_hash, + compression_strategy="kompress", + ) + with contextlib.suppress(Exception): + from ..telemetry import get_toin + + get_toin().record_compression( + tool_signature=signature, + original_count=original_tokens, + compressed_count=compressed_tokens, + original_tokens=original_tokens, + compressed_tokens=compressed_tokens, + strategy="kompress", + ) + return cache_key + except Exception: + return None + + class KompressCompressor(Transform): """Kompress: ModernBERT token compressor. @@ -1549,33 +1590,4 @@ class KompressCompressor(Transform): ) def _store_in_ccr(self, original: str, compressed: str, original_tokens: int) -> str | None: - try: - from ..cache.compression_store import get_compression_store - - signature = _kompress_content_signature(original) - compressed_tokens = len(compressed.split()) - store = get_compression_store() - cache_key = store.store( - original, - compressed, - original_tokens=original_tokens, - compressed_tokens=compressed_tokens, - original_item_count=original_tokens, - compressed_item_count=compressed_tokens, - tool_signature_hash=signature.structure_hash, - compression_strategy="kompress", - ) - with contextlib.suppress(Exception): - from ..telemetry import get_toin - - get_toin().record_compression( - tool_signature=signature, - original_count=original_tokens, - compressed_count=compressed_tokens, - original_tokens=original_tokens, - compressed_tokens=compressed_tokens, - strategy="kompress", - ) - return cache_key - except Exception: - return None + return store_kompress_in_ccr(original, compressed, original_tokens) diff --git a/headroom/transforms/kompress_remote.py b/headroom/transforms/kompress_remote.py new file mode 100644 index 000000000..6c335e5d3 --- /dev/null +++ b/headroom/transforms/kompress_remote.py @@ -0,0 +1,133 @@ +"""Remote Kompress: offload ML compression to a hosted ``/compress`` endpoint. + +Lets a sandboxed proxy — installed WITHOUT the ``[ml]`` extra (no torch/onnx) — +still run Kompress by calling a remote endpoint over HTTP. The class mirrors +:class:`~headroom.transforms.kompress_compressor.KompressCompressor`'s public +surface (``is_ready`` / ``preload`` / ``ensure_background_load`` / ``compress``), +so it is a drop-in at the ContentRouter seam. + +Only the model inference is remote. The CCR store + retrieval marker stay +proxy-local (the endpoint is stateless, ``enable_ccr=False``), so +``headroom_retrieve`` keeps working and original content never persists off-box. + +Enabled by ``HEADROOM_KOMPRESS_ENDPOINT`` (+ optional +``HEADROOM_KOMPRESS_ENDPOINT_TOKEN``) — see ``ContentRouter._get_kompress``. +""" + +from __future__ import annotations + +import logging + +import httpx + +from .kompress_compressor import KompressConfig, KompressResult, store_kompress_in_ccr + +logger = logging.getLogger(__name__) + +# Below this word count local Kompress passes through verbatim (KompressCompressor +# .compress); mirror it so we never pay a round-trip on a trivially small block. +_MIN_WORDS = 10 + +# Accept-any-shrink CCR gate, identical to KompressCompressor.compress: only +# store + mark when the shrink is worth the retrieval marker's own cost. +_CCR_RATIO_GATE = 0.8 + + +class RemoteKompressCompressor: + """Drop-in for KompressCompressor that POSTs to a hosted ``/compress`` endpoint. + + Fails OPEN: any network/HTTP error returns the content verbatim so a flaky + endpoint degrades compression rather than breaking the proxy. + """ + + name = "kompress_compressor" + + def __init__( + self, + endpoint: str, + token: str | None = None, + config: KompressConfig | None = None, + timeout: float = 20.0, + ) -> None: + self.config = config or KompressConfig() + self._url = endpoint.rstrip("/") + "/compress" + self._headers = {"content-type": "application/json"} + if token: + self._headers["authorization"] = f"Bearer {token}" + # httpx.Client is safe to share across the proxy's worker threads. + self._client = httpx.Client(timeout=timeout) + + # Nothing to load locally; short-circuit the router straight to compress(). + def is_ready(self) -> bool: + return True + + def preload(self, *, allow_download: bool = True) -> str: + return "remote" + + def ensure_background_load(self) -> None: + return None + + def _passthrough(self, content: str, n_words: int) -> KompressResult: + return KompressResult( + compressed=content, + original=content, + original_tokens=n_words, + compressed_tokens=n_words, + compression_ratio=1.0, + model_used=self.config.model_id, + ) + + def compress( + self, + content: str, + context: str = "", + content_type: str | None = None, + question: str | None = None, + target_ratio: float | None = None, + *, + allow_download: bool = True, + ) -> KompressResult: + n_words = len(content.split()) + if n_words < _MIN_WORDS: + return self._passthrough(content, n_words) + + try: + resp = self._client.post( + self._url, + headers=self._headers, + json={"content": content, "target_ratio": target_ratio}, + ) + resp.raise_for_status() + data = resp.json() + compressed = data["compressed"] + if not isinstance(compressed, str): + raise TypeError("remote Kompress response field 'compressed' must be a string") + except Exception as e: # fail OPEN — never break the proxy on a bad endpoint + logger.warning("Remote Kompress failed (%s); passing through", e) + return self._passthrough(content, n_words) + + result = KompressResult( + compressed=compressed, + original=content, + original_tokens=int(data.get("original_tokens", n_words)), + compressed_tokens=int(data.get("compressed_tokens", len(compressed.split()))), + compression_ratio=float(data.get("compression_ratio", 1.0)), + model_used=str(data.get("model_used", self.config.model_id)), + ) + + # CCR stays PROXY-LOCAL: endpoint is stateless (enable_ccr=False), so we + # store the mapping + append the retrieval marker here — same policy and + # marker format as KompressCompressor.compress. + if self.config.enable_ccr and result.compression_ratio < _CCR_RATIO_GATE: + cache_key = store_kompress_in_ccr(content, compressed, result.original_tokens) + if cache_key: + result.cache_key = cache_key + result.compressed += ( + f"\n[{result.original_tokens} items compressed to " + f"{result.compressed_tokens}. Retrieve more: hash={cache_key}]" + ) + + return result + + def close(self) -> None: + self._client.close() diff --git a/pyproject.toml b/pyproject.toml index 86ef1a514..01424f235 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -293,6 +293,25 @@ dev = [ all = [ "headroom-ai[proxy,code,ml,memory,relevance,image,reports,otel,evals,voice,html,mcp,spreadsheet]", ] +# Sandbox: a lean proxy with ALL torch-free capability — for running Headroom in +# a locked-down/low-resource sandbox and offloading heavy ML elsewhere. +# +# = [all] MINUS: +# - image (SigLIP/OCR — excluded by request) +# - ml (torch — the PyTorch Kompress backend; ONNX path in [proxy] still +# runs Kompress locally with no torch, or offload it entirely via +# HEADROOM_KOMPRESS_ENDPOINT) +# - voice (excluded by request) +# - memory + evals (both pull sentence-transformers -> torch, i.e. the very +# ML weight a sandbox avoids; evals is a dev/test harness, not a +# runtime feature). Opt back in explicitly if you accept torch: +# pip install headroom-ai[sandbox,memory] +# +# Everything kept here is torch-free: code-aware compression (tree-sitter), +# embedding relevance (fastembed), HTML/spreadsheet ingestion, reports, OTel. +sandbox = [ + "headroom-ai[proxy,code,relevance,reports,otel,html,mcp,spreadsheet]", +] [project.scripts] headroom = "headroom.cli:main" diff --git a/tests/test_transforms/test_kompress_remote.py b/tests/test_transforms/test_kompress_remote.py new file mode 100644 index 000000000..f7d81cc4b --- /dev/null +++ b/tests/test_transforms/test_kompress_remote.py @@ -0,0 +1,111 @@ +import httpx + +from headroom.transforms.content_router import ContentRouter, ContentRouterConfig +from headroom.transforms.kompress_compressor import KompressConfig +from headroom.transforms.kompress_remote import RemoteKompressCompressor + + +def _long_text() -> str: + return " ".join(f"word{i}" for i in range(20)) + + +def _compressor(transport: httpx.BaseTransport) -> RemoteKompressCompressor: + compressor = RemoteKompressCompressor( + "https://kompress.example", + token="secret", + config=KompressConfig(enable_ccr=False), + ) + compressor._client = httpx.Client(transport=transport) + return compressor + + +def test_remote_kompress_posts_content_and_returns_result() -> None: + seen: dict[str, object] = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["url"] = str(request.url) + seen["authorization"] = request.headers.get("authorization") + seen["json"] = request.read().decode() + return httpx.Response( + 200, + json={ + "compressed": "short result", + "original_tokens": 20, + "compressed_tokens": 2, + "compression_ratio": 0.1, + "model_used": "remote-model", + }, + ) + + compressor = _compressor(httpx.MockTransport(handler)) + try: + result = compressor.compress(_long_text(), target_ratio=0.3) + finally: + compressor.close() + + assert seen["url"] == "https://kompress.example/compress" + assert seen["authorization"] == "Bearer secret" + assert '"target_ratio":0.3' in str(seen["json"]).replace(" ", "") + assert result.compressed == "short result" + assert result.original_tokens == 20 + assert result.compressed_tokens == 2 + assert result.compression_ratio == 0.1 + assert result.model_used == "remote-model" + + +def test_remote_kompress_short_input_skips_network() -> None: + called = False + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal called + called = True + return httpx.Response(200, json={"compressed": "unused"}) + + compressor = _compressor(httpx.MockTransport(handler)) + try: + result = compressor.compress("too short") + finally: + compressor.close() + + assert called is False + assert result.compressed == "too short" + assert result.compression_ratio == 1.0 + + +def test_remote_kompress_http_error_fails_open() -> None: + content = _long_text() + compressor = _compressor(httpx.MockTransport(lambda request: httpx.Response(503))) + try: + result = compressor.compress(content) + finally: + compressor.close() + + assert result.compressed == content + assert result.compression_ratio == 1.0 + + +def test_remote_kompress_malformed_success_fails_open() -> None: + content = _long_text() + compressor = _compressor(httpx.MockTransport(lambda request: httpx.Response(200, json={}))) + try: + result = compressor.compress(content) + finally: + compressor.close() + + assert result.compressed == content + assert result.compression_ratio == 1.0 + + +def test_content_router_selects_remote_kompress_from_env(monkeypatch) -> None: + monkeypatch.setenv("HEADROOM_KOMPRESS_ENDPOINT", "https://kompress.example") + monkeypatch.setenv("HEADROOM_KOMPRESS_ENDPOINT_TOKEN", "secret") + + router = ContentRouter(ContentRouterConfig(ccr_inject_marker=False)) + compressor = router._get_kompress() + try: + assert isinstance(compressor, RemoteKompressCompressor) + assert compressor.config == KompressConfig(enable_ccr=False) + assert compressor._url == "https://kompress.example/compress" + assert compressor._headers["authorization"] == "Bearer secret" + finally: + compressor.close()