diff --git a/headroom/pricing/litellm_pricing.py b/headroom/pricing/litellm_pricing.py index 6eb19c30c..0ac08c758 100644 --- a/headroom/pricing/litellm_pricing.py +++ b/headroom/pricing/litellm_pricing.py @@ -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 diff --git a/headroom/providers/openai.py b/headroom/providers/openai.py index 3ebd0c549..cc91d2b71 100644 --- a/headroom/providers/openai.py +++ b/headroom/providers/openai.py @@ -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: diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index f7bbd8eec..8dff6c160 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -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 ``<>`` / ``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, ) diff --git a/headroom/proxy/savings_tracker.py b/headroom/proxy/savings_tracker.py index 5f5bfa599..4b7d82c44 100644 --- a/headroom/proxy/savings_tracker.py +++ b/headroom/proxy/savings_tracker.py @@ -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 diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 56a776d46..3e4afbfa9 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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) diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 44207332f..f317bfbd3 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -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. diff --git a/tests/test_gateway_sidecar_ports.py b/tests/test_gateway_sidecar_ports.py new file mode 100644 index 000000000..9267ef8e5 --- /dev/null +++ b/tests/test_gateway_sidecar_ports.py @@ -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 diff --git a/tests/test_proxy_compress_endpoint.py b/tests/test_proxy_compress_endpoint.py index 3ffef2a19..9ae358341 100644 --- a/tests/test_proxy_compress_endpoint.py +++ b/tests/test_proxy_compress_endpoint.py @@ -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 "< 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)."""