mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(cli): proxy/perf/wrap UX cleanup + perf --hours correctness
Address user-reported UX gaps across the CLI surface: - code-aware: add --code-aware/--no-code-aware (+ HEADROOM_CODE_AWARE_ENABLED env) to the Click CLI. PR #411 had added these only to the orphaned argparse main; the user-facing CLI couldn't reach the flag. Banner status text "remove --no-code-aware to enable" referenced a flag that didn't exist — fix to point at the actual flag/env. Surface code-aware in the click banner and add print_banner=False plumbing to run_server so the click path doesn't print two banners back-to-back. - --mode: hide alias clutter via metavar=[token|cache] and rewrite help to lead with the two real modes. Legacy aliases (token_mode/token_savings/...) still validate. - perf --hours: was documented but ignored. Records are now actually filtered, the report shows the actual time-range covered, and the count of records filtered out (so users can tell when raising --hours helps). - perf TOIN: replace the hash-keyed pattern dump with a strategy-distribution view + recommendation-eligibility from the live store — actionable signal rather than opaque rows. - code-graph: clarify in --help that it indexes cwd / project root. - wrap: spell out supported tools, wrap-vs-proxy distinction, and that `headroom wrap opencode` isn't a thing (use `proxy` directly for opencode; openclaw is not opencode). - mcp: note that mcp__headroom__headroom_retrieve is correct MCP namespacing, not a doubled-prefix bug. Renaming would break the proxy's tool injection. - LLMLingua cleanup: remove [llmlingua] extra from pyproject (no live code uses it). Delete wiki/llmlingua.md and clean retired flag/class references in 6 other wiki pages. Point at [ml] (Kompress) where ML compression is documented. - init -g openclaw: strip mcpServers from existing plugin entries before re-writing — newer openclaw schemas reject it, leaving stale entries from older installs unhealable. Pinned with regression test. Tests: mock_run_server signatures in two existing tests accept **kwargs (needed for the new print_banner plumbing). New test for the openclaw mcpServers strip. Full suite: 4847 passed, 262 skipped, 0 failed.
This commit is contained in:
chopratejas
parent
7aaa4ac48f
commit
265554d4ad
19 changed files with 411 additions and 425 deletions
|
|
@ -76,6 +76,15 @@ def mcp() -> None:
|
|||
3. Claude sees compressed summaries with hash markers
|
||||
4. When Claude needs full details, it calls headroom_retrieve
|
||||
5. The MCP server fetches original content from the proxy
|
||||
|
||||
\b
|
||||
Note on tool naming: MCP clients display tools as
|
||||
`mcp__<server>__<tool>`. Our server is named "headroom" and our
|
||||
tools are named headroom_retrieve / headroom_compress / etc., so
|
||||
Claude Code shows them as `mcp__headroom__headroom_retrieve`. The
|
||||
"headroom" doubling is normal MCP namespacing — not a bug. The
|
||||
proxy's compression markers (and any docs/prompts) reference the
|
||||
bare tool name `headroom_retrieve`.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
|
|
|||
|
|
@ -62,7 +62,11 @@ from .main import main
|
|||
@click.option(
|
||||
"--mode",
|
||||
default=None,
|
||||
metavar="[token|cache]",
|
||||
type=click.Choice(
|
||||
# Canonical modes first; legacy aliases follow for backward compatibility.
|
||||
# `metavar` above hides the alias clutter from --help; users see "[token|cache]"
|
||||
# while internal callers passing "token_mode"/"cost_savings"/etc. still validate.
|
||||
[
|
||||
"token",
|
||||
"cache",
|
||||
|
|
@ -71,12 +75,15 @@ from .main import main
|
|||
"token_savings",
|
||||
"cost_savings",
|
||||
"token_headroom",
|
||||
]
|
||||
],
|
||||
case_sensitive=False,
|
||||
),
|
||||
help=(
|
||||
"Optimization mode: token (prioritize compression) or cache "
|
||||
"(freeze prior turns for prefix-cache stability). "
|
||||
"Legacy aliases are accepted. Default: token. Env: HEADROOM_MODE"
|
||||
"Optimization mode (default: token).\n"
|
||||
" token — prioritize compression; prior turns may be rewritten for max savings.\n"
|
||||
" cache — freeze prior turns to maximise provider prefix-cache hit rate.\n"
|
||||
"Legacy aliases (token_mode, token_savings, token_headroom, cache_mode, "
|
||||
"cost_savings) are still accepted. Env: HEADROOM_MODE."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
|
|
@ -184,11 +191,32 @@ from .main import main
|
|||
envvar="HEADROOM_BUDGET",
|
||||
help="Daily budget limit in USD (env: HEADROOM_BUDGET)",
|
||||
)
|
||||
# Code graph: indexes project + watches files for live reindex via codebase-memory-mcp
|
||||
# Code-aware compression (AST-based, requires `pip install headroom-ai[code]`).
|
||||
# Pair of flags so users can override the env-var default in either direction.
|
||||
# We resolve HEADROOM_CODE_AWARE_ENABLED in the body (not via Click's envvar=),
|
||||
# because Click's envvar handling for paired bool flags is brittle in older
|
||||
# Click versions.
|
||||
@click.option(
|
||||
"--code-aware/--no-code-aware",
|
||||
"code_aware_flag",
|
||||
default=None,
|
||||
help=(
|
||||
"Enable/disable AST-based code compression. Requires the optional "
|
||||
"tree-sitter dependency: pip install headroom-ai[code]. "
|
||||
"Default: disabled. Env: HEADROOM_CODE_AWARE_ENABLED=1 to enable."
|
||||
),
|
||||
)
|
||||
# Code graph: indexes project + watches files for live reindex via codebase-memory-mcp.
|
||||
# Only useful when the proxy is launched from a project root — it indexes the
|
||||
# current working directory.
|
||||
@click.option(
|
||||
"--code-graph",
|
||||
is_flag=True,
|
||||
help="Enable code graph intelligence (indexes project, watches files for live reindex via codebase-memory-mcp)",
|
||||
help=(
|
||||
"Enable code graph intelligence: indexes the current working directory "
|
||||
"and watches files for live reindex via codebase-memory-mcp. Only useful "
|
||||
"when the proxy is launched from a project root."
|
||||
),
|
||||
)
|
||||
# Read lifecycle (ON by default: compresses stale/superseded Read outputs)
|
||||
@click.option(
|
||||
|
|
@ -359,6 +387,7 @@ def proxy(
|
|||
log_file: str | None,
|
||||
log_messages: bool,
|
||||
budget: float | None,
|
||||
code_aware_flag: bool | None,
|
||||
code_graph: bool,
|
||||
no_read_lifecycle: bool,
|
||||
memory: bool,
|
||||
|
|
@ -514,6 +543,17 @@ def proxy(
|
|||
log_full_messages=log_messages
|
||||
or os.environ.get("HEADROOM_LOG_MESSAGES", "").lower() in ("true", "1", "yes", "on"),
|
||||
budget_limit_usd=budget,
|
||||
# Code-aware compression resolution:
|
||||
# 1. Explicit --code-aware / --no-code-aware always wins.
|
||||
# 2. Otherwise read HEADROOM_CODE_AWARE_ENABLED (truthy = on).
|
||||
# 3. Otherwise default off — matches the prior cli/proxy.py behavior so
|
||||
# existing users see no change unless they opt in.
|
||||
code_aware_enabled=(
|
||||
bool(code_aware_flag)
|
||||
if code_aware_flag is not None
|
||||
else os.environ.get("HEADROOM_CODE_AWARE_ENABLED", "").strip().lower()
|
||||
in ("true", "1", "yes", "on")
|
||||
),
|
||||
# Code graph: live file watcher for incremental reindexing
|
||||
code_graph_watcher=code_graph,
|
||||
# Read lifecycle: ON by default (use --no-read-lifecycle to disable)
|
||||
|
|
@ -660,6 +700,13 @@ Memory (Multi-Provider):
|
|||
f"(available: {','.join(_ext_available)})"
|
||||
)
|
||||
|
||||
# Code-aware status line — same logic the inner banner uses, surfaced here
|
||||
# so the click-CLI banner is a complete picture (avoids the dual-banner
|
||||
# confusion this branch retired).
|
||||
from headroom.proxy.server import _get_code_aware_banner_status
|
||||
|
||||
code_aware_line = f" Code-Aware: {_get_code_aware_banner_status(config)}"
|
||||
|
||||
click.echo(f"""
|
||||
╔═══════════════════════════════════════════════════════════════════════╗
|
||||
║ HEADROOM PROXY ║
|
||||
|
|
@ -675,6 +722,7 @@ Starting proxy server...
|
|||
Rate Limit: {"ENABLED" if config.rate_limit_enabled else "DISABLED"}
|
||||
Memory: {memory_status}
|
||||
License: {license_status}
|
||||
{code_aware_line}
|
||||
{extensions_line}
|
||||
{stateless_line}{telemetry_line}
|
||||
{backend_section}
|
||||
|
|
@ -700,11 +748,15 @@ Press Ctrl+C to stop.
|
|||
""")
|
||||
|
||||
try:
|
||||
run_kwargs: dict[str, int] = {}
|
||||
run_kwargs: dict[str, Any] = {}
|
||||
if workers != 1:
|
||||
run_kwargs["workers"] = workers
|
||||
if limit_concurrency != 1000:
|
||||
run_kwargs["limit_concurrency"] = limit_concurrency
|
||||
# Suppress run_server's legacy banner — the click CLI already printed
|
||||
# a richer one above. Direct `python -m headroom.proxy.server` keeps
|
||||
# the legacy banner via run_server's default.
|
||||
run_kwargs["print_banner"] = False
|
||||
run_server(config, **run_kwargs)
|
||||
except KeyboardInterrupt:
|
||||
click.echo("\nShutting down...")
|
||||
|
|
|
|||
|
|
@ -1298,13 +1298,28 @@ def wrap() -> None:
|
|||
the target tool so all API calls route through Headroom automatically.
|
||||
|
||||
\b
|
||||
Supported tools:
|
||||
Supported tools (one Click subcommand per tool):
|
||||
headroom wrap claude # Claude Code (Anthropic)
|
||||
headroom wrap copilot -- --model claude-sonnet-4-20250514
|
||||
headroom wrap codex # OpenAI Codex CLI
|
||||
headroom wrap copilot -- --model claude-sonnet-4-20250514
|
||||
headroom wrap aider # Aider
|
||||
headroom wrap cursor # Cursor (prints config instructions)
|
||||
headroom wrap openclaw # OpenClaw plugin bootstrap
|
||||
|
||||
\b
|
||||
`wrap` vs `proxy`:
|
||||
- `headroom wrap <tool>` — convenience: starts the proxy for you,
|
||||
sets the right env vars, and launches the wrapped CLI.
|
||||
- `headroom proxy` — just the proxy. Use this with any
|
||||
OpenAI/Anthropic-compatible client by setting
|
||||
ANTHROPIC_BASE_URL / OPENAI_BASE_URL yourself. Required for
|
||||
tools without a dedicated `wrap` subcommand
|
||||
(e.g. opencode, Cline, Continue).
|
||||
|
||||
\b
|
||||
Note: `headroom wrap opencode` does NOT exist. For opencode, run
|
||||
`headroom proxy` and point opencode at it via OPENAI_BASE_URL.
|
||||
`openclaw` is a separate tool — different from opencode.
|
||||
"""
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from __future__ import annotations
|
|||
import logging
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from headroom import paths as _paths
|
||||
|
||||
|
|
@ -213,6 +214,31 @@ class PerfReport:
|
|||
toin_records: list[ToinRecord] = field(default_factory=list)
|
||||
log_files_read: int = 0
|
||||
total_lines_parsed: int = 0
|
||||
# Window covered by the report. `requested_hours` is what the caller
|
||||
# asked for; `oldest_kept_ts` / `newest_kept_ts` are the actual
|
||||
# timestamps of the oldest and newest records that survived the
|
||||
# filter (may be narrower if the log doesn't go back that far).
|
||||
# All optional so existing callers keep working.
|
||||
requested_hours: float | None = None
|
||||
oldest_kept_ts: str | None = None
|
||||
newest_kept_ts: str | None = None
|
||||
records_filtered_out: int = 0
|
||||
|
||||
|
||||
# Log timestamps are emitted by Python's `logging` formatter as
|
||||
# `YYYY-MM-DD HH:MM:SS,fff`. We keep the parser permissive so the perf
|
||||
# CLI never throws on a stray malformed line — unparsable records are
|
||||
# just kept (better to over-report than to silently drop data).
|
||||
_LOG_TS_FMT = "%Y-%m-%d %H:%M:%S,%f"
|
||||
|
||||
|
||||
def _parse_log_ts(ts: str | None) -> datetime | None:
|
||||
if not ts:
|
||||
return None
|
||||
try:
|
||||
return datetime.strptime(ts, _LOG_TS_FMT)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def parse_log_files(last_n_hours: float = 168.0) -> PerfReport:
|
||||
|
|
@ -220,15 +246,40 @@ def parse_log_files(last_n_hours: float = 168.0) -> PerfReport:
|
|||
|
||||
Args:
|
||||
last_n_hours: Only include records from the last N hours (default 7 days).
|
||||
Records with un-parseable timestamps are kept (fail-open) — the
|
||||
window in the report header reflects the actual timestamps that
|
||||
survived the filter, so the user can see whether the log went
|
||||
back far enough.
|
||||
|
||||
Returns:
|
||||
PerfReport with all parsed records.
|
||||
"""
|
||||
report = PerfReport()
|
||||
report.requested_hours = last_n_hours
|
||||
|
||||
if not LOG_DIR.exists():
|
||||
return report
|
||||
|
||||
cutoff = datetime.now() - timedelta(hours=last_n_hours) if last_n_hours > 0 else None
|
||||
|
||||
def _within_window(ts_str: str | None) -> bool:
|
||||
# Fail-open: records without a parseable timestamp are kept. The
|
||||
# alternative (silent drop) makes `headroom perf` lie about coverage.
|
||||
if cutoff is None:
|
||||
return True
|
||||
ts = _parse_log_ts(ts_str)
|
||||
if ts is None:
|
||||
return True
|
||||
return ts >= cutoff
|
||||
|
||||
def _track_window(ts_str: str | None) -> None:
|
||||
if not ts_str:
|
||||
return
|
||||
if report.oldest_kept_ts is None or ts_str < report.oldest_kept_ts:
|
||||
report.oldest_kept_ts = ts_str
|
||||
if report.newest_kept_ts is None or ts_str > report.newest_kept_ts:
|
||||
report.newest_kept_ts = ts_str
|
||||
|
||||
# Collect log files: proxy.log, proxy.log.1, proxy.log.2, ...
|
||||
log_files = sorted(LOG_DIR.glob("proxy.log*"), key=lambda p: p.stat().st_mtime)
|
||||
|
||||
|
|
@ -260,9 +311,14 @@ def parse_log_files(last_n_hours: float = 168.0) -> PerfReport:
|
|||
else:
|
||||
# Old comma-separated format
|
||||
transforms = transforms_str.split(",")
|
||||
ts = m.group("ts")
|
||||
if not _within_window(ts):
|
||||
report.records_filtered_out += 1
|
||||
continue
|
||||
_track_window(ts)
|
||||
report.perf_records.append(
|
||||
PerfRecord(
|
||||
timestamp=m.group("ts"),
|
||||
timestamp=ts,
|
||||
request_id=m.group("rid"),
|
||||
model=kv.get("model", ""),
|
||||
num_messages=int(kv.get("msgs", 0)),
|
||||
|
|
@ -283,6 +339,10 @@ def parse_log_files(last_n_hours: float = 168.0) -> PerfReport:
|
|||
m2 = _ROUTER_RE.search(line)
|
||||
if m2:
|
||||
ts = line[:23]
|
||||
if not _within_window(ts):
|
||||
report.records_filtered_out += 1
|
||||
continue
|
||||
_track_window(ts)
|
||||
detail = m2.group("detail")
|
||||
rec = RouterRecord(
|
||||
timestamp=ts,
|
||||
|
|
@ -312,6 +372,10 @@ def parse_log_files(last_n_hours: float = 168.0) -> PerfReport:
|
|||
m3 = _TRANSFORM_RE.search(line)
|
||||
if m3:
|
||||
ts = line[:23]
|
||||
if not _within_window(ts):
|
||||
report.records_filtered_out += 1
|
||||
continue
|
||||
_track_window(ts)
|
||||
report.transform_records.append(
|
||||
TransformRecord(
|
||||
timestamp=ts,
|
||||
|
|
@ -327,6 +391,10 @@ def parse_log_files(last_n_hours: float = 168.0) -> PerfReport:
|
|||
m4 = _TOIN_RE.search(line)
|
||||
if m4:
|
||||
ts = line[:23]
|
||||
if not _within_window(ts):
|
||||
report.records_filtered_out += 1
|
||||
continue
|
||||
_track_window(ts)
|
||||
report.toin_records.append(
|
||||
ToinRecord(
|
||||
timestamp=ts,
|
||||
|
|
@ -357,6 +425,21 @@ def format_report(report: PerfReport) -> str:
|
|||
# Header
|
||||
lines.append("Headroom Performance Report")
|
||||
lines.append("=" * 60)
|
||||
if report.requested_hours is not None:
|
||||
if report.oldest_kept_ts and report.newest_kept_ts:
|
||||
window_str = (
|
||||
f"Window: last {report.requested_hours:g}h "
|
||||
f"(actual data: {report.oldest_kept_ts[:19]} → "
|
||||
f"{report.newest_kept_ts[:19]})"
|
||||
)
|
||||
else:
|
||||
window_str = f"Window: last {report.requested_hours:g}h (no records found in window)"
|
||||
lines.append(window_str)
|
||||
if report.records_filtered_out > 0:
|
||||
lines.append(
|
||||
f"Records outside window: {report.records_filtered_out:,} "
|
||||
"(filtered out — increase --hours to include them)"
|
||||
)
|
||||
lines.append("")
|
||||
|
||||
records = report.perf_records
|
||||
|
|
@ -501,7 +584,7 @@ def format_report(report: PerfReport) -> str:
|
|||
lines.append(" ! Excluded tools dominate — consider compressing stale Read outputs")
|
||||
lines.append("")
|
||||
|
||||
# TOIN status
|
||||
# TOIN status — log-derived counters first, then live-store highlights.
|
||||
if report.toin_records:
|
||||
latest = report.toin_records[-1]
|
||||
lines.append("TOIN Learning")
|
||||
|
|
@ -513,6 +596,17 @@ def format_report(report: PerfReport) -> str:
|
|||
lines.append(" ! 0% retrieval rate — TOIN learning but never used")
|
||||
lines.append("")
|
||||
|
||||
# TOIN highlights — read the live on-disk pattern store and surface
|
||||
# strategy distribution + top high-impact patterns in human-readable
|
||||
# form. Pattern keys are opaque hashes, so the actionable signal is
|
||||
# *which strategies are winning* and *how many patterns have crossed
|
||||
# the recommendation threshold*. Best-effort: if TOIN isn't installed
|
||||
# or the store is empty, we skip the section silently.
|
||||
toin_lines = _format_toin_highlights()
|
||||
if toin_lines:
|
||||
lines.extend(toin_lines)
|
||||
lines.append("")
|
||||
|
||||
# Recommendations
|
||||
recommendations = _generate_recommendations(report)
|
||||
if recommendations:
|
||||
|
|
@ -531,6 +625,88 @@ def format_report(report: PerfReport) -> str:
|
|||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _format_toin_highlights() -> list[str]:
|
||||
"""Render a human-readable TOIN highlights block from the live store.
|
||||
|
||||
Returns an empty list when TOIN is unavailable or has no patterns.
|
||||
Pattern keys (auth_mode, model_family, structure_hash) are opaque
|
||||
hashes so we don't print them as rows — instead we group by the
|
||||
learned ``optimal_strategy`` (a human-readable string like
|
||||
``"lossless:table(240->len=7026)"``) and surface the highest-impact
|
||||
slices via ``avg_token_reduction``.
|
||||
"""
|
||||
try:
|
||||
from headroom.telemetry.toin import get_toin
|
||||
except ImportError:
|
||||
return []
|
||||
|
||||
try:
|
||||
pairs = get_toin().iter_patterns()
|
||||
except Exception: # noqa: BLE001 — perf must never fail on TOIN errors
|
||||
return []
|
||||
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
# Strategy distribution: how many patterns settled on each strategy.
|
||||
strategy_counts: dict[str, int] = {}
|
||||
for _key, pattern in pairs:
|
||||
strategy = pattern.optimal_strategy or "default"
|
||||
strategy_counts[strategy] = strategy_counts.get(strategy, 0) + 1
|
||||
|
||||
# Top patterns by avg token reduction (the high-impact learnings).
|
||||
by_impact = sorted(
|
||||
pairs,
|
||||
key=lambda kp: kp[1].avg_token_reduction,
|
||||
reverse=True,
|
||||
)[:5]
|
||||
|
||||
# How many patterns have enough samples to drive a recommendation.
|
||||
# Falls back to 0 if the threshold attr isn't reachable.
|
||||
try:
|
||||
from headroom.telemetry.toin import get_toin as _get
|
||||
|
||||
threshold = _get()._config.min_samples_for_recommendation
|
||||
except Exception: # noqa: BLE001
|
||||
threshold = 1
|
||||
qualified = sum(1 for _k, p in pairs if p.sample_size >= threshold)
|
||||
|
||||
lines: list[str] = []
|
||||
lines.append("TOIN Highlights (live store)")
|
||||
lines.append("-" * 40)
|
||||
lines.append(
|
||||
f" {qualified}/{len(pairs)} patterns have ≥{threshold} samples "
|
||||
f"(eligible for `python -m headroom.cli.toin_publish`)"
|
||||
)
|
||||
lines.append("")
|
||||
lines.append(" Strategy distribution:")
|
||||
for strategy, count in sorted(strategy_counts.items(), key=lambda kv: kv[1], reverse=True)[:8]:
|
||||
lines.append(f" {count:>4} pattern(s) {strategy}")
|
||||
|
||||
# Only surface patterns with non-trivial impact AND a non-default
|
||||
# strategy — single-digit-token "wins" against the default strategy
|
||||
# are noise, not insight.
|
||||
impact_rows = [
|
||||
(kp[1].avg_token_reduction, kp[1].total_compressions, kp[1].optimal_strategy or "default")
|
||||
for kp in by_impact
|
||||
if kp[1].avg_token_reduction >= 50 and (kp[1].optimal_strategy or "default") != "default"
|
||||
]
|
||||
if impact_rows:
|
||||
lines.append("")
|
||||
lines.append(" Top patterns by avg token reduction:")
|
||||
for avg_red, n, strategy in impact_rows:
|
||||
lines.append(f" {avg_red:>7.0f} tok avg ({n:>3} compression(s)) {strategy}")
|
||||
|
||||
if qualified == 0 and len(pairs) > 0:
|
||||
lines.append("")
|
||||
lines.append(
|
||||
f" ! No pattern has reached {threshold} samples — TOIN is still warming up. "
|
||||
"Recommendations TOML will be empty until traffic grows."
|
||||
)
|
||||
|
||||
return lines
|
||||
|
||||
|
||||
def _generate_recommendations(report: PerfReport) -> list[str]:
|
||||
"""Generate actionable recommendations from the report data."""
|
||||
recs: list[str] = []
|
||||
|
|
|
|||
|
|
@ -35,6 +35,15 @@ def decode_entry_json(raw_value: str | None) -> Any | None:
|
|||
return raw_value
|
||||
|
||||
|
||||
# Keys we know newer openclaw plugin schemas reject when echoed back.
|
||||
# We strip them defensively from `existing_entry` so a stale entry left
|
||||
# over from an older Headroom or older OpenClaw install doesn't cause
|
||||
# `openclaw config set` to fail with "Unrecognized key". The list is
|
||||
# narrow on purpose — anything else is assumed user-managed and
|
||||
# preserved verbatim.
|
||||
_LEGACY_REJECTED_TOP_LEVEL_KEYS: frozenset[str] = frozenset({"mcpServers"})
|
||||
|
||||
|
||||
def build_plugin_entry(
|
||||
*,
|
||||
existing_entry: Any,
|
||||
|
|
@ -46,7 +55,8 @@ def build_plugin_entry(
|
|||
enabled: bool,
|
||||
) -> dict[str, object]:
|
||||
"""Merge managed Headroom plugin settings with any existing entry payload."""
|
||||
base_entry = existing_entry if isinstance(existing_entry, dict) else {}
|
||||
raw_base = existing_entry if isinstance(existing_entry, dict) else {}
|
||||
base_entry = {k: v for k, v in raw_base.items() if k not in _LEGACY_REJECTED_TOP_LEVEL_KEYS}
|
||||
existing_config = base_entry.get("config")
|
||||
next_config = dict(existing_config) if isinstance(existing_config, dict) else {}
|
||||
|
||||
|
|
|
|||
|
|
@ -2598,14 +2598,15 @@ def _get_code_aware_banner_status(config: ProxyConfig) -> str:
|
|||
return "NOT INSTALLED (pip install headroom-ai[code])"
|
||||
else:
|
||||
if is_tree_sitter_available():
|
||||
return "DISABLED (remove --no-code-aware to enable)"
|
||||
return "DISABLED"
|
||||
return "DISABLED (--code-aware or HEADROOM_CODE_AWARE_ENABLED=1 to enable)"
|
||||
return "DISABLED (install headroom-ai[code] to enable)"
|
||||
|
||||
|
||||
def run_server(
|
||||
config: ProxyConfig | None = None,
|
||||
workers: int = 1,
|
||||
limit_concurrency: int = 1000,
|
||||
print_banner: bool = True,
|
||||
):
|
||||
"""Run the proxy server.
|
||||
|
||||
|
|
@ -2613,6 +2614,11 @@ def run_server(
|
|||
config: Proxy configuration
|
||||
workers: Number of worker processes (use N for multi-core scaling)
|
||||
limit_concurrency: Max concurrent connections before 503 response
|
||||
print_banner: When False, skip the legacy ASCII banner. The
|
||||
Click CLI (`headroom proxy`) prints its own startup banner
|
||||
before calling this — printing a second banner here is the
|
||||
"dual banner" UX issue. Direct `python -m headroom.proxy.server`
|
||||
still gets the banner since it has no other startup output.
|
||||
"""
|
||||
if not FASTAPI_AVAILABLE:
|
||||
print("ERROR: FastAPI required. Install: pip install fastapi uvicorn httpx")
|
||||
|
|
@ -2631,7 +2637,8 @@ def run_server(
|
|||
bedrock_region=config.bedrock_region,
|
||||
)
|
||||
|
||||
print(f"""
|
||||
if print_banner:
|
||||
print(f"""
|
||||
╔══════════════════════════════════════════════════════════════════════╗
|
||||
║ HEADROOM PROXY SERVER ║
|
||||
╠══════════════════════════════════════════════════════════════════════╣
|
||||
|
|
|
|||
|
|
@ -75,17 +75,13 @@ proxy = [
|
|||
code = [
|
||||
"tree-sitter-language-pack>=0.10.0",
|
||||
]
|
||||
# ML-based compression with Kompress (ModernBERT)
|
||||
# ML-based compression with Kompress (ModernBERT).
|
||||
# (The legacy [llmlingua] extra was removed in 0.9.x — no live code path used it.
|
||||
# Use [ml] for the supported ML compression dependencies.)
|
||||
ml = [
|
||||
"torch>=2.0.0",
|
||||
"transformers>=4.30.0",
|
||||
]
|
||||
# Legacy ML compression (LLMLingua-2 — use [ml] instead for Kompress)
|
||||
llmlingua = [
|
||||
"llmlingua>=0.2.0",
|
||||
"torch>=2.0.0",
|
||||
"transformers>=4.30.0",
|
||||
]
|
||||
# Memory system (hierarchical memory with vector search)
|
||||
memory = [
|
||||
"hnswlib>=0.8.0",
|
||||
|
|
|
|||
|
|
@ -721,7 +721,7 @@ def _run_cli_capture(args: list[str], env: dict | None = None) -> ProxyConfig:
|
|||
captured: dict[str, ProxyConfig] = {}
|
||||
orig_run = server_mod.run_server
|
||||
|
||||
def _fake_run(config: ProxyConfig): # noqa: D401 - stub
|
||||
def _fake_run(config: ProxyConfig, **_kwargs): # noqa: D401 - stub
|
||||
captured["config"] = config
|
||||
return 0
|
||||
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@ class TestCLIProxyEnvVars:
|
|||
"""HEADROOM_HOST env var should be passed to ProxyConfig."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config):
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
with patch("headroom.proxy.server.run_server", mock_run_server):
|
||||
|
|
@ -48,7 +48,7 @@ class TestCLIProxyEnvVars:
|
|||
"""HEADROOM_PORT env var should be passed to ProxyConfig."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config):
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
with patch("headroom.proxy.server.run_server", mock_run_server):
|
||||
|
|
@ -66,7 +66,7 @@ class TestCLIProxyEnvVars:
|
|||
"""HEADROOM_BUDGET env var should be passed to ProxyConfig."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config):
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
with patch("headroom.proxy.server.run_server", mock_run_server):
|
||||
|
|
@ -84,7 +84,7 @@ class TestCLIProxyEnvVars:
|
|||
"""OPENAI_TARGET_API_URL env var should be passed to ProxyConfig."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config):
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
with patch("headroom.proxy.server.run_server", mock_run_server):
|
||||
|
|
@ -102,7 +102,7 @@ class TestCLIProxyEnvVars:
|
|||
"""GEMINI_TARGET_API_URL env var should be passed to ProxyConfig."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config):
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
with patch("headroom.proxy.server.run_server", mock_run_server):
|
||||
|
|
@ -120,7 +120,7 @@ class TestCLIProxyEnvVars:
|
|||
"""--openai-api-url CLI flag should take precedence."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config):
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
with patch("headroom.proxy.server.run_server", mock_run_server):
|
||||
|
|
@ -137,7 +137,7 @@ class TestCLIProxyEnvVars:
|
|||
"""CLI flag should take precedence over env var."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config):
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
with patch("headroom.proxy.server.run_server", mock_run_server):
|
||||
|
|
@ -155,7 +155,7 @@ class TestCLIProxyEnvVars:
|
|||
"""Without env var or flag, openai_api_url should be None."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config):
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
# Ensure the env var is not set
|
||||
|
|
@ -178,7 +178,7 @@ class TestCLIProxyEnvVars:
|
|||
"""Both OPENAI and GEMINI target URLs can be set via env."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config):
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
with patch("headroom.proxy.server.run_server", mock_run_server):
|
||||
|
|
@ -200,7 +200,7 @@ class TestCLIProxyEnvVars:
|
|||
"""Fast-fail CLI flags should map into ProxyConfig."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config):
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
with patch("headroom.proxy.server.run_server", mock_run_server):
|
||||
|
|
@ -243,7 +243,12 @@ class TestCLIProxyEnvVars:
|
|||
assert result.exit_code == 0, result.output
|
||||
assert captured["config"].max_connections == 200
|
||||
assert captured["config"].max_keepalive_connections == 50
|
||||
assert captured["kwargs"] == {"workers": 4, "limit_concurrency": 250}
|
||||
# Click CLI also passes `print_banner=False` to suppress the legacy
|
||||
# run_server banner (cli/proxy.py prints its own). Assert the
|
||||
# production-scaling keys we care about, not the full kwargs dict.
|
||||
assert captured["kwargs"]["workers"] == 4
|
||||
assert captured["kwargs"]["limit_concurrency"] == 250
|
||||
assert captured["kwargs"].get("print_banner") is False
|
||||
|
||||
def test_production_scaling_cli_flags_override_env_vars(self, runner):
|
||||
captured = {}
|
||||
|
|
@ -278,7 +283,11 @@ class TestCLIProxyEnvVars:
|
|||
assert result.exit_code == 0, result.output
|
||||
assert captured["config"].max_connections == 150
|
||||
assert captured["config"].max_keepalive_connections == 25
|
||||
assert captured["kwargs"] == {"workers": 3, "limit_concurrency": 125}
|
||||
# Click CLI also passes `print_banner=False`. Assert production
|
||||
# scaling keys explicitly rather than the full kwargs dict.
|
||||
assert captured["kwargs"]["workers"] == 3
|
||||
assert captured["kwargs"]["limit_concurrency"] == 125
|
||||
assert captured["kwargs"].get("print_banner") is False
|
||||
|
||||
|
||||
class TestCLIProxyBackend:
|
||||
|
|
@ -288,7 +297,7 @@ class TestCLIProxyBackend:
|
|||
"""--backend litellm-hosted_vllm should be accepted (not rejected)."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config):
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
with patch("headroom.proxy.server.run_server", mock_run_server):
|
||||
|
|
@ -305,7 +314,7 @@ class TestCLIProxyBackend:
|
|||
"""--backend litellm-vertex should be accepted."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config):
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
with patch("headroom.proxy.server.run_server", mock_run_server):
|
||||
|
|
@ -322,7 +331,7 @@ class TestCLIProxyBackend:
|
|||
"""Full vLLM setup: litellm backend + OPENAI_TARGET_API_URL."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config):
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
with patch("headroom.proxy.server.run_server", mock_run_server):
|
||||
|
|
@ -350,7 +359,7 @@ class TestCLIAnyllmProviderEnv:
|
|||
"""HEADROOM_ANYLLM_PROVIDER env var should override the default."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config):
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
with patch("headroom.proxy.server.run_server", mock_run_server):
|
||||
|
|
@ -368,7 +377,7 @@ class TestCLIAnyllmProviderEnv:
|
|||
"""--anyllm-provider flag should still work."""
|
||||
captured_config = {}
|
||||
|
||||
def mock_run_server(config):
|
||||
def mock_run_server(config, **kwargs):
|
||||
captured_config["config"] = config
|
||||
|
||||
with patch("headroom.proxy.server.run_server", mock_run_server):
|
||||
|
|
|
|||
|
|
@ -117,3 +117,34 @@ def test_build_unwrap_entry_disables_plugin_and_removes_managed_keys_only() -> N
|
|||
def test_build_unwrap_entry_handles_non_mapping_input() -> None:
|
||||
# Arrange / Act / Assert
|
||||
assert build_unwrap_entry("not-a-dict") == {"enabled": False, "config": {}}
|
||||
|
||||
|
||||
def test_build_plugin_entry_strips_mcpServers_from_existing_entry() -> None:
|
||||
"""Newer OpenClaw schemas reject `mcpServers` at the plugin-entry root.
|
||||
|
||||
`headroom init -g` was failing with `Config invalid: Unrecognized
|
||||
key: "mcpServers"` because the prior plugin entry in the user's
|
||||
config still had that legacy field, and we were spreading it back in
|
||||
via `**existing_entry`. Pin the strip so we don't regress.
|
||||
"""
|
||||
existing_entry = {
|
||||
"enabled": True,
|
||||
"name": "headroom",
|
||||
"mcpServers": {"some": "stale-block"}, # legacy, must be removed
|
||||
"config": {"keep": "value"},
|
||||
}
|
||||
|
||||
entry = build_plugin_entry(
|
||||
existing_entry=existing_entry,
|
||||
proxy_port=8787,
|
||||
startup_timeout_ms=1500,
|
||||
python_path=None,
|
||||
no_auto_start=False,
|
||||
gateway_provider_ids=None,
|
||||
enabled=True,
|
||||
)
|
||||
|
||||
assert "mcpServers" not in entry
|
||||
assert entry["enabled"] is True
|
||||
assert entry["name"] == "headroom"
|
||||
assert entry["config"]["keep"] == "value"
|
||||
|
|
|
|||
|
|
@ -239,32 +239,18 @@ analysis = {
|
|||
|
||||
---
|
||||
|
||||
#### Transform 4: LLMLingua Compressor (Optional)
|
||||
#### Transform 4: ML Compressor (Optional, Kompress)
|
||||
|
||||
**When to use:** Maximum compression needed and latency is acceptable.
|
||||
|
||||
```python
|
||||
# Opt-in ML-based compression using Microsoft's LLMLingua-2
|
||||
# BERT-based token classifier trained via GPT-4 distillation
|
||||
The proxy ships an opt-in ML compression path backed by **Kompress**
|
||||
(ModernBERT-based token classifier). Install with `pip install
|
||||
'headroom-ai[ml]'`; see `wiki/transforms.md` for current configuration.
|
||||
|
||||
# Before: Long tool output text
|
||||
"The function processUserData takes a user object and validates all fields..."
|
||||
|
||||
# After: Compressed while preserving semantic meaning
|
||||
"function processUserData validates user fields..."
|
||||
```
|
||||
|
||||
**Key characteristics:**
|
||||
- Uses `microsoft/llmlingua-2-xlm-roberta-large-meetingbank` model
|
||||
- Auto-detects content type (code, JSON, text) for optimal compression rates
|
||||
- Stores original in CCR for retrieval if needed
|
||||
- Adds 50-200ms latency per request
|
||||
- Requires ~1GB RAM when loaded
|
||||
|
||||
**Proxy integration (opt-in):**
|
||||
```bash
|
||||
headroom proxy --llmlingua --llmlingua-device cuda
|
||||
```
|
||||
**Note:** The earlier LLMLingua-2 integration (`--llmlingua` flag, the
|
||||
`headroom-ai[llmlingua]` extra, and the `LLMLinguaCompressor` class) was
|
||||
retired and replaced by Kompress. `pip install 'headroom-ai[llmlingua]'`
|
||||
no longer resolves; use `[ml]` instead.
|
||||
|
||||
---
|
||||
|
||||
|
|
@ -418,8 +404,8 @@ def apply(self, messages, ...):
|
|||
# - Compresses to 17 points (preserving spike)
|
||||
# - Factors out constant "host" field
|
||||
|
||||
# Transform 3: LLMLingua (if enabled via --llmlingua)
|
||||
# - ML-based compression on remaining long text
|
||||
# Transform 3: Kompress ML compressor (opt-in via headroom-ai[ml])
|
||||
# - ModernBERT-based compression on remaining long text
|
||||
# - Auto-detects content type for optimal rate
|
||||
# - Stores original in CCR for retrieval
|
||||
|
||||
|
|
@ -1082,7 +1068,7 @@ headroom/
|
|||
│ ├── rolling_window.py # Token limit enforcement (position-based)
|
||||
│ ├── intelligent_context.py # Semantic context management (score-based)
|
||||
│ ├── scoring.py # Message importance scoring
|
||||
│ └── llmlingua_compressor.py # ML-based compression (opt-in)
|
||||
│ └── (legacy llmlingua_compressor.py removed — see [ml] extra for Kompress)
|
||||
│
|
||||
├── cache/ # CCR Architecture - Caching & Storage
|
||||
│ ├── compression_store.py # Phase 1: Store original content
|
||||
|
|
|
|||
|
|
@ -73,24 +73,23 @@ SmartCrusher doesn't use fixed K values. It uses information-theoretic sizing:
|
|||
|
||||
These are kept even if they exceed the K budget.
|
||||
|
||||
## Text Compression (LLMLingua)
|
||||
## ML Text Compression (Kompress, opt-in)
|
||||
|
||||
- **Requires**: `headroom-ai[llmlingua]` — downloads ~2GB model, needs ~1GB RAM
|
||||
- **First call**: 10-30s model load latency (cached globally after)
|
||||
- **Sequence length**: Content chunked at 512 tokens (model limit)
|
||||
- **Content < 100 tokens**: Skipped
|
||||
- **Latency**: Adds overhead that doesn't break even on fast models (GPT-4o Mini, Sonnet). Use for **cost savings**, not speed
|
||||
- **Requires**: `headroom-ai[ml]` — downloads model weights and needs GPU/CPU RAM for inference
|
||||
- **First call**: model-load latency (cached globally after)
|
||||
- **Latency**: Adds overhead that doesn't break even on fast models. Use for **cost savings**, not speed
|
||||
- **Thread safety**: Single global model instance with lock — sequential access under concurrency
|
||||
|
||||
> The earlier LLMLingua-2 integration (`headroom-ai[llmlingua]`) was retired and is no longer installable.
|
||||
|
||||
## Error Handling
|
||||
|
||||
All compressors follow the same principle: **fail gracefully, return original content unchanged**.
|
||||
|
||||
- Invalid JSON → passthrough (no error raised)
|
||||
- AST parse failure in CodeCompressor → falls back to original or LLMLingua
|
||||
- AST parse failure in CodeCompressor → falls back to original
|
||||
- Compression makes output larger → original returned
|
||||
- Missing optional dependencies (tree-sitter, LLMLingua) → passthrough with warning log
|
||||
- **One exception**: LLMLingua out-of-memory during model loading raises `RuntimeError`
|
||||
- Missing optional dependencies (tree-sitter, ML stack) → passthrough with warning log
|
||||
|
||||
Errors are logged at WARNING level and never propagated to callers.
|
||||
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ Universal Compression combines several techniques:
|
|||
|
||||
1. **ML-based Detection** - Automatically detects content type (JSON, code, logs, text) using Magika
|
||||
2. **Structure Preservation** - Keeps keys, signatures, and templates intact via structure masks
|
||||
3. **Intelligent Compression** - Compresses content while preserving meaning with LLMLingua
|
||||
3. **Intelligent Compression** - Compresses content while preserving meaning with the optional ML compressor (Kompress)
|
||||
4. **Reversible via CCR** - Stores originals for retrieval when LLM needs full context
|
||||
|
||||
## Quick Start
|
||||
|
|
@ -51,9 +51,9 @@ result = compressor.compress(content)
|
|||
│ │ │
|
||||
▼ ▼ ▼
|
||||
┌─────────────┐ ┌─────────────┐ ┌─────────────┐
|
||||
│ Magika │ │ Handler │ │ LLMLingua │
|
||||
│ (ML) │ │ (JSON, │ │ (optional) │
|
||||
│ │ │ Code...) │ │ │
|
||||
│ Magika │ │ Handler │ │ Kompress │
|
||||
│ (ML) │ │ (JSON, │ │ (ML, opt- │
|
||||
│ │ │ Code...) │ │ in [ml]) │
|
||||
└─────────────┘ └─────────────┘ └─────────────┘
|
||||
```
|
||||
|
||||
|
|
@ -82,7 +82,9 @@ config = UniversalCompressorConfig(
|
|||
use_magika=True, # Use ML-based detection (requires magika)
|
||||
|
||||
# Compression
|
||||
use_llmlingua=True, # Use LLMLingua for compression
|
||||
# (Note: the legacy `use_llmlingua` flag was retired with the
|
||||
# LLMLingua-2 integration. The optional ML compressor is now Kompress,
|
||||
# installed via `headroom-ai[ml]` and configured separately.)
|
||||
compression_ratio_target=0.3, # Keep 30% of content (70% reduction)
|
||||
min_content_length=100, # Skip content shorter than this
|
||||
|
||||
|
|
|
|||
|
|
@ -59,9 +59,9 @@ headroom proxy --no-ccr-responses
|
|||
# Disable proactive expansion
|
||||
headroom proxy --no-ccr-expansion
|
||||
|
||||
# Enable LLMLingua ML compression
|
||||
headroom proxy --llmlingua
|
||||
headroom proxy --llmlingua --llmlingua-device cuda --llmlingua-rate 0.4
|
||||
# (The earlier --llmlingua flag was retired in 0.9.x and replaced by
|
||||
# Kompress (ModernBERT). See `wiki/transforms.md` for the current
|
||||
# opt-in path via the `[ml]` extra.)
|
||||
```
|
||||
|
||||
### All Options
|
||||
|
|
|
|||
|
|
@ -1,188 +0,0 @@
|
|||
# LLMLingua-2 Integration
|
||||
|
||||
For maximum compression, Headroom integrates with **LLMLingua-2**, Microsoft's BERT-based token classifier trained via GPT-4 distillation. It achieves **up to 20x compression** while preserving semantic meaning.
|
||||
|
||||
## When to Use LLMLingua-2
|
||||
|
||||
| Approach | Best For | Compression | Speed |
|
||||
|----------|----------|-------------|-------|
|
||||
| **SmartCrusher** | JSON tool outputs | 70-90% | ~1ms |
|
||||
| **Text Utilities** | Search/logs | 50-90% | ~1ms |
|
||||
| **LLMLingua-2** | Any text, max compression | 80-95% | ~50-200ms |
|
||||
|
||||
LLMLingua-2 is ideal when you need maximum compression and can tolerate slightly higher latency (e.g., compressing large tool outputs before storage, offline processing).
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
# Adds ~2GB of model weights
|
||||
pip install "headroom-ai[llmlingua]"
|
||||
```
|
||||
|
||||
## Basic Usage
|
||||
|
||||
```python
|
||||
from headroom.transforms import LLMLinguaCompressor
|
||||
|
||||
# Create compressor (model loaded lazily on first use)
|
||||
compressor = LLMLinguaCompressor()
|
||||
|
||||
# Compress any text
|
||||
long_output = "The function processUserData takes a user object and validates..."
|
||||
result = compressor.compress(long_output)
|
||||
|
||||
print(f"Before: {result.original_tokens} tokens")
|
||||
print(f"After: {result.compressed_tokens} tokens")
|
||||
print(f"Saved: {result.savings_percentage:.1f}%")
|
||||
print(result.compressed)
|
||||
```
|
||||
|
||||
## Content-Aware Compression
|
||||
|
||||
LLMLingua-2 automatically adjusts compression based on content type:
|
||||
|
||||
```python
|
||||
from headroom.transforms import LLMLinguaCompressor, LLMLinguaConfig
|
||||
|
||||
# Conservative for code (keep 40% of tokens)
|
||||
config = LLMLinguaConfig(
|
||||
code_compression_rate=0.4, # More conservative
|
||||
json_compression_rate=0.35, # Moderate
|
||||
text_compression_rate=0.25, # Aggressive
|
||||
)
|
||||
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
|
||||
# Auto-detects content type
|
||||
code_result = compressor.compress("def calculate(x): return x * 2")
|
||||
text_result = compressor.compress("This is a verbose explanation...")
|
||||
```
|
||||
|
||||
## Memory Management
|
||||
|
||||
The model uses ~1GB RAM. Unload it when done:
|
||||
|
||||
```python
|
||||
from headroom.transforms import (
|
||||
LLMLinguaCompressor,
|
||||
unload_llmlingua_model,
|
||||
is_llmlingua_model_loaded,
|
||||
)
|
||||
|
||||
compressor = LLMLinguaCompressor()
|
||||
result = compressor.compress(content) # Model loaded here
|
||||
|
||||
# Check if loaded
|
||||
print(is_llmlingua_model_loaded()) # True
|
||||
|
||||
# Free memory when done
|
||||
unload_llmlingua_model() # Frees ~1GB
|
||||
print(is_llmlingua_model_loaded()) # False
|
||||
|
||||
# Next compression will reload automatically
|
||||
```
|
||||
|
||||
## Device Configuration
|
||||
|
||||
```python
|
||||
from headroom.transforms import LLMLinguaConfig, LLMLinguaCompressor
|
||||
|
||||
# Force CPU (slower but works everywhere)
|
||||
config = LLMLinguaConfig(device="cpu")
|
||||
|
||||
# Force GPU (faster but needs CUDA)
|
||||
config = LLMLinguaConfig(device="cuda")
|
||||
|
||||
# Auto-detect (default): uses CUDA > MPS > CPU
|
||||
config = LLMLinguaConfig(device="auto")
|
||||
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
```
|
||||
|
||||
## Use in Pipeline
|
||||
|
||||
```python
|
||||
from headroom.transforms import TransformPipeline, LLMLinguaCompressor, SmartCrusher
|
||||
|
||||
# Combine with other transforms
|
||||
pipeline = TransformPipeline([
|
||||
SmartCrusher(), # First: compress JSON
|
||||
LLMLinguaCompressor(), # Then: ML compression on remaining text
|
||||
])
|
||||
|
||||
result = pipeline.apply(messages, tokenizer)
|
||||
```
|
||||
|
||||
## Proxy Integration
|
||||
|
||||
Enable LLMLingua in the proxy server for automatic ML compression:
|
||||
|
||||
```bash
|
||||
# Enable LLMLingua in proxy (requires: pip install headroom-ai[llmlingua,proxy])
|
||||
headroom proxy --llmlingua
|
||||
|
||||
# With custom settings
|
||||
headroom proxy --llmlingua --llmlingua-device cuda --llmlingua-rate 0.4
|
||||
|
||||
# The proxy shows LLMLingua status at startup:
|
||||
# LLMLingua: ENABLED (device=cuda, rate=0.4)
|
||||
#
|
||||
# If llmlingua is installed but not enabled, you'll see a helpful hint:
|
||||
# LLMLingua: available (enable with --llmlingua for ML compression)
|
||||
```
|
||||
|
||||
## Configuration Reference
|
||||
|
||||
| Option | Default | Description |
|
||||
|--------|---------|-------------|
|
||||
| `device` | `"auto"` | Device to run model on: auto, cpu, cuda, mps |
|
||||
| `code_compression_rate` | `0.4` | Keep 40% of tokens for code |
|
||||
| `json_compression_rate` | `0.35` | Keep 35% of tokens for JSON |
|
||||
| `text_compression_rate` | `0.25` | Keep 25% of tokens for text |
|
||||
| `force_tokens` | `[]` | Tokens to always preserve |
|
||||
| `drop_consecutive` | `True` | Drop consecutive whitespace |
|
||||
|
||||
## Performance Characteristics
|
||||
|
||||
| Metric | Value |
|
||||
|--------|-------|
|
||||
| Model size | ~500MB |
|
||||
| Memory usage | ~1GB RAM |
|
||||
| Cold start | 10-30s (first load) |
|
||||
| Inference | 50-200ms per request |
|
||||
| Compression | 80-95% |
|
||||
|
||||
## Why Opt-In?
|
||||
|
||||
LLMLingua adds significant dependencies and overhead:
|
||||
|
||||
| Aspect | Default Proxy | With LLMLingua |
|
||||
|--------|--------------|----------------|
|
||||
| Dependencies | ~50MB | ~2GB |
|
||||
| Cold start | <1s | 10-30s |
|
||||
| Per-request | ~1-5ms | ~50-200ms |
|
||||
| Compression | 70-90% | 80-95% |
|
||||
|
||||
The default proxy is lightweight and fast. Enable LLMLingua when you need maximum compression and can accept the tradeoffs.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### "Model not found"
|
||||
|
||||
```bash
|
||||
# Ensure llmlingua extra is installed
|
||||
pip install "headroom-ai[llmlingua]"
|
||||
```
|
||||
|
||||
### "CUDA out of memory"
|
||||
|
||||
```python
|
||||
# Force CPU mode
|
||||
config = LLMLinguaConfig(device="cpu")
|
||||
```
|
||||
|
||||
### "Slow compression"
|
||||
|
||||
- Use GPU if available: `device="cuda"`
|
||||
- Batch multiple compressions
|
||||
- Consider using SmartCrusher for JSON (faster, similar results)
|
||||
|
|
@ -202,21 +202,14 @@ Configure additional options in the plist `EnvironmentVariables` section:
|
|||
<key>ANTHROPIC_API_KEY</key>
|
||||
<string>sk-ant-...</string>
|
||||
|
||||
<!-- Optional: Enable LLMLingua compression -->
|
||||
<key>HEADROOM_COMPRESSION_PROVIDER</key>
|
||||
<string>llmlingua</string>
|
||||
|
||||
<!-- Optional: LLMLingua device (auto, cuda, cpu, mps) -->
|
||||
<key>HEADROOM_LLMLINGUA_DEVICE</key>
|
||||
<string>mps</string>
|
||||
</dict>
|
||||
```
|
||||
|
||||
**Note:** LLMLingua requires additional installation:
|
||||
|
||||
```bash
|
||||
pip install headroom-ai[llmlingua]
|
||||
```
|
||||
**Note:** The earlier LLMLingua-2 launch-agent variables
|
||||
(`HEADROOM_COMPRESSION_PROVIDER=llmlingua`, `HEADROOM_LLMLINGUA_DEVICE`,
|
||||
the `headroom-ai[llmlingua]` extra) were retired with the
|
||||
`--llmlingua` flag. For ML compression today, install the `[ml]`
|
||||
extra and follow `wiki/transforms.md`.
|
||||
|
||||
### Crash Recovery
|
||||
|
||||
|
|
@ -683,4 +676,4 @@ A: The LaunchAgent setup is Anthropic-specific. For other providers, see [proxy.
|
|||
|
||||
**Q: Does this work with Apple Silicon (M1/M2/M3)?**
|
||||
|
||||
A: Yes, fully compatible. For LLMLingua compression, use `--llmlingua-device mps` for Apple Silicon acceleration.
|
||||
A: Yes, fully compatible. ML compression (Kompress, opt-in via `headroom-ai[ml]`) auto-detects MPS on Apple Silicon.
|
||||
|
|
|
|||
|
|
@ -303,8 +303,7 @@ curl http://localhost:8787/health
|
|||
{
|
||||
"status": "healthy",
|
||||
"version": "0.1.0",
|
||||
"uptime_seconds": 3600,
|
||||
"llmlingua_enabled": false
|
||||
"uptime_seconds": 3600
|
||||
}
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -117,26 +117,12 @@ headroom proxy --no-intelligent-context
|
|||
headroom proxy --no-intelligent-scoring
|
||||
```
|
||||
|
||||
### LLMLingua Options (ML Compression)
|
||||
### ML Compression — RETIRED `--llmlingua` flag
|
||||
|
||||
| Option | Default | Description |
|
||||
|--------|---------|-------------|
|
||||
| `--llmlingua` | `false` | Enable LLMLingua-2 ML-based compression |
|
||||
| `--llmlingua-device` | `auto` | Device for model: `auto`, `cuda`, `cpu`, `mps` |
|
||||
| `--llmlingua-rate` | `0.3` | Target compression rate (0.3 = keep 30% of tokens) |
|
||||
|
||||
**Note:** LLMLingua requires additional dependencies: `pip install headroom-ai[llmlingua]`
|
||||
|
||||
```bash
|
||||
# Enable LLMLingua with GPU acceleration
|
||||
headroom proxy --llmlingua --llmlingua-device cuda
|
||||
|
||||
# More aggressive compression (keep only 20%)
|
||||
headroom proxy --llmlingua --llmlingua-rate 0.2
|
||||
|
||||
# Conservative compression for code (keep 50%)
|
||||
headroom proxy --llmlingua --llmlingua-rate 0.5
|
||||
```
|
||||
The `--llmlingua` / `--llmlingua-device` / `--llmlingua-rate` flags and
|
||||
the `headroom-ai[llmlingua]` extra were retired and replaced by Kompress
|
||||
(ModernBERT). For the current opt-in path, install `headroom-ai[ml]`
|
||||
and see [transforms.md](transforms.md) and [ARCHITECTURE.md](ARCHITECTURE.md).
|
||||
|
||||
## API Endpoints
|
||||
|
||||
|
|
@ -341,41 +327,14 @@ client = OpenAI(
|
|||
|
||||
## Features
|
||||
|
||||
### LLMLingua ML Compression (Opt-In)
|
||||
### ML Compression (Opt-In, Kompress)
|
||||
|
||||
When enabled, the proxy uses Microsoft's LLMLingua-2 model for ML-based token compression:
|
||||
|
||||
```bash
|
||||
headroom proxy --llmlingua
|
||||
```
|
||||
|
||||
**How it works:**
|
||||
- LLMLinguaCompressor is added to the transform pipeline (before RollingWindow)
|
||||
- Automatically detects content type (JSON, code, text) and adjusts compression
|
||||
- Stores original content in CCR for retrieval if needed
|
||||
|
||||
**Startup feedback:**
|
||||
|
||||
```
|
||||
# When enabled and available:
|
||||
LLMLingua: ENABLED (device=cuda, rate=0.3)
|
||||
|
||||
# When installed but not enabled (helpful hint):
|
||||
LLMLingua: available (enable with --llmlingua for ML compression)
|
||||
|
||||
# When enabled but not installed:
|
||||
WARNING: LLMLingua requested but not installed. Install with: pip install headroom-ai[llmlingua]
|
||||
```
|
||||
|
||||
**Why opt-in?**
|
||||
| Concern | Default Proxy | With LLMLingua |
|
||||
|---------|---------------|----------------|
|
||||
| Dependencies | ~50MB | +2GB (torch, transformers) |
|
||||
| Cold start | <1s | 10-30s (model load) |
|
||||
| Memory | ~100MB | +1GB (model in RAM) |
|
||||
| Overhead | <5ms | 50-200ms per request |
|
||||
|
||||
Enable LLMLingua when maximum compression justifies the resource cost.
|
||||
> The earlier LLMLingua-2 integration documented in this section
|
||||
> (`--llmlingua`, `--llmlingua-device`, `--llmlingua-rate`,
|
||||
> `headroom-ai[llmlingua]`, `LLMLinguaCompressor`) was retired and
|
||||
> replaced by **Kompress** (ModernBERT). Install with `pip install
|
||||
> 'headroom-ai[ml]'`. See [transforms.md](transforms.md) and
|
||||
> [ARCHITECTURE.md](ARCHITECTURE.md) for current configuration.
|
||||
|
||||
### Semantic Caching
|
||||
|
||||
|
|
|
|||
|
|
@ -325,73 +325,16 @@ intelligent_config = IntelligentContextConfig(
|
|||
|
||||
---
|
||||
|
||||
## LLMLinguaCompressor (Optional)
|
||||
## LLMLinguaCompressor — RETIRED
|
||||
|
||||
ML-based compression using Microsoft's LLMLingua-2 model.
|
||||
|
||||
### When to Use
|
||||
|
||||
| Transform | Best For | Speed | Compression |
|
||||
|-----------|----------|-------|-------------|
|
||||
| SmartCrusher | JSON arrays | ~1ms | 70-90% |
|
||||
| Text Utilities | Search/logs | ~1ms | 50-90% |
|
||||
| **LLMLinguaCompressor** | Any text, max compression | 50-200ms | 80-95% |
|
||||
|
||||
### Installation
|
||||
|
||||
```bash
|
||||
pip install "headroom-ai[llmlingua]" # Adds ~2GB
|
||||
```
|
||||
|
||||
### Configuration
|
||||
|
||||
```python
|
||||
from headroom.transforms import LLMLinguaCompressor, LLMLinguaConfig
|
||||
|
||||
config = LLMLinguaConfig(
|
||||
device="auto", # auto, cuda, cpu, mps
|
||||
target_compression_rate=0.3, # Keep 30% of tokens
|
||||
min_tokens_for_compression=100, # Skip small content
|
||||
code_compression_rate=0.4, # Conservative for code
|
||||
json_compression_rate=0.35, # Moderate for JSON
|
||||
text_compression_rate=0.25, # Aggressive for text
|
||||
enable_ccr=True, # Store original for retrieval
|
||||
)
|
||||
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
```
|
||||
|
||||
### Content-Aware Rates
|
||||
|
||||
LLMLinguaCompressor auto-detects content type:
|
||||
|
||||
| Content Type | Default Rate | Behavior |
|
||||
|--------------|--------------|----------|
|
||||
| Code | 0.4 | Conservative - preserves syntax |
|
||||
| JSON | 0.35 | Moderate - keeps structure |
|
||||
| Text | 0.3 | Aggressive - maximum compression |
|
||||
|
||||
### Memory Management
|
||||
|
||||
```python
|
||||
from headroom.transforms import (
|
||||
is_llmlingua_model_loaded,
|
||||
unload_llmlingua_model,
|
||||
)
|
||||
|
||||
# Check if model is loaded
|
||||
print(is_llmlingua_model_loaded()) # True/False
|
||||
|
||||
# Free ~1GB RAM when done
|
||||
unload_llmlingua_model()
|
||||
```
|
||||
|
||||
### Proxy Integration
|
||||
|
||||
```bash
|
||||
# Enable in proxy
|
||||
headroom proxy --llmlingua --llmlingua-device cuda --llmlingua-rate 0.3
|
||||
```
|
||||
The earlier LLMLingua-2 integration (`LLMLinguaCompressor`,
|
||||
`LLMLinguaConfig`, `is_llmlingua_model_loaded`, `unload_llmlingua_model`,
|
||||
the `headroom-ai[llmlingua]` extra, and the `--llmlingua` proxy flag)
|
||||
was retired in 0.9.x and replaced by **Kompress** (ModernBERT).
|
||||
`pip install 'headroom-ai[llmlingua]'` no longer resolves; use the
|
||||
`[ml]` extra instead. The Kompress transform shipped with the proxy
|
||||
runs as Transform 4 in the live-zone pipeline (see
|
||||
[ARCHITECTURE.md](ARCHITECTURE.md)).
|
||||
|
||||
---
|
||||
|
||||
|
|
@ -405,14 +348,14 @@ AST-based compression for source code using tree-sitter.
|
|||
|-----------|----------|-------|-------------|
|
||||
| SmartCrusher | JSON arrays | ~1ms | 70-90% |
|
||||
| **CodeAwareCompressor** | Source code | ~10-50ms | 40-70% |
|
||||
| LLMLinguaCompressor | Any text | 50-200ms | 80-95% |
|
||||
| Kompress (ML) | Any text | 50-200ms | 80-95% |
|
||||
|
||||
### Key Benefits
|
||||
|
||||
- **Syntax validity guaranteed** — Output always parses correctly
|
||||
- **Preserves critical structure** — Imports, signatures, types, error handlers
|
||||
- **Multi-language support** — Python, JavaScript, TypeScript, Go, Rust, Java, C, C++
|
||||
- **Lightweight** — ~50MB vs ~1GB for LLMLingua
|
||||
- **Lightweight** — ~50MB vs ~1GB for the ML compressor
|
||||
|
||||
### Installation
|
||||
|
||||
|
|
@ -436,7 +379,6 @@ config = CodeCompressorConfig(
|
|||
max_body_lines=5, # Lines to keep per function body
|
||||
min_tokens_for_compression=100, # Skip small content
|
||||
language_hint=None, # Auto-detect if None
|
||||
fallback_to_llmlingua=True, # Use LLMLingua for unknown langs
|
||||
)
|
||||
|
||||
compressor = CodeAwareCompressor(config)
|
||||
|
|
@ -553,9 +495,10 @@ print(result.routing_log) # List of routing decisions
|
|||
| SEARCH | Grep/find output | SearchCompressor |
|
||||
| LOG | Log files | LogCompressor |
|
||||
| TEXT | Plain text | TextCompressor |
|
||||
| LLMLINGUA | Any (max compression) | LLMLinguaCompressor |
|
||||
| PASSTHROUGH | Small content | None |
|
||||
|
||||
(The earlier `LLMLINGUA` strategy was retired with the LLMLingua integration; ML compression is now provided by Kompress.)
|
||||
|
||||
### Content Detection
|
||||
|
||||
The router automatically detects content types by analyzing the content itself:
|
||||
|
|
@ -572,7 +515,7 @@ No manual hints required - the router inspects content directly.
|
|||
|
||||
ContentRouter records all compressions to TOIN (Tool Output Intelligence Network) for cross-user learning:
|
||||
|
||||
- **All strategies tracked**: Code, search, logs, text, and LLMLingua compressions are recorded
|
||||
- **All strategies tracked**: Code, search, logs, text, and ML compressions are recorded
|
||||
- **Retrieval feedback**: When users retrieve original content via CCR, TOIN learns which compressions need expansion
|
||||
- **Pattern learning**: TOIN builds signatures for each content type to improve future compressions
|
||||
|
||||
|
|
@ -597,21 +540,9 @@ result = pipeline.transform(messages)
|
|||
print(f"Saved {result.tokens_saved} tokens")
|
||||
```
|
||||
|
||||
### With LLMLingua (Optional)
|
||||
### With ML compression (Optional, Kompress)
|
||||
|
||||
```python
|
||||
from headroom.transforms import (
|
||||
TransformPipeline, SmartCrusher, CacheAligner,
|
||||
RollingWindow, LLMLinguaCompressor
|
||||
)
|
||||
|
||||
pipeline = TransformPipeline([
|
||||
CacheAligner(), # 1. Stabilize prefix
|
||||
SmartCrusher(), # 2. Compress JSON arrays
|
||||
LLMLinguaCompressor(), # 3. ML compression on remaining text
|
||||
RollingWindow(), # 4. Final size constraint (always last)
|
||||
])
|
||||
```
|
||||
The earlier hand-assembled `TransformPipeline([..., LLMLinguaCompressor(), ...])` recipe is no longer supported. ML compression now ships as part of the live-zone pipeline when the `[ml]` extra is installed; see [ARCHITECTURE.md](ARCHITECTURE.md) for the current placement.
|
||||
|
||||
### Recommended Order
|
||||
|
||||
|
|
@ -619,13 +550,13 @@ pipeline = TransformPipeline([
|
|||
|-------|-----------|---------|
|
||||
| 1 | CacheAligner | Stabilize prefix for caching |
|
||||
| 2 | SmartCrusher | Compress JSON tool outputs |
|
||||
| 3 | LLMLinguaCompressor | ML compression (optional) |
|
||||
| 3 | Kompress (ML) | ML compression on remaining text (optional, `[ml]` extra) |
|
||||
| 4 | RollingWindow | Enforce token limits (always last) |
|
||||
|
||||
**Why this order?**
|
||||
- CacheAligner first to maximize prefix stability
|
||||
- SmartCrusher handles JSON arrays efficiently
|
||||
- LLMLingua compresses remaining long text
|
||||
- Kompress compresses remaining long text
|
||||
- RollingWindow truncates only if still over limit
|
||||
|
||||
---
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue