mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
feat(proxy): make /v1/compress usable as a gateway/Kong sidecar (#2458)
## Description Makes the compression-only `POST /v1/compress` endpoint usable as a **network compression sidecar** behind an API gateway (Kong, LiteLLM, ...), and fixes a latent content-detector hang that silently zeroed compression on non-Windows hosts. Motivated by a LiteLLM-sidecar deployment whose team documented five build-time patches; this ports the ones that belong upstream, generalized so they cover any aliasing gateway (not just LiteLLM). Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - **`lossy_inline` compress mode** (`config.mode="lossy_inline"`, alias `"lossless_then_lossy"`): lossless byte/data fold first, then Kompress the folded remainder, with `ccr_inject_marker=False` so every compressor emits **inline, marker-free** output — no `<<ccr:…>>` markers and no CCR store write, so the result is safe to forward straight to a provider with no retrieval round-trip. The mode inherits the deployment's `enable_kompress`. - **`HEADROOM_COMPRESS_ALLOW_REMOTE`** opt-in: drops the loopback dependency on the `/v1/compress` route **only** so an authorized in-network gateway can reach it. Default is unchanged (loopback-only); inbound `HEADROOM_PROXY_TOKEN` auth still applies. - **`HEADROOM_MODEL_ALIAS_MAP`** (gateway-agnostic, fail-soft): one shared resolver in `pricing/litellm_pricing.py` reduces a gateway-aliased model name (e.g. `claude-opus`) to a priced `litellm.model_cost` key, trying the mapped target as-is and with a `bedrock/` / `vertex_ai/` prefix stripped. `proxy/savings_tracker.py` now delegates to it, so the live (`/stats`) and persisted (`/stats-history`) dollar figures price identically. - **`get_context_limit`**: an operator-configured limit (`HEADROOM_MODEL_LIMITS` / `~/.headroom/models.json`) now wins **before** the dynamic LiteLLM lookup, so an aliased name no longer falls through to the 128K default and skews compression. - **fix(content_router): first-call detector watchdog on all platforms.** The native content detector can deadlock on first use (#575, previously flagged Windows-only). The watchdog was `win32`-only, so on macOS/Linux a first-use hang was unbounded → `_detect_content` never returned → the `/v1/compress` executor timeout fired → fail-open → **`tokens_before=0`, silent zero compression**. Now the native detector runs under the watchdog on the first call on every platform; once it returns it is marked verified and the direct fast path is used (zero steady-state overhead). A hang degrades to pure-Python detection with a clear warning. `win32` behavior is unchanged. - Thread `waste_signals` / `pipeline_timing` into the already-present `/v1/compress` outcome record so the guardrail path populates the dashboard panels like the forward-proxy paths. Deliberately **not** ported: the sidecar's LiteLLM-specific `GET /model/info` HTTP fetch (urllib/ssl/threading/TTL). Kong has no such endpoint; the static `HEADROOM_MODEL_ALIAS_MAP` covers any gateway with no network dependency on the pricing path. ## Testing - [x] Unit tests pass (targeted — see output) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ ruff check <changed files> All checks passed! $ mypy <changed source files> Success: no issues found in 6 source files $ pytest tests/test_gateway_sidecar_ports.py tests/test_proxy_compress_endpoint.py -q tests/test_gateway_sidecar_ports.py ........ [ 34%] tests/test_proxy_compress_endpoint.py ............... [100%] ============================= 23 passed in 20.62s ============================== ``` ## Real Behavior Proof - **Environment:** macOS (darwin/arm64), Python 3.12, `.venv`; Kompress offloaded to a Modal endpoint via `HEADROOM_KOMPRESS_ENDPOINT`. - **Exact command / steps:** posted typical tool-output payloads to `POST /v1/compress` (via the FastAPI `TestClient`, loopback) in both `default` and `lossy_inline` modes; separately reproduced the detector hang with `faulthandler.dump_traceback_later`. - **Observed result:** - Real savings through the endpoint (structural/lossless, Kompress off): **JSON 150 records 13,982→9,514 (32.0%)**, **logs 314 lines 12,240→9,549 (22.0%)**, **search 200 hits 5,231→3,471 (33.6%)**. `lossy_inline` emits **zero** CCR markers. - `faulthandler` pinned the pre-fix hang to `content_router.py:_detect_content` → native `_rust_detect`. With the fix, the first call degrades at the 5s watchdog with `"Native content detector hung … using pure-Python detection"` and compression proceeds (previously it hung and the endpoint returned `tokens_before=0`). - Modal Kompress warm latency measured ~0.8s/call; the learned pass compresses prose further (62→56 words on a sample). - **Not tested:** full `pytest` suite (ran the two affected test files only); the native-detector hang was reproduced on a local macOS/arm64 build — the fix's degrade path is verified, but a healthy-native CI Linux run should confirm the fast (verified) path there. ## 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 - The five-item context comes from a downstream LiteLLM sidecar's `PATCHES.md`; item #3 (record an outcome from the guardrail path) was already upstreamed — this PR only adds the missing `waste_signals`/`pipeline_timing` threading. Item #2 (observability read-only exemption when `HEADROOM_PROXY_TOKEN` is set) is not addressed here. - All new config is opt-in and fail-soft; with nothing set, behavior is byte-identical to today.
This commit is contained in:
Tejas Chopra
GitHub
parent
f4070c44cb
commit
1329ed7f1a
8 changed files with 384 additions and 10 deletions
|
|
@ -8,6 +8,9 @@ See: https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_windo
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
|
|
@ -40,15 +43,65 @@ except ImportError:
|
|||
|
||||
_resolved_model_cache: dict[str, str] = {}
|
||||
|
||||
logger = logging.getLogger("headroom.pricing")
|
||||
|
||||
# --- Gateway model-name resolution ---------------------------------------
|
||||
# When Headroom sits behind a gateway (Kong, LiteLLM, ...) that aliases model
|
||||
# names, the raw client name it sees (e.g. "claude-opus") is not a priced key
|
||||
# in litellm.model_cost, so dollar savings read $0. HEADROOM_MODEL_ALIAS_MAP is
|
||||
# an optional, gateway-agnostic, fail-soft static JSON map {client_name: target}
|
||||
# that reduces that name to a priced model_cost key (trying the target as-is and
|
||||
# with a bedrock/ or vertex_ai/ provider prefix stripped). Unset -> behavior is
|
||||
# identical to today's bare-prefix resolution; pricing never breaks.
|
||||
_GATEWAY_PROVIDER_PREFIXES = ("bedrock/", "vertex_ai/")
|
||||
|
||||
|
||||
def _static_alias_map() -> dict[str, str]:
|
||||
raw = os.environ.get("HEADROOM_MODEL_ALIAS_MAP", "").strip()
|
||||
if not raw:
|
||||
return {}
|
||||
try:
|
||||
data = json.loads(raw)
|
||||
except ValueError:
|
||||
logger.debug("invalid HEADROOM_MODEL_ALIAS_MAP JSON", exc_info=True)
|
||||
return {}
|
||||
if not isinstance(data, dict):
|
||||
return {}
|
||||
return {str(k): str(v) for k, v in data.items() if k and v}
|
||||
|
||||
|
||||
def _reduce_to_priced_key(target: str) -> str | None:
|
||||
"""Reduce a gateway target to a priced litellm.model_cost key, or None."""
|
||||
if not LITELLM_AVAILABLE or litellm is None:
|
||||
return None
|
||||
candidates = [target]
|
||||
for prefix in _GATEWAY_PROVIDER_PREFIXES:
|
||||
if target.startswith(prefix):
|
||||
candidates.append(target[len(prefix) :])
|
||||
for candidate in candidates:
|
||||
info = litellm.model_cost.get(candidate)
|
||||
if info and info.get("input_cost_per_token") is not None:
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def resolve_litellm_model(model: str) -> str:
|
||||
"""Resolve model name to one LiteLLM recognizes, adding provider prefix if needed.
|
||||
Results are cached per model name to avoid blocking the event loop
|
||||
with repeated synchronous litellm lookups.
|
||||
|
||||
When HEADROOM_MODEL_ALIAS_MAP is configured, a raw client name / group alias
|
||||
is first reduced to a priced model_cost key; otherwise this falls through to
|
||||
the bare-prefix rules. Shared by the live (cost.py) and persisted
|
||||
(savings_tracker) pricing paths so both figures price identically.
|
||||
"""
|
||||
if model in _resolved_model_cache:
|
||||
return _resolved_model_cache[model]
|
||||
resolved = _resolve_litellm_model_uncached(model)
|
||||
priced: str | None = None
|
||||
alias = _static_alias_map()
|
||||
if alias:
|
||||
priced = _reduce_to_priced_key(alias.get(model, model))
|
||||
resolved = priced if priced is not None else _resolve_litellm_model_uncached(model)
|
||||
_resolved_model_cache[model] = resolved
|
||||
return resolved
|
||||
|
||||
|
|
|
|||
|
|
@ -460,6 +460,16 @@ class OpenAIProvider(Provider):
|
|||
|
||||
Never raises an exception - uses sensible defaults for unknown models.
|
||||
"""
|
||||
# Explicitly configured limits win first. Behind a gateway/alias proxy
|
||||
# (Kong, LiteLLM, ...) Headroom sees the raw client model name
|
||||
# (e.g. "claude-opus"), which litellm.get_model_info can't resolve, so
|
||||
# resolution would fall through to the 128K default + an "Unknown model"
|
||||
# warning and skew compression. Configuring the alias via
|
||||
# HEADROOM_MODEL_LIMITS / ~/.headroom/models.json makes it authoritative
|
||||
# here, before the dynamic LiteLLM lookup. Fail-soft, no network.
|
||||
if model in self._context_limits:
|
||||
return self._context_limits[model]
|
||||
|
||||
# Try LiteLLM first
|
||||
litellm = _get_litellm_module()
|
||||
if litellm is not None:
|
||||
|
|
|
|||
|
|
@ -7939,11 +7939,54 @@ class OpenAIHandlerMixin:
|
|||
with contextlib.suppress(Exception):
|
||||
await websocket.close()
|
||||
|
||||
def _lossy_inline_pipeline(self) -> Any:
|
||||
"""Cached pipeline for ``/v1/compress`` ``config.mode="lossy_inline"``.
|
||||
|
||||
Runs the lossless byte/data fold first, then Kompresses the folded
|
||||
remainder (``lossless_then_lossy``). ``ccr_inject_marker=False`` makes
|
||||
every compressor (Kompress, SmartCrusher, search/log/config) emit inline
|
||||
lossy output with NO ``<<ccr:…>>`` / ``Retrieve more: hash=`` marker and
|
||||
NO CCR store write, so the result is safe to forward straight to a
|
||||
provider with no retrieval round-trip. Derived once from the live OpenAI
|
||||
router's config and reused read-only across requests.
|
||||
|
||||
ponytail: a first-request race just builds it twice — both are
|
||||
equivalent and Kompress weights are cached at module level, so no lock.
|
||||
"""
|
||||
cached = getattr(self, "_lossy_inline_pipeline_cache", None)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
from headroom.transforms.compression_units import find_content_router
|
||||
from headroom.transforms.content_router import ContentRouter
|
||||
from headroom.transforms.pipeline import TransformPipeline
|
||||
|
||||
base = find_content_router(self.openai_pipeline)
|
||||
if base is None: # ponytail: nothing to derive from — use default pipeline
|
||||
return self.openai_pipeline
|
||||
cfg = replace(
|
||||
base.config,
|
||||
lossless=False, # lossy mode (not lossless-only)
|
||||
lossless_then_lossy=True, # fold first, then Kompress the remainder
|
||||
ccr_inject_marker=False, # inline, marker-free everywhere
|
||||
ccr_enabled=False, # no CCR store writes
|
||||
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:
|
||||
"""Compress messages without calling an LLM.
|
||||
|
||||
POST /v1/compress
|
||||
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
|
||||
retrieval round-trip — the mode to use behind a gateway/sidecar.
|
||||
Returns compressed messages + metrics.
|
||||
"""
|
||||
from fastapi.responses import JSONResponse
|
||||
|
|
@ -8040,10 +8083,20 @@ class OpenAIHandlerMixin:
|
|||
)
|
||||
# Extract CompressConfig options from request body
|
||||
compress_config = body.get("config", {})
|
||||
if not isinstance(compress_config, dict):
|
||||
compress_config = {}
|
||||
compress_user_messages = compress_config.get("compress_user_messages", False)
|
||||
target_ratio = compress_config.get("target_ratio")
|
||||
protect_recent = compress_config.get("protect_recent")
|
||||
protect_analysis_context = compress_config.get("protect_analysis_context")
|
||||
# Marker-free lossless-then-lossy mode: safe to forward downstream
|
||||
# with no CCR retrieval round-trip (see _lossy_inline_pipeline).
|
||||
mode = compress_config.get("mode")
|
||||
pipeline = (
|
||||
self._lossy_inline_pipeline()
|
||||
if mode in ("lossy_inline", "lossless_then_lossy")
|
||||
else self.openai_pipeline
|
||||
)
|
||||
|
||||
pipeline_kwargs: dict = {
|
||||
"model_limit": context_limit,
|
||||
|
|
@ -8064,7 +8117,7 @@ class OpenAIHandlerMixin:
|
|||
# until it finished (#718). The executor also enforces a timeout so a
|
||||
# too-large body fails fast instead of hanging forever.
|
||||
result = await self._run_compression_in_executor(
|
||||
lambda: self.openai_pipeline.apply(
|
||||
lambda: pipeline.apply(
|
||||
messages=messages,
|
||||
model=model,
|
||||
**pipeline_kwargs,
|
||||
|
|
@ -8094,6 +8147,12 @@ class OpenAIHandlerMixin:
|
|||
overhead_ms=latency_ms,
|
||||
num_messages=len(messages) if isinstance(messages, list) else 0,
|
||||
transforms_applied=tuple(result.transforms_applied or ()),
|
||||
waste_signals=(
|
||||
result.waste_signals.to_dict()
|
||||
if getattr(result, "waste_signals", None) is not None
|
||||
else None
|
||||
),
|
||||
pipeline_timing=getattr(result, "timing", None) or None,
|
||||
tags=tags,
|
||||
client=client,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -165,11 +165,29 @@ def _normalize_model(value: Any) -> str:
|
|||
|
||||
|
||||
def _resolve_litellm_model(model: str) -> str:
|
||||
"""Resolve model name to one LiteLLM recognizes."""
|
||||
"""Resolve model name to one LiteLLM recognizes.
|
||||
|
||||
Delegates to the shared alias-map-aware resolver in
|
||||
``headroom.pricing.litellm_pricing`` so the persisted /stats-history funnel
|
||||
(PROXY $ SAVED tile + Historical Checkpoints) prices gateway aliases like
|
||||
"claude-opus" identically to the live /stats path. Uses the shared result
|
||||
only when it maps to a priced model_cost key; otherwise falls through to the
|
||||
bare-prefix logic below. Fail-soft: pricing never breaks bookkeeping.
|
||||
"""
|
||||
litellm = _get_litellm_module()
|
||||
if litellm is None:
|
||||
return model
|
||||
|
||||
try:
|
||||
from headroom.pricing.litellm_pricing import resolve_litellm_model
|
||||
|
||||
resolved = resolve_litellm_model(model)
|
||||
info = litellm.model_cost.get(resolved)
|
||||
if info and info.get("input_cost_per_token") is not None:
|
||||
return resolved
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
litellm.cost_per_token(model=model, prompt_tokens=1, completion_tokens=0)
|
||||
return model
|
||||
|
|
|
|||
|
|
@ -4862,8 +4862,21 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
"data": retrieval_data,
|
||||
}
|
||||
|
||||
# Compression-only endpoint (for TypeScript SDK and other HTTP clients)
|
||||
@app.post("/v1/compress", dependencies=[Depends(_require_loopback)])
|
||||
# Compression-only endpoint (for TypeScript SDK and other HTTP clients).
|
||||
# Loopback-only by default (guard added in #1537). An operator can opt in to
|
||||
# network access for an authorized in-network sidecar/gateway (e.g. Kong,
|
||||
# LiteLLM) on a trusted network by setting HEADROOM_COMPRESS_ALLOW_REMOTE=1,
|
||||
# which drops ONLY this route's loopback dependency. Inbound auth
|
||||
# (HEADROOM_PROXY_TOKEN via _security_gate) and network scoping still apply;
|
||||
# all other _require_loopback routes are unaffected. Unset/false preserves
|
||||
# today's loopback-only behavior.
|
||||
_compress_dependencies = (
|
||||
[]
|
||||
if _get_env_bool("HEADROOM_COMPRESS_ALLOW_REMOTE", False)
|
||||
else [Depends(_require_loopback)]
|
||||
)
|
||||
|
||||
@app.post("/v1/compress", dependencies=_compress_dependencies)
|
||||
async def compress_messages(request: Request):
|
||||
return await proxy.handle_compress(request)
|
||||
|
||||
|
|
|
|||
|
|
@ -92,6 +92,7 @@ split_into_sections = _mixed_content.split_into_sections
|
|||
_detect_backend_warned = False
|
||||
_detect_panic_warned = False
|
||||
_detect_native_unhealthy = False # circuit breaker: native detect hung once (#575)
|
||||
_detect_native_verified = False # native detect has returned once -> skip the watchdog
|
||||
|
||||
|
||||
# Shared calibrated fallback estimator (tiktoken cl100k_base ~90% accuracy,
|
||||
|
|
@ -871,6 +872,7 @@ def _detect_content(content: str) -> DetectionResult:
|
|||
`_strategy_from_detection` keys off that field alone.
|
||||
"""
|
||||
global _detect_backend_warned, _detect_panic_warned, _detect_native_unhealthy
|
||||
global _detect_native_verified
|
||||
|
||||
# Detect on the unwrapped payload so a tool-output envelope's tags don't get
|
||||
# the whole result misclassified as HTML/XML (#route-converter corruption).
|
||||
|
|
@ -896,14 +898,19 @@ def _detect_content(content: str) -> DetectionResult:
|
|||
from headroom._core import detect_content_type as _rust_detect
|
||||
|
||||
try:
|
||||
if sys.platform == "win32":
|
||||
# Windows is the only platform where the native detector can deadlock
|
||||
# on first use (#575); bound it with a watchdog so a hang degrades to
|
||||
# the pure-Python detector below. Elsewhere it is the trusted default
|
||||
# hot path — call it directly, with no per-call thread overhead.
|
||||
# The native detector can deadlock on FIRST use (#575 — seen on Windows
|
||||
# and macOS/arm64). Bound it with a watchdog so a hang degrades to the
|
||||
# pure-Python detector; the previous win32-only guard left other
|
||||
# platforms unprotected, so a hung Linux sidecar silently stopped
|
||||
# compressing (every request failed open to passthrough). Watchdog until
|
||||
# the native detector has returned once, then use the direct fast path —
|
||||
# the hang is first-use only, so steady state pays no per-call thread
|
||||
# overhead. win32 keeps watchdogging every call (unchanged).
|
||||
if sys.platform == "win32" or not _detect_native_verified:
|
||||
rust_result = _rust_detect_watchdogged(_rust_detect, content, _detect_timeout_secs())
|
||||
else:
|
||||
rust_result = _rust_detect(content)
|
||||
_detect_native_verified = True # returned without hanging -> trusted hot path
|
||||
# Rust's `content_type` is the lowercase string tag (e.g.
|
||||
# "json_array"); translate to the Python `ContentType` enum so
|
||||
# downstream mapping keys match.
|
||||
|
|
|
|||
134
tests/test_gateway_sidecar_ports.py
Normal file
134
tests/test_gateway_sidecar_ports.py
Normal file
|
|
@ -0,0 +1,134 @@
|
|||
"""Ports of the LiteLLM/Kong sidecar patches (see the sidecar PATCHES.md #1/#4/#5/#6).
|
||||
|
||||
- #1 HEADROOM_COMPRESS_ALLOW_REMOTE opt-in drops the loopback guard on
|
||||
/v1/compress so an authorized in-network gateway (Kong, LiteLLM) can reach it.
|
||||
- #4/#5 HEADROOM_MODEL_ALIAS_MAP reduces a gateway-aliased model name (e.g.
|
||||
"claude-opus") to a priced litellm.model_cost key — one shared resolver for
|
||||
the live (cost.py) and persisted (savings_tracker) price paths.
|
||||
- #6 an operator-configured context limit (HEADROOM_MODEL_LIMITS) wins over the
|
||||
128K default for an aliased name.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("fastapi")
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from headroom.proxy.server import ProxyConfig, create_app
|
||||
|
||||
|
||||
def _clear_pricing_cache() -> None:
|
||||
from headroom.pricing import litellm_pricing as lp
|
||||
|
||||
lp._resolved_model_cache.clear()
|
||||
|
||||
|
||||
def _first_priced_opus_key() -> str:
|
||||
litellm = pytest.importorskip("litellm")
|
||||
for key, val in litellm.model_cost.items():
|
||||
if "opus" in key.lower() and val.get("input_cost_per_token") is not None:
|
||||
return key
|
||||
pytest.skip("no priced opus key in this litellm build")
|
||||
|
||||
|
||||
# ----- #4/#5 pricing: HEADROOM_MODEL_ALIAS_MAP -> priced key -----
|
||||
|
||||
|
||||
def test_alias_map_resolves_gateway_name_to_priced_key(monkeypatch):
|
||||
litellm = pytest.importorskip("litellm")
|
||||
from headroom.pricing.litellm_pricing import resolve_litellm_model
|
||||
|
||||
key = _first_priced_opus_key()
|
||||
monkeypatch.setenv("HEADROOM_MODEL_ALIAS_MAP", json.dumps({"claude-opus": key}))
|
||||
_clear_pricing_cache()
|
||||
|
||||
resolved = resolve_litellm_model("claude-opus")
|
||||
info = litellm.model_cost.get(resolved)
|
||||
assert info and info.get("input_cost_per_token") is not None
|
||||
|
||||
|
||||
def test_alias_map_strips_bedrock_prefix(monkeypatch):
|
||||
pytest.importorskip("litellm")
|
||||
from headroom.pricing.litellm_pricing import resolve_litellm_model
|
||||
|
||||
key = _first_priced_opus_key()
|
||||
monkeypatch.setenv("HEADROOM_MODEL_ALIAS_MAP", json.dumps({"claude-opus": f"bedrock/{key}"}))
|
||||
_clear_pricing_cache()
|
||||
assert resolve_litellm_model("claude-opus") == key
|
||||
|
||||
|
||||
def test_unpriced_alias_falls_through_soft(monkeypatch):
|
||||
from headroom.pricing.litellm_pricing import resolve_litellm_model
|
||||
|
||||
monkeypatch.setenv("HEADROOM_MODEL_ALIAS_MAP", json.dumps({"claude-opus": "not-a-real-model"}))
|
||||
_clear_pricing_cache()
|
||||
# No crash; falls back to bare-prefix resolution (never returns the bogus target).
|
||||
assert resolve_litellm_model("claude-opus") != "not-a-real-model"
|
||||
|
||||
|
||||
def test_unset_env_is_unchanged(monkeypatch):
|
||||
from headroom.pricing.litellm_pricing import resolve_litellm_model
|
||||
|
||||
monkeypatch.delenv("HEADROOM_MODEL_ALIAS_MAP", raising=False)
|
||||
_clear_pricing_cache()
|
||||
assert isinstance(resolve_litellm_model("gpt-4o"), str)
|
||||
|
||||
|
||||
def test_savings_tracker_delegates_to_shared_resolver(monkeypatch):
|
||||
pytest.importorskip("litellm")
|
||||
from headroom.proxy.savings_tracker import _resolve_litellm_model
|
||||
|
||||
key = _first_priced_opus_key()
|
||||
monkeypatch.setenv("HEADROOM_MODEL_ALIAS_MAP", json.dumps({"claude-opus": key}))
|
||||
_clear_pricing_cache()
|
||||
# Persisted funnel prices the alias identically to the live path.
|
||||
assert _resolve_litellm_model("claude-opus") == key
|
||||
|
||||
|
||||
# ----- #6 context limit: configured alias wins over the 128K default -----
|
||||
|
||||
|
||||
def test_configured_context_limit_wins_over_default(monkeypatch):
|
||||
monkeypatch.setenv(
|
||||
"HEADROOM_MODEL_LIMITS",
|
||||
json.dumps({"openai": {"context_limits": {"claude-opus": 200000}}}),
|
||||
)
|
||||
from headroom.providers.openai import OpenAIProvider
|
||||
|
||||
provider = OpenAIProvider()
|
||||
assert provider.get_context_limit("claude-opus") == 200000 # not the 128000 default
|
||||
|
||||
|
||||
# ----- #1 loopback opt-in on /v1/compress -----
|
||||
|
||||
|
||||
def _fast_app():
|
||||
return create_app(
|
||||
ProxyConfig(
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
_BODY = {"messages": [{"role": "user", "content": "hi"}], "model": "gpt-4"}
|
||||
|
||||
|
||||
def test_compress_blocks_non_loopback_by_default(monkeypatch):
|
||||
monkeypatch.delenv("HEADROOM_COMPRESS_ALLOW_REMOTE", raising=False)
|
||||
# A vanilla TestClient presents client.host="testclient" (non-loopback).
|
||||
client = TestClient(_fast_app())
|
||||
assert client.post("/v1/compress", json=_BODY).status_code == 404
|
||||
|
||||
|
||||
def test_compress_allows_non_loopback_with_flag(monkeypatch):
|
||||
monkeypatch.setenv("HEADROOM_COMPRESS_ALLOW_REMOTE", "1")
|
||||
client = TestClient(_fast_app())
|
||||
resp = client.post("/v1/compress", json=_BODY)
|
||||
assert resp.status_code == 200, resp.text
|
||||
|
|
@ -348,6 +348,86 @@ class TestCompressEndpointCompression:
|
|||
assert outcome.tokens_saved == 0
|
||||
|
||||
|
||||
class TestCompressEndpointLossyInlineMode:
|
||||
"""config.mode="lossy_inline" must compress losslessly-then-lossily but emit
|
||||
NO CCR marker / retrieval round-trip, so the output is safe to forward
|
||||
straight to a provider (Kong-sidecar use case)."""
|
||||
|
||||
def _big_tool_message(self):
|
||||
large_data = json.dumps(
|
||||
[
|
||||
{
|
||||
"id": i,
|
||||
"name": f"Item {i}",
|
||||
"description": f"Detailed description for item {i}. "
|
||||
f"Status active, created 2024-01-{(i % 28) + 1:02d}, "
|
||||
f"category=electronics, price={i * 10.99:.2f}, stock={i * 5}.",
|
||||
"tags": ["electronics", "sale", "featured", "new-arrival"],
|
||||
}
|
||||
for i in range(200)
|
||||
]
|
||||
)
|
||||
return [
|
||||
{"role": "user", "content": "What items are available?"},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": large_data},
|
||||
{"role": "user", "content": "Summarize them."},
|
||||
]
|
||||
|
||||
@pytest.fixture
|
||||
def client(self):
|
||||
# disable_kompress keeps the real ONNX model out of the test: marker
|
||||
# suppression is exercised by SmartCrusher (pure-Python) regardless, and
|
||||
# the mode inherits enable_kompress from config so this stays fast.
|
||||
config = ProxyConfig(
|
||||
optimize=True,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
disable_kompress=True,
|
||||
)
|
||||
app = create_app(config)
|
||||
with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as c:
|
||||
yield c
|
||||
|
||||
def test_lossy_inline_emits_no_ccr_markers(self, client):
|
||||
messages = self._big_tool_message()
|
||||
response = client.post(
|
||||
"/v1/compress",
|
||||
json={"messages": messages, "model": "gpt-4", "config": {"mode": "lossy_inline"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Core guarantee: no CCR markers anywhere — no store/retrieval needed.
|
||||
assert data["ccr_hashes"] == []
|
||||
blob = json.dumps(data["messages"])
|
||||
assert "<<ccr:" not in blob
|
||||
assert "Retrieve more: hash=" not in blob
|
||||
assert "Retrieve original: hash=" not in blob
|
||||
|
||||
# Real compression happened. These guard against a fail-open-to-zero
|
||||
# (e.g. a content-detector hang tripping the executor timeout) passing
|
||||
# vacuously as 0 <= 0.
|
||||
assert data["tokens_before"] > 0
|
||||
assert data["tokens_saved"] > 0
|
||||
assert data["tokens_after"] < data["tokens_before"]
|
||||
assert data["tokens_saved"] == data["tokens_before"] - data["tokens_after"]
|
||||
|
||||
def test_lossless_then_lossy_alias(self, client):
|
||||
"""The spelled-out alias selects the same mode."""
|
||||
messages = self._big_tool_message()
|
||||
response = client.post(
|
||||
"/v1/compress",
|
||||
json={
|
||||
"messages": messages,
|
||||
"model": "gpt-4",
|
||||
"config": {"mode": "lossless_then_lossy"},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["ccr_hashes"] == []
|
||||
|
||||
|
||||
class TestCompressEndpointDoesNotBlockLoop:
|
||||
"""/v1/compress must offload to the compression executor so a slow/large
|
||||
payload cannot freeze the single event loop (#718)."""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue