mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(compress): reach the lossless provider seam on the general path and default /v1/compress to marker-free output (#2691)
## Description Two related changes to the compression seams, plus the review fixes for both. Supersedes #2661 and #2662, which are closed in favour of this branch — the fixes are inseparable from the code they fix, so reviewing them together is cheaper than landing two PRs and patching them afterwards. **1. A registered lossless provider now competes on the general path.** The `headroom.transforms.lossless_provider` seam was only ever consulted from `_lossless_compact_excluded`, gated on `DEFAULT_EXCLUDE_TOOLS` (`config.py:216` — `Read/Grep/Glob/Write/Edit/WebSearch/WebFetch`). Gateway traffic carries the caller's own tool names — LiteLLM's `headroom` guardrail (https://docs.litellm.ai/docs/proxy/headroom) posts requests containing tools like `search_docs` / `run_ci` / `fetch_rows` — so a registered provider was structurally unreachable for every gateway/sidecar deployment. The seam existed; nothing could get to it. **2. `POST /v1/compress` is marker-free by default.** A CCR marker is only useful to a caller that also injects the `headroom_retrieve` tool AND can reach `/v1/retrieve`. Neither holds here: tool injection lives in the provider request handlers (`handlers/anthropic.py:1894`), never in `handle_compress`; and every `/v1/retrieve*` route is `Depends(_require_loopback)` (`server.py:4422, 4470, 4749, 4781`) with no remote opt-in — `HEADROOM_COMPRESS_ALLOW_REMOTE` drops the loopback dependency on `/v1/compress` only. So a gateway forwards a `Retrieve more: hash=…` pointer the model cannot follow, and the proxy pays a CCR store write nobody reads. `config.mode="ccr"` opts back in. Closes # ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [x] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made ### Seam — `content_router.py`, `lossless_provider.py` - `_lossless_first` (STAGE 0, every block on every path) consults `get_lossless_provider()` and keeps whichever output is smaller. **Strict no-op when no provider is registered**, which is the default; and because it is best-of rather than authoritative, a provider can never do worse than the built-in folds. - Malformed provider output can no longer escape. Every shape check runs inside the `try`: result must be `None`, or a 2-element tuple/list of two `str`. Anything else is ignored at debug level. (Previously the unpack sat outside the `try`, so a 3-tuple raised `ValueError` up through `TransformPipeline.apply`, which re-raises.) - Empty / whitespace-only candidates are rejected rather than silently replacing block content. - **Providers are never offered diff content.** Diff folding is subtractive with no inverse check and a reflowed hunk breaks `git apply` — the same reason the built-in `diff` fold is restricted at `content_router.py:2481`. - The third-party `kind` label is sanitised against `^[a-z0-9_]{1,32}$` before reaching `transforms_applied` and the per-strategy metric dicts, so a caller-controlled string cannot explode Prometheus label cardinality. `fullmatch`, not `match`: `$` also matches before a trailing newline, which would put a newline in a label. - `set_lossless_provider(provider, *, verifier=None)` — in lossless-only mode, where STAGE 0's output is final and there is no marker to recover from, a registered verifier must confirm the fold or the candidate is dropped. No verifier registered = today's behaviour. `provider=None` clears both. - The provider is invoked once per block, not twice (`_has_lossless_fold` probes `_lossless_first` and discards the result, then STAGE 0 recomputes). Bounded memo, wholesale clear on overflow, no lock — a race costs one redundant fold. The memo keys on the provider registration generation, so registering or clearing a provider after a block was already folded takes effect. - The seam docstring now records that the provider runs on the general path and inside the parallel compression pool, so it must be thread-safe as well as deterministic. ### Route — `handlers/openai.py`, `server.py` - `_derived_compress_pipeline(key, **overrides)` replaces the copy-pasted pipeline-derivation block; `_no_ccr_pipeline` (the new default) and `_lossy_inline_pipeline` both use it. - The default pipeline is **built at startup** and included in `_eager_preload_transforms`, so a fresh pod does not pay ContentRouter construction and compressor load on its first request, inside the compression-executor budget. - An unrecognised `config.mode` returns 400 naming the valid values instead of silently falling back to the default. - Claude-family model names resolve their context limit from the Anthropic provider. Real divergence: `bedrock/anthropic.claude-3-5-sonnet` is 200000 there and 128000 on the OpenAI provider. The tokenizer still comes from the OpenAI pipeline's provider — a separate, larger change, noted in a comment. - Documents why the derived router deliberately does **not** share the base router's compression cache: keys do not encode CCR-marker mode, so sharing would leak marker-laden entries into the marker-free path. ### Behavior change A `/v1/compress` caller that relied on default markers now gets none. The only in-tree caller that can resolve them is the TypeScript SDK (`sdk/typescript/src/client.ts:398 retrieve`, `:422 handleToolCall`); it needs `config: {"mode": "ccr"}` to keep today's behaviour, and landing that SDK default in the same release would leave only gateway callers — for whom markers were never resolvable — seeing a difference. No `HEADROOM_COMPRESS_DEFAULT_MODE` compat env deliberately: a flag nobody sets becomes permanent debt, and the wire-level `mode` already covers the one caller that needs it. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ python -m pytest tests/test_lossless_first_dispatch.py tests/test_lossless_excluded_compaction.py \ tests/test_lossless_mode.py tests/test_lossless_then_lossy.py tests/test_lossless_diff_fold_guard.py \ tests/test_bash_search_lossless_fold.py tests/test_proxy_compress_endpoint.py \ tests/test_ccr_row_drop_store_bridge.py tests/test_gateway_sidecar_ports.py tests/test_compress_api.py \ tests/test_platform_stabilization_functional.py tests/test_proxy_eager_preload_bind.py \ tests/test_proxy_warmup.py tests/test_router_registry_smartcrusher.py -q 189 passed in 32.04s $ ruff check <all 9 changed files> All checks passed! $ ruff format --check <all 9 changed files> 9 files already formatted $ mypy headroom/transforms/content_router.py headroom/transforms/lossless_provider.py \ headroom/proxy/handlers/openai.py headroom/proxy/server.py Success: no issues found in 4 source files ``` New tests cover, one concern each: every malformed provider shape; empty and whitespace-only results; diff content never reaching a provider (call-recording); `kind` sanitisation including the trailing-newline case; the verifier accepting / rejecting / raising; clearing a provider clearing its verifier; single provider invocation per block; memo invalidation on re-registration; unknown and valid `mode` values; the default pipeline existing before any request; and Claude vs OpenAI context-limit resolution with `token_budget` precedence preserved. ## Real Behavior Proof - **Environment:** macOS arm64, Python 3.12.6, proxy built from `_proxy_config_from_env()` with the default `coding` savings profile; Kompress both disabled and offloaded to a remote `kompress-v2-base` endpoint; `HEADROOM_COMPRESSION_TIMEOUT_SECONDS=300`. - **Steps:** `POST /v1/compress` over `TestClient` with OpenAI-shaped payloads under non-excluded tool names (`run_ci`, `list_files`, `code_search`, `fetch_rows`) — a CI log with ANSI escapes and repeated lines, a 160-path listing, a 150-line grep dump, a 150-row JSON array; plus a second payload with a RAG user blob, a 200-row JSON tool result and a 300-line log. Ran with and without a provider registered via `set_lossless_provider`. - **Observed — seam reachability:** | Kompress | no provider registered | provider registered | |---|---|---| | off | 19,284 → 10,265 tokens (46.8%) | 19,284 → **7,879 (59.1%)** | | on (remote) | 19,284 → 9,366 tokens (51.4%) | 19,284 → **7,146 (62.9%)** | Before this change the right-hand column was identical to the left — the registered provider was never called on this payload. - **Observed — marker-free default costs nothing:** | Config | tokens | saved | |---|---|---| | markers on (previous default) | 37,791 → 24,415 | 35.4% | | markers off (new default) | 37,791 → 24,415 | **35.4% — identical** | | `mode="lossy_inline"` | 37,791 → 25,129 | 33.5% | | `--lossless` | 37,791 → 35,100 | 7.1% | - **Observed — memo staleness, before the fix:** registering a provider that folds a grep block to 5 bytes left the block at its 1496-byte built-in fold, and clearing a provider kept serving the provider's output. Both correct after keying on the registration generation. - **Observed — context limit:** `bedrock/anthropic.claude-3-5-sonnet` resolves 200000 via the Anthropic provider, 128000 via the OpenAI provider. - **Not tested:** `tests/test_transforms_content_router.py` was not run — it does not complete on this machine, wedging on its 5th test while that test passes in 5.7s alone. Verified pre-existing before this work: with the diff stashed, the clean tree stalled at the identical test, and it stalled the same way under `HF_HUB_OFFLINE=1 HEADROOM_OFFLINE=1`. The machine was also out of disk at the time, which may be the real cause rather than the suspected native-detector deadlock (#575) — worth a separate issue either way. CI should be the arbiter here. ## 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 - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` ## Additional Notes Two pre-existing test files needed adjusting, both direct consequences rather than scope creep: - `test_platform_stabilization_functional.py::test_v1_compress_success_reports_actual_metrics` patches `openai_pipeline.apply`, which the default mode no longer routes through. **This test was already failing on the marker-free-default commit** before any of the fixes — my original test selection missed it. It now patches the pipeline the route actually uses. - `test_proxy_eager_preload_bind.py` substitutes fake pipelines to control exactly what the preload walks; the eager build injected the real derived router's statuses into an exact-equality assertion. Its shared helper now clears the derived cache, preserving each test's intent without weakening an assertion. Docs unchecked — follow-ups worth doing in the same release: document `config.mode` values in `docs/content/docs/litellm.mdx` and `wiki/proxy.md`; the TS SDK `mode:"ccr"` default; and an operational note that `COMPRESSION_TIMEOUT_SECONDS` defaults to 30 (`helpers.py:687`) while a remote ML endpoint makes one sequential call per unit — on a large payload it trips the executor timeout and the handler fails open, returning `compression_skipped: true` with `tokens_before: 0`, which reads as "nothing to save" rather than "we gave up". Those zeroed counters are misleading and worth a separate fix.
This commit is contained in:
parent
e0ce4b1d48
commit
f2c48e26c6
9 changed files with 883 additions and 41 deletions
|
|
@ -1207,6 +1207,12 @@ RESPONSES_CONTEXT_SEARCH_TIMEOUT_SECONDS = 2.0
|
||||||
# accept) but short enough to bound the damage from a hung peer.
|
# accept) but short enough to bound the damage from a hung peer.
|
||||||
WS_FIRST_FRAME_TIMEOUT_SECONDS = 60.0
|
WS_FIRST_FRAME_TIMEOUT_SECONDS = 60.0
|
||||||
|
|
||||||
|
# Accepted values for ``config.mode`` on POST /v1/compress. Unset/None means
|
||||||
|
# the default marker-free pipeline. Anything else is a 400 rather than a
|
||||||
|
# silent fall-through to the default (a typo or an unsupported mode like
|
||||||
|
# "lossless" would otherwise look like it worked).
|
||||||
|
COMPRESS_MODES = ("ccr", "lossy_inline", "lossless_then_lossy")
|
||||||
|
|
||||||
|
|
||||||
def _extract_codex_handshake_headers(upstream: Any) -> list[tuple[str, str]]:
|
def _extract_codex_handshake_headers(upstream: Any) -> list[tuple[str, str]]:
|
||||||
"""Return the ``x-codex-*`` headers from an upstream WS handshake response.
|
"""Return the ``x-codex-*`` headers from an upstream WS handshake response.
|
||||||
|
|
@ -8252,21 +8258,27 @@ class OpenAIHandlerMixin:
|
||||||
with contextlib.suppress(Exception):
|
with contextlib.suppress(Exception):
|
||||||
await websocket.close()
|
await websocket.close()
|
||||||
|
|
||||||
def _lossy_inline_pipeline(self) -> Any:
|
def _derived_compress_pipeline(self, key: str, **overrides: Any) -> Any:
|
||||||
"""Cached pipeline for ``/v1/compress`` ``config.mode="lossy_inline"``.
|
"""Cached ``/v1/compress`` pipeline derived from the live OpenAI router.
|
||||||
|
|
||||||
Runs the lossless byte/data fold first, then Kompresses the folded
|
``overrides`` are applied to the live ContentRouter's config, so the
|
||||||
remainder (``lossless_then_lossy``). ``ccr_inject_marker=False`` makes
|
derived pipeline inherits every operator setting (Kompress on/off,
|
||||||
every compressor (Kompress, SmartCrusher, search/log/config) emit inline
|
exclusions, profile knobs) and differs only in what the caller needs.
|
||||||
lossy output with NO ``<<ccr:…>>`` / ``Retrieve more: hash=`` marker and
|
|
||||||
NO CCR store write, so the result is safe to forward straight to a
|
The derived router deliberately does NOT share the base router's
|
||||||
provider with no retrieval round-trip. Derived once from the live OpenAI
|
compression cache. ``ContentRouter._cache`` is per-instance, and its
|
||||||
router's config and reused read-only across requests.
|
keys do not encode the CCR-marker mode, so a shared cache would serve
|
||||||
|
marker-laden entries (written by the OpenAI request path) to this
|
||||||
|
marker-free path and vice versa. Sharing would be a correctness bug,
|
||||||
|
not an optimisation — the duplicate cache is the intended trade.
|
||||||
|
|
||||||
ponytail: a first-request race just builds it twice — both are
|
ponytail: a first-request race just builds it twice — both are
|
||||||
equivalent and Kompress weights are cached at module level, so no lock.
|
equivalent and Kompress weights are cached at module level, so no lock.
|
||||||
"""
|
"""
|
||||||
cached = getattr(self, "_lossy_inline_pipeline_cache", None)
|
cache = getattr(self, "_compress_pipeline_cache", None)
|
||||||
|
if cache is None:
|
||||||
|
cache = self._compress_pipeline_cache = {}
|
||||||
|
cached = cache.get(key)
|
||||||
if cached is not None:
|
if cached is not None:
|
||||||
return cached
|
return cached
|
||||||
|
|
||||||
|
|
@ -8277,29 +8289,71 @@ class OpenAIHandlerMixin:
|
||||||
base = find_content_router(self.openai_pipeline)
|
base = find_content_router(self.openai_pipeline)
|
||||||
if base is None: # ponytail: nothing to derive from — use default pipeline
|
if base is None: # ponytail: nothing to derive from — use default pipeline
|
||||||
return self.openai_pipeline
|
return self.openai_pipeline
|
||||||
cfg = replace(
|
pipeline = TransformPipeline(
|
||||||
base.config,
|
transforms=[ContentRouter(replace(base.config, **overrides), observer=self.metrics)],
|
||||||
|
provider=self.openai_provider,
|
||||||
|
)
|
||||||
|
cache[key] = pipeline
|
||||||
|
return pipeline
|
||||||
|
|
||||||
|
def _no_ccr_pipeline(self) -> Any:
|
||||||
|
"""Default ``/v1/compress`` pipeline: compress, but emit no CCR markers.
|
||||||
|
|
||||||
|
Every caller of this route is a gateway/sidecar or SDK client that
|
||||||
|
forwards the returned messages straight to a provider. A
|
||||||
|
``Retrieve more: hash=`` marker is only useful to a caller that also
|
||||||
|
injects the ``headroom_retrieve`` tool AND can reach ``/v1/retrieve``
|
||||||
|
(loopback-only). LiteLLM's `headroom` guardrail — the main consumer —
|
||||||
|
does neither: it swaps ``messages`` and forwards. So markers here are a
|
||||||
|
dangling pointer for the model plus a pointless CCR store write.
|
||||||
|
``config.mode="ccr"`` opts back in for callers that do run the loop
|
||||||
|
(e.g. the TypeScript SDK's ``retrieve``/``handleToolCall``).
|
||||||
|
|
||||||
|
Nothing else changes: same compressors, same aggressiveness. Measured
|
||||||
|
identical token savings to the marker-on default on OpenAI-shaped
|
||||||
|
gateway traffic.
|
||||||
|
"""
|
||||||
|
return self._derived_compress_pipeline(
|
||||||
|
"no_ccr",
|
||||||
|
ccr_inject_marker=False, # no markers in returned content
|
||||||
|
ccr_enabled=False, # no CCR store writes
|
||||||
|
)
|
||||||
|
|
||||||
|
def _lossy_inline_pipeline(self) -> Any:
|
||||||
|
"""Pipeline for ``/v1/compress`` ``config.mode="lossy_inline"``.
|
||||||
|
|
||||||
|
Runs the lossless byte/data fold first, then Kompresses the folded
|
||||||
|
remainder (``lossless_then_lossy``), marker-free like
|
||||||
|
:meth:`_no_ccr_pipeline`. Kept as a distinct mode because the
|
||||||
|
fold-then-Kompress ordering is a different compression posture, not a
|
||||||
|
different CCR setting.
|
||||||
|
"""
|
||||||
|
return self._derived_compress_pipeline(
|
||||||
|
"lossy_inline",
|
||||||
lossless=False, # lossy mode (not lossless-only)
|
lossless=False, # lossy mode (not lossless-only)
|
||||||
lossless_then_lossy=True, # fold first, then Kompress the remainder
|
lossless_then_lossy=True, # fold first, then Kompress the remainder
|
||||||
ccr_inject_marker=False, # inline, marker-free everywhere
|
ccr_inject_marker=False, # inline, marker-free everywhere
|
||||||
ccr_enabled=False, # no CCR store writes
|
ccr_enabled=False, # no CCR store writes
|
||||||
smart_crusher_lossless_only=False, # keep SmartCrusher lossy
|
smart_crusher_lossless_only=False, # keep SmartCrusher lossy
|
||||||
) # enable_kompress inherited: on by default, off if operator disabled it
|
|
||||||
pipeline = TransformPipeline(
|
|
||||||
transforms=[ContentRouter(cfg, observer=self.metrics)],
|
|
||||||
provider=self.openai_provider,
|
|
||||||
)
|
)
|
||||||
self._lossy_inline_pipeline_cache = pipeline
|
|
||||||
return pipeline
|
|
||||||
|
|
||||||
async def handle_compress(self, request: Request) -> JSONResponse:
|
async def handle_compress(self, request: Request) -> JSONResponse:
|
||||||
"""Compress messages without calling an LLM.
|
"""Compress messages without calling an LLM.
|
||||||
|
|
||||||
POST /v1/compress
|
POST /v1/compress
|
||||||
Body: {"messages": [...], "model": "...", "config": {}}
|
Body: {"messages": [...], "model": "...", "config": {}}
|
||||||
``config.mode="lossy_inline"`` (alias ``"lossless_then_lossy"``) selects
|
|
||||||
the marker-free lossless-then-lossy pipeline whose output needs no CCR
|
``config.mode`` selects the pipeline:
|
||||||
retrieval round-trip — the mode to use behind a gateway/sidecar.
|
|
||||||
|
* unset (default) — marker-free; forward the result straight to a
|
||||||
|
provider with no CCR retrieval round-trip (see _no_ccr_pipeline).
|
||||||
|
* ``"ccr"`` — CCR markers + store writes, for a caller that injects the
|
||||||
|
``headroom_retrieve`` tool and can reach loopback ``/v1/retrieve``.
|
||||||
|
* ``"lossy_inline"`` (alias ``"lossless_then_lossy"``) — marker-free,
|
||||||
|
lossless fold first then Kompress the remainder.
|
||||||
|
|
||||||
|
Any other ``config.mode`` value is a 400 (see ``COMPRESS_MODES``).
|
||||||
|
|
||||||
Returns compressed messages + metrics.
|
Returns compressed messages + metrics.
|
||||||
"""
|
"""
|
||||||
from fastapi.responses import JSONResponse
|
from fastapi.responses import JSONResponse
|
||||||
|
|
@ -8389,10 +8443,27 @@ class OpenAIHandlerMixin:
|
||||||
# Allow optional token_budget to override model's context limit
|
# Allow optional token_budget to override model's context limit
|
||||||
# (used by OpenClaw compact() and other callers that need tighter budgets)
|
# (used by OpenClaw compact() and other callers that need tighter budgets)
|
||||||
token_budget = body.get("token_budget")
|
token_budget = body.get("token_budget")
|
||||||
|
# Resolve the context limit against the model's own provider. This
|
||||||
|
# route is OpenAI-shaped, but LiteLLM's `headroom` guardrail passes
|
||||||
|
# Anthropic model names straight through (claude-sonnet-4-5-...,
|
||||||
|
# bedrock/anthropic.claude-3-5-sonnet, anthropic/claude-opus-4), and
|
||||||
|
# the OpenAI provider answers those with its 128K default instead of
|
||||||
|
# 200K+. Substring match is enough here — no model registry.
|
||||||
|
#
|
||||||
|
# Known gap: the TOKENIZER still comes from the OpenAI pipeline's
|
||||||
|
# provider, so token counts remain approximate for Claude models.
|
||||||
|
# Fixing that means selecting the whole pipeline per model family,
|
||||||
|
# which is a larger change and out of scope for this route.
|
||||||
|
model_name = model if isinstance(model, str) else str(model)
|
||||||
|
limit_provider = (
|
||||||
|
self.anthropic_provider
|
||||||
|
if ("claude" in model_name.lower() or "anthropic" in model_name.lower())
|
||||||
|
else self.openai_provider
|
||||||
|
)
|
||||||
context_limit = (
|
context_limit = (
|
||||||
token_budget
|
token_budget
|
||||||
if token_budget and isinstance(token_budget, int)
|
if token_budget and isinstance(token_budget, int)
|
||||||
else self.openai_provider.get_context_limit(model)
|
else limit_provider.get_context_limit(model)
|
||||||
)
|
)
|
||||||
# Extract CompressConfig options from request body
|
# Extract CompressConfig options from request body
|
||||||
compress_config = body.get("config", {})
|
compress_config = body.get("config", {})
|
||||||
|
|
@ -8402,14 +8473,30 @@ class OpenAIHandlerMixin:
|
||||||
target_ratio = compress_config.get("target_ratio")
|
target_ratio = compress_config.get("target_ratio")
|
||||||
protect_recent = compress_config.get("protect_recent")
|
protect_recent = compress_config.get("protect_recent")
|
||||||
protect_analysis_context = compress_config.get("protect_analysis_context")
|
protect_analysis_context = compress_config.get("protect_analysis_context")
|
||||||
# Marker-free lossless-then-lossy mode: safe to forward downstream
|
# Mode selection. Default is marker-free (see _no_ccr_pipeline):
|
||||||
# with no CCR retrieval round-trip (see _lossy_inline_pipeline).
|
# no caller of this route can resolve a CCR marker unless it opts in
|
||||||
|
# with mode="ccr", which restores the full marker + store behaviour.
|
||||||
mode = compress_config.get("mode")
|
mode = compress_config.get("mode")
|
||||||
pipeline = (
|
if mode is not None and mode not in COMPRESS_MODES:
|
||||||
self._lossy_inline_pipeline()
|
return JSONResponse(
|
||||||
if mode in ("lossy_inline", "lossless_then_lossy")
|
status_code=400,
|
||||||
else self.openai_pipeline
|
content={
|
||||||
|
"error": {
|
||||||
|
"type": "invalid_request",
|
||||||
|
"message": (
|
||||||
|
f"Invalid config.mode: {mode!r}. "
|
||||||
|
f"Valid values are: {', '.join(COMPRESS_MODES)} "
|
||||||
|
"(or omit config.mode for the default marker-free mode)."
|
||||||
|
),
|
||||||
|
}
|
||||||
|
},
|
||||||
)
|
)
|
||||||
|
if mode in ("lossy_inline", "lossless_then_lossy"):
|
||||||
|
pipeline = self._lossy_inline_pipeline()
|
||||||
|
elif mode == "ccr":
|
||||||
|
pipeline = self.openai_pipeline
|
||||||
|
else:
|
||||||
|
pipeline = self._no_ccr_pipeline()
|
||||||
|
|
||||||
pipeline_kwargs: dict = {
|
pipeline_kwargs: dict = {
|
||||||
"model_limit": context_limit,
|
"model_limit": context_limit,
|
||||||
|
|
|
||||||
|
|
@ -956,6 +956,21 @@ class HeadroomProxy(
|
||||||
transforms=[*_intercept_prefix, cache_aligner, openai_router],
|
transforms=[*_intercept_prefix, cache_aligner, openai_router],
|
||||||
provider=self.openai_provider,
|
provider=self.openai_provider,
|
||||||
)
|
)
|
||||||
|
# Build the DEFAULT /v1/compress pipeline now, not on first request.
|
||||||
|
# It is a ContentRouter derived from `openai_router` (marker-free), so
|
||||||
|
# a lazy build would land inside the bounded compression executor on a
|
||||||
|
# cold pod's very first gateway request — paying router construction
|
||||||
|
# and, when the ML model runs in-process, model load against the 30 s
|
||||||
|
# compression budget. Building it here also lets startup warmup see it
|
||||||
|
# (`_eager_preload_transforms`). Never fatal: on failure the handler
|
||||||
|
# falls back to the original lazy path.
|
||||||
|
try:
|
||||||
|
self._no_ccr_pipeline()
|
||||||
|
except Exception as exc: # pragma: no cover - defensive
|
||||||
|
logger.debug(
|
||||||
|
"Eager /v1/compress pipeline build failed (%s); deferring to first request",
|
||||||
|
exc,
|
||||||
|
)
|
||||||
|
|
||||||
# Initialize components
|
# Initialize components
|
||||||
self.cache = (
|
self.cache = (
|
||||||
|
|
@ -1580,7 +1595,12 @@ class HeadroomProxy(
|
||||||
eager_status: dict[str, str] = {}
|
eager_status: dict[str, str] = {}
|
||||||
transform_statuses: list[dict[str, str]] = []
|
transform_statuses: list[dict[str, str]] = []
|
||||||
seen_transform_ids: set[int] = set()
|
seen_transform_ids: set[int] = set()
|
||||||
for pipeline in (self.anthropic_pipeline, self.openai_pipeline):
|
# The derived /v1/compress pipelines (built eagerly in __init__) own
|
||||||
|
# their own ContentRouter instances, so dedup-by-id() does not cover
|
||||||
|
# them via the two request pipelines — warm them explicitly or the
|
||||||
|
# first gateway request still pays the compressor load.
|
||||||
|
derived_pipelines = list(getattr(self, "_compress_pipeline_cache", {}).values())
|
||||||
|
for pipeline in (self.anthropic_pipeline, self.openai_pipeline, *derived_pipelines):
|
||||||
for transform in pipeline.transforms:
|
for transform in pipeline.transforms:
|
||||||
if id(transform) in seen_transform_ids:
|
if id(transform) in seen_transform_ids:
|
||||||
continue
|
continue
|
||||||
|
|
|
||||||
|
|
@ -78,7 +78,11 @@ from .content_detector import (
|
||||||
)
|
)
|
||||||
from .content_detector import detect_content_type as _regex_detect_content_type
|
from .content_detector import detect_content_type as _regex_detect_content_type
|
||||||
from .error_detection import content_has_strong_error_indicators
|
from .error_detection import content_has_strong_error_indicators
|
||||||
from .lossless_provider import get_lossless_provider
|
from .lossless_provider import (
|
||||||
|
get_lossless_generation,
|
||||||
|
get_lossless_provider,
|
||||||
|
get_lossless_verifier,
|
||||||
|
)
|
||||||
from .mixed_content import ContentSection, mixed_content_indicators
|
from .mixed_content import ContentSection, mixed_content_indicators
|
||||||
from .relevance_split import build_relevance_query, plan_relevance_split
|
from .relevance_split import build_relevance_query, plan_relevance_split
|
||||||
|
|
||||||
|
|
@ -101,6 +105,17 @@ _detect_native_verified = False # native detect has returned once -> skip the w
|
||||||
_TOKEN_ESTIMATOR = EstimatingTokenCounter()
|
_TOKEN_ESTIMATOR = EstimatingTokenCounter()
|
||||||
|
|
||||||
|
|
||||||
|
# `kind` labels supplied by a third-party lossless provider flow into
|
||||||
|
# `transforms_applied`, `transforms_summary` and the per-strategy metric/timing
|
||||||
|
# dicts (which become Prometheus label values). An unbounded caller-controlled
|
||||||
|
# string there is a label-cardinality explosion, so only a short, lowercase,
|
||||||
|
# identifier-shaped label is accepted; anything else is replaced with
|
||||||
|
# `_PROVIDER_KIND_FALLBACK` rather than raising (a bad label must not fail a
|
||||||
|
# request, and the fold itself is still valid).
|
||||||
|
_PROVIDER_KIND_RE = re.compile(r"^[a-z0-9_]{1,32}$")
|
||||||
|
_PROVIDER_KIND_FALLBACK = "provider"
|
||||||
|
|
||||||
|
|
||||||
def _estimate_tokens(text: str) -> int:
|
def _estimate_tokens(text: str) -> int:
|
||||||
"""Size-proportional token estimate for section ratio decisions.
|
"""Size-proportional token estimate for section ratio decisions.
|
||||||
|
|
||||||
|
|
@ -1700,6 +1715,12 @@ class ContentRouter(Transform):
|
||||||
# (fewer lossy chains, safer); 0 = keep the lossy pass on any improvement.
|
# (fewer lossy chains, safer); 0 = keep the lossy pass on any improvement.
|
||||||
_DEFAULT_LOSSY_MIN_EXTRA_SAVINGS = 0.05
|
_DEFAULT_LOSSY_MIN_EXTRA_SAVINGS = 0.05
|
||||||
|
|
||||||
|
# Entry cap for the `_lossless_first` memo. Sized for "the blocks of a
|
||||||
|
# handful of in-flight requests", not for cross-request reuse — on overflow
|
||||||
|
# the whole dict is dropped rather than evicted entry-by-entry, which keeps
|
||||||
|
# the memo O(1) and unlockable at the cost of an occasional cold start.
|
||||||
|
_LOSSLESS_FIRST_MEMO_MAX = 256
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
config: ContentRouterConfig | None = None,
|
config: ContentRouterConfig | None = None,
|
||||||
|
|
@ -1764,6 +1785,23 @@ class ContentRouter(Transform):
|
||||||
self._relevance_scorer: Any = None
|
self._relevance_scorer: Any = None
|
||||||
self._relevance_scorer_tried: bool = False
|
self._relevance_scorer_tried: bool = False
|
||||||
self._relevance_prewarm_started: bool = False
|
self._relevance_prewarm_started: bool = False
|
||||||
|
# STAGE 0 memo: `(hash(content), len(content), strategy) -> (best, label)`.
|
||||||
|
# `_lossless_first` is a pure function of its arguments but is called
|
||||||
|
# twice for the same block (`_has_lossless_fold` probes it, then STAGE 0
|
||||||
|
# recomputes it), which also means a registered lossless provider gets
|
||||||
|
# invoked twice. Bounded by `_LOSSLESS_FIRST_MEMO_MAX` with a wholesale
|
||||||
|
# clear (no LRU bookkeeping — this is a within-request hit, not a
|
||||||
|
# long-lived cache; `_cache` is the cross-request one).
|
||||||
|
self._lossless_first_memo: dict[
|
||||||
|
tuple[int, int, CompressionStrategy, int], tuple[str, str | None]
|
||||||
|
] = {}
|
||||||
|
# Companion memo for the third-party half of STAGE 0, keyed by content
|
||||||
|
# (plus the provider generation) but NOT by strategy — the provider
|
||||||
|
# contract takes no strategy, so the two probes above, which pass
|
||||||
|
# different strategies for the same block, share one provider
|
||||||
|
# invocation. See `_lossless_provider_result`.
|
||||||
|
self._lossless_provider_memo: dict[tuple[int, int, int], tuple[str, str] | None] = {}
|
||||||
|
|
||||||
# tool_call_id → compact args text, populated by _build_tool_name_map.
|
# tool_call_id → compact args text, populated by _build_tool_name_map.
|
||||||
self._tool_call_args: dict[str, str] = {}
|
self._tool_call_args: dict[str, str] = {}
|
||||||
# tool_call_id → raw shell command (bash-search fold), same population.
|
# tool_call_id → raw shell command (bash-search fold), same population.
|
||||||
|
|
@ -2462,6 +2500,26 @@ class ContentRouter(Transform):
|
||||||
"""
|
"""
|
||||||
from headroom.transforms.lossless_compaction import compact_lossless
|
from headroom.transforms.lossless_compaction import compact_lossless
|
||||||
|
|
||||||
|
# Memo: `_has_lossless_fold` probes this method purely to decide whether a
|
||||||
|
# small block is worth admitting and throws the result away, and the real
|
||||||
|
# STAGE 0 pass then recomputes it. That doubles the built-in folds *and*
|
||||||
|
# the third-party provider call for every block that takes both paths.
|
||||||
|
# The method is a pure function of (content, strategy), so cache it.
|
||||||
|
# Resolved with `getattr` so the method still works on a bare instance
|
||||||
|
# (`object.__new__(ContentRouter)`, used by unit tests that exercise the
|
||||||
|
# folds without paying for a full router init).
|
||||||
|
memo = getattr(self, "_lossless_first_memo", None)
|
||||||
|
if memo is None:
|
||||||
|
memo = self._lossless_first_memo = {}
|
||||||
|
# The provider generation is part of the key: this memo caches a value
|
||||||
|
# that DEPENDS on the registered provider, so a provider registered (or
|
||||||
|
# cleared) after a block was first folded must not be served the old
|
||||||
|
# answer forever.
|
||||||
|
memo_key = (hash(content), len(content), strategy, get_lossless_generation())
|
||||||
|
memoized = memo.get(memo_key)
|
||||||
|
if memoized is not None:
|
||||||
|
return memoized
|
||||||
|
|
||||||
# Apply losslessness to the OUTPUT structure, not to the classification:
|
# Apply losslessness to the OUTPUT structure, not to the classification:
|
||||||
# try the fold implied by the detected strategy first, then the others.
|
# try the fold implied by the detected strategy first, then the others.
|
||||||
# Each compact_lossless call is self-verifying (exact inverse or returns
|
# Each compact_lossless call is self-verifying (exact inverse or returns
|
||||||
|
|
@ -2498,8 +2556,132 @@ class ContentRouter(Transform):
|
||||||
continue
|
continue
|
||||||
if len(cand) < len(best):
|
if len(cand) < len(best):
|
||||||
best, best_label = cand, f"lossless_{kind}"
|
best, best_label = cand, f"lossless_{kind}"
|
||||||
|
|
||||||
|
# A registered lossless provider (headroom.transforms.lossless_provider)
|
||||||
|
# competes on the GENERAL path too, not only on excluded-tool output via
|
||||||
|
# `_lossless_compact_excluded`. Without this an external fold only ever
|
||||||
|
# saw Read/Grep/Glob-style results, so it was inert for gateway traffic
|
||||||
|
# (`/v1/compress`), where tool names are the caller's own. Same contract
|
||||||
|
# (information-preserving + deterministic) and we keep whichever output
|
||||||
|
# is smaller, so a provider can only improve on the built-in folds.
|
||||||
|
#
|
||||||
|
# Diffs are off limits for a provider, for the same reason the built-in
|
||||||
|
# "diff" fold is restricted above: diff folding is subtractive with no
|
||||||
|
# inverse check, and a Kompressed/reflowed hunk breaks `git apply`. The
|
||||||
|
# built-in path can at least reason about its own fold; a third-party one
|
||||||
|
# we cannot, so we simply never offer it diff content.
|
||||||
|
if strategy is not CompressionStrategy.DIFF and not self._looks_like_diff(content):
|
||||||
|
best, best_label = self._apply_lossless_provider(content, best, best_label)
|
||||||
|
|
||||||
|
# Plain dict operations only: this is a pure-function cache, so a racing
|
||||||
|
# double-compute under the parallel compression pool (see
|
||||||
|
# HEADROOM_COMPRESS_WORKERS) costs one redundant fold and nothing else.
|
||||||
|
# A lock here would sit on the hot path of every block for no
|
||||||
|
# correctness gain. Bounded by a wholesale clear rather than LRU
|
||||||
|
# bookkeeping — the memo exists to collapse a within-request double
|
||||||
|
# call, not to be a cross-request cache (`self._cache` is that).
|
||||||
|
if len(memo) >= self._LOSSLESS_FIRST_MEMO_MAX:
|
||||||
|
memo.clear()
|
||||||
|
memo[memo_key] = (best, best_label)
|
||||||
return best, best_label
|
return best, best_label
|
||||||
|
|
||||||
|
def _apply_lossless_provider(
|
||||||
|
self, content: str, best: str, best_label: str | None
|
||||||
|
) -> tuple[str, str | None]:
|
||||||
|
"""Let a registered lossless provider beat ``best``; never raise.
|
||||||
|
|
||||||
|
The provider only wins on a strictly smaller, non-blank candidate — and
|
||||||
|
in lossless-only mode only if an optionally registered verifier confirms
|
||||||
|
the fold is genuinely reversible.
|
||||||
|
"""
|
||||||
|
supplied = self._lossless_provider_result(content)
|
||||||
|
if supplied is None:
|
||||||
|
return best, best_label
|
||||||
|
cand, kind = supplied
|
||||||
|
|
||||||
|
if len(cand) >= len(best):
|
||||||
|
return best, best_label
|
||||||
|
|
||||||
|
# Lossless-only mode has no downstream recovery: STAGE 0's output IS the
|
||||||
|
# answer and there is no CCR marker to retrieve the original from. The
|
||||||
|
# built-in folds are self-verifying; a provider is trusted unless the
|
||||||
|
# operator registered a verifier, in which case it must pass.
|
||||||
|
if self.config.lossless:
|
||||||
|
verifier = get_lossless_verifier()
|
||||||
|
if verifier is not None:
|
||||||
|
try:
|
||||||
|
ok = verifier(content, cand)
|
||||||
|
except Exception: # noqa: BLE001 - a raising verifier means "unverified"
|
||||||
|
logger.debug("lossless verifier raised; rejecting candidate", exc_info=True)
|
||||||
|
return best, best_label
|
||||||
|
if not ok:
|
||||||
|
logger.debug("lossless verifier rejected the provider candidate")
|
||||||
|
return best, best_label
|
||||||
|
|
||||||
|
return cand, f"lossless_{kind}"
|
||||||
|
|
||||||
|
def _lossless_provider_result(self, content: str) -> tuple[str, str] | None:
|
||||||
|
"""Validated + memoized ``(candidate, safe_kind)`` from the provider.
|
||||||
|
|
||||||
|
Everything a third-party provider hands back is treated as hostile: a
|
||||||
|
malformed shape, an empty result, a metric-unsafe ``kind`` label or an
|
||||||
|
outright exception all degrade to ``None`` ("provider ignored"). This
|
||||||
|
runs on the request path and its caller (``TransformPipeline.apply``)
|
||||||
|
re-raises, so nothing here may escape.
|
||||||
|
|
||||||
|
Memoized per ``(hash(content), len(content))`` — the provider contract is
|
||||||
|
``content -> result`` with no other input, so its answer cannot depend on
|
||||||
|
the routing strategy. That matters because one block reaches
|
||||||
|
``_lossless_first`` twice under *different* strategies (the
|
||||||
|
``_has_lossless_fold`` admission probe passes PASSTHROUGH, STAGE 0 passes
|
||||||
|
the detected strategy), which would otherwise invoke third-party code
|
||||||
|
twice per block. Same no-lock rationale as the ``_lossless_first`` memo.
|
||||||
|
"""
|
||||||
|
provider = get_lossless_provider()
|
||||||
|
if provider is None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
memo = getattr(self, "_lossless_provider_memo", None)
|
||||||
|
if memo is None:
|
||||||
|
memo = self._lossless_provider_memo = {}
|
||||||
|
memo_key = (hash(content), len(content), get_lossless_generation())
|
||||||
|
if memo_key in memo:
|
||||||
|
return memo[memo_key]
|
||||||
|
|
||||||
|
result: tuple[str, str] | None = None
|
||||||
|
try:
|
||||||
|
supplied = provider(content)
|
||||||
|
if supplied is None:
|
||||||
|
result = None
|
||||||
|
# Validate defensively *inside* the try: unpacking a 3-tuple / a bare
|
||||||
|
# string / a non-sequence raises, and that exception used to escape
|
||||||
|
# all the way out of the request.
|
||||||
|
elif not isinstance(supplied, tuple | list) or len(supplied) != 2:
|
||||||
|
logger.debug("lossless provider returned a malformed result; ignoring")
|
||||||
|
elif not isinstance(supplied[0], str) or not isinstance(supplied[1], str):
|
||||||
|
logger.debug("lossless provider returned non-str members; ignoring")
|
||||||
|
# An empty / whitespace-only "fold" wins every length comparison and
|
||||||
|
# would silently delete the block's content. That is not lossless.
|
||||||
|
elif not supplied[0].strip():
|
||||||
|
logger.debug("lossless provider returned an empty result; ignoring")
|
||||||
|
else:
|
||||||
|
raw_kind = supplied[1]
|
||||||
|
# fullmatch, not match: Python's `$` also matches just BEFORE a
|
||||||
|
# trailing newline, so `re.match` would let "log\n" through and
|
||||||
|
# put a newline into a metric label.
|
||||||
|
kind = (
|
||||||
|
raw_kind if _PROVIDER_KIND_RE.fullmatch(raw_kind) else _PROVIDER_KIND_FALLBACK
|
||||||
|
)
|
||||||
|
result = (supplied[0], kind)
|
||||||
|
except Exception: # noqa: BLE001 - a broken provider must not break routing
|
||||||
|
logger.debug("lossless provider failed in _lossless_first", exc_info=True)
|
||||||
|
result = None
|
||||||
|
|
||||||
|
if len(memo) >= self._LOSSLESS_FIRST_MEMO_MAX:
|
||||||
|
memo.clear()
|
||||||
|
memo[memo_key] = result
|
||||||
|
return result
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _looks_like_diff(content: str) -> bool:
|
def _looks_like_diff(content: str) -> bool:
|
||||||
"""Cheap structural sniff for unified/git-diff content.
|
"""Cheap structural sniff for unified/git-diff content.
|
||||||
|
|
@ -2508,7 +2690,8 @@ class ContentRouter(Transform):
|
||||||
Kompressing hunks corrupts ``git apply``. This is defense-in-depth beyond the
|
Kompressing hunks corrupts ``git apply``. This is defense-in-depth beyond the
|
||||||
DIFF-strategy and ``lossless_diff``-label checks: a diff can be folded
|
DIFF-strategy and ``lossless_diff``-label checks: a diff can be folded
|
||||||
best under a non-diff label (e.g. blank-line collapse → ``lossless_text``)
|
best under a non-diff label (e.g. blank-line collapse → ``lossless_text``)
|
||||||
or mis-detected, and must still never reach the lossy stage.
|
or mis-detected, and must still never reach the lossy stage. It is also
|
||||||
|
what keeps a third-party lossless provider away from diffs entirely.
|
||||||
"""
|
"""
|
||||||
return (
|
return (
|
||||||
"diff --git " in content
|
"diff --git " in content
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,4 @@
|
||||||
"""Pluggable provider for information-preserving compaction of protected output.
|
"""Pluggable provider for information-preserving compaction of tool output.
|
||||||
|
|
||||||
Excluded ("protected") tool results are kept out of *lossy* compression for
|
Excluded ("protected") tool results are kept out of *lossy* compression for
|
||||||
accuracy; the content router still applies reversible/data-preserving folds to
|
accuracy; the content router still applies reversible/data-preserving folds to
|
||||||
|
|
@ -14,9 +14,39 @@ Contract — ``provider(content: str) -> tuple[compacted: str, kind: str] | None
|
||||||
* The provider MUST be deterministic and depend only on ``content`` (no
|
* The provider MUST be deterministic and depend only on ``content`` (no
|
||||||
cross-message state), so the proxy's prefix cache stays byte-stable across
|
cross-message state), so the proxy's prefix cache stays byte-stable across
|
||||||
turns.
|
turns.
|
||||||
|
* The provider MUST be thread-safe. It is called from the router's parallel
|
||||||
|
compression worker pool (``HEADROOM_COMPRESS_WORKERS``), so several blocks can
|
||||||
|
be handed to it concurrently on different threads.
|
||||||
|
* ``kind`` is a short label that ends up in ``transforms_applied`` and in
|
||||||
|
per-strategy metric/timing keys. Only ``^[a-z0-9_]{1,32}$`` is accepted;
|
||||||
|
anything else is replaced with a fixed fallback label so a third-party string
|
||||||
|
cannot blow up Prometheus label cardinality.
|
||||||
|
* A malformed return (not ``None``, not a 2-tuple of ``str``) or an empty /
|
||||||
|
whitespace-only ``compacted`` is ignored, never fatal.
|
||||||
|
|
||||||
|
Where the provider is called:
|
||||||
|
|
||||||
|
* :meth:`ContentRouter._lossless_compact_excluded` — protected/excluded tool
|
||||||
|
output. Here the provider is *authoritative*: when one is set the router does
|
||||||
|
not run its built-in folds; it falls back to the built-in only if the provider
|
||||||
|
raises.
|
||||||
|
* :meth:`ContentRouter._lossless_first` — the GENERAL compression path, i.e.
|
||||||
|
every block of every provider/harness (including ``/v1/compress`` gateway
|
||||||
|
traffic, where tool names are the caller's own). Here the provider *competes*
|
||||||
|
with the built-in folds and the smaller output wins, so it can only improve on
|
||||||
|
them. Diff content is never handed to the provider — folding a diff can break
|
||||||
|
``git apply``.
|
||||||
|
|
||||||
|
Verification — ``verifier(original: str, compacted: str) -> bool``:
|
||||||
|
|
||||||
|
In lossless-only mode (``--lossless``) the fold IS the final answer: nothing
|
||||||
|
downstream can recover a mistake. The built-in folds are self-verifying
|
||||||
|
(``compact_lossless`` returns its input when it cannot invert), a provider is
|
||||||
|
not. Register an optional ``verifier`` alongside the provider and the router
|
||||||
|
will run it on the provider's candidate in lossless-only mode, rejecting the
|
||||||
|
candidate when it returns falsy *or* raises. With no verifier registered the
|
||||||
|
behaviour is unchanged: the provider is trusted (documented trust).
|
||||||
|
|
||||||
A registered provider is *authoritative*: when one is set the router does not run
|
|
||||||
its built-in folds — it falls back to the built-in only if the provider raises.
|
|
||||||
Default is ``None`` → the router uses its built-in folds, unchanged.
|
Default is ``None`` → the router uses its built-in folds, unchanged.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
|
@ -26,16 +56,47 @@ from collections.abc import Callable
|
||||||
|
|
||||||
#: ``content -> (compacted, kind) | None``.
|
#: ``content -> (compacted, kind) | None``.
|
||||||
LosslessProvider = Callable[[str], "tuple[str, str] | None"]
|
LosslessProvider = Callable[[str], "tuple[str, str] | None"]
|
||||||
|
#: ``(original, compacted) -> is-it-really-lossless``.
|
||||||
|
LosslessVerifier = Callable[[str, str], bool]
|
||||||
|
|
||||||
_provider: LosslessProvider | None = None
|
_provider: LosslessProvider | None = None
|
||||||
|
_verifier: LosslessVerifier | None = None
|
||||||
|
#: Bumped on every registration. Callers that memoize a provider's output must
|
||||||
|
#: include this in their cache key, or a provider registered (or cleared) after
|
||||||
|
#: a block was already folded would be ignored for that block forever. Normal
|
||||||
|
#: deployments register once at extension install, but tests and any hot-reload
|
||||||
|
#: path re-register at will.
|
||||||
|
_generation: int = 0
|
||||||
|
|
||||||
|
|
||||||
def set_lossless_provider(provider: LosslessProvider | None) -> None:
|
def set_lossless_provider(
|
||||||
"""Register (or clear, with ``None``) the lossless compaction provider."""
|
provider: LosslessProvider | None,
|
||||||
global _provider
|
*,
|
||||||
|
verifier: LosslessVerifier | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Register (or clear, with ``None``) the lossless compaction provider.
|
||||||
|
|
||||||
|
``verifier`` is optional and keyword-only so existing single-argument calls
|
||||||
|
keep working. Passing ``provider=None`` clears *both* the provider and any
|
||||||
|
previously registered verifier — a verifier without a provider is dead
|
||||||
|
state, and leaving one behind would silently apply to the next provider.
|
||||||
|
"""
|
||||||
|
global _provider, _verifier, _generation
|
||||||
_provider = provider
|
_provider = provider
|
||||||
|
_verifier = verifier if provider is not None else None
|
||||||
|
_generation += 1
|
||||||
|
|
||||||
|
|
||||||
def get_lossless_provider() -> LosslessProvider | None:
|
def get_lossless_provider() -> LosslessProvider | None:
|
||||||
"""Return the registered provider, or ``None`` if the built-in should run."""
|
"""Return the registered provider, or ``None`` if the built-in should run."""
|
||||||
return _provider
|
return _provider
|
||||||
|
|
||||||
|
|
||||||
|
def get_lossless_generation() -> int:
|
||||||
|
"""Registration counter — include it in any key that memoizes provider output."""
|
||||||
|
return _generation
|
||||||
|
|
||||||
|
|
||||||
|
def get_lossless_verifier() -> LosslessVerifier | None:
|
||||||
|
"""Return the registered verifier, or ``None`` when the provider is trusted."""
|
||||||
|
return _verifier
|
||||||
|
|
|
||||||
|
|
@ -385,6 +385,9 @@ def test_v1_compress_then_v1_retrieve_resolves_marker_hash() -> None:
|
||||||
]
|
]
|
||||||
req = {
|
req = {
|
||||||
"model": "gpt-4o",
|
"model": "gpt-4o",
|
||||||
|
# /v1/compress is marker-free by default (no gateway caller can resolve a
|
||||||
|
# marker); mode="ccr" is the opt-in for callers that run the retrieve loop.
|
||||||
|
"config": {"mode": "ccr"},
|
||||||
"messages": [
|
"messages": [
|
||||||
{"role": "user", "content": "Get items"},
|
{"role": "user", "content": "Get items"},
|
||||||
{
|
{
|
||||||
|
|
|
||||||
|
|
@ -9,8 +9,17 @@ its byte-exact fold — the lossless floor is never discarded by a later lossy
|
||||||
stage.
|
stage.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from headroom.transforms.content_router import ContentRouter, ContentRouterConfig
|
from headroom.transforms.content_router import (
|
||||||
|
CompressionStrategy,
|
||||||
|
ContentRouter,
|
||||||
|
ContentRouterConfig,
|
||||||
|
)
|
||||||
from headroom.transforms.lossless_compaction import search_unheading
|
from headroom.transforms.lossless_compaction import search_unheading
|
||||||
|
from headroom.transforms.lossless_provider import (
|
||||||
|
get_lossless_provider,
|
||||||
|
get_lossless_verifier,
|
||||||
|
set_lossless_provider,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def _grep_block() -> str:
|
def _grep_block() -> str:
|
||||||
|
|
@ -105,6 +114,44 @@ def test_has_lossless_fold_admits_small_block_below_size_floor():
|
||||||
assert router._has_lossless_fold("def f():\n return 1\n") is False
|
assert router._has_lossless_fold("def f():\n return 1\n") is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_registered_provider_competes_on_the_general_path():
|
||||||
|
# A registered lossless provider must be consulted for ANY block, not only
|
||||||
|
# excluded-tool output — that gate made external folds inert for gateway
|
||||||
|
# traffic (/v1/compress), where tool names are the caller's own. The smaller
|
||||||
|
# output wins, so a provider can only improve on the built-in folds.
|
||||||
|
from headroom.transforms.lossless_provider import set_lossless_provider
|
||||||
|
|
||||||
|
block = _grep_block()
|
||||||
|
baseline, _, _ = _compress(block, lossless=True)
|
||||||
|
better = "SHORTER-THAN-ANY-BUILTIN-FOLD\n"
|
||||||
|
assert len(better) < len(baseline)
|
||||||
|
|
||||||
|
try:
|
||||||
|
set_lossless_provider(lambda content: (better, "plugin"))
|
||||||
|
out, was, tr = _compress(block + "\n", lossless=True) # fresh cache key
|
||||||
|
assert was is True
|
||||||
|
assert out == better
|
||||||
|
assert tr == ["router:tool_result:lossless_plugin"]
|
||||||
|
|
||||||
|
# A provider that loses to the built-in fold is ignored, not adopted.
|
||||||
|
set_lossless_provider(lambda content: (content + "padding" * 100, "plugin"))
|
||||||
|
out2, _, tr2 = _compress(block + "\n\n", lossless=True)
|
||||||
|
assert tr2 == ["router:tool_result:lossless_search"]
|
||||||
|
assert len(out2) < len(block)
|
||||||
|
|
||||||
|
# A raising provider must not break routing.
|
||||||
|
set_lossless_provider(_raise)
|
||||||
|
out3, was3, _ = _compress(block + "\n\n\n", lossless=True)
|
||||||
|
assert was3 is True
|
||||||
|
assert len(out3) < len(block)
|
||||||
|
finally:
|
||||||
|
set_lossless_provider(None)
|
||||||
|
|
||||||
|
|
||||||
|
def _raise(content: str):
|
||||||
|
raise RuntimeError("broken provider")
|
||||||
|
|
||||||
|
|
||||||
def test_lossless_mode_non_foldable_is_lossless_noop_not_ratio_too_high():
|
def test_lossless_mode_non_foldable_is_lossless_noop_not_ratio_too_high():
|
||||||
# In lossless-only mode, code with no byte-lossless fold is left verbatim.
|
# In lossless-only mode, code with no byte-lossless fold is left verbatim.
|
||||||
# That is NOT a rejected compression, so it must not be bucketed as
|
# That is NOT a rejected compression, so it must not be bucketed as
|
||||||
|
|
@ -129,3 +176,226 @@ def test_lossless_mode_non_foldable_is_lossless_noop_not_ratio_too_high():
|
||||||
assert was is False
|
assert was is False
|
||||||
assert rc.get("lossless_noop", 0) >= 1
|
assert rc.get("lossless_noop", 0) >= 1
|
||||||
assert rc.get("ratio_too_high", 0) == 0
|
assert rc.get("ratio_too_high", 0) == 0
|
||||||
|
|
||||||
|
|
||||||
|
# ── Hostile / third-party provider input ──────────────────────────────────────
|
||||||
|
# The provider is arbitrary out-of-tree code called on the request path, and
|
||||||
|
# `_lossless_first`'s caller (`TransformPipeline.apply`) re-raises. Nothing a
|
||||||
|
# provider returns may fail a request or silently destroy a block.
|
||||||
|
|
||||||
|
_SHORT = "SHORTER-THAN-ANY-BUILTIN-FOLD\n"
|
||||||
|
|
||||||
|
_MALFORMED_RESULTS = [
|
||||||
|
(_SHORT, "plugin", "extra"), # 3-tuple -> would raise on unpack
|
||||||
|
_SHORT, # bare string -> would unpack into 2 chars, or raise
|
||||||
|
(_SHORT, 42), # non-str kind
|
||||||
|
(42, "plugin"), # non-str candidate
|
||||||
|
{"compacted": _SHORT}, # non-tuple entirely
|
||||||
|
17, # not even iterable
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def test_malformed_provider_result_is_ignored_not_raised():
|
||||||
|
block = _grep_block()
|
||||||
|
for i, bad in enumerate(_MALFORMED_RESULTS):
|
||||||
|
try:
|
||||||
|
set_lossless_provider(lambda content, bad=bad: bad)
|
||||||
|
# Distinct content per shape so nothing is served from a cache/memo.
|
||||||
|
out, was, tr = _compress(block + "\n" * (i + 1), lossless=True)
|
||||||
|
finally:
|
||||||
|
set_lossless_provider(None)
|
||||||
|
# No exception, and the built-in fold is still the answer.
|
||||||
|
assert was is True, bad
|
||||||
|
assert tr == ["router:tool_result:lossless_search"], bad
|
||||||
|
assert len(out) < len(block), bad
|
||||||
|
|
||||||
|
|
||||||
|
def test_empty_provider_result_is_rejected():
|
||||||
|
# "" beats every candidate on length, so an unchecked empty result would
|
||||||
|
# win and silently delete the block's content.
|
||||||
|
block = _grep_block()
|
||||||
|
for i, blank in enumerate(("", " ", "\n\t \n")):
|
||||||
|
try:
|
||||||
|
set_lossless_provider(lambda content, blank=blank: (blank, "plugin"))
|
||||||
|
out, was, tr = _compress(block + "x" * (i + 1) + "\n", lossless=True)
|
||||||
|
finally:
|
||||||
|
set_lossless_provider(None)
|
||||||
|
assert was is True, repr(blank)
|
||||||
|
assert tr == ["router:tool_result:lossless_search"], repr(blank)
|
||||||
|
assert out.strip(), repr(blank)
|
||||||
|
|
||||||
|
|
||||||
|
def test_provider_is_never_offered_diff_content():
|
||||||
|
# Diff folding is subtractive with no inverse check and a reflowed hunk
|
||||||
|
# breaks `git apply`, so a third-party fold must never see a diff — the
|
||||||
|
# same reason the built-in "diff" fold is strategy-gated.
|
||||||
|
diff = (
|
||||||
|
"diff --git a/x b/x\nindex 1111111..2222222 100644\n--- a/x\n+++ b/x\n@@ -1 +1 @@\n-a\n+b\n"
|
||||||
|
)
|
||||||
|
router = ContentRouter(ContentRouterConfig(lossless=True))
|
||||||
|
baseline_diff = router._lossless_first(diff, CompressionStrategy.DIFF)
|
||||||
|
baseline_pass = router._lossless_first(diff, CompressionStrategy.PASSTHROUGH)
|
||||||
|
|
||||||
|
seen: list[str] = []
|
||||||
|
|
||||||
|
def recording(content: str):
|
||||||
|
seen.append(content)
|
||||||
|
return (_SHORT, "plugin")
|
||||||
|
|
||||||
|
try:
|
||||||
|
set_lossless_provider(recording)
|
||||||
|
fresh = ContentRouter(ContentRouterConfig(lossless=True))
|
||||||
|
# Both the DIFF strategy and diff-*shaped* content under another
|
||||||
|
# strategy must bypass the provider entirely.
|
||||||
|
assert fresh._lossless_first(diff, CompressionStrategy.DIFF) == baseline_diff
|
||||||
|
assert fresh._lossless_first(diff, CompressionStrategy.PASSTHROUGH) == baseline_pass
|
||||||
|
assert seen == []
|
||||||
|
finally:
|
||||||
|
set_lossless_provider(None)
|
||||||
|
|
||||||
|
|
||||||
|
def test_provider_kind_label_is_sanitized():
|
||||||
|
# `kind` reaches transforms_applied and the per-strategy metric/timing keys
|
||||||
|
# (Prometheus label values), so an unbounded caller-controlled string is a
|
||||||
|
# cardinality bomb. Only ^[a-z0-9_]{1,32}$ survives.
|
||||||
|
block = _grep_block()
|
||||||
|
# "log\n" is the subtle one: Python's `$` matches before a trailing newline,
|
||||||
|
# so a `re.match`-based check would let a newline into a metric label.
|
||||||
|
bogus = ["UPPER", "has space", "punct!", "x" * 200, "", "kebab-case", "log\n"]
|
||||||
|
for i, kind in enumerate(bogus):
|
||||||
|
try:
|
||||||
|
set_lossless_provider(lambda content, kind=kind: (_SHORT, kind))
|
||||||
|
out, was, tr = _compress(block + "\n" * (i + 1), lossless=True)
|
||||||
|
finally:
|
||||||
|
set_lossless_provider(None)
|
||||||
|
assert was is True, kind
|
||||||
|
assert out == _SHORT, kind
|
||||||
|
assert tr == ["router:tool_result:lossless_provider"], kind
|
||||||
|
|
||||||
|
# A clean label is preserved verbatim.
|
||||||
|
try:
|
||||||
|
set_lossless_provider(lambda content: (_SHORT, "log_fold_2"))
|
||||||
|
_out, _was, tr = _compress(block + "\t", lossless=True)
|
||||||
|
assert tr == ["router:tool_result:lossless_log_fold_2"]
|
||||||
|
finally:
|
||||||
|
set_lossless_provider(None)
|
||||||
|
|
||||||
|
|
||||||
|
def test_lossless_mode_verifier_gates_the_provider_candidate():
|
||||||
|
# In --lossless mode STAGE 0's output IS the answer: no CCR marker, no
|
||||||
|
# retrieval. The built-in folds self-verify; a provider only does if the
|
||||||
|
# operator registered a verifier.
|
||||||
|
block = _grep_block()
|
||||||
|
provider = lambda content: (_SHORT, "plugin") # noqa: E731
|
||||||
|
|
||||||
|
def _run(verifier, suffix):
|
||||||
|
try:
|
||||||
|
set_lossless_provider(provider, verifier=verifier)
|
||||||
|
return _compress(block + suffix, lossless=True)
|
||||||
|
finally:
|
||||||
|
set_lossless_provider(None)
|
||||||
|
|
||||||
|
# No verifier -> documented trust, unchanged behaviour.
|
||||||
|
_out, was, tr = _run(None, "\n")
|
||||||
|
assert was is True
|
||||||
|
assert tr == ["router:tool_result:lossless_plugin"]
|
||||||
|
|
||||||
|
# Verifier says yes -> adopted.
|
||||||
|
out, was, tr = _run(lambda original, compacted: True, "\n\n")
|
||||||
|
assert out == _SHORT
|
||||||
|
assert tr == ["router:tool_result:lossless_plugin"]
|
||||||
|
|
||||||
|
# Verifier says no -> rejected, built-in fold stands.
|
||||||
|
out, was, tr = _run(lambda original, compacted: False, "\n\n\n")
|
||||||
|
assert tr == ["router:tool_result:lossless_search"]
|
||||||
|
assert len(out) < len(block)
|
||||||
|
|
||||||
|
# Verifier raises -> "unverified" -> rejected, never fatal.
|
||||||
|
def _boom(original, compacted):
|
||||||
|
raise RuntimeError("verifier exploded")
|
||||||
|
|
||||||
|
out, was, tr = _run(_boom, "\n\n\n\n")
|
||||||
|
assert tr == ["router:tool_result:lossless_search"]
|
||||||
|
assert len(out) < len(block)
|
||||||
|
|
||||||
|
|
||||||
|
def test_memo_does_not_go_stale_when_the_provider_changes():
|
||||||
|
# The STAGE 0 memo caches a value that depends on the registered provider,
|
||||||
|
# so its key includes the registration generation. Without that, a provider
|
||||||
|
# registered (or cleared) after a block was already folded is ignored for
|
||||||
|
# that exact block forever — invisible in production (extensions register at
|
||||||
|
# startup) but wrong, and a trap for tests and any hot-reload path.
|
||||||
|
from headroom.transforms.content_router import CompressionStrategy
|
||||||
|
|
||||||
|
router = ContentRouter(ContentRouterConfig(lossless=True))
|
||||||
|
block = _grep_block()
|
||||||
|
tiny = "TINY\n"
|
||||||
|
|
||||||
|
try:
|
||||||
|
set_lossless_provider(None)
|
||||||
|
baseline, baseline_label = router._lossless_first(block, CompressionStrategy.SEARCH)
|
||||||
|
assert baseline_label == "lossless_search"
|
||||||
|
|
||||||
|
# Register AFTER the block was already folded once.
|
||||||
|
set_lossless_provider(lambda content: (tiny, "plugin"))
|
||||||
|
out, label = router._lossless_first(block, CompressionStrategy.SEARCH)
|
||||||
|
assert out == tiny, "provider registered after first fold was ignored (stale memo)"
|
||||||
|
assert label == "lossless_plugin"
|
||||||
|
|
||||||
|
# Clearing it must take effect too.
|
||||||
|
set_lossless_provider(None)
|
||||||
|
out, label = router._lossless_first(block, CompressionStrategy.SEARCH)
|
||||||
|
assert out == baseline, "cleared provider still served from the memo"
|
||||||
|
assert label == "lossless_search"
|
||||||
|
finally:
|
||||||
|
set_lossless_provider(None)
|
||||||
|
|
||||||
|
|
||||||
|
def test_clearing_the_provider_also_clears_the_verifier():
|
||||||
|
# A verifier with no provider is dead state that would silently start
|
||||||
|
# gating the *next* provider someone registers.
|
||||||
|
try:
|
||||||
|
set_lossless_provider(lambda content: None, verifier=lambda o, c: True)
|
||||||
|
assert get_lossless_verifier() is not None
|
||||||
|
set_lossless_provider(None)
|
||||||
|
assert get_lossless_provider() is None
|
||||||
|
assert get_lossless_verifier() is None
|
||||||
|
finally:
|
||||||
|
set_lossless_provider(None)
|
||||||
|
|
||||||
|
|
||||||
|
def test_provider_is_called_once_per_block_not_twice():
|
||||||
|
# `_has_lossless_fold` (the small-block admission probe) and STAGE 0 both
|
||||||
|
# run `_lossless_first` over the same block, under *different* strategies.
|
||||||
|
# Memoizing the provider's answer by content collapses that to one call
|
||||||
|
# into third-party code.
|
||||||
|
calls: list[str] = []
|
||||||
|
|
||||||
|
def counting(content: str):
|
||||||
|
calls.append(content)
|
||||||
|
return None
|
||||||
|
|
||||||
|
small = "\n".join(f"pkg/mod/long_filename.py:{n}:value = {n}" for n in range(1, 8)) + "\n"
|
||||||
|
try:
|
||||||
|
set_lossless_provider(counting)
|
||||||
|
router = ContentRouter(ContentRouterConfig(lossless=True))
|
||||||
|
assert router._has_lossless_fold(small) is True
|
||||||
|
assert len(calls) == 1
|
||||||
|
out, was = router._compress_block_content(
|
||||||
|
small,
|
||||||
|
hash(small),
|
||||||
|
"",
|
||||||
|
1.0,
|
||||||
|
1.0,
|
||||||
|
None,
|
||||||
|
[],
|
||||||
|
{},
|
||||||
|
[],
|
||||||
|
"tool_result",
|
||||||
|
"tool",
|
||||||
|
True,
|
||||||
|
)
|
||||||
|
assert was is True
|
||||||
|
assert len(calls) == 1
|
||||||
|
finally:
|
||||||
|
set_lossless_provider(None)
|
||||||
|
|
|
||||||
|
|
@ -68,7 +68,10 @@ def test_v1_compress_success_reports_actual_metrics(monkeypatch) -> None:
|
||||||
markers_inserted=["marker-1"],
|
markers_inserted=["marker-1"],
|
||||||
)
|
)
|
||||||
|
|
||||||
monkeypatch.setattr(proxy.openai_pipeline, "apply", fake_apply)
|
# The default /v1/compress mode runs a marker-free pipeline derived from
|
||||||
|
# `openai_pipeline`, not `openai_pipeline` itself, so patch the one the
|
||||||
|
# route actually uses. It is built eagerly at create_app() time.
|
||||||
|
monkeypatch.setattr(proxy._compress_pipeline_cache["no_ccr"], "apply", fake_apply)
|
||||||
|
|
||||||
with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as client:
|
with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as client:
|
||||||
response = client.post(
|
response = client.post(
|
||||||
|
|
|
||||||
|
|
@ -428,6 +428,211 @@ class TestCompressEndpointLossyInlineMode:
|
||||||
assert response.json()["ccr_hashes"] == []
|
assert response.json()["ccr_hashes"] == []
|
||||||
|
|
||||||
|
|
||||||
|
class TestCompressEndpointModeValidation:
|
||||||
|
"""``config.mode`` must be validated, not silently ignored.
|
||||||
|
|
||||||
|
Before this, ``mode: "lossless"`` (a mode that does not exist) or any typo
|
||||||
|
fell through to the default pipeline and the caller got a 200 describing a
|
||||||
|
compression posture it never asked for.
|
||||||
|
"""
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"bad_mode",
|
||||||
|
["lossless", "CCR", "ccr ", "no_ccr", "", 7, ["ccr"]],
|
||||||
|
)
|
||||||
|
def test_unknown_mode_returns_400(self, client, bad_mode):
|
||||||
|
response = client.post(
|
||||||
|
"/v1/compress",
|
||||||
|
json={
|
||||||
|
"messages": [{"role": "user", "content": "hello"}],
|
||||||
|
"model": "gpt-4",
|
||||||
|
"config": {"mode": bad_mode},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
assert response.status_code == 400
|
||||||
|
data = response.json()
|
||||||
|
assert data["error"]["type"] == "invalid_request"
|
||||||
|
message = data["error"]["message"]
|
||||||
|
# The message must tell the caller what IS valid.
|
||||||
|
for valid in ("ccr", "lossy_inline", "lossless_then_lossy"):
|
||||||
|
assert valid in message
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"config",
|
||||||
|
[
|
||||||
|
{}, # mode unset -> default marker-free pipeline
|
||||||
|
{"mode": None}, # explicit null is the same as unset
|
||||||
|
{"mode": "ccr"},
|
||||||
|
{"mode": "lossy_inline"},
|
||||||
|
{"mode": "lossless_then_lossy"},
|
||||||
|
],
|
||||||
|
ids=["unset", "null", "ccr", "lossy_inline", "lossless_then_lossy"],
|
||||||
|
)
|
||||||
|
def test_valid_modes_return_200(self, client, config):
|
||||||
|
response = client.post(
|
||||||
|
"/v1/compress",
|
||||||
|
json={
|
||||||
|
"messages": [{"role": "user", "content": "hello"}],
|
||||||
|
"model": "gpt-4",
|
||||||
|
"config": config,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert isinstance(response.json()["messages"], list)
|
||||||
|
|
||||||
|
|
||||||
|
class TestCompressDefaultPipelineBuiltAtStartup:
|
||||||
|
"""The default (marker-free) /v1/compress pipeline must exist before the
|
||||||
|
first request.
|
||||||
|
|
||||||
|
It used to be built lazily, so a cold pod paid ContentRouter construction —
|
||||||
|
and in-process ML model load — inside the bounded compression executor on
|
||||||
|
its first real gateway request.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def test_default_pipeline_exists_before_any_request(self):
|
||||||
|
config = ProxyConfig(
|
||||||
|
optimize=True,
|
||||||
|
cache_enabled=False,
|
||||||
|
rate_limit_enabled=False,
|
||||||
|
cost_tracking_enabled=False,
|
||||||
|
)
|
||||||
|
# No TestClient / no request: only create_app().
|
||||||
|
proxy = create_app(config).state.proxy
|
||||||
|
|
||||||
|
cache = getattr(proxy, "_compress_pipeline_cache", None)
|
||||||
|
assert cache, "default /v1/compress pipeline was not built at startup"
|
||||||
|
assert "no_ccr" in cache
|
||||||
|
# It must be a DERIVED pipeline, not the shared request pipeline.
|
||||||
|
assert cache["no_ccr"] is not proxy.openai_pipeline
|
||||||
|
|
||||||
|
def test_startup_warmup_covers_the_derived_router(self):
|
||||||
|
"""The eager compressor preload must walk the derived pipeline too."""
|
||||||
|
config = ProxyConfig(
|
||||||
|
optimize=True,
|
||||||
|
cache_enabled=False,
|
||||||
|
rate_limit_enabled=False,
|
||||||
|
cost_tracking_enabled=False,
|
||||||
|
)
|
||||||
|
proxy = create_app(config).state.proxy
|
||||||
|
|
||||||
|
derived = proxy._compress_pipeline_cache["no_ccr"]
|
||||||
|
derived_ids = {id(t) for t in derived.transforms}
|
||||||
|
|
||||||
|
seen: list[int] = []
|
||||||
|
for pipeline in (proxy.anthropic_pipeline, proxy.openai_pipeline):
|
||||||
|
seen.extend(id(t) for t in pipeline.transforms)
|
||||||
|
# Precondition: the derived router is NOT reachable via the request
|
||||||
|
# pipelines, so dedup-by-id() cannot have covered it implicitly.
|
||||||
|
assert derived_ids - set(seen)
|
||||||
|
|
||||||
|
_status, transform_statuses = proxy._eager_preload_transforms()
|
||||||
|
# Base router + derived router both report a status dict.
|
||||||
|
assert len(transform_statuses) >= 2
|
||||||
|
|
||||||
|
|
||||||
|
class TestCompressContextLimitByModelFamily:
|
||||||
|
"""LiteLLM's guardrail forwards Anthropic model names through this
|
||||||
|
OpenAI-shaped route; the context limit must come from the Anthropic
|
||||||
|
provider for those, not the OpenAI provider's 128K default."""
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _spy_providers(proxy, monkeypatch):
|
||||||
|
anthropic_calls: list[str] = []
|
||||||
|
openai_calls: list[str] = []
|
||||||
|
|
||||||
|
def anthropic_limit(model):
|
||||||
|
anthropic_calls.append(model)
|
||||||
|
return 987_654
|
||||||
|
|
||||||
|
def openai_limit(model):
|
||||||
|
openai_calls.append(model)
|
||||||
|
return 123_456
|
||||||
|
|
||||||
|
monkeypatch.setattr(proxy.anthropic_provider, "get_context_limit", anthropic_limit)
|
||||||
|
monkeypatch.setattr(proxy.openai_provider, "get_context_limit", openai_limit)
|
||||||
|
return anthropic_calls, openai_calls
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _spy_pipeline(proxy, monkeypatch):
|
||||||
|
"""Capture the kwargs handed to the default (marker-free) pipeline."""
|
||||||
|
seen: dict = {}
|
||||||
|
|
||||||
|
def fake_apply(**kwargs):
|
||||||
|
seen.update(kwargs)
|
||||||
|
return SimpleNamespace(
|
||||||
|
messages=kwargs["messages"],
|
||||||
|
tokens_before=10,
|
||||||
|
tokens_after=10,
|
||||||
|
transforms_applied=[],
|
||||||
|
transforms_summary={},
|
||||||
|
markers_inserted=[],
|
||||||
|
)
|
||||||
|
|
||||||
|
monkeypatch.setattr(proxy._compress_pipeline_cache["no_ccr"], "apply", fake_apply)
|
||||||
|
return seen
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"model",
|
||||||
|
[
|
||||||
|
"claude-sonnet-4-5-20250929",
|
||||||
|
"bedrock/anthropic.claude-3-5-sonnet",
|
||||||
|
"anthropic/claude-opus-4",
|
||||||
|
"CLAUDE-Sonnet-4-5", # case-insensitive
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_claude_models_use_anthropic_context_limit(self, client, monkeypatch, model):
|
||||||
|
proxy = client.app.state.proxy
|
||||||
|
anthropic_calls, openai_calls = self._spy_providers(proxy, monkeypatch)
|
||||||
|
seen = self._spy_pipeline(proxy, monkeypatch)
|
||||||
|
|
||||||
|
response = client.post(
|
||||||
|
"/v1/compress",
|
||||||
|
json={"messages": [{"role": "user", "content": "hello"}], "model": model},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert anthropic_calls == [model]
|
||||||
|
assert openai_calls == []
|
||||||
|
assert seen["model_limit"] == 987_654
|
||||||
|
|
||||||
|
def test_openai_models_still_use_openai_context_limit(self, client, monkeypatch):
|
||||||
|
proxy = client.app.state.proxy
|
||||||
|
anthropic_calls, openai_calls = self._spy_providers(proxy, monkeypatch)
|
||||||
|
seen = self._spy_pipeline(proxy, monkeypatch)
|
||||||
|
|
||||||
|
response = client.post(
|
||||||
|
"/v1/compress",
|
||||||
|
json={"messages": [{"role": "user", "content": "hello"}], "model": "gpt-4o"},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert openai_calls == ["gpt-4o"]
|
||||||
|
assert anthropic_calls == []
|
||||||
|
assert seen["model_limit"] == 123_456
|
||||||
|
|
||||||
|
def test_token_budget_still_overrides_for_claude_models(self, client, monkeypatch):
|
||||||
|
"""token_budget precedence must survive the provider routing."""
|
||||||
|
proxy = client.app.state.proxy
|
||||||
|
anthropic_calls, openai_calls = self._spy_providers(proxy, monkeypatch)
|
||||||
|
seen = self._spy_pipeline(proxy, monkeypatch)
|
||||||
|
|
||||||
|
response = client.post(
|
||||||
|
"/v1/compress",
|
||||||
|
json={
|
||||||
|
"messages": [{"role": "user", "content": "hello"}],
|
||||||
|
"model": "claude-sonnet-4-5-20250929",
|
||||||
|
"token_budget": 4096,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert seen["model_limit"] == 4096
|
||||||
|
# Neither provider is consulted when the caller pins a budget.
|
||||||
|
assert anthropic_calls == []
|
||||||
|
assert openai_calls == []
|
||||||
|
|
||||||
|
|
||||||
class TestCompressEndpointDoesNotBlockLoop:
|
class TestCompressEndpointDoesNotBlockLoop:
|
||||||
"""/v1/compress must offload to the compression executor so a slow/large
|
"""/v1/compress must offload to the compression executor so a slow/large
|
||||||
payload cannot freeze the single event loop (#718)."""
|
payload cannot freeze the single event loop (#718)."""
|
||||||
|
|
@ -477,6 +682,9 @@ class TestCompressEndpointDoesNotBlockLoop:
|
||||||
json={
|
json={
|
||||||
"messages": [{"role": "user", "content": "hello world"}],
|
"messages": [{"role": "user", "content": "hello world"}],
|
||||||
"model": "gpt-4",
|
"model": "gpt-4",
|
||||||
|
# mode="ccr" routes to `openai_pipeline` (the default is a
|
||||||
|
# derived marker-free pipeline); this test patches that one.
|
||||||
|
"config": {"mode": "ccr"},
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -36,7 +36,14 @@ def _make_proxy(*, optimize: bool):
|
||||||
image_optimize=False,
|
image_optimize=False,
|
||||||
subscription_tracking_enabled=False,
|
subscription_tracking_enabled=False,
|
||||||
)
|
)
|
||||||
return create_app(config).state.proxy
|
proxy = create_app(config).state.proxy
|
||||||
|
# Every test here substitutes fake pipelines to control exactly what the
|
||||||
|
# preload walks. The proxy also eagerly builds the default /v1/compress
|
||||||
|
# pipeline (a derived ContentRouter, warmed alongside the request
|
||||||
|
# pipelines), which would inject real transform statuses into those
|
||||||
|
# assertions — drop it so the fakes remain the only input.
|
||||||
|
proxy._compress_pipeline_cache = {}
|
||||||
|
return proxy
|
||||||
|
|
||||||
|
|
||||||
class _FastTransform:
|
class _FastTransform:
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue